Neural Network Quantization by Learning Low-Loss Subspaces

Traditional neural network quantization often relies on Quantization-Aware Training (QAT), which uses non-differentiable rounding operations and surrogate gradients (like the Straight-Through Estimator) to simulate low-precision behavior. These methods can be unstable and sensitive to initialization. This paper asks whether one can instead construct a quantized model by identifying a low-loss subspace within the full-precision (FP) parameter space, where the model is inherently robust to discretization.

Instead of optimizing a single point in the weight space, the authors propose Quantization by Learning Subspaces (QLS). They parameterize the weights as a linear path between two trainable endpoints, Θ(1) and Θ(2). During training, they apply a regularization term that forces the distance between these endpoints to be at least as large as the quantization scale factor. This ensures that the midpoint of the subspace—the average of the two endpoints—is "quantization-friendly." Because the training objective is optimized entirely in full-precision without rounding, it avoids the biases and instabilities associated with surrogate gradients. After training, the model is quantized by simply rounding the midpoint of the learned subspace.

• Performance: The QLS method produces quantized models with accuracy comparable to, and often exceeding, standard QAT.
• No Discretization During Training: The method avoids explicit rounding and the use of the Straight-Through Estimator, leading to more stable optimization.
• Cross-Bit Adaptability: Models trained with QLS are less sensitive to changes in bit-width and can be deployed at different precision levels without needing to retrain.