Low-Dimensional High-Leverage Subspace Optimization: Beyond Full-Parameter Coupled Training for Neural Network Quantization

Research Question and Approach

Low-bit quantization often leads to significant accuracy degradation in compact neural networks, primarily because standard training paradigms treat all parameters as equally important, ignoring the heterogeneity of different parameter subspaces. The authors investigate whether specific, low-dimensional parameter subsets can be leveraged to absorb quantization errors more effectively than full-parameter joint training. They identify the affine parameters (scale and shift) in normalization layers (Batch Normalization and RMSNorm) as a critical 'response-control' subspace.

To exploit this, the authors propose Normalization Affine Preconditioning (NAP). Unlike conventional Post-Training Quantization (PTQ) or Quantization-Aware Training (QAT), NAP freezes the backbone weights and focuses exclusively on fine-tuning these normalization affine parameters under the target fake-quantization graph. This approach is applied in three ways: as a pre-conditioning step before PTQ, as a post-QAT adaptation stage for saturated models, and as an alternating optimization scheme that decouples feature learning from numerical calibration.

The study demonstrates that normalization affine parameters act as a high-leverage subspace capable of canceling the structured channel-wise affine component of quantization distortion. Experiments on ImageNet, CIFAR-100, and large language models (e.g., Qwen2.5-3B) show that NAP consistently recovers performance in severely collapsed low-bit settings.

Notably, NAP outperforms saturated full-parameter QAT with significantly lower tuning costs and fewer trainable parameters. The authors provide a theoretical analysis showing that while these affine parameters can perfectly compensate for channel-wise scale and shift errors, they are bounded by irreducible non-linear residuals (rounding and clipping). The distillation-guided NAP is interpreted as a local projection of the teacher-student logit mismatch onto the restricted normalization-affine subspace.