仅注意力Transformer的对照研究

# 仅注意力Transformer的受控研究
来源: https://arxiv.org/html/2607.18363
Henry Ndubuaku, Karen Mosoyan, Jakub Mroz, Noah Cylich, Satyajit Kumar, Parkirat Sandhu, Roman Shemet, Justin H Lee

###### 摘要

前馈网络占据Transformer非嵌入参数的三分之二，然而该架构尚未经历同时控制参数、计算量和深度的必要性检验。我们预训练了仅注意力的解码器Transformer（简单注意力网络，SANs），与标准Transformer进行对比，这些标准Transformer在参数数量、训练FLOPs和深度（各分支深度为2到48层）上分别匹配，每个分支进行学习率扫描，使用最多105B token的数据（68B token的推理密集语料的1.5个epoch），总参数范围为6M到87M。直接删除前馈层代价高昂：在深度匹配时，标准Transformer领先0.47 nats；在训练FLOPs匹配时，领先0.26 nats——此时注意力将计算花费在无参数的二次项上，因此在相同成本下携带更少参数。将释放的预算重新分配到注意力深度上缩小了差距：在参数数量匹配时，差异为0.006 nats，即损失值的0.27%，在干净种子对上的可重复性达到万分之一，并在分别训练的5B、30B和105B预算下持续缩小，且在固定31.5B token预算下，跨越29倍的非嵌入参数范围仍保持在接近0.02 nats。三项独立测量将剩余差距定位到参数化记忆：基于token区域、基于任务类型和零样本基准的损失分析。仅注意力模型在基于上下文的答案上表现更好，在必须从权重中获取知识时表现更差。权重谱揭示了原因：路由矩阵（Q/K）在训练的前四分之一阶段结晶，而内容矩阵在稳定阶段积累秩；删除前馈层将这种积累转移到注意力输出投影。QK归一化（而非前馈层或残差门控）使得48层仅注意力堆栈保持可训练性。该差距是集中的，而非分散的：低上下文查询预测的每个token赤字是总体赤字的五倍，但仅占语料损失的8%；token加权分解在2%以内重现其总体值；在最大预算下，赤字完全定位在该样本上。一项预先注册的测试证实了这一解释：它预测了在知识密集型网络文本上0.02到0.05 nats的差距，并且在fineweb-edu上训练的一对匹配模型测量值为0.040。在所测试的范围内，注意力完成了其余工作。

## 引言

前馈网络（FFN）自Transformer引入以来就一直是其固定组成部分（Vaswani et al. 2017 (https://arxiv.org/html/2607.18363#bib.bib33)），在现代解码器配置中，它占据约三分之二的非嵌入参数。大量可解释性文献认为这些层充当模型的参数化记忆：训练事实的键值存储（Geva et al. 2021 (https://arxiv.org/html/2607.18363#bib.bib10), 2022 (https://arxiv.org/html/2607.18363#bib.bib9)）、可编辑事实关联的所在（Meng et al. 2022 (https://arxiv.org/html/2607.18363#bib.bib19)）以及可识别知识神经元的家园（Dai et al. 2022 (https://arxiv.org/html/2607.18363#bib.bib2)）。然而，反向实验——完全移除FFN并测量实际损失了什么——从未在受控条件下进行。最接近的尝试是将同等容量折叠回注意力中作为持久记忆向量（Sukhbaatar et al. 2019 (https://arxiv.org/html/2607.18363#bib.bib31)）；探测和编辑研究定位了功能但未删除它；而纯注意力的经典理论（Dong, Cordonnier, and Loukas 2021 (https://arxiv.org/html/2607.18363#bib.bib4)）分析了没有残差连接的设置，没有人在实际训练中使用过这种设置。

这个实验比看起来更困难，因为移除FFN会同时扰动三个量：参数数量、每个token的计算量以及非线性组合步骤的数量。因此，无论固定哪个轴，单一控制对比都是混淆的。我们通过三次单独的匹配进行必要性检验，每次固定一个轴（表1 (https://arxiv.org/html/2607.18363#Sx3.T1)），每个分支都有自己的学习率扫描，并带有边界扩展协议，以确保没有分支在另一个分支偏好的超参数下竞争；并使用一个优化器等价的对照组来校准实验的噪声底限。所有预测都在相应测量之前注册，其中一项关于知识密集文本的差距是在运行启动之前以数值形式注册的。

结果是一个单调序列，逐步解析了这个问题。在深度匹配的情况下直接删除前馈层，验证损失增加0.47 nats。在训练FLOPs匹配的情况下，注意力的二次项会挤占参数，标准Transformer仍然领先0.26 nats。但在参数数量匹配的情况下，将释放的预算重新分配到注意力深度，差距缩小到一个小而高度可重复的数值：在干净种子对上为+0.0055和+0.0054 nats（损失值的0.27%，一致到10^{-4}），并且在分别训练的5B、30B和105B token预算下单调减小，同时在固定31.5B token预算下，跨越29倍的非嵌入参数范围保持在接近0.02 nats。FFN的参数很重要；但在该分布上，其函数形式在很大程度上并不重要。

剩余的差距并非分散。通过对token区域、任务类型以及零样本基准的损失分解，赤字集中体现在低上下文预测上：即上下文提供很少信息用于路由的token，因此只有参数化知识有所帮助。在最大预算下，仅注意力模型在分解样本的每个答案区域（包括记忆练习）上都领先，仅在低上下文查询token上落后；在基准测试中，答案位于已提供段落中的任务更倾向于SAN且差距不断增大，而需要分布外回忆的任务则始终在每个预算下更偏好FFN模型。权重谱分析提供了机制：在每个模型中，路由矩阵（Q/K）在训练的前四分之一阶段从谱上结晶，而内容矩阵在稳定阶段积累稳定秩；移除FFN将这种积累转移到注意力输出投影——这是仅注意力模型中唯一的写入残差流的路径。

我们的贡献：(i) 首次对Transformer FFN进行受控必要性检验，包含参数、FLOP和深度匹配、每个分支的学习率公平性以及校准的噪声底限；(ii) 将剩余差距三重定位到低上下文预测，并通过数值预先注册的分布外测试（预测fineweb-edu上为0.02–0.05 nats；测量值为0.040）确认；(iii) 确定了QK归一化（而非前馈层或残差门控）是保持深度仅注意力堆栈可训练性的组件；(iv) 提供权重谱解释：路由早期结晶，内容在稳定阶段积累，存储空间在FFN移除后重新定位。

## 相关工作

作为记忆的FFN。Geva等人（2021 (https://arxiv.org/html/2607.18363#bib.bib10)）将FFN层识别为键值记忆；Geva等人（2022 (https://arxiv.org/html/2607.18363#bib.bib9)）完善了该机制；ROME（Meng等人 2022 (https://arxiv.org/html/2607.18363#bib.bib19)）和知识神经元（Dai等人 2022 (https://arxiv.org/html/2607.18363#bib.bib2)）将可编辑事实定位于中间层FFN。这些是探测和编辑结果；我们的结果是有因果关系的补充，通过删除表明损失的能力正是参数化存储。注意力权重也存储关联（Elhage等人 2021 (https://arxiv.org/html/2607.18363#bib.bib6)），因此存储是一种具有因果证据的趋势，而非排他性位置；我们的分解测量了该趋势的大小，而不是假设它。Pires等人（2023 (https://arxiv.org/html/2607.18363#bib.bib26)）发现FFN在层间冗余；He和Hofmann（2024 (https://arxiv.org/html/2607.18363#bib.bib12)）通过信号传播论证移除了其他块组件。

仅注意力模型。Sukhbaatar等人（2019 (https://arxiv.org/html/2607.18363#bib.bib31)）移除了FFN但将其容量恢复为持久记忆向量，留下了必要性问题的开放性。机械可解释性文献将小型仅注意力Transformer作为一个可分析模型类别进行研究（Elhage等人 2021 (https://arxiv.org/html/2607.18363#bib.bib6); Olsson等人 2022 (https://arxiv.org/html/2607.18363#bib.bib24)），建立了QK/OV路由与内容分解，我们的谱动力学在现实深度下确认了这一点。Dong、Cordonnier和Loukas（2021 (https://arxiv.org/html/2607.18363#bib.bib4)）证明了无残差注意力堆栈的秩崩溃；我们的48层结果和秩测量限定了该分析适用于训练实践的范围，并识别出QK归一化（Henry等人 2020 (https://arxiv.org/html/2607.18363#bib.bib13); Dehghani等人 2023 (https://arxiv.org/html/2607.18363#bib.bib3)）作为有效的稳定器。

对偶消融。仅MLP架构（Tolstikhin等人 2021 (https://arxiv.org/html/2607.18363#bib.bib32); Liu等人 2021 (https://arxiv.org/html/2607.18363#bib.bib17)）映射了无注意力的方向；MetaFormer（Yu等人 2022 (https://arxiv.org/html/2607.18363#bib.bib34)）认为块框架比混合器更重要。我们为语言建模提供了缺失的仅注意力方向：在匹配参数下删除通道MLP代价为0.006 nats，具有特定的、可预测的残差轮廓。

组件和优化器。我们的标量残差门属于Highway/ReZero/LayerScale家族（Srivastava, Greff, and Schmidhuber 2015 (https://arxiv.org/html/2607.18363#bib.bib30); Bachlechner等人 2020 (https://arxiv.org/html/2607.18363#bib.bib1)）；夹层归一化遵循Shleifer, Weston, and Ott (2021 (https://arxiv.org/html/2607.18363#bib.bib29))；零中心范数增益遵循OLMo Team (2024 (https://arxiv.org/html/2607.18363#bib.bib23))；门控注意力输出由Qiu等人（2025 (https://arxiv.org/html/2607.18363#bib.bib28)）研究，逐层残差重加权由Kimi Team（2026 (https://arxiv.org/html/2607.18363#bib.bib16)）研究。我们使用Muon（Jordan等人 2024 (https://arxiv.org/html/2607.18363#bib.bib15); Liu等人 2025 (https://arxiv.org/html/2607.18363#bib.bib18)）训练，并与AdamW对照组对比。

数据范围。我们在SYNTH（PleIAs 2025 (https://arxiv.org/html/2607.18363#bib.bib27)）上预训练，这是一个推理密集的合成语料，相对于计算最优（Hoffmann等人 2022 (https://arxiv.org/html/2607.18363#bib.bib14)）有意识地过度训练，属于小型模型合成预训练谱系（Eldan and Li 2023 (https://arxiv.org/html/2607.18363#bib.bib5); Gunasekar等人 2023 (https://arxiv.org/html/2607.18363#bib.bib11)）；重复行为将Muennighoff等人（2023 (https://arxiv.org/html/2607.18363#bib.bib21)）扩展到合成数据。思维链表达能力结果（Merrill and Sabharwal 2024 (https://arxiv.org/html/2607.18363#bib.bib20); Feng等人 2023 (https://arxiv.org/html/2607.18363#bib.bib7)）启发跟踪格式化的训练范围，但需注意这些构造使用了MLP。

## 实验设置

架构。SAN块是一个预归一化注意力块，删除了FFN，确切来说：

ZCN(z)=(1+γ)⊙z/RMS(z), γ初始化为0
u = ZCN(x)
q,k,v = W_q u, W_k u, W_v u (GQA: 8Q, 4KV)
q,k = RoPE(ZCN_h(q), ZCN_h(k))
A = softmax(qk^T/√d_h + M)
y = x + σ(g) · W_o (A v)

其中M是因果的、文档边界的掩码，作用在打包的2048 token行上；g是每个子层的标量（初始化为0，因此每个分支开始时为半强度）；绑定嵌入；在输出头之前有最后的ZCN。控制臂插入一个SwiGLU FFN：y' = y + σ(g_2) · W_down SwiGLU(W_gate ũ, W_up ũ)，其中d_ff=4d，在注意力之后；两个分支共享一个实现。SAN中每个位置的非线性总和是注意力softmax和归一化：没有学习到的每位置特征映射，这正是整个实验操作。四个简要结果精确说明了该操作移除了什么。

###### 命题 1（条件线性性）。

对于任何固定的行随机注意力模式A，从归一化上下文(û_1, ..., û_T)到该层更新的映射是线性的。

###### 证明。

它是线性映射W_v、v ↦ A v和W_o的复合，由常数σ(g)缩放。∎

###### 命题 2（单纯形传输）。

每个头上，位置i的投影前更新位于{W_v û_j: j ≤ i}的凸包中。

###### 证明。

Softmax行非负且和为1。∎

SAN层选择和传输上下文中的内容；它不能合成上下文不支持的表征。这是全文使用的"基于上下文"的形式含义。

###### 引理 1（标量瓶颈）。

在序列长度1时，整个深度L网络是线性映射 x_L = [∏_ℓ (I + B̃_ℓ/ρ_ℓ)] x_0，其中B̃_ℓ = σ(g_ℓ) W_o W_v diag(1+γ_ℓ)是常数，ρ_ℓ = RMS(x_ℓ)：一个由恰好L个标量调制的线性映射。

###### 证明。

A=[1]且位置0处的RoPE是恒等映射，因此每个块计算 x_{ℓ+1} = (I + B̃_ℓ/ρ_ℓ) x_ℓ；对层进行归纳。∎

所有每位置非线性通过这个L标量瓶颈流动；跨token注意力是唯一的逃逸路径，而定价这一限制正是实验所做的工作。

###### 定理 1（初始化时流方差有界）。

在实现的初始化下（精确RMS归一化，门控σ(0)=1/2，输出投影独立同分布零均值，方差s^2/(2N)），E‖x_N‖^2 = E‖x_0‖^2 + Θ(1) 在深度N上一致有界。

###### 证明。

交叉项E⟨x_ℓ, b_ℓ⟩消失，因为W_o是零均值且与上游一切独立，因此二阶矩相加。每个分支输入经过RMS归一化，因此其二阶矩有界且与流幅度无关；行随机A不能增加最大范数（命题2）；1/(2N)的输出方差给出E‖b_ℓ‖^2 = Θ(1/N)，N层上的总和缩放到Θ(1)。∎

每个假设在后续中承担重要角色。