【人工智能】纠正线性注意力中的幅度忽略问题-论文阅读

Arxiv：https://arxiv.org/abs/2507.00698。2025年7月1日-v1，2025年8月2日-v3。

作者团队来自中国科学院自动化研究所和中国科学院大学人工智能学院。

1. 论文动机（为什么要做这件事）

背景：Transformer 的核心是 Softmax Attention，它建模能力强但复杂度为\(O(N^2)\)，在视觉任务中代价太大；为此出现了多种降低复杂度的方法，其中 Linear Attention 通过把相似度写成\(\phi(Q)\phi(K)^T\)的形式，将复杂度降为线性\(O(N)\)。然而，实际中线性注意力性能明显落后于 Softmax Attention（尤其在保留局部性/“尖锐”注意分布方面表现欠佳）。

论文的动机是：找出线性注意力性能差的根本原因并修正，使其既保留线性复杂度又能获得更接近 Softmax 的建模行为。

具体发现：线性注意力在数学推导上完全忽略了 Query 的幅值（magnitude）信息，只保留了方向成分。结果是：当 Query 放大时，Softmax 会让注意分布变得更“尖”（更集中到原来高分的 keys 上），但线性注意力不会随 Query 幅值变化，注意分布保持平滑、缺乏动态变尖的能力，这导致局部性弱、表现欠佳。论文以理论推导和实验证明这一点（比如把 Softmax 的 Q 替换为\(Q/|Q|\)会让 Softmax 的行为变得像线性注意力且性能下降）。

2. 核心方法（MALA：Magnitude-Aware Linear Attention）

目标：在不破坏线性复杂度的前提下，让线性注意力感知并利用 Query 的幅值，从而让注意分布随幅值适度“变尖”，更接近 Softmax 的行为，但避免 Softmax 过度的指数放大。

主要改动（数学上）：

①仍基于核分解\(\mathrm{Sim}(Q,K)\approx \phi(Q)\phi(K)^T\)，但注意权重不再用原来那种除法归一（division-based normalization），而是使用带比例因子 β 和偏置 γ 的线性调整并做加性归一化：

\(\mathrm{Attn}(Q_i, K_j) = \beta\,\phi(Q_i)\phi(K_j)^T – \gamma\)

其中 β 和 γ 都依赖于\(\phi(Q_i)\)与所有\(\phi(K_m)\)的和，具体定义见论文；同时论文保证所有注意分数加和为 1（通过构造使得\(\sum_j \mathrm{Attn}=1\)）。

②通过分析证明：当\(|\phi(Q)|\)被放大（因子\(a>1\)）时，MALA 中的注意分数比率会随之增大（即变得更尖），虽然增长速率是分式（fractional）而非 Softmax 的指数增长。论文给出 β、γ 的更新形式并证明新的比率\(p_m>p\)（详见补充证明）。这意味着 MALA 在保留线性复杂度的同时，恢复了 Query 幅值对注意分布的影响，使注意既有“集中”能力又不过度尖锐。

③实现细节：论文在视觉模型中使用\(\phi(x)=\mathrm{ELU}(x)+1\)（保证非负），并将上述 MALA 替换模型中所有的 Softmax/线性注意块，构建出基于 MALA 的视觉骨干 MAViT。

3. 实验设置与主要结果（论文直接报告的数值）

论文做了非常全面的横向比较：图像分类（ImageNet-1K）/目标检测与实例分割（COCO 框架下 Mask R-CNN / Cascade）/语义分割（ADE20K）/NLP 基准/语音识别/图像生成扩散，并有消融分析。下面把重点数字列出来（均来源于论文表格）：

ImageNet-1K（从头训练，300 epoch）：MAViT 在不同规模均优于同级别的线性注意力或许多 Softmax/其它骨干。MAViT-L（98M params, 16.1G FLOPs）Top-1 = 86.0%（优于 MILA 等其他线性方法）。

目标检测 / 实例分割（COCO 下 Mask R-CNN / Cascade Mask R-CNN）：在 “3×+MS” schedule 下，MAViT-L 在 Cascade Mask R-CNN 上报告 AP_b = 56.0, AP_m = 48.4，优于许多同级别 Transformer/CNN/线性模型。

语义分割（ADE20K）：UperNet 框架：MAViT-L 达到 mIoU = 53.6%（比论文中其他线性方法与部分 Transformer 更好）。

推理效率：在 A100 上的低分辨率 (224×224, batch 64) 与高分辨率 (512×2048, batch1) 任务中，MAViT 展示了很好的吞吐量—在论文绘图中其在精度与吞吐量的折中上优于多种基线（图 5 和图 6）。

NLP（语言理解基准）：论文训练了 0.3B MALA 模型（在 15B tokens 上），在 LMB / PIQA / Hella / Wino 四项都有竞争力的结果（表格给出与其它方法的对比，MALA 表现强劲或持平）。

语音识别（基于 Conformer 替换 Attention）：与原始 Conformer（Softmax）比较：MALA 在 WER 上优于 Linear Attention 并略优于或接近原始 Conformer（表中显示 MALA 在 test-clean/test-other 上显著改善）。表格给出例：Conformer(S) WER without LM test-clean = 2.7，Linear Attn = 3.4，MALA = 2.4（数字来自论文表 7）。

扩散式图像生成（Diffusion）：在 DiT/DiG/DiC 的比较表中，基于 MALA 的生成模型给出更低 FID（论文中示例：MALA FID=49.62、IS=32.18 且吞吐率更高）。

对比消融（线性注意力比较）：在 DeiT-T / Swin-T / Swin-S 三套设置下，MALA 均优于多种线性注意力变体（如 InLine Attn、Focused Linear 等），例如在 Swin-S 设置中 MALA 达到 85.3%（比 InLine +1.7）。

4. 论文里对比的主要算法（baseline）

论文在不同任务中与多类方法比较，涵盖三大类：

①Softmax 基的 Transformer：如 Swin、DeiT、RMT、SG-Former 等。

②线性注意力 / 相关改进：Flatten Transformer（Focused Linear）、InLine Attention、MILA、RAVLT、SOFT 等。论文重点把 MALA 与这些“线性家族”比较以证明修正有效。

③其它架构：传统 CNN（ConvNeXt、InternImage、MogaNet 等）、Mamba/Visual SSM 相关（Vmamba、MSVMamba）等也被列为参考。

5. 用到的数据集与评测标准（论文明确提及）

ImageNet-1K（图像分类，训练细节：300 epochs, batch 1024 等）。

COCO（用于目标检测、实例分割，使用 Mask R-CNN / Cascade Mask R-CNN，1× 和 3×+MS 两套 schedule）。论文在表格里以 AP 等指标给出对比。

ADE20K（语义分割，SemanticFPN 和 UperNet 框架，80K / 160K iterations）。

NLP 训练语料：论文提到训练 0.3B MALA 模型在 15B tokens 上（用于 LMB/PIQA/Hella/Wino 基准评测）。

语音识别：使用 Conformer 设置并报告 WER（表格列出 testclean/testother，虽论文未在正文显式写出“LibriSpeech”字样，但测试项与 LibriSpeech 常用拆分一致）。

图像生成 / Diffusion：论文在与 DiT/DiG/DiC 比较时使用常见 diffusion 训练设置（400K iterations 等），并报出 FID/IS/throughput 等。

6. 消融、稳定性与设计要点（论文已验证的关键点）

β 与 γ 必不可少：论文消融显示把 β 或 γ 去掉会导致性能急剧下降（去掉 γ 甚至会出现 NaN），把 β/γ 设为可学习参数也会显著劣化，说明论文设计的解析解式 β/γ 对数值稳定性与归一化至关重要。

核函数选择不敏感：使用\(\phi=\mathrm{ELU}+1\)、ReLU 或 exp 等不同非负核，对最终性能影响不大（表 10 显示差异微小）。所以非负性比具体形式更重要。

MALA 与 Softmax 的差别：MALA 的注意比率随幅值放大而“增大”，但比率增长是分式（温和）而非 Softmax 的指数，论文认为这让分布“既有集中性又不过度尖锐”，从而在多数任务上表现优异（图示与理论证明同步）。

7. 论文指出/或由实验可推断的改进空间（可做的事与原因）

下面列出论文里暗示的、或者基于实验结果合理的改进方向，并说明为什么这些方向值得尝试（便于你做后续研究或复现的 Ti 风格思路）：

1. 更严格分析 β/γ 的数值边界与稳定性

论文指出 β/γ 的解析构造对稳定性重要，但当\(\phi(Q)\)很小时分母可能带来数值问题（论文中已遇到 NaN 的情形在移除 γ 时）。可以做更细致的数值分析、截断策略或带自适应下限的归一化来提高健壮性（这是做实际工程部署时常遇到的问题）。

2. 理论上比较 MALA 的分式增长与 Softmax 指数增长的泛化/鲁棒性

论文认为温和增长比指数增长更好，但缺少严格的泛化或噪声鲁棒性证明。可以从理论或合成实验研究在输入噪声、对抗扰动或长序列依赖时两种增长形式的差异。

3. 与局部卷积/混合结构的联合设计

线性注意仍容易弱化局部感受野（论文也讨论线性注意“过于平滑”导致局部性感知弱）。虽然 MALA 改善了这一点，但进一步结合轻量卷积或 local-window 机制可能提升小目标/高分辨率任务性能（论文在比较中已有相关模型作参考）。

4. 扩展到更长的序列与多模态任务，测试伸缩性

论文主要在视觉密集任务（高分辨率）与常规 NLP 任务证明效果，进一步验证在极长序列（比如长文本、长语音片段）或跨模态（视频 + 文本）任务的行为值得。线性复杂度是天然优势，但 MALA 的 β/γ 计算是否在极长序列与分布偏移下仍稳健需检验。

5. 可学习但受约束的 β/γ 形式

论文尝试直接把 β/γ 设为可学习参数效果不好，但或许可以设计受解析约束的可学习变体（比如把解析表达式作为初值并允许小幅可学习修正），在保持数值/归一性质的同时增加模型表达力。实验表明完全可学习会劣化，说明约束形式必须谨慎设计。

6. 更细粒度的注意分布可视化与解释性研究

论文给出若干注意图与统计，但可以进一步对“为什么 MALA 在某些场景下明显优于 Softmax/线性注意力”做定量解析（如不同类别的 tokens、空间尺度对注意分布变化的响应函数），这有助于理解何种类型任务最受益。

7. 图像生成和扩散模型的更多训练策略与更广数据集评估

论文在 diffusion 上显示良好结果，但扩展实验（像更大尺度的图像、更多迭代/采样策略）能更完整地评估 MALA 在生成任务的潜力与局限。

8. 总结

论文提出的核心诊断是明确且具解释力的：线性注意力“丢失”了 Query 的幅值信息，导致注意分布不会随幅值发生聚集；基于这个诊断，提出的 MALA 用解析构造的比例与偏置（β, γ）把幅值效应“恢复”回去，同时保留线性复杂度，结果在分类、检测、分割、NLP、语音、扩散生成等多个任务上都表现优异，并在消融中证明 β/γ 设计关键。