Записи в полях: Улучшенный шаблон вывода для извлечения длинного контекста

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

August 27, 2024

Авторы: Melisa Russak, Umar Jamil, Christopher Bryant, Kiran Kamble, Axel Magnuson, Mateusz Russak, Waseem AlShikh

cs.AI

Аннотация

В данной статье мы представляем Writing in the Margins (WiM), новый шаблон вывода для крупных языковых моделей, разработанный для оптимизации обработки длинных входных последовательностей в задачах, ориентированных на извлечение информации. Этот подход использует сегментированное предзаполнение кэша ключ-значение для выполнения вывода по сегментам, что обеспечивает эффективную обработку обширных контекстов вместе с генерацией и классификацией промежуточной информации ("маржей"), направляющей модель на выполнение конкретных задач. Этот метод незначительно увеличивает вычислительную нагрузку, значительно повышая производительность стандартных моделей без необходимости тонкой настройки. В частности, мы наблюдаем, что WiM обеспечивает в среднем улучшение точности на 7.5% для навыков рассуждения (HotpotQA, MultiHop-RAG) и более чем на 30.0% увеличение F1-оценки для задач агрегации (CWE). Кроме того, мы показываем, как предложенный шаблон вписывается в интерактивный дизайн извлечения, предоставляя конечным пользователям текущие обновления о ходе обработки контекста и выделяя интеграцию соответствующей информации в окончательный ответ. Мы предоставляем нашу реализацию WiM с использованием библиотеки Hugging Face Transformers по адресу 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.