动态线性注意力 — HF Daily Papers

基于信息感知的自适应状态合并，提升线性注意力的长上下文建模能力

前置知识

标准自注意力

Transformer 中的核心机制，计算查询 Q、键 K、值 V 三者的关系：$O = \text{softmax}(QK^\top \odot M)V$，其中 M 是因果掩码。其时间和内存复杂度为 $O(T^2)$，T 是序列长度，这使得它在处理长文本时计算开销巨大

论文的核心动机就是解决标准注意力的二次复杂度瓶颈，理解它是理解线性注意力改进的基础

线性注意力

通过移除 softmax 归一化来降低复杂度的注意力机制，形式为 $O = (QK^\top \odot M)V$。它可以通过结合律重写为递归形式：$S_t = S_{t-1} + v_t k_t^\top$，$o_t = S_t q_t$，从而实现 $O(T)$ 的时间复杂度和常量内存

DLA 是在线性注意力基础上的改进，需要先理解其基本原理和实现方式

多状态线性注意力

将历史上下文分段并总结为多个记忆状态的线性注意力扩展方法，如 Log-Linear Attention 使用 Fenwick 树分解将前缀分成最多 $L = \lceil \log_2(t+1) \rceil + 1$ 个不同的桶，每个桶对应一个状态 $S^{(\ell)}_t = \sum_{s \in B^{(\ell)}_t} v_s k_s^\top$

这是本文要改进的现有方法，理解其固定分层调度的局限性是理解 DLA 优势的关键

状态信息分数

DLA 引入的度量，衡量当前 token 相对于最近记忆状态的信息变化量：$I_t = \frac{\|S_t - S_{t-1}\|_F}{\|S_{t-1}\|_F + \epsilon}$，其中 $S_t = \phi(k_t) v_t^\top$ 是新 token 的状态，$\|\cdot\|_F$ 是 Frobenius 范数

这是 DLA 信息感知动态状态合并的核心机制，决定了何时创建新状态

研究动机

大语言模型扩展到长上下文设置面临根本性挑战，标准自注意力的二次计算和内存复杂度 $O(T^2)$ 严重限制了其在长序列上的应用。线性注意力通过移除 softmax 归一化实现了 $O(T)$ 复杂度，但现有的多状态线性注意力方法（如 Log-Linear Attention）存在关键缺陷：它们依赖固定的状态合并策略（如每 K 个 token 分块合并）或基于规则的边界，无法适应动态变化的 token 重要性。这种内容无关的固定策略会不可逆地将关键 token 吸收到粗粒度摘要中，特别是在语义转换点附近的 token 被过早合并，导致误差在长序列上累积。实验显示，随着上下文长度增加，现有方法的性能出现明显下降

本文的目标是本文提出 Dynamic Linear Attention（DLA），目标是构建一个信息感知的线性注意力框架，通过自适应地确定状态粒度来改进长上下文下的线性注意力。DLA 旨在实现两个目标：一是基于 token 级别的信息变化动态确定状态边界，在语义转换点附近保持高分辨率表示，而在稳定区域积极进行摘要；二是通过容量受限的内存建模维护固定大小的按时间顺序排列的状态缓存，确保预测的计算和内存成本

与已有工作不同的是，本文的独特切入角度是从理论角度证明固定合并策略在非平稳序列上的次优性，并提出基于信息感知的动态状态合并方案。与 Log-Linear Attention 等方法依赖固定的时间分层调度不同，DLA 在线评估每个 incoming token 的表示变化，仅在信息变化量小于阈值时才合并，从而近似最小化总结偏差的主要项。这种贪婪在线策略能够在保持计算效率的同时，更好地适应长序列中非均匀的信息密度分布

核心方法

DLA 的整体思路是通过信息感知的动态状态构建和容量受限的内存建模来实现自适应的线性注意力。在推理过程中，DLA 按顺序处理 token，对每个新 token 计算轻量级的状态信息分数来衡量其相对于最近记忆状态的信息变化量。信息变化量低的 token 被合并到当前状态，而表现出显著漂移的 token 则启动新状态。这种机制在语义转换点周围实现细粒度建模，同时在稳定 token 跨度上积极进行摘要。为了限制内存和计算，DLA 维护容量受限的状态缓存，当缓存达到最大大小时，选择性地合并相邻的低信息状态，在保持时间顺序的同时最小化信息损失

DLA 的核心创新点是信息感知的动态状态合并机制，它与传统多状态线性注意力方法的本质区别在于：传统方法（如 Log-Linear Attention）使用固定的分层调度或预定义的块边界来确定哪些历史 token 应该被合并到摘要状态中，这种设计在很大程度上与序列的语义演化无关。而 DLA 通过计算状态信息分数 $I_t$ 在线评估表示漂移，基于 token 级别的信息变化自适应地确定状态边界。具体来说，DLA 使用边界指示器 $b_t = \mathbb{1}[I_t \geq \tau]$ 决定是否启动新状态，其中 $\tau$ 是可学习的阈值。这种信息感知的机制使得 DLA 能够在非平稳序列中更好地适应动态出现的语义转换，避免过早将重要转换吸收到粗粒度摘要中

方法步骤详情

DLA 的完整工作流程包含两个主要步骤。首先是信息感知的动态状态合并：对于每个新 token t，计算其状态 $S_t = \phi(k_t) v_t^\top$，然后计算状态信息分数 $I_t = \frac{\|S_t - S_{t-1}\|_F}{\|S_{t-1}\|_F + \epsilon}$，其中 $\|\cdot\|_F$ 是 Frobenius 范数，$\epsilon$ 是数值稳定常数。如果 $I_t \geq \tau$，则启动新状态 $S_{cur}^t = S_t$ 并追加到记忆缓存；否则将新 token 合并到当前状态 $S_{cur}^t = S_{cur}^{t-1} + S_t$。在预训练期间使用软门控使决策可微，推理时切换到硬分段策略。其次是容量受限的内存建模：维护状态缓存 $M = \{(S_i, n_i, \bar{I}_i)\}_{i=1}^m$，其中 $S_i$ 是第 i 个记忆状态，$n_i$ 是该状态总结的 token 数量，$\bar{I}_i$ 是聚合信息分数。当缓存未满时直接插入新状态；当缓存达到容量 K 时，触发压缩步骤，选择信息密度最低的相邻状态对 $(i^\star, i^\star+1) = \arg\min_i \frac{\bar{I}_i + \bar{I}_{i+1}}{n_i + n_{i+1}}$ 进行合并。最后在解码步骤 t，通过线性注意力形式计算输出 $o_t = \sum_{i=1}^m \lambda_{t,i} q_t^\top S_i$，其中 $\lambda_{t,i}$ 是预训练期间学习的查询相关权重

技术新颖性

DLA 的技术新颖性体现在三个方面：一是理论层面，本文提供了固定合并策略次优性的严格证明。定理 3.1 证明了块级总结引起的偏差由块内异质性控制，其上界为 $B(\pi; q) = \|q\|_2 \sum_i \sqrt{|C_i|} \sqrt{\sum_{t \in C_i} \|u_t - \bar{u}_i\|_2^2}$。推论 3.2 进一步证明，对于非平稳序列，任何固定策略的偏差边界都严格大于与语义变化点对齐的自适应策略。二是算法层面，DLA 提出的状态信息分数 $I_t$ 提供了一种轻量级的信息变化度量，能够在线评估表示漂移而无需额外的可学习参数。三是工程层面，容量受限的内存建模通过选择性合并相邻低信息状态，在保持时间顺序的同时实现稳定的推理成本，这种设计在实际部署中具有重要意义

**Figure 2**: Standard linear attention (left) vs. log-linear attention (mid) vs. dynamic linear attention (right)

实验结果

DLA 在 16 个数据集上的全面评估显示了其显著优势。在常识推理任务上，DLA 在两个骨干网络上都一致性地超越基线方法。对于 Mamba-2-780M，DLA 在 LAMBADA、PIQA、HellaSwag、WinoGrande、ARC-e、ARC-c、OpenBookQA 和 CommonsenseQA 八个数据集上的平均准确率达到 34.2%，相比 Log-Linear 变体的 31.0% 提升了 8%，相比基线 Mamba-2 的 31.8% 提升了 7.6%。值得注意的是，DLA 在 Mamba-2 上的表现甚至超越了参数量相当的完整注意力 Transformer（32.9%），证明了 DLA 能够缩小甚至超越线性注意力与全注意力之间的准确率差距。对于 Gated DeltaNet-1.3B，DLA 的平均准确率达到 34.8%，相比 Log-Linear 变体的 32.6% 提升了 6%。在上下文检索任务上，DLA 表现出更显著的提升，Mamba-2 变体在 SQuAD、TriviaQA、SWDE、FDA、NQ 和 Drop 六个数据集上的平均准确率从 6.3% 提升到 13.7%，相对提升 44%。在长上下文建模任务上，DLA 在 RULER 的单针任务（S-NIAH-1/2/3）和多针任务（MK-NIAH-1、MQ-NIAH、MV-NIAH）上表现出色，特别是 Mamba-2 在 S-NIAH-3 任务上从 10.6% 提升到 37.1%，相对提升高达 250%。在 LongBench 上，DLA 在单文档问答、多文档问答、摘要和少样本学习等多个子任务上都取得了一致的性能提升

**Table 1**: Performance comparison of DLA and baseline methods on zero-shot commonsense reasoning tasks

**Table 2**: Performance on in-context retrieval benchmarks

**Table 3**: Evaluation results of single-needle tasks (S-NIAH-1–3) and multi-needle tasks (MK-1, MQ, MV) on RULER (4K context)

**Table 4**: Performance on LongBench datasets with different types of tasks

**Table 5**: Ablation study of Mamba-2 and Gated DeltaNet with different variants

**Table 6**: Ablation study of Mamba-2 DLA variant with different memory budget k and merge boundary $\tau$

**Figure 3**: Throughput and runtime memory consumption of vanilla, Log-Linear, and DLA variants of Mamba-2 (780M)

查看结构化数据

任务	指标	本文	基线	提升
常识推理（8个数据集平均）	准确率	34.2% (Mamba-2), 34.8% (Gated DeltaNet)	31.0% (Mamba-2 Log-Linear), 32.6% (Gated DeltaNet Log-Linear)	8% (Mamba-2), 6% (Gated DeltaNet)
上下文检索（6个数据集平均）	准确率	13.7% (Mamba-2)	6.3% (Mamba-2 Log-Linear)	44% 相对提升
RULER S-NIAH-3（多针检索）	准确率	37.1% (Mamba-2)	10.6% (Mamba-2 Log-Linear)	250% 相对提升
RULER MQ-NIAH	准确率	42.5% (Mamba-2)	25.2% (Mamba-2 Log-Linear)	67% 相对提升
LongBench 单文档问答	EM/F1分数	11.1 (Mamba-2)	6.8 (Mamba-2 Log-Linear)	32% 相对提升

局限与改进

作者在论文中承认了一些局限性。首先是超参数敏感性分析显示，虽然 DLA 在一定范围内对内存预算 k 和合并边界 $\tau$ 具有鲁棒性，但极端值可能会影响性能。实验表明，当 k 从 30 降低到 20 或增加到 40 时，性能变化相对较小，但这并不意味着在所有场景下都能保持这种鲁棒性。其次是理论分析假设 token 的贡献可以独立表示为加性项，而在实际模型中，多层注意力和非线性激活可能会引入更复杂的交互，这可能导致理论偏差边界在实际场景中不够紧。此外，从技术观察来看，DLA 的信息感知机制依赖于状态信息分数 $I_t$ 的计算，这需要额外的矩阵范数运算，虽然成本相对较低，但在极端长序列（如 1M token）场景下可能会成为瓶颈。最后，DLA 的容量受限设计要求预定义固定的状态缓存大小 K，这在不同任务和数据分布下可能需要重新调优

独立分析的弱点

DLA 的一个潜在弱点是其在极端非平稳序列上的表现。虽然理论上 DLA 的贪婪在线策略能够近似最小化偏差边界，但在实践中，这种策略可能过于局部化，无法捕捉全局最优的分割方案。例如，在某些情况下，早期的一个错误合并决策可能会在后续阶段被放大，而 DLA 的不可逆合并机制无法纠正这种错误。一个改进方向是引入回溯机制或两阶段策略：第一阶段使用贪心的在线算法生成候选状态分割，第二阶段使用更昂贵的离线优化（如动态规划）重新调整状态边界。另一个弱点是 DLA 的状态信息分数仅依赖于当前 token 与最近状态的表示变化，而没有考虑更远期历史的信息模式。在某些场景下，一个 token 的重要性可能需要通过更长期的上下文来判断。改进方向可以引入多尺度信息分数，考虑不同时间窗口内的表示变化。此外，DLA 的容量约束通过合并相邻低信息状态实现，这在某些语义组织上可能不够灵活。例如，两个低信息状态可能对应不同的语义主题，合并它们可能会破坏语义连贯性。一个改进方向是引入语义感知的合并策略，在决定合并时不仅考虑信息密度，还考虑语义相似性

未来方向

作者没有明确提及具体的未来工作方向，但基于论文的贡献和当前大语言模型的发展趋势，可以延伸出多个研究方向。首先是探索 DLA 在更大规模模型上的效果，当前实验在 780M 和 1.3B 参数规模上进行，而在 7B、13B 或更大的模型上，DLA 的性能优势和效率收益可能会有不同的表现。其次是研究 DLA 与其他高效注意力机制的组合，如将 DLA 的信息感知机制与稀疏注意力或局部注意力结合，以进一步提高长上下文建模的效率和质量。另一个有前途的方向是将 DLA 应用于多模态场景，如图像-文本建模或视频理解，这些场景中信息的时空分布可能更加复杂，信息感知的动态状态构建可能会带来更大的优势。此外，从理论角度，可以进一步研究 DLA 的理论性质，如收敛性分析、稳定性分析和在更广泛的非平稳序列类上的理论保证。最后，DLA 的动态内存管理思想可以扩展到其他架构，如将状态信息分数的概念应用于 Transformer 的 KV cache 管理，实现更智能的缓存策略

复现评估

论文的复现性评估需要考虑多个方面。从代码和开源情况来看，论文没有明确提及代码开源计划，这意味着独立研究者可能需要基于论文中的算法描述和数学公式从头实现。从数据集来看，所有评估数据集都是公开的标准基准，包括常识推理数据集（LAMBADA、PIQA、HellaSwag、WinoGrande、ARC-e、ARC-c、OpenBookQA、CommonsenseQA）、上下文检索数据集（SWDE、SQuAD、FDA、TriviaQA、Drop、NQ）和长上下文数据集（RULER、LongBench），这些都可以通过标准的评估框架（如 LM-Evaluation-Harness）获取。从算力需求来看，预训练使用了 50B tokens，在 4 张 NVIDIA A100 GPU 上进行学术规模的语言建模预训练，这对大多数学术实验室来说是可访问的。从实现细节来看，论文提供了相对完整的算法伪代码（Algorithm 1 和 Algorithm 2）以及超参数设置（状态缓存容量 K=30，与 Log-Linear Attention 的最大状态数相同），这为实现提供了清晰的指导。总体而言，虽然有代码开源缺失的问题，但论文提供了足够的实验细节和算法描述，有经验的独立研究者应该能够复现主要结果，尽管预训练部分可能需要大量计算资源

动态线性注意力 Dynamic Linear Attention