Improving the Straight-Through Estimator with Zeroth-Order Information

TL;DR

This paper introduces FOGZO, a new gradient estimation method that combines the benefits of STE and zeroth-order approaches, improving training efficiency and accuracy in quantized neural networks.

Contribution

The paper proposes FOGZO, a novel method that reduces STE bias and computational cost by guiding zeroth-order gradients with first-order information.

Findings

FOGZO improves accuracy by 1-8% on DeiT models.

FOGZO achieves 1-22 perplexity point improvement on LLaMA.

FOGZO reduces computation by up to 796x compared to n-SPSA.

Abstract

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 Pre-Training. 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 models with up to 0.3 billion parameters. For the same loss, FOGZO yields a 796$\times$ reduction in computation versus n-SPSA for a 2-layer MLP on MNIST. Code is available at https://github.com/1733116199/fogzo.