Improving the Straight-Through Estimator with Zeroth-Order Information
Ningfeng Yang, Tor Aamodt
We study the problem of training neural networks with quantized parameters.   Learning low-precision quantized parameters by enabling computation of gradients via the Straight-Through Estimator (STE) can be challenging.  While the STE enables back-propagation, which is a first-order method, recent works have explored the use of zeroth-order (ZO) gradient descent for fine-tuning.  We note that the STE provides high-quality biased gradients and ZO gradients are unbiased but can be expensive.  We thus propose First-Order-Guided Zeroth-Order Gradient Descent (FOGZO) that reduces STE bias while reducing computations relative to ZO methods.   Empirically we show FOGZO improves the tradeoff between quality and training time in quantization aware training (QAT).   Specifically, versus STE at the same number of iterations we show a 1-8\% accuracy improvement for DeiT Tiny/Small, 1-2\% accuracy improvement on Resnet 18/50, and 1-22 perplexity point improvement for LLaMA 9m/20m.   We also show that for the same loss FOGZO yields a 796$\times$ reduction in computation versus n-SPSA for a 2-layer MLP on MNIST.