여백에 쓰기: 긴 맥락 검색을 위한 더 나은 추론 패턴

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)이며 집계 작업에 대한 F1 점수가 30.0% 이상 향상되는 것을 관찰합니다 (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.