Stochastic KV Routing：实现自适应深度缓存共享
2026/05/05 08:00阅读量 3
苹果机器学习研究团队提出随机KV路由方法，通过训练时随机选择层间注意力，使大语言模型能够在深度维度上共享KV缓存，从而显著降低推理时的内存占用，同时保持或提升模型性能。
事件概述
Transformer语言模型在自回归生成中依赖KV缓存来避免重复计算，但其内存占用显著影响服务成本。现有研究多从时间维度进行压缩或淘汰，而苹果团队提出从深度维度进行优化——通过层间缓存共享来减少内存需求。
核心方法
提出“随机跨层注意力”训练方案：在训练过程中，每一层随机选择关注自身的KV状态或前一层（或更早层）的KV状态。这种随机过程使模型对不同的深度缓存共享策略具有鲁棒性，部署时可根据硬件约束灵活调整。
评估结果
该方法在预训练或微调阶段均可应用，适用于多种模型家族。对于数据受限环境下的更大模型，该方法具有类似正则化的效果，在显著降低缓存内存占用的同时，往往能够保持甚至提升性能。