arXiv:2505.16950cs.LGcs.AI2025-05被引 4

通过周期性重写记忆键值对,提升大模型的推理能力。

Bottlenecked Transformers: Periodic KV Cache Consolidation for Generalised Reasoning

  • 用辅助模块在推理时定期重写键值缓存,模拟大脑记忆巩固机制。
  • 在数学推理任务上最高提升6.6个百分点,优于普通Transformer和加暂停令牌的方法。
  • 适合研究长序列推理、记忆机制或想优化大模型思考过程的研究者。

Transformer 类大语言模型展现出随推理计算量增长而增强的推理能力,主要体现在基于令牌空间的“思维链”中。现有方法多集中在潜空间计算(ALSC),包括令牌引导的潜空间演化、残差/激活调控和记忆(KV)压缩。本文提出一种未被充分探索的方向:记忆巩固与再巩固。受大脑记忆机制启发,该方法在推理过程中对新生成和已回忆的键值段进行原位重写。本文从信息瓶颈理论出发,论证了通过键值缓存重写实现记忆再巩固有助于提升模型泛化能力。为此,我们提出瓶颈化Transformer(Bottlenecked Transformer),其在骨干模型外增设一个缓存处理器,该处理器在换行标记处周期性地执行非因果、原位的键值重写,合并近期条目并选择前k个最相关的旧条目进行再巩固。我们在数学推理基准上评估该架构,结果表明,相比基础Transformer和暂停令牌增强基线,性能持续提升,在部分任务/模型上最高达+6.6个百分点。

原文摘要 · Abstract (English)

Transformer LLMs have been shown to exhibit strong reasoning ability that scales with inference-time compute, most prominently through token-space "thinking" chains of thought. A growing line of work pushes extra computation into the model's latent space, which we term Auxiliary Latent-Space Computation (ALSC). Existing ALSC methods largely fall into three buckets: (i) token-mediated latent rollouts, (ii) residual/activation steering, and (iii) memory (KV) compression. An underexplored alternative is memory consolidation/reconsolidation, two processes in the brain that are responsible for stabilising newly formed memory traces, and, upon recall, transiently rendering established traces plastic such they can integrate new contextual information before restabilising. In Transformer LLMs, this can be seen as analogous to performing in-place rewrites of new KV segments, and rewrites of recalled past segments. In this work, we give a theoretical justification as to why memory (re)consolidation via KV cache rewrites is beneficial for improved reasoning. We do this through the lens of Information Bottleneck (IB) theory, which posits that model generalisation emerges from an optimal balance between input information compression and retention of predictive information in latent representations. We then introduce the Bottlenecked Transformer, which augments a backbone LLM with a Cache Processor, an auxiliary Transformer that performs periodic, non-causal, in-place KV rewrites at newline-delimited reasoning step boundaries. The Processor consolidates recently written KV entries and reconsolidates a small, top-k attention-selected set of prior entries. We evaluate our Bottlenecked Transformer architecture on math reasoning benchmarks. Our model sees consistent performance gains over vanilla Transformers and pause-token augmented baselines, with gains of up to +6.6pp for selected tasks/backbones.

推理增强记忆机制键值缓存信息瓶颈

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。