Écriture dans les marges : Meilleur motif d'inférence pour la récupération de longs contextes

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

Résumé

Dans cet article, nous présentons Writing in the Margins (WiM), un nouveau schéma d'inférence pour les grands modèles de langage conçu pour optimiser le traitement de longues séquences d'entrée dans des tâches orientées vers la récupération. Cette approche exploite le pré-remplissage segmenté du cache clé-valeur pour effectuer une inférence par segment, ce qui permet un traitement efficace de contextes étendus ainsi que la génération et la classification d'informations intermédiaires ("marges") qui guident le modèle vers des tâches spécifiques. Cette méthode augmente légèrement la charge de calcul tout en améliorant significativement les performances des modèles prêts à l'emploi sans nécessiter de fine-tuning. Plus précisément, nous observons que WiM apporte une amélioration moyenne de 7,5 % en termes de précision pour les compétences de raisonnement (HotpotQA, MultiHop-RAG) et une augmentation de plus de 30,0 % du score F1 pour les tâches d'agrégation (CWE). De plus, nous montrons comment le schéma proposé s'intègre dans une conception interactive de récupération qui fournit aux utilisateurs finaux des mises à jour continues sur l'avancement du traitement du contexte, et souligne l'intégration d'informations pertinentes dans la réponse finale. Nous mettons à disposition notre implémentation de WiM en utilisant la bibliothèque Hugging Face Transformers sur 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.