Escribiendo en los Márgenes: Mejor Patrón de Inferencia para la Recuperación de Contextos Largos

Writing in the Margins: Better Inference Pattern for Long Context Retrieval

August 27, 2024

Autores: Melisa Russak, Umar Jamil, Christopher Bryant, Kiran Kamble, Axel Magnuson, Mateusz Russak, Waseem AlShikh

cs.AI

Resumen

En este artículo, presentamos Writing in the Margins (WiM), un nuevo patrón de inferencia para Modelos de Lenguaje Grandes diseñado para optimizar el manejo de secuencias de entrada largas en tareas orientadas a la recuperación. Este enfoque aprovecha el precargado segmentado de la caché de clave-valor para realizar inferencia por segmentos, lo que permite el procesamiento eficiente de contextos extensos junto con la generación y clasificación de información intermedia ("márgenes") que guían al modelo hacia tareas específicas. Este método aumenta mínimamente la carga computacional mientras mejora significativamente el rendimiento de modelos listos para usar sin necesidad de ajustes finos. Específicamente, observamos que WiM proporciona un aumento promedio del 7.5% en precisión para habilidades de razonamiento (HotpotQA, MultiHop-RAG) y más de un 30.0% en el puntaje F1 para tareas de agregación (CWE). Además, mostramos cómo el patrón propuesto encaja en un diseño interactivo de recuperación que brinda a los usuarios finales actualizaciones continuas sobre el progreso del procesamiento del contexto, y señala la integración de información relevante en la respuesta final. Publicamos nuestra implementación de WiM utilizando la biblioteca Hugging Face Transformers en https://github.com/writer/writing-in-the-margins.

English

In this paper, we introduce Writing in the Margins (WiM), a new inference pattern for Large Language Models designed to optimize the handling of long input sequences in retrieval-oriented tasks. This approach leverages the chunked prefill of the key-value cache to perform segment-wise inference, which enables efficient processing of extensive contexts along with the generation and classification of intermediate information ("margins") that guide the model towards specific tasks. This method increases computational overhead marginally while significantly enhancing the performance of off-the-shelf models without the need for fine-tuning. Specifically, we observe that WiM provides an average enhancement of 7.5% in accuracy for reasoning skills (HotpotQA, MultiHop-RAG) and more than a 30.0% increase in the F1-score for aggregation tasks (CWE). Additionally, we show how the proposed pattern fits into an interactive retrieval design that provides end-users with ongoing updates about the progress of context processing, and pinpoints the integration of relevant information into the final response. We release our implementation of WiM using Hugging Face Transformers library at https://github.com/writer/writing-in-the-margins.