Research map / Optimization / Optimizers
Improving the Straight-Through Estimator with Zeroth-Order Information
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes FOGZO, a zeroth-order gradient estimator that uses the straight-through estimator gradient as a biased source (mixed with random perturbations, with a decaying mixing ratio) to reduce variance and cost of n-SPSA, improving quantization-aware pre-training over STE.
Problem
Training neural networks with quantized parameters: the STE provides cheap but biased gradients that can point in wrong directions at low precision, while unbiased zeroth-order (n-SPSA) gradients are too expensive and high-variance for a fixed compute budget.
Approach
Combine the STE gradient (normalized, randomly sign-flipped) with random perturbations into a single perturbation direction for a two-point finite-difference estimator with n=1, using a linearly decayed mixing ratio beta; the STE-biased term is automatically suppressed when the STE direction misaligns with the true smoothed gradient.
Key results
- FOGZO improves DeiT-Tiny on ImageNet-100 with identity STE from 62.72 to 70.06 accuracy at same iterations.
- With LSQ + identity STE, FOGZO lowers LLaMA-9M perplexity on 328M C4 tokens from 104.31 to 99.74 (2-bit weights).
- FOGZO reaches the same loss as n-SPSA with 796x less computation on a 2-layer MLP on MNIST.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use FOGZO instead of STE with LSQ on quantized weights | DeiT-Tiny and DeiT-Small · ImageNet-100 | accuracy: 72.76±0.62 -> 72.92±0.60; 79.55±0.86 -> 80.06±0.48 | Helped |
| Use FOGZO instead of STE with LSQ on quantized weights | DeiT-Tiny · ImageNet-1K | accuracy: 63.19±0.30 -> 63.38±0.52 | Helped |
| Use FOGZO instead of STE with LSQ on quantized weights | LLaMA-9M · C4, 328M tokens | perplexity: 104.31±6.54 -> 99.74±4.42 | Helped |
| Use FOGZO instead of STE with LSQ on quantized weights | LLaMA models · C4, 1.3B–13B tokens | perplexity: 59.21±1.22 -> 58.93±0.90; 50.85±0.76 -> 50.61±0.64 | Helped |
| Use FOGZO instead of STE with LSQ on quantized weights | ResNet-18 · ImageNet-100 | accuracy: 80.23±0.84 -> 80.04±1.14 | No clear effect |
| Use FOGZO instead of STE with LSQ on quantized weights | ResNet-18 · ImageNet-1K | accuracy: 66.51±0.30 -> 66.13±0.50 | No clear effect |
| Use FOGZO instead of STE with LSQ on quantized weights | ResNet-50 · ImageNet-100 | accuracy: 82.81±0.82 -> 83.67±0.26 | Helped |
| Add FOGZO to QuEST | LLaMA with 2-bit weights and activations · C4, 3B tokens | perplexity: 37.75 -> 37.37 | Helped |
| Add FOGZO to LSQ | LLaMA with 2-bit weights and activations · C4, 3B tokens | perplexity: 39.06 -> 37.38 | Helped |
| Use FOGZO rather than STE at matched training time | LLaMA-30M · C4 | perplexity: 38.69 -> 38.67; 38.43 -> 38.38; 38.25 -> 37.93 | Helped |
| Set beta to 1, removing the unbiased perturbation contribution (ablation) | FOGZO · MNIST, 2-layer MLP | training loss | Hurt |
| Set epsilon scale c to 1 rather than sweeping other scales (ablation) | FOGZO · MNIST, 2-layer MLP | training loss | No clear effect |
Surprising findings
- At beta=1 (purely STE-biased direction, no unbiased term), FOGZO cannot outperform the STE, indicating the unbiased component is needed.
- Training the first part of the run with pure STE and only the last 10-30% with FOGZO yields lower perplexity than STE at equal wall-clock time.
Limitations
- Suffers precision loss when weights are stored in 16-bit because small updates are hard to apply, so it is incompatible with BF16 weights.
- Only partially zeroth-order; extra cost of two forward passes per iteration over STE.
- Largest models tested are about 300M parameters; scaling to large LLM pre-training is not verified.
Experimental setup
Datasets: MNIST ImageNet-1K ImageNet-100 C4
Benchmarks: MNIST ImageNet-1K ImageNet-100 C4
Metrics: top-1 accuracy perplexity training loss FLOPs
Baselines: identity STE tanh STE ApproxSign STE n-SPSA LSQ QuEST
Builds on: straight-through estimator n-SPSA MeZO LSQ QuEST
Closest papers by meaning
- Fine-tuning Quantized Neural Networks with Zeroth-order Optimization (ICLR 2026 · similarity 0.89)
- DeepZero: Scaling Up Zeroth-Order Optimization for Deep Model Training (ICLR 2024 · similarity 0.88)
- Refining Adaptive Zeroth-Order Optimization at Ease (ICML 2025 · similarity 0.87)
- Sparse MeZO: Less Parameters for Better Performance in Zeroth-Order LLM Fine-Tuning (NeurIPS 2025 · similarity 0.86)
- Powering Up Zeroth-Order Training via Subspace Gradient Orthogonalization (NeurIPS 2026 · similarity 0.86)
- One-Step Forward and Backtrack: Overcoming Zig-Zagging in Loss-Aware Quantization Training (AAAI 2024 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.