Noise amplification limits memory efficient reinforcement learning updates
The Cost of Mismatch: Noise Amplification in Zeroth-Order Reinforcement Learning
Machine Learning
Summary
This paper studies a technique called zeroth-order optimization, used to update large language models in reinforcement learning. The authors find that if the data used in two gradient estimates doesn't match, noise in the updates becomes much larger, which harms learning. They propose a method called ZO-BOLT that matches the data in these estimates and reduces this noise, improving learning reliability without increasing memory needs. Their experiments show this approach consistently improves performance on reasoning tasks.
What this means in practice
- •For machine learning engineers: Improve memory-efficient fine-tuning of large language models with reinforcement learning by reducing update noise using ZO-BOLT.
- •For ai product teams: Deploy improved reinforcement learning updates that maintain small runtime memory footprints for better model adaptation in products.
Authors
Lianmin Chen, Junbin Qiu, Chenxing Wei, Yao Shu, Kun He
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.