카트리지: 자기 학습을 통한 경량 및 범용 장문맥 표현

초록

대형 언어 모델은 종종 대규모 텍스트 코퍼스(예: 코드베이스, 법률 문서, 채팅 기록 등)를 기반으로 한 질의에 답변하기 위해 사용되며, 이때 전체 코퍼스를 컨텍스트 창에 배치하고 컨텍스트 내 학습(ICL)을 활용합니다. 현재 모델은 100K-1M 토큰의 컨텍스트를 지원하지만, 이 설정은 KV 캐시의 메모리 소비가 입력 길이에 따라 증가하기 때문에 서비스 비용이 높습니다. 우리는 이를 대체할 방법으로, 각 코퍼스에 대해 더 작은 KV 캐시를 오프라인에서 훈련시키는 방식을 탐구합니다. 추론 시에는 이 훈련된 KV 캐시를 로드하여 응답을 디코딩하며, 이를 카트리지(Cartridge)라고 부릅니다. 중요한 점은, 동일한 코퍼스를 참조하는 모든 질의에 대해 카트리지 훈련 비용을 분산시킬 수 있다는 것입니다. 그러나 코퍼스에 대해 다음 토큰 예측으로 카트리지를 훈련하는 단순한 접근법은 ICL과 경쟁력이 없음을 발견했습니다. 대신, 우리는 코퍼스에 대한 합성 대화를 생성하고 컨텍스트 증류 목표로 카트리지를 훈련시키는 자기 학습(self-study)이라는 훈련 방법을 제안합니다. 자기 학습으로 훈련된 카트리지는 ICL의 기능을 복제하면서도 서비스 비용이 훨씬 저렴합니다. 도전적인 장문 컨텍스트 벤치마크에서, 자기 학습으로 훈련된 카트리지는 ICL 성능을 유지하면서 메모리 사용량을 38.6배 줄이고 처리량을 26.4배 높였습니다. 또한 자기 학습은 모델의 유효 컨텍스트 길이를 확장하며(예: MTOB에서 128k에서 484k 토큰으로), 놀랍게도 재훈련 없이 추론 시 카트리지를 조합할 수 있게 합니다.

English

Large language models are often used to answer queries grounded in large text corpora (e.g. codebases, legal documents, or chat histories) by placing the entire corpus in the context window and leveraging in-context learning (ICL). Although current models support contexts of 100K-1M tokens, this setup is costly to serve because the memory consumption of the KV cache scales with input length. We explore an alternative: training a smaller KV cache offline on each corpus. At inference time, we load this trained KV cache, which we call a Cartridge, and decode a response. Critically, the cost of training a Cartridge can be amortized across all the queries referencing the same corpus. However, we find that the naive approach of training the Cartridge with next-token prediction on the corpus is not competitive with ICL. Instead, we propose self-study, a training recipe in which we generate synthetic conversations about the corpus and train the Cartridge with a context-distillation objective. We find that Cartridges trained with self-study replicate the functionality of ICL, while being significantly cheaper to serve. On challenging long-context benchmarks, Cartridges trained with self-study match ICL performance while using 38.6x less memory and enabling 26.4x higher throughput. Self-study also extends the model's effective context length (e.g. from 128k to 484k tokens on MTOB) and surprisingly, leads to Cartridges that can be composed at inference time without retraining.

카트리지: 자기 학습을 통한 경량 및 범용 장문맥 표현

Cartridges: Lightweight and general-purpose long context representations via self-study

초록

Support