Adaptive Memory Decay for Log-Linear Attention
作者: Yaxita Amin, Helen Zichen Li, Mengfan Zhang, Samet Ayhan
分类: cs.LG, cs.AI
发布日期: 2026-05-07
备注: 19 pages, 13 figures. Preprint
💡 一句话要点
提出自适应记忆衰减机制，优化对数线性注意力模型的长程上下文建模能力
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 对数线性注意力 长程依赖 自适应记忆 序列建模 Fenwick树 上下文压缩
📋 核心要点
- 现有对数线性注意力模型采用固定的衰减参数λ，无法根据输入内容动态调整记忆权重，导致在处理复杂序列时记忆分配僵化。
- 提出一种基于轻量级MLP的自适应衰减机制，根据输入内容实时生成逐Token、逐层级的衰减因子，实现对记忆的动态控制。
- 实验表明，该方法在关联召回、选择性复制及语言建模任务中均优于基线模型，尤其在长程记忆场景下表现出显著的性能提升。
📝 摘要（中文）
序列模型在记忆容量与计算效率之间面临权衡。Transformer虽具备强大的上下文建模能力，但计算复杂度呈二次方增长；线性注意力和状态空间模型通过将上下文压缩为固定大小的隐藏状态实现线性时间复杂度，却限制了长程信息的召回。对数线性注意力机制通过Fenwick树层次结构组织记忆，使隐藏状态随序列长度对数级增长，计算复杂度为对数线性。然而，其固定的衰减参数λ忽略了输入内容差异，导致记忆分配僵化。本文提出通过轻量级双层MLP从输入中学习λ，实现针对不同Token和不同层级的自适应衰减。利用Softplus激活函数避免了层间竞争，在保持对数线性复杂度的同时，显著提升了长程记忆任务的性能。
🔬 方法详解
问题定义：现有对数线性注意力模型依赖固定的衰减参数λ，这种全局统一的权重分配方式无法捕捉序列中不同信息的重要性差异，限制了模型在长程依赖任务中的召回精度。
核心思路：引入输入依赖的动态衰减机制。通过感知输入内容，让模型能够根据当前Token的语义特征，灵活调整Fenwick树各层级的记忆保留程度，从而实现对关键信息的选择性存储。
技术框架：模型在原有的对数线性注意力架构基础上，增加了一个轻量级的双层MLP模块。该模块以当前输入Token的嵌入向量作为输入，输出对应于Fenwick树各层级的衰减参数λ。
关键创新：核心创新在于将静态的超参数转化为动态的预测变量。通过引入Softplus激活函数替代Softmax，使得各层级衰减因子的学习相互独立，有效消除了层间竞争，确保了模型对不同时间尺度信息的捕捉能力。
关键设计：采用双层MLP进行参数映射，保证了计算开销的极小化。通过Softplus函数确保衰减因子为正，并结合Fenwick树的层次结构，在保持对数线性计算复杂度的前提下，实现了对记忆衰减的精细化控制。
🖼️ 关键图片
📊 实验亮点
在关联召回、选择性复制及语言建模等基准测试中，自适应衰减机制均表现出优于固定衰减基线的性能。特别是在长程记忆任务中，当基线模型的衰减参数因序列过长而失效或崩溃时，该方法仍能保持稳定的召回精度，证明了其在处理极端长序列时的鲁棒性与有效性。
🎯 应用场景
该技术适用于长文档理解、长序列时间序列预测及大规模语言模型推理等场景。通过优化长程记忆的存储效率，能够显著提升模型在处理超长上下文时的信息召回能力，降低计算资源消耗，在金融分析、法律文档审查及长程对话系统等领域具有广阔的应用前景。
📄 摘要（原文）
Sequence models face a fundamental tradeoff between memory capacity and computational efficiency. Transformers achieve expressive context modeling at quadratic cost, while linear attention and state-space models run in linear time by compressing context into a fixed-size hidden state, inherently limiting recall. Log-linear attention navigates this tradeoff by organizing memory across a Fenwick tree hierarchy, growing its hidden state logarithmically with sequence length at log-linear compute cost. However, its memory decay parameter λ is fixed and independent of the input, assigning uniform weights across all hierarchy levels regardless of the content, which introduces unnecessary rigidity. We propose learning λ directly from the input via a lightweight two-layer MLP, producing per-token, per-level decay that adapts to content rather than position. A softplus activation lets each Fenwick tree level scale independently, avoiding the inter-level competition that softmax introduces. This modification preserves log-linear complexity exactly and adds negligible parameter overhead. We evaluate on associative recall, selective copying, and language modeling, finding that input-dependent decay consistently outperforms the baseline, with the largest gains in long-range memory settings where baseline λ degrades or collapses entirely.