用自学习训练轻量级上下文卡片,大幅降低长文本推理成本
Cartridges: Lightweight and general-purpose long context representations via self-study
- 通过自学习生成合成对话,训练小型上下文缓存(Cartridge)
- 在长文本任务上性能媲美原生模型,内存减少38.6倍,吞吐提升26.4倍
- 支持推理时组合使用,有效上下文长度可扩展至484k tokens
大型语言模型常通过将整个文本语料(如代码库、法律文档或聊天历史)放入上下文窗口,并利用上下文学习(ICL)来回答问题。尽管当前模型支持10万至100万词元的上下文,但服务成本高昂,因为键值缓存(KV cache)内存随输入长度线性增长。本文提出一种替代方案:预先在每个语料上离线训练一个更小的KV缓存,推理时加载该缓存(称为Cartridge)进行解码。关键优势在于,训练成本可被同一语料的所有查询分摊。然而,直接用语料进行下一项预测训练的朴素方法不如ICL。因此,本文提出自学习(self-study)训练策略:生成关于语料的合成对话,并以上下文蒸馏为目标训练Cartridge。实验表明,自学习训练的Cartridge能复现ICL功能,且服务成本显著降低。在多个长上下文基准测试中,其性能与ICL相当,同时内存占用减少38.6倍,吞吐量提升26.4倍。自学习还使模型有效上下文长度从128k扩展至484k(MTOB数据集),并意外发现推理时无需重训即可组合多个Cartridge。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。