跳转至

Normalized Architectures are Natively 4-Bit

会议: NeurIPS 2026
arXiv: 2605.06067
代码: https://github.com/anonymous452026/ngpt-nvfp4
领域: 模型压缩
关键词: 归一化Transformer、NVFP4量化、点积信噪比、超球面约束、训练动态性

一句话总结

本文证明将权重与表征约束在单位超球面上的 nGPT 架构天生具备 4 位量化鲁棒性,无需随机阿达马变换与动态逐张量缩放即可在 Blackwell 硬件上实现端到端稳定的 NVFP4 训练,其鲁棒性根源于超球面约束促使信号项在点积累加中产生相干正相关放大,而非抑制单点量化噪声。

研究背景与动机

随着大语言模型参数规模迈向数千亿乃至万亿级别,采用 4 位浮点格式(如 NVIDIA Blackwell 架构原生支持的 NVFP4)进行预训练已成为压缩显存占用、提升计算吞吐的关键路径。然而超低比特训练极度脆弱,标准 Transformer 在 4 位精度下极易出现激活值异常离群点与梯度方差爆炸,导致训练过程发散或下游性能剧烈退化。工业界和学术界目前的主流应对策略是将量化鲁棒性视为后处理算法挑战,通过引入随机阿达马变换(Randomized Hadamard Transforms, RHT)打散特征离群值、执行细粒度动态逐张量缩放计算(amax reduction)以及随机舍入(Stochastic Rounding)来维持训练稳定性。

这些补丁式的干预手段带来了显著的系统额外开销,包括额外的矩阵变换延迟、归约同步瓶颈以及复杂的超参数微调成本。更关键的矛盾在于,传统思路普遍假设量化失效的根源在于局部量化噪声过大,因而所有优化都围绕着压低离群误差、调整量化网格或在敏感层回退高精度展开。这种将模型结构与量化算术割裂的视角,回避了一个根本性命题:能否从网络架构的内在几何特性出发,让模型天生具备低比特量化免疫力,从而彻底抛弃高开销的量化干预技巧?

本文的切入视角正是打破“后验算法修补”的定势,重新审视将所有隐藏状态与权重参数约束在单位超球面上的归一化 Transformer(nGPT)。研究发现,超球面几何约束赋予了网络原生的 4 位浮点算术鲁棒性。核心 idea:量化鲁棒性本质上由点积累加中的信号相干性决定,nGPT 的超球面几何约束阻断了模型依赖少数无界极端坐标的捷径,迫使优化器在数千维度上建立微弱但一致的正向协同,使得真实信号随维度增长产生相干相长累加,在点积求和中以更大幅度甩开不相干的随机量化噪声。

方法详解

整体框架

标准 Transformer 每一层线性投影在量化时均面临激活异常值冲击与离群点失真,而归一化架构(nGPT)通过将权重与输入特征投影到单位超球面上,重构了底层矩阵乘法(GEMM)的数据分布。在端到端 NVFP4 训练流程中,输入词元向量进入网络后,在各层注意力机制与前馈网络中始终保持归一化状态;所有输入投影在送入低精度张量核心计算前均经过单位球投影,使得动态缩放因子可固定预设,并完全剔除输入侧的 RHT 混杂变换。

在计算 \(Y = WX\) 时,标量输出 \(y = \sum_{k=1}^D w_k x_k\) 对应 \(D\) 维向量的点积。由于权重与激活均有界且模长归一,单点量化误差与基线相当,但其点积输出信噪比(SNR)却高出约 7 dB。随后输出通过超球面缩放因子进行重标定,无需重量级量化干预即可保持全网络端到端梯度的数值稳定。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入词元嵌入序列"] --> B["超球面归一化约束<br/>权重与激活全输入投影至单位超球面"]
    B --> C["点积相干信号累加<br/>弱正相关促使信号相长叠加并大幅提升SNR"]
    C --> D["原生 NVFP4 训练管线<br/>剔除 RHT 与动态缩放并采用固定缩放与随机舍入"]
    D --> E["稳定低比特梯度回传与下游输出"]

关键设计

1. 超球面归一化约束:阻断单坐标主导并消除离群点依赖

传统 Transformer 倾向于在少数特征通道上形成幅值极大的尖锐离群点(outliers),这些无界坐标主导了点积计算,但在 4 位浮点受限的动态范围内会遭遇严重截断与精度损失。nGPT 针对此缺陷,在架构上显式剔除了所有的 RMSNorm 与 LayerNorm 层,转而在注意力的 Query、Key、Value 投影以及 MLP 的两级前向矩阵乘前,对输入和权重施加严格的 \(\ell_2\) 归一化,即约束在单位超球面 \(\mathbb{S}^{D-1}\) 上。本文进一步扩展了官方 nGPT 的实现,将注意力输出投影与 MLP 第二层投影的输入也纳入归一化范畴,实现了真正的全输入超球面约束。由于每个维度的数值上限被几何严格限制,模型无法通过拉大单个坐标的绝对幅值来改变输出,从而迫使优化过程将表征能力分散至所有隐藏维度,从源头上根除了因局部特征幅值爆炸所引发的量化溢出与精度坍塌。

2. 点积相干信号累加:弱正相关驱动的信噪比阶跃

学界直觉通常认为低精度鲁棒性源于压低了单点量化噪声,但层级结构分析表明,在单权重 \(w_k\)、单激活 \(x_k\) 及其逐元素乘积 \(s_k = w_k x_k\) 阶段,nGPT 与标准 GPT 的量化信噪比完全一致(元素级约 19 dB,乘积级约 16.5 dB)。真正的质变发生在全求和阶段 \(S = \sum_{k=1}^D s_k\),nGPT 的点积 SNR 达到 26 dB,相比乘积阶段产生了 +9.4 dB 的累加增益(标准 GPT 仅有 +2.2 dB)。其根本机制在于,超球面约束诱导不同维度的逐元素乘积项之间产生了微弱但系统性的正相关性:

\[\bar{\rho}_s = \frac{1}{D-1}\left(\frac{\text{Var}(\sum_{k=1}^D s_k)}{\sum_{k=1}^D \text{Var}(s_k)} - 1\right)\]

在 NVFP4 量化下,标准 GPT 的乘积项相关系数极小(\(\bar{\rho}_s \approx 9.31 \times 10^{-5}\)),而 nGPT 的相关系数高出 14 倍(\(\bar{\rho}_s \approx 1.32 \times 10^{-3}\))。与此相对,两者的量化噪声相关性均接近于零(\(\bar{\rho}_n \sim 10^{-6}\))。在点积方差分解中,零相关的量化噪声呈现无相干随机游走(\(\text{Var}(N) \approx D \sigma_n^2\)),而带有正相关性的信号方差则表现为 \(\text{Var}(S) \approx D \sigma_s^2 [1 + (D-1)\rho_s]\)。由于 \((D-1)\rho_s\) 随隐藏维度 \(D\) 线性放大,有效信号以相长干涉(constructive accumulation)的方式快速增长,点积信噪比随之产生倍增倍率:

\[\frac{\text{SNR}_{\text{nGPT}}(D)}{\text{SNR}_{\text{GPT}}(D)} \approx \frac{1 + D\,\bar{\rho}_{\text{nGPT}}}{1 + D\,\bar{\rho}_{\text{GPT}}}\]

当维度处于过渡区间(\(755 \ll D \ll 10738\))时,nGPT 的信噪比相对增益随模型宽度线性上升,在主流 \(D=4096\) 规模下稳步扩大,最终在极端大模型尺度饱和于约 14.2 倍。

3. 原生 NVFP4 训练管线:算子解耦与学习率直接迁移

基于超球面约束带来的高点积信噪比与平坦损失曲面,模型对权重摄动与超参数波动的容忍度显著提高。在标准 Transformer 的 NVFP4 训练管线中,为了避免发散,必须插入两道高延迟屏障:在激活张量送入 GEMM 前通过阿达马矩阵变换打乱空间结构以平滑离群值,以及跨张量扫描绝对最大值(amax)以计算每块动态缩放系数。nGPT 架构因为激活张量天然位于单位球面上且方差有界,直接将动态缩放因子替换为固定标量,并完全剔除 RHT 模块,仅保留梯度的随机舍入(SR)。此外,分析表明同时扰动全网络参数时,GPT 损失的恶化速度比 nGPT 快 3.5 倍;这种内在平坦性使得 nGPT 在 BF16 精度下搜索出的最佳学习率可以直接 1:1 迁移至 NVFP4 训练,彻底免除了标准低比特量化训练中通常面临的 16 倍学习率漂移和昂贵的二次超参数网格搜索。

损失函数 / 训练策略

在预训练目标上,模型沿用标准自回归因果语言建模交叉熵损失。优化器选用 AdamW(\(\beta_1=0.9, \beta_2=0.95\)),调度策略采用 Warmup-Stable-Decay (WSD)。不同于标准 GPT 的繁琐优化技巧,超球面归一化架构使得训练策略大幅精简: - 权重衰减与预热移除:在 1.2B 密集模型和 MoE 模型中,nGPT 均完全去除了权重衰减(weight decay 设为 0)和学习率预热(Warmup samples 设为 0),标准 GPT 则需 1M 步预热与 0.1 的衰减系数; - 学习率设定:nGPT 的基准学习率设置为标准模型的一半(例如 1.2B 模型中 nGPT 采用 \(0.6 \times 10^{-3}\),GPT 采用 \(1.2 \times 10^{-3}\)),且在 BF16 与 NVFP4 下保持完全一致; - 量化范围:对所有层的核心矩阵乘法(GEMM)全面开启 NVFP4 精度计算,而将词表嵌入(Embedding)、非线性激活与小批量批处理 GEMM 保留在 BF16。在 3B/30B 超大 MoE 实验中,沿用主流方案将最后的 15% 顶层维持在高精度以保护 logits 分布。

实验关键数据

主实验

在 NVIDIA Blackwell 硬件与 Nemotron-Pretraining-Datasets 数据集上,作者系统评估了 1.2B 参数量密集模型在 1T tokens 训练下的最终预训练 Loss 及各项主流下游零样本/少样本基准性能。

配置 训练精度 Loss ↓ HellaSwag ↑ PIQA ↑ WinoGrande ↑ MMLU (5-shot) ↑ MMLU-Pro (5-shot) ↑ GSM8K (8-shot) ↑
GPT 基线 BF16 1.47 60.54 73.72 58.17 45.23 13.98 33.66
nGPT 归一化 BF16 1.46 61.57 74.27 59.51 45.50 14.34 36.24
GPT 基线 (含RHT/缩放) NVFP4 1.52 58.86 73.78 58.64 42.56 13.06 29.49
nGPT 归一化 (免RHT/免缩放) NVFP4 1.50 59.93 74.37 58.64 43.57 14.59 35.56

从结果可以看出,在 NVFP4 精度下,去除了 RHT 和动态 per-tensor scaling 的 nGPT 不仅训练 Loss 优于标准 GPT(1.50 vs 1.52),而且在 GSM8K(+6.07%)、MMLU-Pro(+1.53%)及 HellaSwag(+1.07%)等高难度推理和常识任务上全面超越基线,甚至在 GSM8K 上逼近其自身 BF16 的水准。

消融实验

为了验证 nGPT 架构优势在不同模型拓扑与训练开销上的普适性,作者在单层吞吐、量化干预依赖度以及混合 Mamba-Transformer MoE 架构(400M/600M 及 3B/30B)上进行了消融与对比。

架构 / 规模 量化配置 训练开销 (RHT / Scale / SR) 训练 Token 量 相对量化误差 (Loss gap) GB200 单层加速比 (vs BF16)
1.2B Dense (GPT) NVFP4 包含 RHT / 包含 Scale / SR 1.0T +3.40% 约 2.4× (大维度下)
1.2B Dense (nGPT) NVFP4 免 RHT / 免 Scale / SR 1.0T +2.74% 3.3× - 3.6×
MoE 400M/600M (GPT) NVFP4 包含 RHT / 包含 Scale / SR 36B 显著正误差 ~2.2×
MoE 400M/600M (nGPT) NVFP4 免 RHT / 免 Scale / SR 36B 显著更低相对误差 ~3.0×
MoE 3B/30B (GPT) NVFP4 包含 RHT / 包含 Scale / SR 500B 持续存在正向 gap ~2.5×
MoE 3B/30B (nGPT) NVFP4 免 RHT / 免 Scale / SR 500B ~0% 相对误差 (无损) ~3.5×

关键发现

  • 量化相对误差在超大 MoE 上归零:在 3B 激活 / 30B 总参数的超大混合 Mamba-Transformer MoE 上训练 500B tokens,nGPT-NVFP4 的相对误差曲线贴合在 0% 附近,实现了与 BF16 训练几乎无损的对齐,证明该架构对复杂稀疏路由与循环状态同样具有普适鲁棒性;
  • Blackwell 硬件端吞吐实现 3.3×–3.6× 飞跃:在 GB200 GPU 上的单层微基准测试中,随着隐藏层宽度 \(D\) 增大,GEMM 耗时占据绝对主导。nGPT 剔除动态 amax 规约和 RHT 之后,成功打通了极简 NVFP4 张量核心加速路径,达到 BF16 基线 3.6 倍的执行速度;
  • 学习率超参敏感度近乎免疫:在跨越两个数量级的学习率扫描中,nGPT 在 BF16 与 NVFP4 下的验证集 BPB(bits-per-byte)曲线均呈现极宽的平坦底部,最佳学习率完全重合;而标准 GPT 在量化后最佳学习率突变了 16 倍,若不重新细致扫参会导致发散或剧烈性能退化。

亮点与洞察

  • 颠覆了量化鲁棒性的优化归因:传统认知一味追求减小单点截断噪声,本文用严格的数据分解证明,nGPT 的优势根本不在于单点噪声抑制,而在于几何约束使点积信号呈现弱正相关相长累加,以结构化信号增长超越非相干噪声;
  • 量化自适应性由架构内在赋予:无需引入复杂的动态量化调度器或可微缩放估算器,仅通过超球面归一化这一单纯的数学投影,就从根本上消除了对阿达马变换等高开销前处理的依赖;
  • 极具工程落地价值的超参迁移性:超低精度训练最耗费算力的是超参数网格搜索,nGPT 展现出的 BF16 到 NVFP4 最优学习率零漂移特性,能够大幅缩短新模型进入低精度预训练的试错周期。

局限与展望

  • 理论机制的深层因果尚待闭环:论文发现了超球面约束必然伴随正向相关性 \(\bar{\rho}_s\) 产生这一经验现象,但未能从微分几何或随机矩阵理论层面给出从超球投影到点积正相关的严密解析证明;
  • 预训练 Token 预算相对有限:在 3B/30B MoE 实验中仅推进至 500B tokens,距离生产级开源前沿模型普遍消耗的 15T–20T tokens 仍有差距,万亿规模超长周期的漂移现象有待进一步验证;
  • 部分非线性算子尚未量化:当前实现中注意力内部 Batched-GEMM 与非线性激活仍运行在 BF16,未来若实现真正的“All-NVFP4”端到端全管线,可能需要更深入的流形算子融合。

相关工作与启发

  • vs 标准 NVFP4 预训练方案 (Abecassis et al., 2025):基线方法依赖高复杂度的 RHT 离群点平滑与动态张量 amax 缩放来防止 4 位训练发散;本文通过 nGPT 超球面约束从架构层面消除离群点,彻底剥离了 RHT 和动态缩放,实现了更低相对误差与更高硬件吞吐;
  • vs 自适应块缩放与量化估算器 (Cook et al., 2025 / Wang et al., 2025):现有改进工作主要集中于微调量化网格与反向梯度方差估计;本文证明模型几何与优化动态本身是更本质的变量,架构级归一化可以大幅简化底层量化算法的工程实现。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次从架构与超球面几何角度揭示 4 位量化鲁棒性本质,彻底打破了仅靠算法打补丁的思维定势
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 1.2B 密集模型与高达 30B MoE 架构,并在 Blackwell GB200 硬件上完成了端到端真实吞吐与下游评测
  • 写作质量: ⭐⭐⭐⭐⭐ 逻辑严密,由浅入深,从信噪比拆解、相关性度量到宽度缩放定律步步推导,论证令人信服
  • 价值: ⭐⭐⭐⭐⭐ 对未来万亿参数规模大模型的 4 位甚至更低比特原生硬件预训练具有极其深远的指导意义与实用价值