PARQ: Piecewise-Affine Regularized Quantization

We develop a novel optimization method for quantization-aware training (QAT). Specifically, we show that convex, piecewise-affine regularization (PAR) can effectively induce neural network weights to cluster towards discrete values. We minimize PAR-regularized loss functions using an aggregate proximal stochastic gradient method (AProx) and prove that it enjoys 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.

Authors

• Lisa Jin
• Jianhao Ma
• Zechun Liu
• Andrey Gromov
• Aaron Defazio
• Lin Xiao

Keywords

• quantization-aware training
• proximal gradient method
• convex regularization
• stochastic gradient method
• last-iterate convergence