📊 statistics

Variance-Aware Baselines and Adaptive Learning Rates for Reinforcement Learning with Verifiable Rewards

This paper establishes a theoretical framework for variance-aware baselines and adaptive learning rates in reinforcement learning with verifiable rewards (RLVR), deriving an optimal gradient-weighted baseline and a signal-to-noise ratio-based learning-rate schedule that collectively form the OBLR-PO method to significantly improve policy optimization performance.

Original authors: Zixun Huang, Jiayi Sheng, Zeyu Zheng

Published 2026-07-31

📖 7 min read🧠 Deep dive

CC BY 4.0

Original authors: Zixun Huang, Jiayi Sheng, Zeyu Zheng

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very smart, but slightly chaotic, robot how to solve complex math puzzles. You don't have a teacher standing over its shoulder correcting every mistake; instead, you only get a simple "Yes" or "No" at the end of its attempt. This is the world of Reinforcement Learning with Verifiable Rewards (RLVR). It's a popular way to train giant AI brains to reason better. The robot tries a solution, gets a score, and then tweaks its brain to do better next time.

However, there's a catch: the robot's learning process is like trying to steer a ship through a foggy storm using a compass that jitters wildly. To make sense of the "Yes" or "No" scores, the robot needs two things: a baseline (a reference point to know if a score is actually good or just average) and a learning rate (how big of a step it should take based on that score). Currently, engineers pick these settings by guessing or following old rules of thumb, which often leaves the robot stumbling in the fog, learning too slowly or taking steps so wild it forgets everything it just learned.

This paper, titled "Variance-Aware Baselines and Adaptive Learning Rates for Reinforcement Learning with Verifiable Rewards," acts like a new, high-tech navigation system for that ship. The authors, Zixun Huang, Jiayi Sheng, and Zeyu Zheng, decided to stop guessing. They built a mathematical map to figure out exactly how to choose the best reference point and the perfect step size for the robot. They found that the old way of averaging scores is too clumsy. Instead, they discovered that the robot should weigh its experiences based on how much "effort" (mathematically, the gradient magnitude) each experience required. Furthermore, they realized the robot should take giant, confident steps when the signal is clear and tiny, cautious steps when the signal is noisy. By combining these two insights, they created a new method called OBLR-PO (Optimal Baseline and Learning-Rate Policy Optimization). When they tested this on a math-solving AI called Qwen3-4B-Base, it learned faster and solved more problems correctly than the previous best methods, proving that a little bit of math can go a long way in taming the chaos of AI training.

The Problem: Steering in the Fog

Think of training an AI like teaching a dog to fetch. If the dog brings back a stick, you say "Good!" If it brings back a rock, you say "No." But what if the dog brings back a stick that is half-buried in mud? Is that a "Good" or a "No"? In the AI world, this is the baseline problem. You need to know what a "normal" performance looks like so you can tell if a specific attempt was a true success or just a fluke.

Most current methods use a simple average. They say, "Last time the dog got a stick, a rock, and a shoe. The average is a 'medium' reward, so let's compare today's stick to that." But this is like judging a marathon runner by the average speed of everyone in the city, including people walking to the grocery store. It doesn't account for how hard the runner actually tried.

Then there's the learning rate. This is the size of the step the AI takes after learning something. If the AI is 100% sure it did well, it should take a big step forward. If it's confused and the data is messy, it should take a tiny, careful step. Most methods use a fixed step size, like a robot that walks at the same speed whether it's on a smooth sidewalk or a pile of rubble. This often leads to the AI either learning too slowly or stumbling and falling over.

The Solution: A Smarter Compass

The authors of this paper asked two simple questions: "What is the perfect reference point to measure success?" and "How big should the next step be?"

They started by looking at the math behind the AI's "brain updates." They found that the old way of calculating the baseline was missing a crucial piece of information: how much each specific example influenced the learning. Some examples are easy; others are hard. The "hard" ones carry more weight in the learning process.

The New Baseline: The Weighted Average
Instead of a simple average, the authors proposed a variance-optimal baseline. Imagine you are grading a class of students. A simple average adds up all the scores and divides by the number of students. But the new method says, "Let's give more weight to the students who studied the hardest." In the AI's case, it weighs the rewards based on the "gradient magnitude"—a fancy way of saying "how much effort this specific example required to learn."

They also added a safety net called KL regularization. This is like a "don't forget your roots" rule. It prevents the AI from changing its personality too drastically while it learns. The authors showed that when you mix this "don't forget your roots" rule with the "weighted effort" baseline, you get a much more stable and accurate way to measure success.

The New Learning Rate: The Signal-to-Noise Ratio
For the step size, the authors introduced a concept called the Signal-to-Noise Ratio (SNR). Imagine you are trying to hear a friend's voice in a crowded room.

High Signal: Your friend is shouting clearly. You can trust what they say and react immediately (take a big step).
High Noise: The room is loud, and you can barely hear them. You should pause, listen carefully, and maybe just nod (take a tiny step).

The paper proves mathematically that the best learning rate is directly tied to this SNR. If the AI's gradient (the direction it wants to move) is clear and strong, the learning rate goes up. If the gradient is messy and noisy, the learning rate shrinks. This is a huge shift from the old method of just picking a fixed number and sticking with it.

The Result: OBLR-PO

The authors combined these two ideas into a single, powerful method called OBLR-PO.

The Baseline: It uses a "gradient-weighted" average of the rewards, meaning it pays extra attention to the examples that were hardest to learn.
The Learning Rate: It dynamically adjusts the step size based on how clear the learning signal is.

They tested this on a 4-billion-parameter AI model (Qwen3-4B-Base) trained on math problems. The results were impressive.

On its own, the new baseline made the AI more accurate than standard methods like PPO, ReMax, and RLOO.
On its own, the new learning rate rule improved performance across all the different methods it was tested on.
Together, they created the strongest performer. On benchmarks like GSM8K (a math reasoning test), the new method achieved 87.19% accuracy, beating the previous best of 85.60% (GRPO) and 83.93% (RLOO).

Why It Matters

This paper doesn't just suggest a tweak; it provides a theoretical proof that these changes are optimal under certain conditions. The authors showed that their method reduces the "noise" in the learning process, making the AI's path to intelligence smoother and more direct.

They also demonstrated that these two improvements are modular. You can use the new learning rate rule with any existing AI training method, and it will likely make it better. You can also use the new baseline with other methods. But when you put them together, you get the full power of the new system.

In short, the authors took the chaotic, guesswork-heavy process of training AI to solve math problems and replaced it with a precise, mathematically grounded strategy. They showed that by listening to the "effort" of each example and adjusting the "step size" based on how clear the signal is, we can teach AI to reason better, faster, and more reliably. It's a reminder that sometimes, the best way to move forward is to stop guessing and start measuring the noise.

Technical Summary: Variance-Aware Baselines and Adaptive Learning Rates for RLVR

Problem Statement
Reinforcement Learning with Verifiable Rewards (RLVR) has become a dominant paradigm for post-training large language models (LLMs) to enhance reasoning capabilities. However, the design of two critical components in policy-gradient updates—the baseline used to construct the gradient estimator and the learning rate schedule—remains largely heuristic. Existing methods (e.g., PPO, GRPO, RLOO) utilize baselines derived from learned value functions, group-normalized rewards, or leave-one-out reward averages without theoretically accounting for how individual samples influence gradient magnitude. Similarly, fixed or predefined learning rates fail to adapt to the changing reliability (signal-to-noise ratio) of stochastic gradients during training. This lack of theoretical grounding limits the understanding of the statistical properties of policy-gradient estimators and their interaction with optimization dynamics.

Methodology and Theoretical Framework
The authors develop a rigorous theoretical framework under a KL-regularized policy-optimization objective. The core of their approach involves analyzing the bias, variance, and convergence properties of the gradient estimator to derive principled alternatives to heuristic designs.

Unbiasedness and Variance Characterization:
The authors first establish the unbiasedness of the KL-regularized gradient estimator. They derive an exact expression for the variance of this estimator, which includes a specific KL cross-covariance term arising from the interaction between the policy gradient and the KL regularization term. This derivation reveals that the total variance depends not only on the reward variance but also on the magnitude of the policy gradients (score norms) for each sample.
Optimal Learning Rate (SNR-Adaptive):
By deriving an upper bound on the optimization loss, the authors optimize this bound to determine the ideal learning rate schedule. They prove that the optimal learning rate is governed by the Signal-to-Noise Ratio (SNR) of the policy gradient. Specifically, the learning rate should scale with the ratio of the squared expected gradient norm (signal) to the trace of the covariance matrix (noise). This results in an adaptive rule where updates are larger when the gradient is reliable and smaller when it is noisy.
Variance-Optimal Baseline:
Minimizing the derived variance expression leads to a novel baseline design. The authors prove that the variance-optimal baseline is not a simple average of rewards (as in standard leave-one-out methods) but a gradient-weighted estimator of the KL-regularized reward. In this formulation, samples with larger gradient norms (higher influence on the policy) are weighted more heavily when computing the baseline. This accounts for the fact that high-variance samples contribute disproportionately to the estimator's noise.

Key Contributions
The paper presents four primary contributions:

Theoretical Framework: A unified analysis establishing the unbiasedness, exact variance expressions (including KL cross-covariance), and an optimization-loss upper bound for KL-regularized policy optimization.
SNR-Adaptive Learning Rate: A derived learning-rate schedule that adapts to gradient reliability. This module is agnostic to the specific baseline used and can be integrated into existing estimators.
Variance-Optimal Baseline: A characterization of the optimal baseline as a gradient-weighted average of the KL-regularized reward, offering a principled alternative to reward-based baselines.
OBLR-PO Method: The combination of these two components into a single algorithm, Optimal Baseline and Learning-Rate Policy Optimization (OBLR-PO).

Experimental Results
The authors validate their framework by post-training Qwen3-4B-Base on mathematical reasoning tasks (AMC23, GSM8K, MATH-500, OlympiadBench) using the DeepMath-103K dataset.

SNR-Adaptive Learning Rate: When applied independently to existing methods (RLOO, ReMax, GRPO), the adaptive learning rate consistently improved both training reward and final Pass@1 accuracy compared to fixed learning rates. Notably, it improved GRPO performance despite GRPO not strictly satisfying the paper's theoretical assumptions, suggesting robustness.
Variance-Optimal Baseline: Using the gradient-weighted baseline with a fixed learning rate outperformed standard baselines (PPO, ReMax, RLOO, GRPO) across all benchmarks, achieving the lowest gradient noise-to-signal ratio ( $1/\widehat{SNR}$ ).
OBLR-PO: Combining both components yielded the strongest overall performance. On GSM8K, MATH-500, and AMC23, OBLR-PO achieved the highest Pass@1 accuracy (e.g., 87.19% on GSM8K, surpassing the next best method by over 1.5 percentage points). On OlympiadBench, the variance-optimal baseline alone remained slightly superior, indicating a minor trade-off with the adaptive schedule on that specific dataset.

Significance and Claims
The paper claims to provide a principled, theoretical foundation for two heuristics that have traditionally been tuned empirically. By linking the learning rate to the SNR and the baseline to the gradient-weighted reward, the authors offer a method to control both the magnitude and direction variance of policy updates. The significance lies in the modularity of the approach: the adaptive learning rate can be "dropped in" to existing algorithms for immediate gains, while the variance-optimal baseline offers a new standard for estimator design. The authors conclude that while their guarantees rely on an upper bound and smoothness assumptions, the empirical results demonstrate that these theoretical insights translate to tangible improvements in LLM reasoning performance.

The Problem: Steering in the Fog

The Solution: A Smarter Compass

The Result: OBLR-PO

Why It Matters

Technical Summary: Variance-Aware Baselines and Adaptive Learning Rates for RLVR

More like this