LOTION: Smoothing the Optimization Landscape for Quantized Training¶

会议: NeurIPS 2026
arXiv: 2510.08757
领域: 模型压缩
关键词: 神经网络量化, 量化感知训练, 随机舍入, 损失函数平滑, 二阶曲率正则

一句话总结¶

针对低比特量化训练中量化算子分段常数导致的梯度消失与直通估计器（STE）不稳定性，LOTION 提出在随机舍入噪声的期望下直接平滑量化损失函数，证明其保持原始量化全局极小值并导出等价于对角高斯-牛顿曲率正则的连续可微目标，在 150M/300M 语言模型 INT4/FP4 训练中显著优于传统 QAT。

研究背景与动机¶

大语言模型在现代加速器上的部署正受到内存带宽与访存开销的严重制约，参数在显存层级中的搬运延迟促使低比特量化技术（如 INT4、FP4、INT8）成为实际推理落地的标准配置。然而，将量化算子直接引入神经网络优化循环会引入本质上的非连续性：前向传播将连续全精度权重映射到有限码本（codebook），导致目标函数的导数几乎处处为零，仅在量化阈值跃迁边界处呈现未定义的阶跃。后训练量化（PTQ）尽管绕开了非连续优化，但在 4-bit 等激进低比特宽度下精度严重劣化；量化感知训练（QAT）则普遍依赖直通估计器（Straight-Through Estimator, STE），在反向传播中将量化算子强行替代为恒等映射。

这种经验性的梯度松弛策略存在根本性的理论与优化矛盾。STE 缺乏收敛性保证，在 INT4 或 FP4 等超低比特场景下极易诱发严重的梯度方差与震荡，导致训练轨迹迅速陷入次优平台甚至完全发散。现有改进工作多停留在经验性的梯度缩放、步长参数学习或退火噪声注入，往往需要复杂的超参数调优，且未能真正解决量化损失表面本质非连续的根本症结。优化目标本身的分段常数特性使得常规一阶与二阶凸优化和非光滑优化收敛理论彻底失效。

本文的切入角度是跳出反向传播阶段经验修剪梯度的传统思维，受 Nesterov 平滑与随机平滑理论启发，直接在损失函数层面构造连续代理。核心 idea：利用无偏随机舍入噪声在期望下平滑离散量化损失曲面，在严格保持原量化目标全局极小值的前提下建立处处可微的平滑替代函数，并将其二阶展开解析映射为对角高斯-牛顿曲率正则项，实现具备收敛理论保障且计算极轻量的全精度平滑优化。

方法详解¶

整体框架¶

LOTION（Low-precision Optimization via sTochastic-noIse smOothiNg）的核心思想是将不可微且分段常数的离散优化问题转化为连续可微的随机扰动期望优化问题。在整套训练管线中，模型权重在参数空间保持全精度（FP32）维护与更新，无需依赖启发式 STE 梯度替换。针对前向传播中的非连续映射，LOTION 定义参数在无偏随机舍入分布下的期望损失作为优化目标。通过二阶泰勒展开与正定高斯-牛顿矩阵近似，该平滑损失被解析推导为一个标准全精度任务损失与一个数据驱动的对角曲率惩罚项之和。曲率惩罚强度由参数所在量化区间的舍入方差及局部 Hessian 谱共同决定，从而惩罚处于损失曲面尖锐方向上的量化误差。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["全精度权重 w (FP32)"] --> B["细粒度共享缩放因子构建<br/>计算分块绝对最大值尺度 s_B 与残差 Δ_i"]
    B --> C["无偏随机舍入平滑算子<br/>期望平滑保持原始量化全局极小值"]
    C --> D["高斯-牛顿二阶展开正则化<br/>经验 Fisher 估计对角曲率 g_ii"]
    D --> E["联合目标更新: L(w) + 0.5 * λ * Σ g_ii * σ_i^2"]
    E --> F["标准优化器反向传播 (Adam/SGD)"]
    F -->|无 STE 梯度替换| A
    E --> G["评估/推理部署: 确定性 RTN 或随机舍入 RR 量化"]

关键设计¶

1. 细粒度共享缩放因子构建：无溢出绝对最大值自适应网格

传统仿射量化通常需要零点（zero-point）与截断（clipping）参数，在权重量化过程中会引入额外的计算与显存开销。LOTION 采用细粒度共享缩放的有符号对称整数量化框架（如 DeepSeek FP8 / LLM.int8() 采用的分块机制）。权重参数被划分为均匀小块 \(B\)（最小可退化为单个元素），每个分块计算基于绝对最大值的浮点缩放因子 \(s_B\)，将块内权重缩放到无符号整型表示范围 \([-(2^{n-1}-1), 2^{n-1}-1]\) 内：

\[s_B = \frac{\max_{i \in B} |w_i|}{2^{n-1}-1}, \quad z_i' = \frac{w_i}{s_B}, \quad z_i = \lfloor z_i' \rceil, \quad [\text{cast}(w)]_i = s_B z_i\]

该构建天然保证 \(|w_i| \le (2^{n-1}-1)s_B\)，因而在量化映射中无需任何显式截断。由于缩放尺度 \(s_B\) 随权重优化动态移动，量化网格与网络参数在训练动态中保持紧密耦合，为后续舍入噪声方差的推导奠定了自适应尺度的几何基础。

2. 无偏随机舍入平滑算子：保持原始量化全局极小值的连续代理

为了根除分段常数损失造成的零梯度问题，LOTION 将非连续目标函数 \(\mathcal{L}(\text{cast}(w))\) 替换为其在随机舍入（Randomized Rounding, RR）分布下的期望：\(\mathcal{L}_{\text{smooth}}(w) = \mathbb{E}_{q \sim \text{RR}(w)} [\mathcal{L}(q)]\)。定义归一化浮点数值 \(z_i' = w_i / s_B\) 与其量化下界距离 \(\Delta_i = z_i' - \lfloor z_i' \rfloor \in [0, 1]\)，随机舍入算子独立作用于每个坐标：

\[R_i = \begin{cases} \lfloor z_i' \rfloor, & \text{概率 } \lceil z_i' \rceil - z_i' = 1 - \Delta_i \\ \lceil z_i' \rceil, & \text{概率 } z_i' - \lfloor z_i' \rfloor = \Delta_i \end{cases}, \quad \varepsilon_i = s_B (R_i - z_i')\]

该舍入算子满足严格的三大性质：(1) 无偏性：\(\mathbb{E}[\varepsilon_i] = 0\)，即 \(\mathbb{E}_{q \sim \text{RR}(w)}[q] = w\)；(2) 连续且局部有界：在 Wasserstein-\(2\) 度量下关于 \(w\) 连续；(3) 网格点保真性：若 \(w \in Q\) 满足 \(\text{cast}(w) = w\)，则 \(\text{RR}(w)\) 以概率 1 落在 \(w\) 本身。作者在理论上证明了两个核心引理：首先，平滑损失 \(\mathbb{E}_{q \sim \text{RR}(w)}[\mathcal{L}(q)]\) 关于 \(w\) 连续且几乎处处可微，使标准一阶与二阶优化器具备传统非光滑凸/非凸收敛保证；其次，\(\min_{w \in \mathbb{R}^d} \mathbb{E}_{q \sim \text{RR}(w)} [\mathcal{L}(q)] = \min_{w \in \mathbb{R}^d} \mathcal{L}(\text{cast}(w))\)，证明了平滑过程绝对不会偏移或破坏原始离散量化曲面的全局最优解。

3. 高斯-牛顿二阶展开正则化：感知曲率的自适应脊惩罚

直接计算高维网络上的期望扰动需要昂贵的蒙特卡洛多样本前向传递。LOTION 创新性地对平滑损失实施二阶泰勒展开。由于随机舍入噪声的均值为零（\(\mathbb{E}[\varepsilon] = 0\)），一阶梯度项直接抵消，从而导出紧凑的二阶解析形式：

\[\mathcal{L}_{\text{smooth}}(w) = \mathcal{L}(w) + \frac{1}{2} \operatorname{tr}(H(w) \Sigma_\varepsilon(w)) + \mathcal{O}(\mathbb{E}[\|\varepsilon\|^3])\]

由于各坐标舍入相互独立，噪声协方差矩阵为严格对角阵 \(\Sigma_\varepsilon = \operatorname{diag}(\sigma_1^2, \dots, \sigma_d^2)\)，其中各分量的方差具备精确闭式解 \(\sigma_i^2 = s_{B(i)}^2 \Delta_i (1 - \Delta_i)\)。针对非凸深度网络中全 Hessian 矩阵可能出现的负曲率问题，LOTION 采用外积形式的正半定高斯-牛顿（Gauss-Newton）矩阵 \(G(w) = \nabla_w f^\top \nabla_f^2 \ell \nabla_w f\) 替代全 Hessian，得到代理目标：

\[\mathcal{L}_{\text{GN}}(w) = \mathcal{L}(w) + \frac{1}{2} \sum_{i=1}^d g_{ii}(w) s_{B(i)}^2 \Delta_i (1 - \Delta_i)\]

该设计在物理意义上等价于向原始任务损失注入了一个加权 \(\ell_2\) 脊正则化项。其惩罚权重由两部分相乘决定：一是曲率项 \(g_{ii}(w)\)，二是舍入方差 \(s_B^2 \Delta_i (1 - \Delta_i)\)。当参数处于尖锐极小值方向（\(g_{ii}\) 极大）或处于量化区间正中央（\(\Delta_i = 0.5\)，舍入不确定性与方差达到峰值 \(0.25 s_B^2\)）时，正则项施加最强惩罚；反之，若参数自然靠近整数格点（\(\Delta_i \to 0\) 或 \(1\)）或处于平坦盆地（\(g_{ii} \to 0\)），惩罚趋于零。这使得网络自动向着低曲率、容忍量化误差的平坦区域演化。

损失函数 / 训练策略¶

在大规模语言模型（如 Transformer / OLMo 架构）的实际训练中，精确计算高斯-牛顿对角线开销过大。LOTION 采用类 Adam 的经验 Fisher 信息矩阵对角线作为实用近似，即通过反向传播梯度的平方滑动平均累积估计 \(g_{ii}(w) \approx \widehat{\mathbb{E}}[(\nabla_{w_i} \ell)^2]\)。

为降低计算图负担并提升吞吐，实践中对经验 Fisher 项阻断反向传播梯度（stop-gradient），并引入全局正则化超参数 \(\lambda\) 进行缩放：

\[\mathcal{L}_{\text{LOTION}}(w) = \mathcal{L}(w) + \frac{\lambda}{2} \sum_{i=1}^d \operatorname{sg}[g_{ii}(w)] \cdot s_{B(i)}^2 \Delta_i (1 - \Delta_i)\]

由于舍入方差项的计算仅涉及分块最大值尺度与逐元素取模距离，计算完全可向量化并行，计算额外开销几乎可以忽略不计。全流程在 FP32 下使用标准优化器（AdamW / SGD）进行常规反向更新，无需任何量化正向/反向模拟算子；在训练结束或周期性验证时，直接使用确定性四舍五入（Round-to-Nearest, RTN）或随机舍入（RR）固化为低比特权重部署。

实验关键数据¶

主实验¶

论文在两类主流设置下评估 LOTION：一是在合成高维二次损失及双层线性网络（\(d=12000\)）上验证优化动力学；二是在 C4 预训练数据集上基于 OLMo 架构预训练 150M 与 300M 参数的语言模型，采用 Chinchilla 最优计算预算（tokens 数量为参数量的 20 倍）以及 5× Chinchilla 预算（100× 参数量），系统对比 PTQ、QAT、RAT（Rounding-Aware Training，以随机舍入作为前向模拟并依赖 STE 反向）与 LOTION 在 INT4、INT8 及 FP4 格式下的量化验证损失。

下表列出 150M 参数语言模型在 C4 验证集上的最终交叉熵损失对比（评估分别采用随机舍入 RR 与四舍五入 RTN）：

方法	评估量化方式	INT4 验证损失	INT8 验证损失
PTQ	RR	4.686	3.092
PTQ	RTN	3.864	3.085
RAT (STE 梯度)	RR	3.608	3.090
QAT (STE 梯度)	RTN	3.315	3.079
LOTION (本文)	RR	3.295	3.078
LOTION (本文)	RTN	3.276	3.076

在 300M 参数语言模型上的最终交叉熵验证损失对比如下：

方法	评估量化方式	INT4 验证损失	INT8 验证损失
PTQ	RR	3.9745	3.0045
PTQ	RTN	3.6062	2.9992
QAT (STE 梯度)	RTN	3.2230	2.9972
LOTION (本文)	RR	3.1772	2.9959
LOTION (本文)	RTN	3.1031	2.9905

消融实验与机制分析¶

论文在合成高维二次优化任务（\(d=12000\)，Hessian 谱服从幂律衰减 \(\lambda_i \propto 1/i^{1.1}\)）上进行了严格的机制对比实验。此时理论全局最优为目标权重 \(w^*\)，PTQ 直接量化 \(w^*\)，QAT 与 RAT 依赖直通估计器，LOTION 采用精确二阶对角展开式进行平滑优化。

优化算法	评估指标 (INT4 验证损失)	优化收敛状态与机理解析
LOTION (RR)	0.13988	损失单调平滑下降，找到比解析目标更抗量化的全精度解
LOTION (RTN)	0.14419	确定性舍入下依然保持最低损失，显著收敛于平坦极小值
PTQ (RTN)	0.20566	直接量化理论最优全精度解 \(w^*\)，因忽略离散网格而落入次优
RAT (RR)	0.33230	随机前向+STE，Lemma 3 证明其梯度期望退化为普通梯度无正则效果
PTQ (RR)	0.40449	随机量化 \(w^*\)，误差方差未经曲率自适应加权
QAT (RTN)	0.79181	损失轨迹剧烈锯齿状震荡，过早陷入高方差局部停滞平台

关键发现¶

跳出 STE 的虚假收敛平台：在 150M 模型的 5× Chinchilla 扩展实验中，QAT 在训练中后期由于 STE 梯度噪声累积而陷入损失平台期无法继续下降，而 LOTION 的量化验证损失保持稳定持续下降，验证了连续平滑曲面在长程训练中的可优化优势。
超低比特与现代 FP4 格式的强泛化性：在非均匀的现代 FP4 格式量化中，尽管浮点表示缓解了近零区域的截断误差，PTQ 依然很早停滞；LOTION 依然显著超越 QAT，证实了曲率感知平滑在非均匀微小尺度下的普适性。
RAT 理论退化证伪：论文在理论上证明了在二次损失下，基于随机舍入模拟前向并配合 STE 的算法（RAT）其梯度期望恰好等于原始全精度梯度（\(\mathbb{E}_\varepsilon [\nabla \mathcal{L}(w+\varepsilon)] = \nabla \mathcal{L}(w)\)），因而在机理上无法提供曲率正则化收益，实验结果中 RAT 性能（0.3323）大幅劣于 LOTION（0.1398）。

亮点与洞察¶

全局极小值保真理论：LOTION 从数学上严格证明了无偏且局部有界的随机舍入扰动期望损失不仅是连续可微函数，而且其在全精度空间内的全局极小值与原始离散分段常数量化损失的全局极小值严格重合，从根本上消除了代理函数带来的目标漂移风险。
随机扰动与二阶平滑的解析对齐：论文揭示了参数化随机舍入与高斯-牛顿曲率正则之间的对偶等价性，将原本需要昂贵多样本前向采样的平滑算法简化为一个极轻量的逐元素解析惩罚项，单次计算复杂度仅为 \(\mathcal{O}(d)\)。
即插即用且解耦前向：该方法完全在全精度域内运行，训练阶段彻底剥离了模拟量化与假量化算子（FakeQuantize），不改变模型的前向图拓扑与内存排布，具备极高的工程可迁移性。

局限与展望¶

边界处导数非完全光滑：随机舍入诱导的平滑损失曲面在量化区间的离散分界点处导数依然存在不连续或未定义的情况，未来可探索结合连续核密度平滑以实现全域任意阶可微。
激活值量化的扩展空白：当前 LOTION 主要针对权重参数的量化感知训练，未将输入激活值动态量化纳入平滑框架。将自适应随机平滑拓展到变长序列的激活值分块量化是实现端到端低比特计算训练的重要方向。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙建立随机舍入与高斯-牛顿曲率正则的理论等价性，彻底避开 STE 启发式修补。
实验充分度: ⭐⭐⭐⭐☆ 涵盖高维合成动力学分析以及 150M/300M LLM 预训练，INT4/INT8/FP4 覆盖全面。
写作质量: ⭐⭐⭐⭐⭐ 理论推导清晰严谨，从二次动机过渡到通用高斯-牛顿近似的逻辑自洽通透。
价值: ⭐⭐⭐⭐☆ 为低比特甚至 4-bit 以下大模型训练提供了兼具收敛保证与高吞吐易部署特性的新范式。