Scrivere nei margini: Miglioramento dei modelli di inferenza per il recupero di contesti lunghi

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

August 27, 2024

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

cs.AI

Abstract

In questo articolo, presentiamo Writing in the Margins (WiM), un nuovo schema di inferenza per Large Language Models progettato per ottimizzare la gestione di lunghe sequenze di input in compiti orientati al recupero. Questo approccio sfrutta il precaricamento a blocchi della cache chiave-valore per eseguire un'inferezza a livello di segmento, che consente un elaborazione efficiente di contesti estesi insieme alla generazione e classificazione di informazioni intermedie ("margini") che guidano il modello verso compiti specifici. Questo metodo aumenta marginalmente l'onere computazionale mentre migliora significativamente le prestazioni dei modelli standard senza la necessità di un raffinamento. In particolare, osserviamo che WiM fornisce un miglioramento medio del 7,5% in accuratezza per le abilità di ragionamento (HotpotQA, MultiHop-RAG) e oltre un aumento del 30,0% nel punteggio F1 per i compiti di aggregazione (CWE). Inoltre, mostriamo come lo schema proposto si inserisca in un design di recupero interattivo che fornisce agli utenti finali aggiornamenti continui sul progresso dell'elaborazione del contesto e individua l'integrazione delle informazioni rilevanti nella risposta finale. Rilasciamo la nostra implementazione di WiM utilizzando la libreria Hugging Face Transformers su 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.