Maximum Likelihood Reinforcement Learning

AI summary

Overview

• Research area: Reinforcement learning for correctness-based tasks, particularly LLM reasoning, connecting RL objectives to maximum likelihood estimation.
• Technical level: Intermediate — the core idea is intuitive, but the paper uses population-level gradient derivations and Maclaurin expansions.
• Scope: The paper argues that standard expected-reward RL is only a first-order approximation of maximum likelihood on binary-correctness tasks, and introduces MaxRL, a compute-indexed family of objectives that interpolates toward exact maximum likelihood as sampling compute grows.

What This Paper Is About

Many modern tasks — navigation, program synthesis, structured prediction, multi-step reasoning in large language models (LLMs) — end with a binary correct/incorrect verdict from an external verifier after a stochastic generation process. This means the model implicitly defines a probability of success for each input, i.e., an implicit likelihood over correctness, but RL is used instead because intermediate sampling is non-differentiable and blocks direct maximum likelihood training. The paper's goal is to close that gap: it shows that standard RL optimizes only a first-order approximation of the likelihood objective and proposes a family of sample-based objectives that recover maximum likelihood as more sampling compute is spent.

Key Contributions

1. Formalizes correctness-based RL as a latent-generation maximum likelihood problem, where decoding pushes the latent policy forward to an output distribution, and proves that standard RL optimizes only the first-order approximation of the maximum likelihood objective.
2. Introduces a compute-indexed family of objectives, built from a Maclaurin expansion in pass@k and fail@k events, that interpolates between expected reward (truncation level T = 1) and exact maximum likelihood (T → ∞).
3. Analyzes a simple on-policy estimator whose expected gradient exactly matches the compute-indexed approximation of the likelihood objective, meaning extra sampling improves the optimized objective itself rather than only reducing variance. The estimator is a one-line change to standard RL implementations: it normalizes by the per-task mean reward instead of the number of samples (RLOO) or the reward standard deviation (GRPO).
4. Empirically demonstrates Pareto dominance over RLOO and GRPO across all tested models and tasks, including a controlled image classification setting, a procedurally generated maze environment, GSM8K, and billion-parameter mathematical reasoning models.

Main Findings

• RL is a first-order approximation of ML: The maximum likelihood gradient equals an infinite harmonic mixture of pass@k gradients, and the standard RL gradient keeps only the leading pass@1 term. Maximum likelihood reweights by 1/p, emphasizing hard, low-success inputs.
• Truncation level controls the objective: For truncation level T, MaxRL's gradient is the sum from k = 1 to T of (1/k) times the pass@k gradient. T = 1 recovers RL, T → ∞ recovers maximum likelihood, and the practical estimator in Algorithm 1 is unbiased for the MaxRL gradient of order T = N − 1 rather than T = N.
• Estimator difference is just normalization: REINFORCE averages score functions and normalizes by total samples N; MaxRL normalizes by successful samples K. They are unbiased for fundamentally different objectives. More compute under REINFORCE reduces variance of a fixed objective, whereas under MaxRL it raises the approximation order.
• Unifying weight-function view: At the population level, RL has weight w(p) = 1, GRPO has w(p) = 1/sqrt(p(1−p)), MaxRL(T) has w(p) = (1 − (1−p)^T)/p, and ML has w(p) = 1/p. GRPO upweights low-pass-rate inputs approximately as 1/sqrt(p) but its weighting function inverts for sufficiently large pass rates, increasing as p → 1, unlike likelihood-based objectives.
• ImageNet controlled comparison: On ImageNet with a ResNet-50-based classifier, REINFORCE with a standard baseline fails to achieve meaningful improvements even with a very high per-input sampling budget, while exact maximum likelihood training yields steady gains on Pass@1 and Pass@k. MaxRL, trained on the same sparse successful trajectories as REINFORCE, tracks exact maximum likelihood as rollout counts increase.
• Infinite-data maze regime: Using 1 million unique 17×17 mazes, a held-out set of 256 mazes, a ~3M-parameter transformer, 9,000 RL steps, batch size 256, and rollouts varied from 4 to 128, RLOO at the highest compute budget (128 rollouts per prompt) fails to match MaxRL at the lowest budget (4 rollouts) across all pass@k metrics. GRPO at the highest compute trails MaxRL trained with only 16 rollouts per prompt at every pass@k evaluation. There is almost no gap between pass@k and pass@1 for GRPO and RLOO despite no repeated data, whereas MaxRL shows a substantial gap.
• Data-scarce GSM8K regime: Training SmolLM2-360M-Instruct on GSM8K for up to 50 epochs, RLOO and GRPO reach peak pass@1 faster (around 10 epochs) but MaxRL overtakes them at approximately 30 epochs and keeps increasing to the end of training. Only MaxRL consistently exceeds the base model in pass@k; RLOO and GRPO show massive pass@k degradation.
• Large reasoning models: Training Qwen3-1.7B-Base and Qwen3-4B-Base on POLARIS-53K (approximately 50,000 mathematical reasoning prompts, 256 prompts per batch, 16 rollouts per prompt, 1000 RL steps) and evaluating on AIME 2025, BeyondAIME, MATH-500, and Minerva, MaxRL Pareto-dominates GRPO with higher pass@1 and improved pass@k. MaxRL improves pass@k relative to both the pretrained base model and the GRPO checkpoint in 7 out of 8 evaluation settings. The paper reports up to 20× test-time scaling efficiency gains over GRPO with a perfect verifier, and describes a 2.3×–19.2× speedup range in the corresponding figure caption.
• Different optimization dynamics: MaxRL produces higher gradient norms on harder prompts and lower gradient norms on easier prompts (analyzed with Qwen2.5-1.5B-Instruct on MATH-500), matching cross-entropy behavior in the image classification setting. MaxRL also generates at least one correct rollout for a larger fraction of prompts during training (measured out of 128, 16, and 16 rollouts for SmolLM-360M-Instruct, Qwen3-1.7B-Base, and Qwen3-4B-Base respectively). GRPO and RLOO drive the per-prompt pass-rate distribution to the extremes of 0 and 1, consistent with distribution sharpening, whereas MaxRL maintains a broad pass-rate distribution.

Methodology in Plain English

The researchers start from a simple observation: if a task ends in success or failure, a model implicitly assigns a probability to success for each input. Maximizing the log of that probability is the natural objective, but it cannot be differentiated through the sampling process, so RL is used instead.

They then write the log of the pass rate as a Maclaurin series in failure events, giving an infinite sum of fail@k terms. Differentiating shows that the maximum likelihood gradient is a weighted mixture of pass@k gradients, where the first term is exactly what standard RL optimizes. Cutting the infinite sum at a finite term T produces a family of objectives: small T behaves like RL, large T behaves like maximum likelihood.

To optimize these objectives, they use a theorem showing that the maximum likelihood gradient equals the expected score function conditioned on success. Instead of the expensive rejection-sampling approach that existing work uses to obtain success-conditioned samples, they sample from the ordinary policy and average score functions only over the successful rollouts, normalizing by the number of successes K rather than the total number of samples N. They add a zero-mean control variate (the unconditional average score) to reduce variance, and show that the resulting on-policy algorithm differs from standard methods only in how the advantage is normalized.

They validate this across four settings: an image classification task where exact maximum likelihood is available as cross-entropy, a procedurally generated data-rich maze environment, a data-scarce fixed GSM8K dataset trained for many epochs, and billion-parameter Qwen3 models trained on mathematical reasoning. In all cases they train on-policy and compare training objectives rather than algorithms, with RLOO and GRPO as primary baselines.

Why This Matters

The paper argues that a widely used class of RL objectives is optimizing the wrong quantity at the population level on binary-correctness tasks, and that a minimal implementation change recovers a better approximation of likelihood as compute grows. The 20× test-time scaling efficiency claim suggests practical inference savings, not just training gains. The related-work section notes that supervised learning and RL are complementary paradigms, with supervised training described as stable, sample-efficient, and well-calibrated within the training distribution.

Real-world applications where a strong verifier exists, per the paper:

• Mathematical reasoning — the paper's own large-scale experiments evaluate on AIME 2025, BeyondAIME, MATH-500, and Minerva, and train on POLARIS-53K.
• Programming — the paper cites programming as a setting that admits a strong verifier.
• Formal proof assistance — the paper cites Lean as a domain with a strong verifier.
• Navigation and structured prediction — both listed among tasks commonly handled with RL despite defining an implicit likelihood over successes.

Industry relevance: The method is a one-line change to standard RL implementations and targets settings where an external verifier produces binary feedback, so it can be dropped into existing LLM post-training pipelines. Reported benefits include less pass@k degradation over long training and fewer samples needed at inference when filtering with a perfect verifier.

Future Directions

• Investigating distribution sharpening — the paper conjectures that GRPO's inverted weighting at high pass rates may contribute to Pass@k degradation when datasets contain many overly easy inputs, and explicitly leaves a detailed analysis to future work.
• Deeper study of optimization dynamics — the authors state that MaxRL differs in significant ways from other RL objectives and leave a deeper investigation of these differences to future work.
• Improving MaxRL gradient estimators — the paper notes that any improvement in pass@k gradient estimators directly translates into improved estimators for the truncated maximum likelihood objective.
• Non-binary rewards and control variates — the appendix is described as examining how MaxRL and GRPO differ in distribution sharpening and non-binary reward settings, and the control variate baseline is noted as having no theoretical guarantees despite empirically reducing gradient variance.

Target Audience

Researchers and practitioners working on reinforcement learning for LLM post-training and reasoning, especially those who implement policy-gradient methods like RLOO and GRPO and want to understand how objective choice affects pass@1, pass@k, and inference efficiency. It also suits theoretically inclined readers interested in the connection between likelihood-based objectives and RL, and engineers looking for low-cost modifications to existing RL training pipelines in verifier-backed domains.

Authors’ abstract

Reinforcement learning (RL) is the method of choice for training models in setups where the objective function can only be evaluated by sampling from the model. Our key observation is that when the feedback is terminal and binary, models implicitly induce a likelihood over correct rollouts. Maximum likelihood would be the natural framework in such settings, but RL is used instead as a workaround to the non-differentiability. We prove that the standard, expected-reward RL formulation is only a first-order approximation of the likelihood. To remedy this mismatch, we introduce Maximum Likelihood Reinforcement Learning (MaxRL), a compute-indexed family of sample-based objectives that interpolate between expected-reward RL and maximum likelihood as sampling compute is scaled. The resulting objective is a one-line change to standard RL implementations. MaxRL Pareto-dominates existing methods in all tested models and tasks, achieves up to $20\times$ gains in test-time scaling efficiency over GRPO, and scales more favorably with additional training data and compute.