Title: Improving the Straight-Through Estimator with Zeroth-Order Information
Venue: NeurIPS 2025
Accepted: yes
Keywords: quantization

straight-through estimator

quantization-aware training
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](https://github.com/1733116199/fogzo).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the challenges of training neural networks with quantized parameters by proposing First-Order-Guided Zeroth-Order Gradient Descent (FOGZO), a hybrid gradient estimator. FOGZO combines the Straight-Through Estimator (STE), which is computationally efficient but biased, with zeroth-order (ZO) methods, which are theoretically sound but computationally expensive. The authors theoretically analyze their approach, claiming reduced bias compared to STE and improved efficiency relative to ZO methods, supported by empirical evaluations across diverse neural network architectures.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\nIntegrating STE with ZO methods offers a relatively novel methodological perspective. \n\nThe paper provides detailed mathematical analyses demonstrating how FOGZO interpolates between biased STE gradients and unbiased ZO gradients. The paper evaluates multiple neural architectures on relevant datasets, providing some empirical evidence supporting their claims.\n\nImproving quantization-aware training is relevant given the increasing demand for efficient neural networks in resource-constrained environments.\n\n# Weaknesses\n\nThe paper is very hard to follow. The authors overuse fuzzy adverbs of frequency such as mostly, often, some, etc.\n\nWhile experiments do indicate improvements over STE, the practical significance of some results is relatively limited (e.g. only 0.5 to 2% improvement on popular models like ResNet). Given that, despite being more efficient than pure zeroth-order methods, FOGZO still incurs a substantial computational increase (~50%) over STE, it's practical relevance might be limited. Moreover, the approach currently applies strictly to weight quantization, potentially reducing it's relevance in practice even further.\n\nThe experimental results seem to be limited to individual runs. Repreating each experiment multiple times and reporting mean and std would make results more robutst. Right now, the instances where FOGZO does improve over STE could be outliers, particularily in experiments where the difference to the baseline is not that large.\n\nQuantization in Graph Neural Networks is not discussed, despite these architectures having gained importance over the recent years.\n\n# Conclusion\n\nThis is a borderline paper to me. If the authors sufficiently address the weak empirical part, I might consider changing my score from borderline reject (3) to borderline accept (4), but not more given the hard-to-follow write-up and expected limited impact.\n\nQUESTIONS:\nPlease address the above-mentioned weaknesses.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces ​​FOGZO​​, a novel method for QAT that synergistically combines the computational efficiency of the STE  with the unbiased gradients of ZO methods. FOGZO demonstrates significant empirical improvements over STE and ZO baselines across multiple architectures (ResNet, DeiT, LLaMA) and datasets (ImageNet, C4), with minimal computational overhead. The paper is well-written, technically sound, and contribute a new direction for optimizing STE.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths​:\n- The combination of STE and ZO gradients​​ via a biased-unbiased gradient mixing strategy is novel and new for me.\n- The theoretical analysis is detailed, showing the effective of FOGZO.\n- Comprehensive experiments​​ on vision and language models show consistent gains of FOGZO.\n- The paper is well-written and easy to follow.\n\nWeaknesses:\n- The searching for hyperparameter beta is time-cost.\n- Although FOGZO reduces the time cost compared to ZO based method, but it is still time cost, narrowing the use range for large models.\n- It doesn't work good with weight-activation quantization.\n\nQUESTIONS:\n1. Could perturbing activations (not just weights) during ZO steps resolve the performance gap? \n2. Or can your maintain the STE for activation scale training to see how FOGZO works with weight-activation quantization?\n3. Could you compared your methods with EWGS? This is a paper using second-order information to suppress the STE, maybe both the efficiency and performance.\n4. Does FOGZO’s gradient variance reduction amplify benefits from optimizers like AdamW or Lion?\n5. Could you share the time cost and memory usage of your methods compared to STE?\n6. I saw you added hardmard transform in your llama models, would this method be helpful with hardmard transform or can this improve the performance after transformed?\n7. Could you share the FLOPs for different step-size of ZOs?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes FOGZO, a novel method that combines the Straight-Through Estimator (STE) and Zeroth-Order (ZO) gradient descent for training quantized neural networks. While STE is efficient but biased, and ZO methods are unbiased but computationally expensive, FOGZO leverages STE as a bias source to guide low-cost ZO updates, suppressing STE’s flaws. The method introduces a mixing parameter to balance biased and unbiased gradients. Experiments show FOGZO improves training accuracy and efficiency across various models, including ResNet, DeiT, and LLaMA, outperforming both STE and ZO baselines.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. Introduces FOGZO, a first-order-guided zeroth-order gradient estimator that leverages the biased yet efficient STE to guide the unbiased zeroth-order optimization.\n2. Provides theoretical analysis showing that FOGZO can suppress incorrect STE gradients while preserving computational efficiency.\n3. Performed numerical experiments to demonstrate the theoretical findings. \n\nWeakness: \n1. Increased the training costs and training time. \n2. More hyperparameters are introduced, which may require more careful tuning. \n3. Experiments are primarily conducted on small-scale models. Gains with SOTA QAT methods, like LSQ, are marginal for most models except DeiT-Tiny.\n\nQUESTIONS:\n1. Figure 1 is hard to understand. Why are reciprocals of loss and flops used? They make readers hard to clearly and easily figure out which line is better. \n2. The chosen $\\beta$ value is very close to 1, implying that the STE part dominates the ZO part. It is better to call it ZO Guided STE? In addition, as $\\beta$ is very close to 1 and $1-\\beta$ is very close to 0. Does it cause any numerical issues during training, especially in FP16/BF16 training?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe straight through estimator (STE) is a commonly used work-around to circumvent the non-differentiability of rounding operators, and is frequently used in quantization aware training of neural networks. This paper proposes blending the STE with zero-order gradient estimation, essentially by taking their weighted sum as the gradient estimate. In particular, the method relies on perturbing the current parameter estimate in a random direction in order to obtain the zero-th order estimate. The proposed method is given a heuristic derivation,  as are the probability distributions associated with different straight through estimators,  and the method is numerically tested by training neural networks of various sizes and architectures.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \n- The paper is generally well written and generally well motivated. \n- The arguments for the proposed algorithm were convincing (albeit heuristic). \n\nWeaknesses: \n- The theoretical components of the paper are rather heuristic. In other words there are lots of $\\approx$ symbols hiding approximation errors that are never quantified. \n- In some places, the paper suffers from awkward writing.\n- The plots were a little hard to read due to some unusual presentation choices.\n- The appendix is written a bit too casually.\n\nQUESTIONS:\n- Section 2.2: $u$ is a vector, so when you say it is unit variance, do you mean it is isotropic, i.e., having an identity as its covariance matrix? \n- Eq (4) looks a little different from the definitions used in the n-SPSA papers [45, 46]. Is the difference cosmetic or substantial? \n- \"While some believe the STE is completely unjustifiable\" -- who? what are there arguments? \n- Eq (5) the multiplication by $s_i$ seems to suggest that you have a 50/50 chance of biasing your update in the negative direction to that of $\\hat{g}$. Is there an explanation for why this is useful (beyond the need for making $v_i$ 0-mean)?\n- Section 3.2: I am hesitant to call this section \"analysis\" given that it relies on the heuristic approximations (indicated by the $\\approx$ symbol in equation (7). Perhaps, \"heuristic derivation\", or \"heuristic motivation\" are better? \n- Eq (8): third = sign, there is no need for the expectation here as all sources of randomness have been taken care of. \n- Line 187, \"are set\" instead of \"is set\"? (similar subject-verb agreement issues elsewhere, like line 195). \n- Section 3.3: This section is a bit awkwardly written. One needs to read all the way to the equations to understand what is meant by (implicit) smoothness and perturbation of an STE. \n- Eq (9): is it assumed (without loss of generality) that $x>0$. \n- hardtanh should be defined as it's definition is used in the derivation.\n- Figure 1: I found this figure confusing to read. Why use the reciprocals of the number of flops and tranining loss? Also, I would have found an explicit reporting of $n$ helpful. \n- Figure 1: How is the gradient norm computed? for a fixed iteration, over all the iterations? \n- What was the range of n used in section 4.1? \n- I appreciated the discussion of activation quantization, even if it was a stumbling block for the method.\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-02"}