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

2025-10-12

AI Summary

Straight-through estimators (STE) are widely adopted in quantized neural network training due to the discrete, non-differentiable nature of the objective, yet the mechanistic role of hyperparameters—such as bit-width and quantization range—in governing STE’s learning dynamics remains poorly understood. Method: We rigorously derive, in the high-dimensional limit, that STE dynamics converge to a deterministic differential equation. We analyze fixed points to quantify the asymptotic bias relative to unquantized linear models and extend the framework to nonlinear settings—incorporating non-convex optimization and SGD theory—to model proxy-gradient dynamics for both weights and inputs. Results: Our analysis reveals that quantization range primarily governs the duration of the generalization error plateau, while bit-width controls convergence speed. We observe and explain the characteristic plateau-and-drop behavior in generalization error evolution. This work establishes the first analytically tractable, high-dimensional dynamical systems framework for efficient quantized training.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Search and Optimization: Learning to Search
• Reasoning under Uncertainty: Stochastic Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Analyzes how quantization hyperparameters affect neural network learning dynamics
• Models STE training convergence using high-dimensional ordinary differential equations
• Quantifies generalization error plateau and deviation from unquantized models

Innovation

• STE enables backpropagation via surrogate gradients
• Quantization hyperparameters influence learning dynamics plateau
• High-dimensional STE converges to deterministic differential equation