瓶颈变换器:周期性KV缓存整合的通用推理
机器学习
2026-03-26 v4 人工智能
信息论
math.IT
摘要
变换器大语言模型(LLM)被证明在推理能力方面表现出强大的推理能力,这种能力随推理时间计算量而扩展,最突出体现在token空间的“思考”链式反应。日益增长的工作将额外的计算推送到模型的潜在空间,我们称之为辅助潜在空间计算(ALSC)。现有的ALSC方法大致分为三类:token中介的潜在滚动,残差/激活引导,以及内存(KV)压缩。一种较少探索的替代方案是记忆整合/重巩固,这两种大脑过程负责稳定新形成的记忆痕迹,并在回想时暂时使 established traces 具有可塑性,以整合新的上下文信息后再稳定。 在变换器LLM中,这可以看到对新KV片段的原地重写,以及对被调召回的过去片段的重写。在本工作中,我们通过信息瓶颈(IB)理论提供了记忆(重新)整合通过KV缓存重写对改进推理有益的理论依据。该理论认为,模型的通用性从输入信息压缩与保持预测信息在潜在表示之间的最佳平衡中产生。我们引入瓶颈变换器,通过在换行符分隔的推理步骤边界上执行周期性、非因果的KV重写来增强主干LLM。处理器整合最近写入的KV条目,并对小的、基于注意力选择的top-k先前条目进行重巩固。我们在数学推理基准测试中评估了瓶颈变换器架构。我们的模型在vanilla变换器和暂停标记增强的基线上看到持续的性能提升,提升幅度最高可达+6.6pp。
引用
@article{arxiv.2505.16950,
title = {Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning},
author = {Adnan Oomerjee and Zafeirios Fountas and Haitham Bou-Ammar and Jun Wang},
journal= {arXiv preprint arXiv:2505.16950},
year = {2026}
}