Learning in Compact Spaces with Approximately Normalized Transformer¶

会议: NeurIPS 2025
arXiv: 2505.22014
代码: github.com/automl/anGPT
领域: LLM效率
关键词: 近似归一化, 高维集中度, 紧凑空间, 无权重衰减, 收敛加速

一句话总结¶

提出 anGPT（近似归一化 Transformer），利用高维空间中向量范数的集中现象，用简单标量乘法替代逐层精确归一化，在消除权重衰减和学习率预热的同时实现了相比 GPT+（含 QK-norm）40% 的收敛加速，仅增加 3% 运行时开销。

研究背景与动机¶

领域现状：深度 Transformer 训练依赖大量正则化和归一化技术（LayerNorm / RMSNorm / QK-Norm）来保证稳定性。这些技术需要调节的超参数众多，包括权重衰减系数、学习率预热步数等。
现有痛点：nGPT 将所有参数和表示严格约束在超球面上，虽然大幅加速收敛，但引入了额外的归一化层，导致训练运行时增加约 9%、推理增加约 3%。精确归一化的计算开销在大模型上不可忽视。
核心矛盾：归一化带来的收敛加速 vs. 归一化本身的计算开销之间存在 trade-off。理想情况是保留归一化的好处（稳定输入尺度、防止残差流范数膨胀），同时避免昂贵的逐样本归一化操作。
本文要解决什么：(a) 能否用更廉价的方式获得"近似归一化"的效果？(b) 能否同时消除权重衰减和学习率预热？
切入角度：高维空间中的集中度现象（Concentration of Measure）——当维度 \(d\) 足够大时，Lipschitz 函数在超球面上的值高度集中于其期望值附近。因此向量范数可以用一个与输入无关的常数因子来近似。
核心 idea：用输入无关的标量归一化因子 \(\nu\) 替代逐样本的精确归一化，配合参数范数约束（\(\|w\|_2 \leq 1\)），构建"紧凑空间"中的 Transformer，无需额外归一化层即可获得归一化的收敛优势。

方法详解¶

整体框架¶

anGPT 基于标准 GPT 架构（SwiGLU + RoPE + QK-Norm），做了三个层面的修改：(1) 用 LERP 残差更新替代经典残差连接；(2) 移除 RMSNorm 层，换成标量近似归一化因子 \(\nu\)；(3) 约束所有权重矩阵行向量的 L2 范数 \(\leq 1\)（而非 \(= 1\)）。输入 token embedding 后，经过 L 层 MHA + MLP block，最终通过可学习缩放向量 \(s_Z\) 和 head 线性层产出 logits。

关键设计¶

近似归一化因子 \(\nu\)
做什么：为每个线性映射和激活函数的输出提供一个与输入无关的标量缩放因子，使输出向量范数近似为 1
核心思路：利用 Lipschitz 函数在高维球面上的集中定理。对于输出 \(g(x)\)，其范数 \(\|g(x)\|_2\) 以指数概率集中于期望值 \(\mathbb{E}[\|g(x)\|_2]\) 附近，偏差概率 \(\leq 2\exp(-cdt^2/\|g\|_{\text{Lip}}^2)\)。因此可预先计算 \(\nu = 1/\sqrt{\mathbb{E}[\|g(x)\|_2^2]}\) 作为固定常数
具体因子：QKV 投影 \(\nu_{qkv} = \sqrt{d_m/d_h}\)；输出投影 \(\nu_p = \sqrt{d_h/d_m}\)；上采样 \(\nu_{uz} = \sqrt{d_m/(4d_m)} = 0.5\)；下采样 \(\nu_d = \sqrt{4d_m/d_m} = 2\)；SiLU 激活 \(\nu_{acf} \approx 3.74\)（Monte Carlo 估计）
设计动机：避免了 nGPT 中每步都需要精确计算 \(\|x\|_2\) 的开销。这些因子是常数，推理时可以直接吸收进权重矩阵
LERP 残差更新 + 近似归一化
做什么：用线性插值 \(h \leftarrow h + \alpha(x - h)\) 替代经典残差 \(h \leftarrow h + x\)，并附加归一化因子
核心思路：残差更新后的范数期望为 \(\mathbb{E}[\|(1-\alpha)h + \alpha x\|^2] = 1 - 2\alpha + 2\alpha^2\)（假设 \(h, x\) 独立、单位范数、零均值），因此归一化因子为 \(\nu(\alpha) = 1/\sqrt{\alpha^2 + (1-\alpha)^2}\)。该因子本身是 \(\alpha\) 的函数，每步需要更新但计算量极小
设计动机：nGPT 在 LERP 后还加了一个精确 norm 操作，anGPT 用 \(\nu(\alpha)\) 标量乘法替代，节省计算量
参数范数约束（\(\|w\|_2 \leq 1\)）
做什么：约束权重矩阵每行的 L2 范数不超过 1（而非严格等于 1）
核心思路：nGPT 要求权重严格归一化（\(\|w\|_2 = 1\)），但这排除了零向量，限制了表达能力。anGPT 放宽为上界约束，允许权重自由学习但保持有界
设计动机：(a) 消除了权重衰减的必要性（参数已经有界）；(b) 保证了前面推导的归一化因子仍然有效（因为 \(\|Wx\|\) 的上界由 \(\|w\|\) 控制）；(c) 训练初始化时将权重归一化，消除了学习率预热的需求
参数重参数化
做什么：对所有可学习的缩放参数（\(\alpha_A, \alpha_M, s_Z\)）进行重参数化以统一优化动力学
核心思路：实际优化代理参数 \(\hat{s}_a\)，通过 \(s_a = (s_{a,\text{init}}/s_{a,\text{scale}}) \cdot \hat{s}_a\) 映射回真实参数，确保所有存储参数量级可比（约为 \(s_{a,\text{scale}}\)），让 Adam 的自适应学习率机制在不同参数类型间均匀工作

损失函数 / 训练策略¶

标准交叉熵语言建模损失
使用 Adam（anGPT/nGPT）或 AdamW（GPT+），cosine 退火调度
anGPT 不需要权重衰减和学习率预热；\(s_{a,\text{init}} = 0.01\)
数据集：SlimPajama（627B tokens），分词器：GPT-NeoX（50K），上下文 2048

实验关键数据¶

主实验：0.5B 模型收敛加速¶

设置	GPT+	nGPT	anGPT
含 QK-norm，平均收敛加速	1.0×	1.29×	1.40×
无 QK-norm，平均收敛加速	1.0×	1.8×	2.0×
训练步耗时（A100×4）	0.1416s	0.1552s (+9.6%)	0.1455s (+2.75%)
推理额外开销	—	~3%	0%（因子吸收进参数）

250M 和 1.0B 模型上也观察到一致的 ~40% 收敛加速。

消融实验（0.5B, 10B tokens, OpenWebText）¶

配置	相对 GPT+(无QK) 提升
GPT+ 无 QK-norm	基线
GPT+ 含 QK-norm	+2.1%
nGPT	+3.0%
nGPT + 参数范数约束（替代严格归一化）	+3.4%
anGPT（完整配置）	+3.8%

关键发现¶

缩放定律一致性：anGPT 和 GPT+ 的缩放指数几乎相同（32M→1B），说明架构改进是"常数加速"而非改变缩放行为
更大批大小：anGPT 在大模型上可使用更大 batch size，有利于多机并行加速
残差归一化因子最关键：消融显示移除残差 LERP 的归一化因子会导致训练不稳定；其他因子 2% 以内的估计误差不影响结果
下游任务：1B 模型在 6 个 benchmark 上一致优于 GPT+，提升 3%-22%
与 DyT 对比：DyT（Dynamic Tanh）替换归一化层反而掉了 6.5%，LN scaling 无明显增益，anGPT 提升 1.3%

亮点与洞察¶

"维度的祝福"：高维向量范数的集中性使得用常数因子近似逐样本归一化成为可能。这个数学视角非常优雅——把维度从"诅咒"变成了"祝福"
解耦稳定化与损失优化：传统 LayerNorm 的可学习参数 \(\gamma\) 同时承担稳定激活和参与梯度下降两个角色，anGPT 将两者解耦——\(\nu\) 负责稳定，模型参数专注优化损失
超参数简化：消除权重衰减和预热后，需要调的超参只剩学习率和批大小，scaling law 推导也更高效
推理零开销：常数归一化因子可以在部署前乘进权重矩阵，推理时无任何额外计算

局限性 / 可改进方向¶

仅在 GPT 架构上验证，未测试 Vision Transformer、扩散模型等其他架构
最大训练规模为 7× Chinchilla（0.5B×70B tokens），更大规模下是否仍然有效未知
注意力矩阵的归一化因子推导假设了 dense attention，但实际 attention 通常是稀疏的。FlashAttention 不提供有效非零比例信息，是一个开放问题
低精度（FP8）训练适配是一个很有前景的方向：紧凑空间中的有界激活和权重天然适合低精度表示

评分¶

新颖性: ⭐⭐⭐⭐ 高维集中度→近似归一化的理论切入角度新颖，但整体架构改动是 nGPT 的自然延伸
实验充分度: ⭐⭐⭐⭐⭐ 多尺度模型(32M-1B)、消融、scaling law、下游评估、与 DyT/LN-scaling 对比，非常全面
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰，Table 1 的三种架构对比一目了然
价值: ⭐⭐⭐⭐ 实用价值高——消除超参、降低开销；但需要更大规模验证