Improving the Straight-Through Estimator with Zeroth-Order Information

AI summary

Overview

Research area: Quantization-Aware Training (QAT) for deep neural networks, specifically gradient estimation for networks with quantized (low-precision) parameters.

Technical level: Advanced. The paper assumes familiarity with backpropagation, the Straight-Through Estimator, zeroth-order/derivative-free optimization, and quantization schemes.

Scope (one sentence): The paper proposes FOGZO, a gradient estimator that mixes the biased first-order gradient from the Straight-Through Estimator with zeroth-order finite-difference samples from n-SPSA, and evaluates it on 2-layer MLPs, ResNets, DeiT/ViT models, and small LLaMA models under 1-bit, 2-bit, and weight-activation quantization.

What This Paper Is About

Training a network whose weights are rounded to low precision is hard because the rounding function has zero gradient almost everywhere, so ordinary backpropagation cannot update the weights. The standard fix, the Straight-Through Estimator (STE), works well in practice but is theoretically questionable and has been shown to produce parameter oscillations at rounding boundaries at low precision; the theoretically sound alternative, zeroth-order (ZO) estimation such as n-SPSA, is unbiased but requires many extra loss evaluations and converges slowly under a realistic compute budget. The paper's goal is to build a method that keeps the STE's cheapness while using ZO information to correct the STE's occasional wrong-direction gradients.

Key Contributions

1. FOGZO (First-Order-Guided Zeroth-Order gradient descent): a gradient estimator that normalizes the STE gradient to a unit vector, randomly signs it (50% each way) to make it symmetric and zero-mean, mixes it with an isotropic ZO perturbation using a mixing ratio β, and then applies the n-SPSA finite-difference formula to the mixed sample.
2. A linear β decay schedule: β is annealed linearly from 1 to a hyperparameter β_min during training, motivated by the observation that STE failures tend to occur in later training stages while learning rates (and hence tolerance to gradient variance) are small.
3. A principled way to set the ZO hyperparameters: the paper argues that every STE implicitly defines a smoothed surrogate of the rounding/sign operator, and solves for the implicit smoothness ε̄ and perturbation distribution p̄(u) for the identity, hardtanh, tanh, and ApproxSign STEs (for hardtanh this yields a uniform p̄(u) = U(−√3, √3) with ε̄ = 1/√3; for identity, ε̄ = 1/(2√3) with p̄(u) = U(−√3, √3); for tanh, ε̄ = π/√12; for ApproxSign, ε̄ = 1/√6), setting FOGZO's ε = α ε̄ for a quantization scale α.
4. Empirical demonstration that FOGZO improves over a variety of STEs in Quantization-Aware Pre-Training, integrates with LSQ and QuEST, and can be scheduled so that training time matches STE exactly.

Main Findings

• Reported headline gains: versus STE at the same number of iterations, the abstract reports 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.
• Compute reduction: for the same loss, FOGZO yields a 796× reduction in computation versus n-SPSA for a 2-layer MLP on MNIST.
• β = 1 is not optimal: on the 2-layer MLP, when β = 1 FOGZO completely relies on its source of bias and cannot outperform the STE; at β = 0.999 with n = 1, FOGZO outperforms the STE while requiring only 2 additional forward passes per iteration. The paper concludes the optimal β is close to, but not exactly, one.
• ε scale of 1 is a good default: sweeping ε = c·α·ε̄ with n = 1 and β = 0.999, the authors report c = 1 is one optimal configuration, although larger values of c are equally good.
• n-SPSA's cost: on the MLP sweep with n ∈ {1, 4, 10, 100, 1000, 7960}, n-SPSA can outperform the identity STE but only with orders of magnitude more computation, and underperforms at low n because of large gradient variance.
• Works across STEs and architectures (Table 1, 2-bit and 1-bit weights, fixed α): DeiT-Tiny on Imagenet-100 improves from 62.72 (Identity-STE) to 70.06 (Identity-FOGZO), from 41.98 (tanh) to 46.8 (tanh-FOGZO), and from 30.16 (ApproxSign) to 39.61 (ApproxSign-FOGZO); ResNet-18 on Imagenet-100 improves from 79.92 to 80.42, 74.68 to 75.02, and 68.62 to 70.91 respectively; LLaMA-9m on 328M C4 tokens goes from 109.95 to 105.64, 123.97 to 121.51, and 159.17 to 137.
• Scaling up model and dataset size with LSQ (Table 2, 2-bit weights, learnable α): DeiT-Small on Imagenet-100 improves from 2.62 ± 0.012 loss / 79.55 ± 0.86 accuracy to 2.57 ± 0.016 / 80.06 ± 0.48; DeiT-Tiny on Imagenet-1K from 4.26 ± 0.015 / 63.19 ± 0.30 to 4.25 ± 0.025 / 63.38 ± 0.52; LLaMA-9m on 3.28B C4 tokens from 4.41 ± 0.020 / 82.54 ± 1.72 to 4.38 ± 0.020 / 79.94 ± 1.18; ResNet-50 on Imagenet-100 from 0.43 ± 0.0182 / 82.81 ± 0.82 to 0.39 ± 0.0086 / 83.67 ± 0.26. Not every cell improves: for ResNet-18 on Imagenet-100 the accuracy is 80.23 ± 0.84 (STE) versus 80.04 ± 1.14 (FOGZO), and on Imagenet-1K it is 66.51 ± 0.30 versus 66.13 ± 0.50.
• Weight-activation quantization (Table 3, 2-bit weights and 2-bit activations): with QuEST, perplexity improves from 37.75 to 37.37 (95M total parameters, n = 30M, 3B tokens), 32.28 to 32.06 (125M, n = 50M, 5B tokens), 26.63 to 26.45 (200M, n = 100M, 10B tokens), and 22.90 to 22.72 (300M, n = 200M, 20B tokens); with LSQ, 39.06 to 37.38 (95M) and 33.28 to 32.41 (125M). The authors state FOGZO can shrink the perplexity difference between these methods.
• Matching training time (Table 4, LLaMA-30m, one Nvidia RTX 5090): at 3.3 hours, STE with 3.174B C4 tokens reaches 38.69 perplexity versus 38.67 for 90% STE + 10% FOGZO using 3B tokens; at 3.5 hours, STE with 3.348B tokens reaches 38.43 versus 38.38 for 80% STE + 20% FOGZO with 3B tokens; at 3.7 hours, STE with 3.522B tokens reaches 38.25 versus 37.93 for 70% STE + 30% FOGZO with 3B tokens. The paper states this also implies higher data efficiency, since FOGZO runs use less training data.

Methodology in Plain English

The starting point is that rounding weights to a small set of values has no usable gradient, so the STE pretends the rounding operation is a smooth function (identity, hardtanh, tanh, ApproxSign, etc.) during the backward pass. That gives a cheap but biased gradient. The other option, n-SPSA, perturbs the model weights in random directions, measures how the loss changes in the positive and negative directions, and averages those measurements; this is unbiased but each sample needs two forward passes, and averaging enough samples to cut the variance requires far more computation than the STE.

FOGZO combines them in one estimator. Each step it computes the STE gradient g, normalizes it to a unit vector ĝ, and randomly flips its sign so the biased component is symmetric and zero-mean. It then forms a mixed direction v_i = √β · s_i ĝ + √(1−β_i) · u_i, where u_i is a random perturbation, and plugs v_i into the same finite-difference formula n-SPSA uses. The heuristic derivation in the paper shows that, for large n, the expected FOGZO gradient is approximately a linear interpolation between the unbiased smoothed gradient and its projection onto the one-dimensional direction ĝ — a rank-1 compression, similar in spirit to MeZO. The practical consequence: when the STE direction agrees with the true smoothed gradient the bias term contributes fully; when the STE points orthogonally to it, the bias term contributes nothing, effectively suppressing bad STE gradients.

Because β controls how much the estimator trusts the STE, it is decayed linearly from 1 down to β_min. For deep models, the authors always use n = 1, so FOGZO costs one backward pass plus two extra forward passes per iteration relative to STE. In the training-time study they additionally run the first r% of training in pure STE mode (r ∈ {70, 80, 90}) and only turn on the ZO sampling afterwards. FOGZO's ε and perturbation distribution are not grid-searched; instead they are derived by treating each STE surrogate as the expected value of the true rounding/sign operator under some perturbation, and solving for that perturbation.

Why This Matters

Quantization is one of the main levers for reducing the compute and memory cost of deploying neural networks, and QAT methods produce more accurate low-precision models than post-training quantization can, especially at 2-bit or 1-bit precision. If FOGZO's gains hold up, it offers a way to get better low-precision models without abandoning the STE and without paying the heavy cost of zeroth-order training, which matters most in Quantization-Aware Pre-Training where the model starts from random initialization and training runs are long.

Real-world applications implied by the paper:

• Edge and embedded deployment: the paper motivates lower-precision models by noting that large hardware deployments are expensive to change, so edge applications benefit most from very low precision.
• Quantization-aware pre-training pipelines: the method slots into existing QAT code as an extra backward/forward step and was tested with the publicly available QuEST source code and with LSQ.
• Memory- and bandwidth-limited inference: 2-bit weight and 2-bit activation models reduce memory traffic, which is relevant for serving many small language models.
• Small language model training on limited hardware: the training-time experiment was run on a single Nvidia RTX 5090, and FOGZO was evaluated on LLaMA variants from 9M to 200M non-embedding parameters.

Industry relevance: the method is designed to be a drop-in change to the backward pass of an existing quantization-aware training loop, requires no extra memory because parameters are perturbed and restored in place (as in MeZO), and the scheduling trick (run mostly STE, finish with FOGZO) lets practitioners trade data for compute while holding wall-clock training time fixed.

Future Directions

• Scaling to larger models: the evaluation is limited by resource constraints to ResNet-18/50, DeiT-Tiny/Small, ViT-Tiny/Small, and LLaMA variants up to 200M non-embedding parameters (300M total, described in the abstract as up to 0.3 billion); whether the gains persist at much larger scales is not established by the reported experiments.
• Reducing the extra cost: FOGZO adds a fixed overhead of 2 forward passes per iteration over the STE, and the paper notes that a mild increase in training cost is acceptable only in some settings.
• Adaptive mixing: the paper keeps the β decay linear and reports that the optimal β_min sits close to but below 1 while β_min = 1 can never beat the STE; an adaptive or alignment-based β schedule is a natural open question raised by the derivation.
• Implicit smoothing for other operators and STEs: the paper derives the implicit (ε̄, p̄(u)) pair for identity, hardtanh, tanh, and ApproxSign, and notes other formulations of the dequantization function f (zero-point offsets, outlier smoothing matrices such as the Hadamard matrix) as possible extensions.
• Note on the source text: the provided paper content is truncated inside the Related Works and Limitations section ("Most existing QAT meth…"), so the authors' own stated limitations are not available in full here.

Target Audience

Researchers and engineers working on quantization-aware training, low-precision model deployment, and gradient estimation for non-differentiable objectives, as well as practitioners running quantization-aware pre-training who want better accuracy at 2-bit or 1-bit precision without paying zeroth-order training costs. Readers should already be comfortable with backpropagation, the STE, and zeroth-order optimization; the "Sufficiently good" framing and theoretical caveats assume some background in gradient estimator bias-variance tradeoffs.

Authors’ 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.