High-Dimensional Learning Dynamics of Quantized Models with Straight-Through Estimator

Quantized neural network training optimizes a discrete, non-differentiable objective. The straight-through estimator (STE) enables backpropagation through surrogate gradients and is widely used. While previous studies have primarily focused on the properties of surrogate gradients and their convergence, the influence of quantization hyperparameters, such as bit width and quantization range, on learning dynamics remains largely unexplored. We theoretically show that in the high-dimensional limit, STE dynamics converge to a deterministic ordinary differential equation. This reveals that STE training exhibits a plateau followed by a sharp drop in generalization error, with plateau length depending on the quantization range. A fixed-point analysis quantifies the asymptotic deviation from the unquantized linear model. We also extend analytical techniques for stochastic gradient descent to nonlinear transformations of weights and inputs.

Modern machine learning models are often very large, which makes them expensive to run on devices with limited memory or computing power. Quantization is a common technique that reduces this cost by representing model values with fewer bits, but training such low-precision models is difficult because the usual gradient-based methods no longer apply directly. This paper studies a widely used workaround called the straight-through estimator and analyzes how training behaves when both model weights and inputs are quantized. We show that the number of bits and the quantization range can strongly affect learning speed, final accuracy, and stability. Our results help explain when quantization may slow learning, when it can improve stability, and how quantization settings should be chosen more carefully.