Intrinsically Interpretable Attention via Sparse Post-Training¶

会议: NeurIPS 2026
arXiv: 2512.05865
领域: 可解释性
关键词: 稀疏注意力, 机械可解释性, 事后训练, 激活补丁, 交叉层转码器

一句话总结¶

提出一种基于拉格朗日约束优化的稀疏注意力事后训练方法，在保持原有语言模型预训练损失几乎无损的前提下，将注意力连接边削减至 0.4% 以下，诱导模型形成紧凑、模块化的内生可解释计算回路。

研究背景与动机¶

大语言模型虽然在广泛的复杂推理与生成任务中展现出卓越能力，但其内部计算机制长期处于高度不透明的状态。机械可解释性（Mechanistic Interpretability）试图通过反向工程还原神经网络内部的特征与回路，从而剖析行为机制。然而，现有的反向工程手段往往受制于模型本身的计算冗余：在标准 Softmax 注意力驱动下，模型倾向于学到极其分散、弥漫全序列的计算方式。即使是在最简单的两数相加或指代识别任务中，单个语义特征或计算步骤也会被离散地切分并分散到数十乃至上百个注意力头与跨层通路中，导致提取出的计算回路极度庞大且充斥着组合爆炸的非直接依赖，为精细的归因分析设置了难以逾越的障碍。

这一困境的核心矛盾在于：模型在常规预训练目标下，没有任何先验或内在驱动力去寻求简洁、紧凑的算法实现。现有的稀疏注意力工作（如局部滑动窗口、固定步长跨步或 Top-\(k\) 截断）主要为了缓解自注意力的二次方计算复杂度，依赖强先验人工规则或刚性的硬截断；一旦在预训练模型之后直接施加硬 Top-\(k\) 截断，会导致模型性能严重崩溃，且固定的 \(k\) 值无法自适应不同注意力头与上下文差异。而在另一侧，传统的 \(L_1\) 软正则化又无法使 Softmax 权重真正归零，微小但非零的注意力连接依然在持续传递信息，无法斩断弥漫式通信。

本文转换了思考视角：将稀疏性视作一种内生可解释性的结构先验，而不是单纯的推理加速手段。作者主张通过事后微调（Post-Training），引导已有预训练模型在保留原有预测能力的同时重构内部连接。核心 idea：引入基于 Gumbel-Softmax 参数化的硬二值注意力门控，并结合 GECO 对偶拉格朗日约束优化框架，在严格限制交叉熵损失不劣于预训练基线的条件下自适应最大化注意力稀疏度，从而强制模型用极度精简、人类可理解的局部计算回路完成信息路由。

方法详解¶

整体框架¶

方法的核心是将标准 Transformer 的稠密全连接自注意力机制改造为可学习的硬稀疏注意力，并通过受约束的对偶优化进行事后训练。输入 token 嵌入经过线性投影得到查询、键和值向量；随后模型并不直接计算全连接的注意力加权，而是通过可微采样生成一个离散的二值门控矩阵，对注意力权重执行逐元素掩码过滤。只有被门控激活的边才允许发生实际的数值信息交互，其余边彻底置零。在优化阶段，算法将稀疏度惩罚项作为主最小化目标，同时将预训练模型的预测交叉熵损失设定为硬约束上界，通过交替梯度更新自适应调整正则化强度。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 Token 嵌入"] --> B["Q / K / V 线性投影"]
    B --> C["可微二值注意力门控<br/>Bernoulli 采样与 Gumbel-Softmax 重参数化"]
    C --> D["Splash Attention 稀疏消息传递<br/>硬二值掩码过滤非零边并与 Softmax 融合"]
    D --> E["GECO 对偶拉格朗日约束优化<br/>动态平衡稀疏惩罚与预测交叉熵阈值"]
    E --> F["输出内生可解释稀疏注意力回路"]

关键设计¶

1. 可微二值注意力门控：彻底阻断 token 间冗余信息流 针对标准 Softmax 注意力下微小权重仍持续传导信息的痛点，本文在键值投影之上构建了一层独立的伯努利门控分布。对于任意位置的查询向量 \(q_i\) 与键向量 \(k_j\)，定义边激活概率为逻辑 Sigmoid 响应 \(\sigma(q_i^T k_j)\)，并据此采样二值门控变量 \(A_{ij} \sim \text{Bern}(\sigma(q_i^T k_j))\)。在前向传播中，门控矩阵 \(A\) 以逐元素相乘的形式直接作用于标准 Softmax 权重矩阵：

\[\text{SparseAttn}(Q, K, V) = \left[ A \odot \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) \right] V\]

由于 \(A_{ij} \in \{0, 1\}\) 是严格离散的二值选择，未被选中的注意力连接输出严格为零，从而完全截断了跨位置的消息传递，实现了严格的 \(L_0\) 稀疏化。为了使离散采样过程能够端到端反向传播，算法采用 Gumbel-Softmax 直通估计器（Straight-Through Estimator），在反向传播时将硬采样的梯度平滑回传至 \(q_i\) 与 \(k_j\) 的投影参数。在结构形式上，该设计与原始 Transformer 保持完全兼容，可以直接无缝加载开源预训练底座的全部权重，无需对网络主干拓扑做结构性重构。

2. GECO 对偶拉格朗日约束优化：在无损性能边界内自适应逼近极限稀疏 传统的固定系数稀疏惩罚往往会导致超参数敏感与性能断崖式下跌，且无法保证不同模型和层级的性能保留。为此，本文采用源自 VAE 正则化的 GECO（Generalized Categorical Expectation-Maximization Optimization）约束优化框架。优化目标被严谨表述为一个带不等式约束的凸松弛问题：在保证下一个 token 预测的交叉熵损失 \(CE\) 不超过预设阈值 \(\tau\) 的前提下，最小化全模型各层期望激活边数 \(\sum_l \mathbb{E}[|A_l|] = \sum_l \sum_{i,j} \sigma(q_i^T k_j)\)：

\[\max_{\lambda > 0} \min_{\theta} \left[ \sum_{l} \mathbb{E}\big[|A_l|\big] + \lambda (CE - \tau) \right]\]

在训练循环中，网络参数 \(\theta\) 与对偶拉格朗日乘子 \(\lambda\) 采用交替梯度步优化。这种动态对偶博弈构成了一个自适应调节器：只要当前模型的交叉熵低于目标阈值 \(\tau\)，\(\lambda\) 就会自动减小，从而放大稀疏惩罚在梯度更新中的相对比重，强力剥离多余的计算连接；一旦模型遭遇性能波动或损失越过容忍线，\(\lambda\) 迅速拉升，迫使优化焦点重新收敛至语言建模任务本身。这种自适应调度彻底摆脱了人工退火超参数的盲目试错，确保模型能够在性能无损的边界上平稳逼近极限稀疏度。

3. Splash Attention 融合算子与高效后训练策略：消除显存与计算开销瓶颈 虽然理论上门控仅增加了逐元素乘法，但朴素的 PyTorch 实现需要物化庞大的注意力矩阵与门控掩码，极易触发显存瓶颈。针对这一工程挑战，作者基于 Helion 领域特定语言定制了高性能 GPU 融合内核 Splash Attention（Sparse Flash Attention）。该算子继承了 FlashAttention-2 的分块计算（Tiling）与在线 Softmax 机制，在前向计算过程中实时原位计算 \(q_i^T k_j\) 与门控采样，并将 Gumbel-Softmax 采样、直通梯度回传以及每层期望边数统计完全融合成单一内核，在 A100 GPU 上实现了接近 45% 的 MFU（Model FLOPs Utilization）。对于 7B 参数规模的大模型（如 OLMo-7B），作者进一步结合了高秩 LoRA 微调（\(r=400, \alpha=800\)）与知识蒸馏损失（引入基座模型的 KL 散度监督），在单张 H100 GPU 的学术级算力下即可完成端到端后训练，为在大规模基础模型上推行内生可解释性打通了工程可行性。

一个完整示例¶

以归因图分析中的典型逆序反义词任务为例：“The opposite of ‘large’ is”。模型最终需要在最后一个 token 位置准确输出 “small”。

在标准密集 GPT-2 模型中，这一推理过程被深度纠缠在超过 40 个分散于各层的注意力头中：位于位置 5 的 “large” 词元特征会向各个中间层的残差流弥散广播，并在多个位置形成错综复杂的间接交互，人类分析者极难厘清信息究竟如何从“反义”关系流向“尺寸”属性。

经过稀疏事后训练后，整个回路发生了极具物理意义的聚类与收缩：通过交叉层转码器（CLT）提取的归因图中，仅保留了 82 个显著特征，明确解耦为四个功能簇：反义关系簇（位于位置 3 “opposite”）、尺寸量级簇（位于位置 5 “large”）、括号语法簇以及最终的目标 Logit 簇。当分析 “large” 簇如何将信息传递至目标 “small” 时，80% 的绝对归因贡献仅由5 个位于高层的稀疏注意力头（如 L8-H6）并行承担。这些头直接执行从位置 5（Key）到位置 8（Query）的定向路由。进一步分解其 Key-Query 交互可知：位置 8 的 Query 向量受语法与最终预测需求驱动，而位置 5 的 Key 向量则由编码了“large”与“opposite”的语义特征激活。这 5 个注意力头在此扮演了清晰的“信息搬运与线性映射”管道角色，使得特征层级的交互与宏观注意力回路在同一视图下获得了完美的机械解释。

损失函数 / 训练策略¶

整个事后训练的联合优化损失函数可表示为带有蒸馏辅助的受约束拉格朗日形式：

\[\mathcal{L}(\theta, \lambda) = \sum_{l=1}^{L} \sum_{i,j} \sigma(q_{i,l}^T k_{j,l}) + \lambda \big( \mathcal{L}_{\text{CE}}(\theta) - \tau \big) + \beta \mathcal{L}_{\text{KD}}(\theta, \theta_{\text{base}})\]

其中 \(\tau\) 为目标交叉熵阈值（GPT-2 设为 3.50，基线为 3.48；OLMo-7B 设为 2.29，基线为 2.25），\(\mathcal{L}_{\text{KD}}\) 为基座模型与稀疏微调模型之间的 KL 散度蒸馏项，用于稳定大模型的表示分布漂移。优化器采用 Adam，学习率配置余弦退火衰减策略（峰值 \(1 \times 10^{-5}\)，最低 \(1 \times 10^{-6}\)），预热步数为 1,000 步；对偶变量 \(\lambda\) 的学习率分别设置为 0.1（GPT-2）与 0.01（OLMo-7B）。

实验关键数据¶

主实验¶

主实验在小模型（GPT-2 124M，全量微调）与大模型（OLMo-7B，LoRA 微调）上全面评估了稀疏化程度、预测损失保持情况，以及激活补丁（Activation Patching）下解释 90% 模型预测行为所需的计算回路规模。

模型	参数量	预训练数据集	验证集 CE（基线 / 稀疏）	活跃注意力边比例	解释 90% 行为所需头数减少倍数	解释 90% 行为所需边数减少倍数
GPT-2	124M	OpenWebText	3.48 / 3.50 (±0.01)	0.22%	1.4× ~ 4.5×	5.4× ~ 97×
OLMo	7B	Dolma-v1	2.25 / 2.29 (±0.01)	0.44%	显著精简	最多达 2 个数量级
LLaMA-3.2	1B	Dolma 3 LongMiNo	检索精度持平基线	~0.4%	长上下文检索验证	Passkey / Needle 精度无损

在典型的 Copy 任务中，密集 GPT-2 需要多达 61 个注意力头才能解释 90% 的复制行为，且注意力图斑驳杂乱；而稀疏 GPT-2 仅需 9 个注意力头即可完整复现该能力，且这些头呈现出教科书般清晰的固定相对偏移感应头（Induction Head）对角线模式。

消融实验¶

消融与机理分析实验重点考查了注意力的层级分布异质性、BOS 标记的注意力汇聚（Attention Sink）现象，以及交叉层转码器（CLT）归因图中的组件中介复杂度。

评估维度 / 指标	密集基线模型 (Dense)	稀疏后训练模型 (Sparse)	机制变化与影响说明
高层 BOS Sink 比例 (L6-L11)	57% ~ 74%	0.00%	彻底消除了无语义的惰性注意力沉底现象
CLT 归因达到 90% 所需 K-Q 对	基准 (1.0×)	减少 16.1×	大幅降低特征间中介分析的组合计算空间
CLT 归因达到 90% 所需注意力头	基准 (1.0×)	减少 3.4×	绝大多数非活跃头归因贡献严格为 0
归因图构建耗时	数小时 / 图	数分钟 / 图	预先剪除 99.6% 零边，算力消耗呈数量级骤降

关键发现¶

注意力连接削减至 0.4% 仍能维持原模型通用性能：实验证明绝大部分注意力权重是高度冗余的信息噪声，将其压制到不足 0.5% 的极度稀疏状态完全不损害预训练预测质量。
天然消除惰性 BOS Attention Sink：在密集模型中，由于 Softmax 必须分配概率质量，深层注意力往往有 60%~70% 惰性地沉淀在无意义的起始 token 上；而引入二值硬门控后，深层层级的 BOS 注意力直接归零，强迫注意力分配回归真实的跨 token 语义交互。
稀疏回路不仅是原回路的子图裁剪，而是发生重组：对比密集与稀疏模型选出的 Top-\(k\) 头，重合率仅为 40%~50%，边重合度仅 5%~10%。这表明稀疏压力促使模型在内部重组了计算拓扑，而非被动过滤。

亮点与洞察¶

将稀疏性重新定义为可解释性先验而非工程加速手段：摆脱了以往仅将注意力稀疏化用于长文本省显存的传统定势，系统性证明了局部稀疏化能够引发全局回路的级联精简。
自适应拉格朗日约束机制的稳健性：采用 GECO 框架让正则化强度随损失阈值动态博弈，不仅避免了稀疏化导致的性能崩溃，还实现了全自动的收敛调度，极为优雅。
架起回路分析与特征级转码器的桥梁：将稀疏注意力与交叉层转码器（CLT）结合，解决了特征间信息传递受困于上百个纠缠头归因的难题，让归因图的单条因果边可以直接对应到极少数物理注意力头。

局限与展望¶

目前仅聚焦于后训练（Post-Training）阶段：尚未探究如果在从头预训练阶段就全面引入该稀疏归纳偏置，是否会孕育出结构完全不同、模块化程度更彻底的计算拓扑。
仅针对自注意力维度展开稀疏化：模型的 MLP 块以及特征层级依然存在维度叠置（Superposition），未来可与稀疏混合专家（MoE）、权重稀疏化或正交性正则化等方法正交结合。
超长文本前沿场景的探索尚浅：虽然在 8K 长度的 LLaMA 探索中展现了潜力，但在 64K~128K 等超长上下文窗口下的端到端可解释性表现仍有待进一步验证。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙地将稀疏性约束用作可解释性归纳偏置，打通了回路发现与特征归因。
实验充分度: ⭐⭐⭐⭐⭐ 涵盖从玩具加法、GPT-2 到 7B OLMo 与 LLaMA，结合了激活补丁、CLT 转码器及定量图表。
写作质量: ⭐⭐⭐⭐⭐ 逻辑闭环紧密，动机、数学推导与实验洞见层层递进。
价值: ⭐⭐⭐⭐⭐ 为解决大语言模型“黑盒”提供了极具实用性与扩展性的模型内生改造新范式。