Optimizing neural networks for quantized objectives is fundamentally challenging because the quantizer is piece-wise constant, yielding zero gradients everywhere except at quantization thresholds where the derivative is undefined. Most existing methods deal with this issue by relaxing gradient computations with techniques like Straight Through Estimators (STE) and do not provide any guarantees of convergence. In this work, taking inspiration from Nesterov smoothing, we approximate the quantized loss surface with a continuous loss surface. In particular, we introduce LOTION, \textbf{L}ow-precision \textbf{O}ptimization via s\textbf{T}ochastic-no\textbf{I}se sm\textbf{O}othi\textbf{N}g, a principled smoothing framework that replaces the raw quantized loss with its expectation under unbiased randomized-rounding noise. In this framework, standard optimizers are guaranteed to converge to a local minimum of the loss surface. Moreover, when using noise derived from stochastic rounding, we show that the global minima of the original quantized loss are preserved. We empirically demonstrate that this method outperforms standard QAT on synthetic testbeds and on 150M- and 300M- parameter language models.
翻译:为量化目标优化神经网络具有根本性挑战,因为量化器是分段常数函数,除量化阈值处导数未定义外,其余位置梯度处处为零。现有方法大多通过直通估计器(STE)等技术松弛梯度计算来处理此问题,且无法提供任何收敛性保证。本研究受Nesterov平滑方法启发,通过连续损失曲面逼近量化损失曲面。具体而言,我们提出LOTION(基于随机噪声平滑的低精度优化框架),该原理性平滑框架将原始量化损失替换为无偏随机舍入噪声下的期望损失。在此框架中,标准优化器可保证收敛至损失曲面的局部极小值。进一步地,当使用随机舍入生成的噪声时,我们证明原始量化损失的全局极小值得以保持。实验结果表明,该方法在合成测试集及1.5亿/3亿参数语言模型上均优于标准量化感知训练(QAT)。