⏶5
Cartridges: 通过自学实现轻量级通用长上下文表示
发表
由
Sabri Eyuboglu 提交
作者:
Sabri Eyuboglu, Ryan Ehrlich, Simran Arora, Neel Guha, Dylan Zinsley, Emily Liu, Will Tennien, Atri Rudra, James Zou, Azalia Mirhoseini, Christopher Re
摘要
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.
当我们向语言模型中放入大量文本(例如整个代码仓库)作为上下文时,由于KV缓存的大小,生成成本会急剧上升。如果我们在离线状态下为我们的文档训练一个更小的KV缓存呢?通过一种名为“自学习”的测试时训练方法,我们表明这个简单的想法可以在保持质量的同时将吞吐量提高26倍。