点击蓝字
关注我们
AI TIME欢迎每一位AI爱好者的加入！
作者简介
闻其帅，北京邮电大学硕士生
内容简介
注意力机制已在众多领域取得了显著的经验性成效，但其背后的优化目标仍不明确。此外，自注意力机制的二次复杂度也逐渐成为应用瓶颈。尽管可解释性与高效性本是两个相互促进的目标，但以往的研究通常割裂地探索两者。本文作者提出了一个统一优化目标，并通过算法展开推导出一种本质可解释且高效的注意力机制。具体来说，作者将所提出的优化目标的梯度下降优化过程构建为收缩-广播自注意力（Contract-and-Broadcast Self-Attention, CBSA）算子的前向计算过程。通过收缩少量的代表性 token，该机制实现了对输入 token 向低维结构的高效压缩。该机制不仅能通过固定代表性 token 的数量取得线性复杂度，还能通过改变代表性 token 的数量和结构，推广出其他注意力机制。通过大量实验，作者证明了该机制在视觉任务上相比于黑盒注意力机制有着可比的性能和额外的优势。本文为可解释性与高效性的融合，以及探索注意力机制的统一形式提供了新的思路。
论文地址：https://arxiv.org/pdf/2509.16875
代码链接：https://github.com/QishuaiWen/CBSA
论文解读
本文的题目是“通向可解释高效注意力机制：以收缩局部来压缩全局”。此处提及的“可解释”，是指基于确切的优化目标推导并设计注意力机制；而“压缩”则是将非结构化的表征分布向低维语义结构转化。
在介绍本文的核心方法前，首先简要回顾高效性和可解释性的发展历程。一方面，在高效注意力领域，各类高效机制不断涌现，但设计思路始终围绕早期形成的几个方向展开。工程上，可以混合使用这些机制，或尝试构建统一的计算形式。另一方面，关于可解释性，随着注意力机制的应用深化，很多难以解释的现象开始显现，而当前针对这些现象的解释却缺乏一个统一的框架。
事实上，存在一类先验可解释方法将注意力机制的本质归结为特定的优化目标，如能量最小化和去噪等。通过验证这些优化目标的优化过程与注意力机制的计算过程的对应性，可以实现对注意力机制的先验解释。
然而，现有先验解释仍停留在解释层面，却未能发挥实际的指导作用。比如，在解释了注意力机制后，如何进一步简化它，或进一步解释现有的高效注意力机制，本应顺理成章。但目前，可解释性与高效性的协同研究仍十分匮乏。对高效机制的探索依旧由经验和实验驱动，更难以对应用中观察到的各类现象进行系统解释。因此，本研究旨在以可解释的方式，推导新型高效注意力，并为理解现有的各种高效注意力机制提供新的视角。
本研究采用了“压缩即智能”的第一性原理，并通过“最大化码率约减”（MCR²）这一优化目标将该追求具体化。下面将简要介绍该优化目标。首先，通过编码率（coding rate）及其约减情况来刻画表征分布的压缩情况。直观来看，编码率计算了表征分布可被多少个半径为 ε 的小球覆盖。
整个最大化码率约减目标包含三项内容：扩张项，压缩项和稀疏项。其中，稀疏项作为常见的正则项，有利于促进神经元的后验可解释性。压缩项将输入表征向一组低维子空间压缩，上图中的合成数据集实验展示了文中推导出的一个算子在三维空间中将不同类别的输入 token 分别压缩到不同一维子空间。至于扩张项的作用是拉开子空间间的距离，避免子空间重叠导致不可分或相互干扰。
马毅老师团队提出的“白盒”Transformer发现，该目标的压缩项的一步梯度下降过程，可被近似地写成 Softmax Attention 的形式。二者的核心区别在于：标准注意力机制中的Wvalue、WK、Wquery三个投影矩阵，在该框架下统一为单一投影矩阵UK——即表征需向UK张成的低维子空间压缩。由此形成“多头子空间自注意力（MSSA）”，尽管其效率略高于标准注意力机制，但仍保持平方复杂度。
基于此，本研究期望从最大化码率约减目标出发，实现优化过程的高效化，进而推导得到高效的注意力机制。核心思路如下：在输入 token（Z）之外，引入一组数量更少的代表性 token（记为Q），通过 Q 代理原本输入 token Z 的压缩过程，将对 N 个输入 token 的压缩转化为对 M 个代表性 token 的压缩（其中 M 远小于 N ）。
在有限维空间中，压缩模式可通过少量代表性 token 实现等价替代或近似。引入的代表性 token Q（绿色标记）独立于输入 token（红色标记），其收缩过程可更高效地计算。随后通过 Q 与原始 token 之间的线性关系（图中虚线所示），将 Q 的收缩量传递至每个原始 token，最终实现“以局部收缩压缩全局（compressing all by contracting a few）”。
对新压缩项执行一步梯度下降后，其计算过程的梯度部分可拆分为两项理解：第一项为收缩项，这里假设代表性 token 选取主成分分析中的主方向（彼此正交且捕获不同大小的投影方差），其计算特点为：方差小的投影方向会产生较大的收缩量（红色箭头所示），导致该方向被抹除；方差大的投影方向收缩量较小，得以保留；第二项为广播项，本质是原始 token Z 与代表性 token Q 之间的线性关系表达矩阵（维度为m×N），能够将收缩项计算得到的 d×m 维收缩量传递至所有原始 token。基于这两项核心设计，该注意力机制被命名为“收缩-广播自注意力（Contract-and-Broadcast Self-Attention, CBSA）”。
接下来将详细说明代表性 token Q 及其与原始 token 的关系矩阵 A 的计算方式：首先通过对输入 token 进行池化操作，将样本维度从 N 降至 m，完成代表性 token 的初始化（也可通过可学习参数实现初始化）；初始化后，利用 Q 对原始输入 Z 进行查询，构建交叉注意力，实现代表性 token 的抽取——该交叉注意力过程中的注意力矩阵可直接作为线性系数矩阵 A 。
原始收缩项包含求逆操作，为降低计算开销，通过矩阵级数的冯·诺依曼展开，将逆矩阵近似为 Gram 矩阵；结合工程应用需求，在 Gram 矩阵外引入 Softmax 函数，最终转化为自注意力操作，即通过自注意力实现收缩过程。下面右图展示了不同序列长度下，该机制相较于现有可解释注意力机制的效率提升，其复杂度已降至线性水平。
值得注意的是，若移除收缩步骤，CBSA 将退化为 Agent Attention 。实验表明，这种简化在多数情况下可提升性能，但收缩步骤在理论层面（尤其是对不同注意力机制的统一）具有关键意义：通过调整代表性 token 的数量与结构，可推导出多种现有注意力机制。例如，令输入 token 自身作为代表性 token ，可推导出 Softmax Attention；令输入 token 完备且正交（如前文提及的主方向），可得到线性注意力；令正交的代表性 token 与输入无关（固定不变），则线性注意力退化为通道注意力。由此可见，各类注意力机制的本质差异，源于代表性 token 的数量与结构不同所导致的压缩模式差异。
具体来看不同注意力机制的压缩模式：通道注意力沿固定坐标轴进行压缩，例如将黄色椭圆形分布的输入 token 压缩至固定坐标轴；线性注意力可实现数据自适应压缩（如沿主成分方向压缩），这种动态特性对上下文学习及缓解叠加效应（superposition）至关重要；Softmax Attention 对每个 token 单独计算压缩量，具备更强的灵活性。而本文提出的 CBSA ，旨在以更低的计算成本近似 Softmax Attention 的压缩模式。
最后简要介绍实验结果，主要体现可解释性带来的三大优势：
其一，全局压缩由局部压缩驱动。对比输入 token 与代表性 token 经注意力机制处理后的编码率减少量，二者变化趋势基本一致，证实全局压缩确实由局部压缩带动。且代表性 token 的编码率减少量显著更大，这一现象可解释为“杠杆效应”——以大的局部改变量撬动全局的微妙优化。
其二，在简单有监督训练下涌现分割特性，该现象的内在机制仍需进一步研究。
其三，对参数扰动具备强鲁棒性。由于注意力机制中的投影矩阵UK本质是子空间基底，对其施加随机噪声扰动不会显著改变张成的子空间，因此对最终性能影响极小。实验显示，当噪声方差为1时，其他方法的性能已降至接近零，而本文方法仍保持50%的准确率。
在图像分类与语义分割任务中的实验结果如下：
图像分类任务（ImageNet-1K）中，本文提出的 CBT-S（small）模型仅使用30%的参数量与40%的计算量，便实现了与 ViT-S 相当的分类准确率。为验证缩放效果，将 ImageNet-21K 预训练模型的注意力机制（forward函数）替换为 CBSA 后进行微调，结果显示 CBSA 与线性注意力效果接近；若效仿 Agent Attention 移除收缩步骤，性能将显著超越线性注意力。
语义分割任务（ ADE20K数据集）中，以预训练 ViT 为基础模型，本文方法的性能优于基于 Softmax Attention 的其他模型。
未来展望：未来的研究方向或将脱离混合注意力机制，转向数学层面可统一的高效注意力机制。
本期文章由支昕整理
往期精彩文章推荐
关于AI TIME
AI TIME源起于2019年，旨在发扬科学思辨精神，邀请各界人士对人工智能理论、算法和场景应用的本质问题进行探索，加强思想碰撞，链接全球AI学者、行业专家和爱好者，希望以辩论的形式，探讨人工智能和人类未来之间的矛盾，探索人工智能领域的未来。
迄今为止，AI TIME已经邀请了2000多位海内外讲者，举办了逾800场活动，超1000万人次观看。
我知道你
在看
提出观点，表达想法，欢迎
留言
点击阅读原文查看作者直播回放！
