Schrijven in de marges: Een beter inferentiepatroon voor het ophalen van lange contexten

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

August 27, 2024

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

cs.AI

Samenvatting

In dit artikel introduceren we Writing in the Margins (WiM), een nieuw inferentiepatroon voor Large Language Models dat is ontworpen om de verwerking van lange invoerreeksen in retrievalgerichte taken te optimaliseren. Deze benadering maakt gebruik van de chunked prefill van de key-value cache om segmentgewijze inferentie uit te voeren, wat efficiënte verwerking van uitgebreide contexten mogelijk maakt, samen met de generatie en classificatie van tussentijdse informatie ("margins") die het model naar specifieke taken leidt. Deze methode verhoogt de rekenkosten marginaal, terwijl de prestaties van standaardmodellen aanzienlijk worden verbeterd zonder dat fine-tuning nodig is. Specifiek observeren we dat WiM een gemiddelde verbetering van 7,5% in nauwkeurigheid biedt voor redeneervaardigheden (HotpotQA, MultiHop-RAG) en een toename van meer dan 30,0% in de F1-score voor aggregatietaken (CWE). Daarnaast laten we zien hoe het voorgestelde patroon past in een interactief retrievalontwerp dat eindgebruikers voortdurend op de hoogte houdt van de voortgang van contextverwerking, en de integratie van relevante informatie in het uiteindelijke antwoord aanwijst. We hebben onze implementatie van WiM vrijgegeven met behulp van de Hugging Face Transformers-bibliotheek op 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.