🤖 machine learning

Maximum Likelihood Reinforcement Learning

This paper introduces Maximum Likelihood Reinforcement Learning (MaxRL), a novel framework that bridges the gap between expected-reward RL and maximum likelihood by offering a compute-indexed objective that Pareto-dominates existing methods and significantly improves test-time scaling efficiency.

Original authors: Fahim Tajwar, Guanning Zeng, Yueer Zhou, Yuda Song, Daman Arora, Yiding Jiang, Jeff Schneider, Ruslan Salakhutdinov, Haiwen Feng, Andrea Zanette

Published 2026-08-21

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Fahim Tajwar, Guanning Zeng, Yueer Zhou, Yuda Song, Daman Arora, Yiding Jiang, Jeff Schneider, Ruslan Salakhutdinov, Haiwen Feng, Andrea Zanette

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

In the world of artificial intelligence, there are two dominant ways to teach a computer how to think. The first is like a strict teacher who provides the correct answer for every question, allowing the student to compare their work against a key and adjust their understanding accordingly. This method, known as supervised learning, has powered the most impressive image recognizers and language tools we use today. The second approach is more like a child learning to walk: there is no answer key, only the environment itself. The learner tries a step, falls, tries again, and eventually learns to balance by feeling the consequences of their actions. This is reinforcement learning, a method designed for situations where the path to a solution is not a straight line of data, but a series of choices that lead to a final outcome. For years, when AI systems faced problems where the only feedback was a simple "yes" or "no" at the very end—such as solving a math problem or navigating a maze—researchers relied on this second method. They treated the problem as a game of maximizing rewards, assuming that if the AI could just get enough "yes" signals, it would learn to succeed.

A new study challenges this long-held assumption, suggesting that the standard way of teaching these systems is actually a rough approximation of a much more powerful, but previously inaccessible, principle. The researchers, working across several universities, discovered that when an AI model generates a correct answer, it is implicitly creating a probability of success. In a perfect world, the best way to train the model would be to maximize this probability directly, a concept known as maximum likelihood. However, because the process of generating an answer often involves unpredictable steps that cannot be easily measured by a computer's internal math, scientists have been forced to use the "game" approach of reinforcement learning instead. The new work proves that this standard approach is only a first-order guess at the true goal. It misses the subtle, critical information hidden in the rare moments when the model struggles but eventually succeeds. By developing a new training framework called Maximum Likelihood Reinforcement Learning, or MaxRL, the team has found a way to bridge this gap, allowing the system to learn from its successes with a precision that was previously impossible.

The core of the discovery lies in how the computer weighs its mistakes and its victories. In the traditional method, every time the model produces a correct answer, it receives a reward, and the system adjusts its internal settings to make that outcome more likely. However, this method treats a correct answer the same way regardless of how difficult the problem was. If the model solves a simple puzzle and a complex one with equal ease, the traditional system gives them equal credit. The new research shows that this is inefficient. The true goal is to maximize the likelihood of being correct, which mathematically demands that the system pay much more attention to the hard problems where success is rare. The standard method fails to do this because it only looks at the average success rate. The new framework, MaxRL, changes the calculation by looking at the entire history of attempts. It asks: "If we tried this problem many times, how often would we get it right?" By analyzing the frequency of success across many attempts, the system can infer the true difficulty of the problem and adjust its learning accordingly.

To test this idea, the researchers built a system that could simulate this more sophisticated learning process. They created a family of objectives that could be tuned like a dial. At one end of the dial, the system behaves exactly like the standard reinforcement learning methods used today. At the other end, it behaves like the ideal maximum likelihood trainer, which is theoretically perfect but usually impossible to run. In the middle, the dial allows the system to use more computing power to get a better, more accurate approximation of that ideal. The team found that as they increased the amount of computing power dedicated to generating multiple attempts at each problem, the system's performance improved dramatically. It wasn't just that the system became more stable; the very nature of what it was learning changed. It began to focus intensely on the difficult problems that the old methods ignored, leading to a much deeper understanding of the task.

The results of this approach were striking across a variety of tests. In a controlled environment where the researchers could compare the new method directly against the theoretical ideal, the new system closely matched the performance of the perfect trainer as they added more computing power. In contrast, the standard methods hit a wall, failing to make significant progress even when given massive amounts of data. When the team moved to more complex, real-world scenarios like navigating mazes and solving mathematical problems, the advantage became even clearer. On tasks involving mathematical reasoning, the new method achieved performance gains that were up to twenty times more efficient than the current state-of-the-art methods. This means that to reach the same level of skill, the new system required far fewer attempts and less computing time. Perhaps most importantly, the new method did not suffer from the "overfitting" problem that often plagues these systems, where a model learns to memorize the training data so well that it fails to generalize to new, unseen problems. While other methods saw their ability to generate diverse, correct solutions degrade over time, the new framework maintained a healthy variety of correct answers, suggesting it was truly learning the underlying logic rather than just memorizing patterns.

The implications of this work extend beyond just making AI smarter; it changes how we think about the limits of machine learning. For a long time, the difficulty in training these systems on complex reasoning tasks was blamed on the algorithms themselves or the lack of data. This study suggests that the bottleneck was actually the objective function—the mathematical goal the system was trying to achieve. By shifting the goal from simply maximizing the average reward to maximizing the likelihood of correctness, the researchers unlocked a new level of efficiency. The system is no longer just guessing at the right answer; it is learning to understand the probability of success in a way that mirrors how humans learn from rare, hard-won victories. The researchers note that this approach works best when there is a clear way to verify if an answer is correct, such as in math or coding, but the principles could potentially apply to other areas where the outcome is binary. As artificial intelligence continues to tackle more complex and abstract problems, the ability to learn efficiently from sparse feedback will be crucial. This new framework provides a path forward, showing that by refining the way we define success, we can teach machines to learn not just faster, but better.

Technical Summary: Maximum Likelihood Reinforcement Learning (MaxRL)

Problem Statement

Reinforcement Learning (RL) is the standard paradigm for training models in settings where the objective function is non-differentiable and can only be evaluated via sampling, such as tasks with binary success/failure outcomes (e.g., program synthesis, mathematical reasoning, navigation). In these scenarios, the model induces an implicit likelihood of correctness, $p_\theta(y^*(x) | x)$ , for a given input $x$ . While Maximum Likelihood (ML) is the principled objective for optimizing such probabilities in differentiable settings, it is traditionally intractable here due to the non-differentiable stochastic generation process. Consequently, practitioners resort to standard RL, which optimizes the expected reward (the expected pass rate, $E[p_\theta]$ ), rather than the log-likelihood ( $E[\log p_\theta]$ ).

The authors observe that standard RL and ML induce fundamentally different optimization behaviors. Specifically, the gradient of the ML objective involves an inverse-probability reweighting ( $1/p_\theta$ ), which places significantly higher emphasis on "hard" inputs with low success probabilities. In contrast, standard RL gradients ( $\nabla p_\theta$ ) treat all inputs linearly, often underweighting difficult cases. This mismatch suggests that RL is merely a workaround for non-differentiability rather than the optimal objective for correctness-based tasks.

Methodology: Maximum Likelihood Reinforcement Learning (MaxRL)

The paper introduces Maximum Likelihood Reinforcement Learning (MaxRL), a framework that bridges the gap between standard expected-reward RL and exact Maximum Likelihood by leveraging increased sampling compute.

Theoretical Foundation

The authors derive a population-level expansion of the Maximum Likelihood objective using a Maclaurin series in terms of failure events. For a pass rate $p$ , the log-likelihood is:
$\log p = -\sum_{k=1}^{\infty} \frac{(1-p)^k}{k} = -\sum_{k=1}^{\infty} \frac{\text{fail@}k}{k}$
Differentiating this yields the ML gradient as an infinite harmonic mixture of pass@k gradients:
$\nabla_\theta J_{ML} = \sum_{k=1}^{\infty} \frac{1}{k} \nabla_\theta \text{pass@}k$
Standard RL corresponds to the first-order approximation ( $k=1$ ), optimizing only $\nabla_\theta \text{pass@}1$ . MaxRL proposes truncating this series at a finite order $T$ , creating a family of objectives $J_{MaxRL}^{(T)}$ that interpolate between RL ( $T=1$ ) and exact ML ( $T \to \infty$ ).

Gradient Estimation

To estimate these truncated objectives without adaptive sampling (which is computationally expensive for low pass rates), the authors propose a simple on-policy estimator.

Sampling: For a given input, sample $N$ trajectories $z_1, \dots, z_N$ from the policy $m_\theta$ .
Success Conditioning: Identify successful trajectories where $f(z_i) = y^*(x)$ . Let $K$ be the count of successes.
Estimator: The gradient estimator averages the score functions ( $\nabla_\theta \log m_\theta(z_i|x)$ ) only over the successful trajectories, normalized by $K$ :
$\hat{g}_N(x) = \frac{1}{K} \sum_{i=1}^N r_i S_i \quad (\text{if } K \ge 1)$
where $r_i$ is the binary reward and $S_i$ is the score.
Variance Reduction: A control variate (the unconditional average score) is subtracted to reduce variance, resulting in the final estimator used in Algorithm 1.

Key Insight: Theorem 2 establishes that this estimator is an unbiased estimator for the truncated ML objective of order $T=N$ (or $T=N-1$ with baseline subtraction). Crucially, increasing the number of rollouts $N$ does not just reduce the variance of a fixed objective; it increases the approximation order of the objective itself, moving the optimization target closer to the true Maximum Likelihood.

Implementation

MaxRL requires only a one-line modification to standard on-policy RL implementations (like RLOO or GRPO). Instead of normalizing the advantage by the total number of samples or the standard deviation, the advantage is normalized by the mean reward (or the count of successful samples), effectively reweighting the gradient signal to match the ML objective.

Key Contributions

Formalization: The paper formalizes correctness-based RL as a latent-generation Maximum Likelihood problem, proving that standard RL optimizes only the first-order approximation of the ML objective.
Compute-Indexed Objectives: It introduces a family of objectives that interpolate between expected reward and exact ML via a Maclaurin expansion in pass@k events, controlled by the sampling compute budget.
Estimator-Objective Equivalence: It provides a simple, unbiased on-policy estimator where increasing the rollout count $N$ directly improves the fidelity of the approximated likelihood objective, rather than merely reducing variance.
Empirical Validation: Extensive experiments demonstrate that MaxRL Pareto-dominates existing methods (RLOO, GRPO) across various scales and data regimes.

Experimental Results

The authors evaluate MaxRL across four distinct regimes:

Controlled Differentiable Setting (ImageNet): When exact ML (Cross-Entropy) is available, MaxRL converges to it as the number of rollouts increases. In contrast, standard REINFORCE fails to make progress from low initial pass rates even with massive sampling budgets.
Infinite Data Regime (Maze Navigation): In a data-rich setting with 1M unique mazes, MaxRL scales significantly better with compute than GRPO and RLOO. While other methods plateau, MaxRL continues to improve, showing a substantial gap between pass@1 and pass@k, indicating retained diversity and learning potential.
Data-Scarce Regime (GSM8K): On a fixed dataset (GSM8K) trained for many epochs, MaxRL demonstrates superior resistance to overfitting. While RLOO and GRPO suffer from "pass@k degradation" (loss of diversity) and plateau early, MaxRL sustains improvement over many epochs, achieving higher peak performance while maintaining output diversity.
Large-Scale Reasoning (Qwen3 Models): Training 1.7B and 4B parameter models on mathematical reasoning tasks, MaxRL Pareto-dominates GRPO. It achieves higher pass@1 scores while simultaneously improving pass@k (coverage). Notably, MaxRL yields up to 20x test-time scaling efficiency gains over GRPO when using a perfect verifier, as it requires fewer samples to achieve the same coverage.

Significance and Claims

The paper claims that MaxRL offers a principled alternative to standard RL for binary-reward tasks. Its primary significance lies in reframing the limitations of foundation model RL not as inherent to the optimization or sampling process, but as a consequence of objective choice.

Optimization Dynamics: MaxRL concentrates learning signals on harder prompts (low pass rates), unlike GRPO and RLOO which can drive distributions to extremes (0 or 1).
Scaling: MaxRL scales more favorably with both additional data and compute. In data-scarce settings, it avoids the diversity collapse often seen in prolonged RL training.
Efficiency: By aligning the training objective with the test-time goal (maximizing the probability of finding any correct solution among $k$ samples), MaxRL drastically improves inference efficiency.

The authors conclude that MaxRL is a promising direction for scaling RL in correctness-based domains, offering a path to recover the benefits of Maximum Likelihood training in non-differentiable settings through increased sampling compute.