The Cost of Mismatch: Noise Amplification in Zeroth-Order Reinforcement Learning

AI summary

Overview

• Research area: Machine learning — memory-efficient post-training of large language models, specifically zeroth-order (gradient-free) optimization applied to reinforcement learning with verifiable rewards (RLVR).
• Technical level: Advanced (the core of the paper is a variance-theoretic and convergence analysis with explicit constants; the empirical section is approachable).
• Scope: The paper explains why zeroth-order gradient estimates are unreliable when applied directly to on-policy RL, proves the difference between "mismatched-data" and "matched-data" evaluation regimes, and proposes ZO-BOLT, a matched-data zeroth-order method that avoids the problem.

What This Paper Is About

Zeroth-order optimization estimates gradients using only forward passes, which makes LLM post-training far cheaper in memory than backpropagation. The authors ask why this seemingly ideal tool has not been reliably applied to on-policy RLVR, and trace the answer to a structural detail: whether the two function evaluations inside a finite difference are performed on the same data or on independently sampled data. They show that on-policy RL forces independent (mismatched) evaluations, which amplifies noise catastrophically, and that a fixed-dataset objective called Boltzmann Targeted SFT (BOLT) allows shared (matched) evaluations that cancel the noise.

Key Contributions

1. A generic variance-theoretic framework for noisy zeroth-order optimization. Lemma 1 gives an exact decomposition of the estimator error into a noise-free term and an oracle-noise term, showing that the oracle-noise term depends entirely on the correlation between the two paired evaluations.
2. A characterization of the two data regimes. Proposition 1 proves that on-policy RL under the mismatched-data regime produces a 1/μ² oracle-noise amplification, while Proposition 2 proves that BOLT under the matched-data regime yields a bound with no 1/μ² term, leaving only a minibatch-level variance term. These are unified in an oracle-noise separation theorem (Theorem 1) and turned into convergence guarantees with explicit constants (Theorem 2).
3. Empirical validation of the predicted noise amplification. Controlled perturbation-scale μ sweeps confirm the 1/μ² trend for independently sampled evaluations, and matched-versus-mismatched end-to-end comparisons confirm the optimization consequences.
4. ZO-BOLT, a practical matched-data algorithm. Evaluated on three model variants and two mathematical reasoning benchmarks, it improves over the pretrained baseline while keeping a near-inference memory footprint. Code is released at https://github.com/lianmin-chen/the_cost_of_mismatch.

Main Findings

• Mismatched data causes 1/μ² amplification. In on-policy RL, the evaluation at θ and the evaluation at θ + μu use rollouts from different policies, so their noise is uncorrelated and adds rather than cancels. The resulting bound is 2 d σ_roll² / μ².
• The amplification is doubly severe. It is linear in dimension d and quadratic in 1/μ. The paper notes that even with d = 10⁶ in a parameter-efficient configuration, the bound reaches 2 × 10¹² σ_roll² at μ = 10⁻³.
• Matched data cancels the noise. For BOLT, using the same minibatch B for both evaluations gives a bound of (√((d+2)σ_mb²/B) + Lμ√(d(d+2)(d+4)))², which contains no 1/μ² factor. The remaining O(d³) term is argued to be negligible when μ²d²L² ≪ σ_mb²/B.
• A noise floor appears only in the mismatched case. For on-policy RL, shrinking μ to reduce smoothing bias simultaneously inflates the noise floor, creating a lower bound on achievable gradient norm that is independent of the number of steps T. For BOLT the bound is controllable through (μ, B, η, T).
• Trajectory reuse is not a free fix. Reusing trajectories replaces the true on-policy objective with a frozen-distribution surrogate, and the distribution-shift term vanishes only when P_{θ+μu} = P_θ. Importance weighting can correct this, but the second moment of trajectory-level weights can grow exponentially with sequence length T. PPO/GRPO-style frozen-distribution surrogates obtain matched data only by changing the optimization target.
• Measured oracle-noise variance confirms the theory. In Table 1 (Qwen2.5-1.5B, 1.09M trainable LoRA parameters, GSM8K), the dimension-normalized variance V(μ) for matched BOLT stays roughly flat at 2.12 × 10⁻² from μ = 10⁻⁴ to μ = 10⁻³ and rises only to 3.31 × 10⁻² at μ = 10⁻², whereas mismatched BOLT goes from 2.77 × 10⁶ at μ = 10⁻⁴ to 2.77 × 10² at μ = 10⁻², and on-policy RL goes from 2.44 × 10⁶ to 2.30 × 10².
• End-to-end, matching wins. On GSM8K with Qwen2.5-1.5B and LoRA, matched ZO-BOLT reaches a best accuracy of 63.78 (+11.54 over the 52.24 baseline), while mismatched ZO-BOLT reaches 49.35 (−2.89). Matched ZO-GRPO reaches 59.00 (+6.76) versus 53.08 (+0.84) for mismatched ZO-GRPO.
• ZO-BOLT improves over pretrained models in all six settings. On Qwen2.5-1.5B it reaches 63.78 (+11.54) on GSM8K and 32.84 (+3.02) on MATH; on Qwen2.5-7B, 85.73 (+9.80) and 58.00 (+7.70); on DeepSeek-R1-Distill-Qwen-1.5B, 66.89 (+5.92) and 52.78 (+0.50). Recovery of the first-order BOLT gain is 98.2, 79.9, 92.2, 69.2, 60.4 and 21.0 percent respectively.
• Memory stays near inference level. ZO-BOLT peaks at 4,201 MB on Qwen2.5-1.5B (only +920 MB over the 3,281 MB inference-only footprint) and 16,014 MB on Qwen2.5-7B (+1,114 MB over 14,900 MB), reducing peak GPU memory by 59.5% and 43.7% relative to first-order BOLT at the two scales.

Methodology in Plain English

Zeroth-order optimization works by nudging the parameters in a random direction, measuring the loss, and using the difference between two such measurements to estimate a gradient. The authors start from a simple but overlooked observation: in a fixed-dataset setting, both measurements can be taken on the same minibatch, but in on-policy RL each measurement needs fresh rollouts from the current policy, so the two measurements see different data.

They write the estimator error as an exact sum of a clean term and a noise term, and show that the noise term depends on how correlated the noise is between the two measurements. If the two measurements are independent, the noise variances add, and dividing by the small perturbation scale μ causes a 1/μ² blow-up. If the two measurements share data, the leading noise cancels. They then plug these two variance bounds into a standard descent argument for Gaussian-smoothed nonconvex optimization to get convergence rates whose constants differ sharply between the regimes.

To test the theory, they measure the noise variance directly across five perturbation scales for three configurations: matched BOLT, mismatched BOLT, and on-policy RL. They then run matched-versus-mismatched training comparisons for both ZO-BOLT and ZO-GRPO. Finally, they build ZO-BOLT by applying zeroth-order optimization to BOLT's fixed weighted supervised fine-tuning objective, adding symmetric finite differences, ZoAR query reuse, and R-AdaZO, and evaluate it on GSM8K and MATH with LoRA (r = 8, α = 16).

Why This Matters

The paper identifies a concrete, structural failure mode rather than a tuning issue: no learning-rate adjustment can fix a noise term that scales as 1/μ². This reframes how zeroth-order methods should be paired with RL-style objectives and gives a principled reason to prefer fixed-data, RLVR-aligned objectives.

Real-world applications:

• Memory-constrained LLM fine-tuning: post-training reasoning models on consumer or single-GPU setups where first-order optimizers exceed available memory.
• Mathematical and verifiable-reward reasoning training: improving models on benchmarks such as GSM8K and MATH without holding activations for backpropagation.
• On-device or edge model adaptation: keeping peak memory close to inference level so adaptation fits within deployment budgets.
• API-only or black-box model optimization: settings where gradients are inaccessible and only forward evaluations are available.

Industry relevance: the reported memory reductions (59.5% and 43.7% relative to first-order BOLT at the 1.5B and 7B scales, with overheads of only +920 MB and +1,114 MB over inference) speak directly to the cost of post-training large models. The result that matched-data ZO-BOLT recovers 98.2% of the first-order BOLT gain on Qwen2.5-1.5B GSM8K suggests zeroth-order RLVR can be a practical substitute rather than a compromise.

Future Directions

• Extending matched-data guarantees to settings without a fixed dataset. The paper shows that trajectory reuse and importance weighting cannot resolve the mismatch for the true on-policy objective; whether some other construction can achieve matched evaluations without changing the target remains open.
• Improving recovery on harder benchmarks. ZO-BOLT recovers only 21.0% of the first-order BOLT gain on DeepSeek-R1-Distill-Qwen-1.5B with MATH, and the paper's own results show weaker recovery on MATH generally; closing this gap is a natural next step.
• Characterizing when surrogate objectives are acceptable. Since PPO/GRPO-style frozen-distribution surrogates obtain matched data only by changing the objective, quantifying the trade-off between the surrogate gap and the noise reduction would clarify when such surrogates are preferable.
• Scaling and combining with other zeroth-order accelerations. The paper already combines symmetric finite differences, ZoAR query reuse, and R-AdaZO; whether further variance-reduction or query-reuse techniques interact differently with the matched-data bound is unexplored in the reported content.

Target Audience

Researchers and engineers working on memory-efficient LLM post-training, zeroth-order or gradient-free optimization, and RLVR/reasoning-model training. The theoretical sections suit readers comfortable with nonconvex optimization, Gaussian smoothing, and variance arguments; the empirical section is useful to practitioners who want a concrete recipe and memory numbers before adopting a method. Readers looking for a purely applied tutorial will find the theory-heavy middle demanding, but the matched-versus-mismatched framing is accessible and directly actionable.

Authors’ abstract

Zeroth-order (ZO) optimization has emerged as a promising approach for memory-efficient large language model (LLM) post-training, yet its reliable application to reinforcement learning with verifiable rewards (RLVR) remains largely unexplored. Motivated by the distinct data regimes of on-policy RL and fixed-data optimization, we develop a generic variance-theoretic framework to investigate what limits the reliability of ZO gradient estimation. Our analysis reveals that estimation reliability is structurally determined by whether the two finite-difference evaluations are performed on matched data. On-policy reinforcement learning (RL) under the mismatched-data regime produces a $1/μ^2$ oracle-noise amplification that sets an irreducible noise floor as the perturbation scale $μ$ decreases, whereas Boltzmann Targeted SFT (BOLT) under the matched-data regime eliminates this amplification via noise cancellation. We further derive convergence guarantees with explicit constants that cleanly separate the two regimes. These results naturally lead to ZO-BOLT, which applies ZO optimization to BOLT's RLVR-aligned weighted-SFT objective. Controlled perturbation-scale sweeps and matched-versus-mismatched end-to-end comparisons validate the predicted effect of data matching. Experiments across three model variants and two mathematical reasoning benchmarks further show that ZO-BOLT consistently improves over the pretrained baseline while retaining a near-inference memory footprint. The code is publicly available at https://github.com/lianmin-chen/the_cost_of_mismatch.