arXiv:2504.05646cs.LGcs.AI2025-04被引 18

提出新型RNN记忆压缩机制,实现次二次复杂度的高效序列建模。

Lattice: Learning to Efficiently Compress the Memory

  • 利用键值矩阵低秩特性,将缓存压缩为固定数量的记忆槽。
  • 在语言建模与关联回忆任务中,内存占用显著降低且性能更优。
  • 采用正交更新机制,仅存储新信息,减少对已有记忆的干扰。

注意力机制虽革新了序列学习,但存在二次计算复杂度问题。本文提出新型循环神经网络(RNN)机制Lattice,通过挖掘键值(K-V)矩阵的内在低秩结构,将缓存高效压缩为固定数量的记忆槽,实现次二次复杂度。该压缩被形式化为在线优化问题,并基于单步梯度下降推导出动态记忆更新规则。其核心创新在于正交更新:每个记忆槽仅接收与其当前状态正交的信息,确保仅纳入新颖、非冗余数据,从而最小化对已有信息的干扰。我们设计了高效的正交更新计算,并通过分块并行近似进一步保障训练可扩展性。实验表明,Lattice在不同上下文长度与模型规模下,均优于强基线,在语言建模与关联回忆任务中表现出更优的内存效率,且内存大小显著减小。

原文摘要 · Abstract (English)

Attention mechanisms have revolutionized sequence learning but suffer from quadratic computational complexity. This paper introduces \model, a novel recurrent neural network (RNN) mechanism that leverages the inherent low-rank structure of K-V matrices to efficiently compress the cache into a fixed number of memory slots, achieving sub-quadratic complexity. We formulate this compression as an online optimization problem and derive a dynamic memory update rule based on a single gradient descent step. The resulting recurrence features a state- and input-dependent gating mechanism, offering an interpretable memory update process. The core innovation is the orthogonal update: each memory slot is updated exclusively with information orthogonal to its current state, hence incorporating only novel, non-redundant data to minimize interference with previously stored information. We derive an efficient computation for this orthogonal update rule and further approximate it with chunk-wise parallelization to ensure training scalability. Empirically, Lattice outperforms strong baselines on language modeling and associative recall tasks across diverse context lengths and model sizes, achieving superior memory efficiency with significantly reduced memory sizes.

序列建模记忆压缩RNN

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