Cartridges:通过自学实现的轻量级通用长上下文表示
计算与语言
2025-06-16 v3 人工智能
机器学习
摘要
大型语言模型常通过将整个语料库(例如代码库、法律文件或聊天记录)放入上下文窗口并利用情境学习(in-context learning, ICL)来回答查询。虽然当前模型支持10万至100万 token 的上下文,但此 setup 代价高昂,因为KV 缓存的内存消耗随输入长度而扩展。我们探索了另一种方法:在每个语料库上离线训练较小的KV缓存。在推理时,我们加载训练好的KV缓存,称为 Cartridge。关键在于,训练 Cartridge 的成本可以跨所有引用同一语料库的查询进行摊销。然而,我们发现仅在语料库上进行下一标记预测训练 Cartridge的做法并不具备与ICL相竞争的力量。相反,我们提出了自学(self-study)方法,即生成关于语料库的合成对话,并以情境蒸馏目标训练 Cartridge。我们发现,采用自学训练的 Cartridge能够复制ICL的功能,同时服务成本显著降低。在具有挑战性的长情境基准测试中,采用自学训练的Cartridge在使用38.6倍更少内存的情况下匹配ICL性能,并实现26.4倍的吞吐量提升。自学还能够扩展模型的有效情境长度(例如将MTOB上的情境长度从128k提升至484k),令人惊讶的是,训练好的Cartridge在推理时可以无需重新训练地进行组合。
引用
@article{arxiv.2506.06266,
title = {Cartridges: Lightweight and general-purpose long context representations via self-study},
author = {Sabri Eyuboglu and Ryan Ehrlich and Simran Arora and Neel Guha and Dylan Zinsley and Emily Liu and Will Tennien and Atri Rudra and James Zou and Azalia Mirhoseini and Christopher Re},
journal= {arXiv preprint arXiv:2506.06266},
year = {2025}
}