Beyond Linearity in Attention Projections: The Case for Nonlinear Queries¶

会议: ICLR 2026 Workshop (GRaM)
arXiv: 2603.13381
代码: GitHub
领域: 其他
关键词: nonlinear query, attention projection, identity prior, bottleneck MLP, transformer architecture

一句话总结¶

基于 \(W_Q\) 代数冗余性的理论发现，将线性 Query 投影替换为非线性残差形式 \(Q(X)=(X+f_\theta(X))/2\)，在不增加参数的情况下超越 +12.5% 参数的基线模型。

研究背景与动机¶

代数冗余性发现：Karbevski & Mijoski (2024) 证明 Transformer 中存在重参数化不变性——对任意可逆矩阵 \(\Theta\)，将 \((X, W_Q, W_K, W_V, W_O)\) 映射为 \((X\Theta, \Theta^{-1}W_Q, ...)\) 后 MHA 输出不变。取 \(\Theta = W_Q\) 可使 \(W_Q \to I\)，表明 \(W_Q\) 的线性参数完全可被相邻层吸收——代数上冗余

实验验证：设置 \(W_Q = I\) 的模型与标准基线性能相当，且在 3× 更低 weight decay 下仍稳定，证明恒等映射是 query 路径的良好先验

核心推理：既然线性参数冗余（可被吸收），若要在 query 路径有效分配参数，就必须引入非线性——非线性是不可被吸收的

信息瓶颈视角：从单个 token \(x\) 生成 q/k/v/残差四个向量全为 \(x\) 的线性函数，构成信息瓶颈。非线性 query 部分解耦了这个瓶颈

为何选择 Query：GQA 共享 \(W_K/W_V\)，只有 \(W_Q\) 可安全替换而不破坏共享结构

方法详解¶

整体框架¶

方法只动 Transformer 注意力里的 Query 路径一处：把原本的线性投影 \(Q=XW_Q\) 替换成非线性残差形式 \(Q(X)=(X+f_\theta(X))/2\)，其中 \(f_\theta\) 是一个瓶颈 MLP（bottleneck MLP）。同一个 token 表示 \(X\) 在这里分两路走——一路进 \(f_\theta\) 算出非线性增量，另一路作为恒等项 \(X\) 直接相加，二者取平均后得到 Query；与此同时，Key、Value 仍走标准线性投影，两条路最后在多头注意力里汇合。这样既绕开了"线性 \(W_Q\) 代数上可被吸收"的冗余陷阱，又把新增的非线性容量约束在与原始 \(W_Q\) 同阶的参数预算内，做到参数中性。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400, 'subGraphTitleMargin': {'top': 8, 'bottom': 16}}}}%%
flowchart TD
    X["输入 token X<br/>(上一 block 输出)"]
    subgraph Q["非线性 Query 投影"]
        direction TB
        MLP["瓶颈 MLP f_θ<br/>RMSNorm→GELU<br/>先降维 d→d/2 再升回 d"]
        RES["残差合成<br/>Q=(X+f_θ(X))/2"]
        MLP --> RES
    end
    X -->|"增量项: 算 f_θ(X)"| MLP
    X -->|"恒等项 X 直通"| RES
    X -->|"标准线性投影"| KV["保留 K/V 线性<br/>K=XW_K, V=XW_V"]
    RES --> ATT["多头注意力<br/>softmax(QKᵀ/√d_k)·V"]
    KV --> ATT
    ATT --> OUT["注意力层输出"]

关键设计¶

1. 非线性 Query 投影：把"可被吸收"的线性参数换成不可吸收的非线性

动机来自上文的代数冗余性——既然 \(W_Q\) 的线性部分总能被相邻层吸收，往这里堆线性参数就是浪费。于是把 Query 改成 \(Q(X)=(X+f_\theta(X))/2\)，残差里的恒等项 \(X\) 把投影锚定在已被实验证明是良好先验的 \(W_Q=I\) 上，保证梯度直通、训练稳定；增量项 \(f_\theta(X)\) 则提供线性投影给不了的、无法被重参数化消去的有效容量。系数 \(1/2\) 沿用 Karbevski & Mijoski 的建议，把残差与增量的幅度压平，防止激活幅度膨胀。

2. 瓶颈 MLP 结构 \(f_\theta\)：在参数预算内塞进非线性

增量项用一个先降维再升维的瓶颈 MLP 实现：\(f_\theta(X)=\text{LN}(\text{GELU}(\text{RMSNorm}(X)\cdot W_1)\cdot W_2)\)，其中 \(W_1\in\mathbb{R}^{d\times r}\)、\(W_2\in\mathbb{R}^{r\times d}\)，瓶颈维度取 \(r=d/2\)。这样两个矩阵的参数量合计 \(2dr=d^2\)，恰好与被替换掉的原始 \(W_Q\) 同阶；外层的 RMSNorm 与 LN 只额外引入 \(O(d)\) 量级参数（占比 <0.1%），可忽略。中间的 GELU 是整条路径真正引入非线性的地方，也是参数得以"有效分配"的关键。

3. 只改 Query、保留 K/V 线性：兼容现代注意力变体

Key 与 Value 保持标准线性投影不动，这一限制并非偶然——在 GQA/MQA 中 \(W_K/W_V\) 是被多个 query head 共享的，若改成非线性会破坏共享结构，而 \(W_Q\) 是唯一可以安全替换的投影。正因为改动局限在 Query 一侧，该设计天然兼容 GQA/MQA、RoPE、MoE 等现代架构组件，可作为即插即用的替换。

实验关键数据¶

主实验（GPT-3 Small, ~124M, OpenWebText, 60k steps ≈ 29B tokens）¶

模型	非嵌入参数	Val Loss (59k)	相对提升
Baseline	85M	2.956	0
MLP 4.75（宽MLP, +12.5%参数）	96M	2.927	0.98%
MLP 4.75 (scaled LR)	96M	2.928	0.94%
Res. GELU (本文)	85M	2.919	1.24%
Res. GELU (最优超参)	85M	2.915	1.40%

训练稳定性¶

配置	结果	说明
Baseline, WD=0.05	20k步前发散	标准模型不稳定
Res. GELU, WD=0.03, LR=3e-3	稳定到60k	可承受5×更高LR

关键发现¶

训练远超 Chinchilla 最优（29B tokens vs 2.5B optimal），确保改进不是因 token 不足产生的"水分"
所有模型在固定随机种子下看到完全相同的训练和验证数据，控制变量严格
非线性变体在 warmup 阶段增益最大，中期减小，末期最佳变体增益回升
作者明确表示 1.40% 可能是下界而非上界——超参搜索非常有限

亮点与洞察¶

理论驱动架构修改：从代数冗余性出发，逻辑链完整（\(W_Q\) 冗余→线性无用→引入非线性）
参数中性改进：不增加参数即超越 +12.5% 参数的模型——说明参数效率而非容量才是瓶颈
训练稳定性双赢：非线性版本不仅性能更好，还能在更激进的超参数（低 WD、高 LR）下保持稳定
代码和 checkpoint 完全开源

局限与展望¶

仅在 ~124M 单一规模验证，未测试大模型（是否在大模型上冗余性依然成立？）
未进行多种子实验（通过固定数据顺序和长时间训练缓解）
未测量推理速度——非线性引入串行依赖（瓶颈 MLP 必须先于注意力完成）
超参数搜索极不充分，各种归一化变体和激活函数未系统探索
未评估下游任务表现（仅报告预训练 val loss）

评分¶

新颖性: ⭐⭐⭐⭐ 理论驱动的架构修改，方向新颖，但改动量较小
实验充分度: ⭐⭐⭐ 单一规模、单一数据集，但控制变量极为严谨
写作质量: ⭐⭐⭐⭐ Workshop 篇幅内逻辑清晰，数学简洁
价值: ⭐⭐⭐⭐ 如能在大规模模型上验证将有重要意义