Title: LOTION: Smoothing the Optimization Landscape for Quantized Training
Primary area: foundation or frontier models, including LLMs
Keywords: optimization, large language models, precision, quantization, quantization aware training

Abstract: 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, Low-precision Optimization via sTochastic-noIse smOothiNg, 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.

TL;DR: We introduce LOTION, a principled smoothing framework that replaces the raw, discontinuous quantized loss with its expectation under unbiased stochastic-rounding noise