LOTION: Smoothing the Optimization Landscape for Quantized Training (2510.08757v1)

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, \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 LLMs.

• The paper introduces LOTION, a loss-smoothing framework that replaces the discontinuous quantized loss with its expectation under randomized rounding to preserve global minima.
• It employs curvature-aware regularization via Hessian or Gauss-Newton approximations, stabilizing training compared to traditional QAT and PTQ methods.
• LOTION demonstrates consistent performance gains on synthetic tasks, two-layer networks, and large language models, particularly at aggressive low bit-width quantization.

LOTION: Smoothing the Optimization Landscape for Quantized Training

Introduction and Motivation

Quantization is a critical technique for reducing the memory and computational footprint of neural networks, especially for LLMs where inference is bandwidth-bound. However, optimizing neural networks directly for quantized objectives is fundamentally challenging due to the discontinuous, piecewise-constant nature of quantizers, which results in vanishing gradients almost everywhere. Existing approaches, such as Quantization-Aware Training (QAT) with Straight-Through Estimators (STE), lack theoretical convergence guarantees and often become unstable at low bit-widths (e.g., INT4). Post-Training Quantization (PTQ) avoids instability but typically underperforms QAT at aggressive quantization levels.

This work introduces LOTION (Low-precision Optimization via sTochastic-noIse smOothiNg), a principled loss-smoothing framework inspired by Nesterov smoothing. LOTION replaces the discontinuous quantized loss with its expectation under unbiased randomized-rounding noise, yielding a continuous, almost-everywhere differentiable objective. This approach preserves the global minima of the original quantized loss and enables the use of standard optimizers with convergence guarantees.

Methodology: Randomized Rounding and Loss Smoothing

The core idea of LOTION is to smooth the non-differentiable quantized loss by considering the expected loss under a randomized rounding distribution. For a neural network with weights $w$, the quantized loss is $\mathcal{L}(\mathrm{cast}(w))$, where $\mathrm{cast}$ is the quantization operator. LOTION defines the smoothed loss as:

$$L_{\mathrm{smooth}}(w) = \mathbb{E}_{q \sim \mathrm{RR}(w)}[\mathcal{L}(q)]$$

where $\mathrm{RR}(w)$ is a randomized rounding distribution satisfying three properties: unbiasedness ($\mathbb{E}[q] = w$), continuity, and idempotence on quantized points.

This smoothing yields two key theoretical properties:

• Continuity: $L_{\mathrm{smooth}}(w)$ is continuous, unlike the original quantized loss.
• Minima Preservation: The global minima of $L_{\mathrm{smooth}}(w)$ coincide with those of the original quantized loss.

Quadratic Loss Case

For quadratic losses, the expectation under randomized rounding admits a closed-form:

$$\mathcal{L}_{\text{smooth}}(w) = \mathcal{L}(w) + \frac{1}{2} \operatorname{tr}(H \Sigma_\varepsilon)$$

where $H$ is the Hessian and $\Sigma_\varepsilon$ is the covariance of the rounding noise. This introduces a curvature-aware regularizer, penalizing sharp directions in the loss landscape.

Figure 2: LOTION smooths the discontinuous quantized loss surface, yielding a continuous surrogate that shares the same minima as the quantized loss.

General Losses and Gauss-Newton Regularization

For general (non-quadratic) losses, a second-order Taylor expansion leads to a similar regularization term, where the Hessian is replaced by the Gauss-Newton matrix $G(w)$ for stability:

$$L_{\mathrm{GN}}(w) = L(w) + \frac{1}{2} \sum_{i=1}^d g_{ii}(w) s_{B(i)}^2 \Delta_i (1 - \Delta_i)$$

Here, $g_{ii}(w)$ is the $i$-th diagonal of the Gauss-Newton matrix, $s_{B(i)}$ is the block scale, and $\Delta_i$ is the fractional distance to the nearest quantization bin. This regularizer is efficiently computable using empirical Fisher approximations, as in Adam or Sophia.

Experimental Results

Synthetic Linear Regression

LOTION was evaluated on a high-dimensional linear regression task with a power-law spectrum for the input covariance, mimicking neural network Hessians. LOTION consistently achieved lower quantized validation loss than both QAT and PTQ, with QAT exhibiting unstable, jagged loss curves and early plateaus.

Figure 3: LOTION achieves lower INT4 quantized validation loss than QAT and PTQ on a synthetic linear regression task.

Two-Layer Linear Networks

Scaling the hidden dimension $k$ in a two-layer linear network, the quantized loss for all methods decreases as $k$ increases, but LOTION maintains a consistent advantage over QAT and PTQ across all model sizes.

Figure 4: LOTION outperforms QAT and PTQ in quantized loss as the hidden dimension $k$ increases in a two-layer linear network.

LLMs

LOTION was applied to 150M- and 300M-parameter LLMs trained on C4 with Chinchilla-optimal token budgets. At INT4 and INT8 precision, LOTION achieved lower post-quantization validation loss than both QAT and PTQ. Notably, the performance gap persisted even with extended training budgets (5x Chinchilla), with QAT plateauing while LOTION continued to improve.

Figure 5: On a 150M-parameter model, LOTION achieves lower INT4 quantized validation loss than QAT, with continued improvement as training progresses.

Figure 6: On a 300M-parameter model, LOTION outperforms QAT and PTQ at both INT4 and INT8 precision.

FP4 Quantization

LOTION also demonstrated superior performance over QAT and PTQ under FP4 quantization, a non-uniform format favored for its ability to represent small values and outliers. The performance gap, while reduced compared to INT4, remained significant.

Figure 1: LOTION maintains a performance advantage over QAT and PTQ under FP4 quantization.

Implementation Considerations

• Randomized Rounding: For each parameter, compute the fractional distance to the nearest quantization bins and sample accordingly. This is trivially parallelizable and incurs negligible overhead.
• Regularizer Computation: The curvature-aware regularizer can be computed using the empirical Fisher information, requiring only accumulation of squared gradients.
• Scalability: All experiments, including those on 300M-parameter models, were performed with simulated quantization in FP32, making the approach hardware-agnostic and compatible with standard training pipelines.
• Hyperparameters: The regularization strength $\lambda$ can be tuned, but the method is otherwise parameter-free and robust across quantization formats and optimizers.

Theoretical and Practical Implications

LOTION provides a theoretically grounded alternative to STE-based QAT, offering:

• Convergence Guarantees: Standard optimizers converge to stationary points of the smoothed loss.
• Minima Preservation: The global minima of the quantized loss are preserved.
• Curvature Awareness: The regularizer penalizes sharp minima, improving quantization robustness.

Empirically, LOTION consistently outperforms QAT and PTQ, especially at low bit-widths and in large-scale settings. The method is extensible to other quantization formats and potentially to activation quantization, which remains an open direction.

Limitations and Future Directions

While LOTION smooths the loss surface, it does not fully eliminate non-smoothness at quantization bin boundaries. Exploring alternative noise distributions that preserve global minima while yielding even smoother objectives is a promising avenue. Extending the framework to activation quantization and mixed-precision regimes could further enhance efficiency.

Conclusion

LOTION introduces a principled, loss-level smoothing framework for quantized training, leveraging randomized rounding to yield a continuous, curvature-aware objective. The method provides both theoretical guarantees and strong empirical performance across synthetic and large-scale language modeling tasks, outperforming existing QAT and PTQ baselines, particularly at aggressive quantization levels. LOTION represents a robust foundation for future research in quantization-friendly neural network optimization.