We develop a principled method for quantization-aware training (QAT) of large-scale machine learning models. Specifically, we show that convex, piecewise-affine regularization (PAR) can effectively induce the model parameters to cluster towards discrete values. We minimize PAR-regularized loss functions using an aggregate proximal stochastic gradient method (AProx) and prove that it has last-iterate convergence. Our approach provides an interpretation of the straight-through estimator (STE), a widely used heuristic for QAT, as the asymptotic form of PARQ. We conduct experiments to demonstrate that PARQ obtains competitive performance on convolution- and transformer-based vision tasks.
@article{arxiv.2503.15748,
title = {PARQ: Piecewise-Affine Regularized Quantization},
author = {Lisa Jin and Jianhao Ma and Zechun Liu and Andrey Gromov and Aaron Defazio and Lin Xiao},
journal= {arXiv preprint arXiv:2503.15748},
year = {2025}
}