Escrevendo nas Margens: Melhor Padrão de Inferência para Recuperação de Contexto Longo

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

Resumo

Neste artigo, apresentamos a Escrita nas Margens (WiM), um novo padrão de inferência para Modelos de Linguagem Grandes projetado para otimizar o tratamento de sequências de entrada longas em tarefas orientadas para recuperação. Esta abordagem aproveita o preenchimento segmentado do cache chave-valor para realizar inferência por segmento, o que possibilita o processamento eficiente de contextos extensos juntamente com a geração e classificação de informações intermediárias ("margens") que orientam o modelo em direção a tarefas específicas. Este método aumenta marginalmente a sobrecarga computacional, ao mesmo tempo em que melhora significativamente o desempenho de modelos prontos para uso sem a necessidade de ajustes finos. Especificamente, observamos que o WiM proporciona um aumento médio de 7,5% na precisão para habilidades de raciocínio (HotpotQA, MultiHop-RAG) e mais de 30,0% no aumento do escore F1 para tarefas de agregação (CWE). Além disso, demonstramos como o padrão proposto se encaixa em um design de recuperação interativa que fornece aos usuários finais atualizações contínuas sobre o progresso do processamento de contexto e destaca a integração de informações relevantes na resposta final. Disponibilizamos nossa implementação do WiM utilizando a biblioteca Hugging Face Transformers em 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.