Research Topic
Could changing how an AI model weighs its successful training attempts, a single-line modification, dramatically improve how it learns to solve math and coding problems?
Project Summary
Training AI models to reason through math and code problems is deceptively hard. Standard reinforcement learning methods teach a model by rewarding correct answers, but they only optimize for "what's the chance I get this right on the first try?" That turns out to be a first-order approximation of a deeper, more principled objective: maximum likelihood.
Fahim Tajwar, Andrea Zanette, and collaborators at Carnegie Mellon University identified a surprisingly simple fix. Their framework, Maximum Likelihood Reinforcement Learning (MaxRL), shows that normalizing the learning signal by the number of successful attempts rather than the total number of attempts causes the training objective to converge to exact maximum likelihood as more compute is applied. The practical change is a single line in the advantage calculation.
The mathematical basis comes from a Maclaurin expansion of the log-likelihood objective. Standard RL (and GRPO) optimizes only pass@1. MaxRL optimizes a harmonic mixture of successively higher pass@k gradients, meaning harder problems where correct solutions are rare receive stronger learning signals rather than being effectively ignored.
The experiments ran across two ACCESS GPU resources: NCSA Delta and NCSA DeltaAI. The paper is structured around four compute tiers, designed to both validate MaxRL progressively and make the results accessible to research groups with different amounts of compute available:
- ImageNet (image classification): Runnable on a single GPU. A single experiment takes 10–15 hours depending on GPU type (A100 vs. H200/GH200).
- Maze navigation: Requires autoregressive generation from a transformer, making it significantly more compute-intensive. One full node of 4×GH200 for 48–72 hours.
- GSM8K math reasoning (SmolLM2-360M): Requires 2 nodes of 4×GH200, or 1 node of 8×H200 on Delta, for up to 96 hours (two 48-hour jobs run in succession).
- Large-scale LLM training (Qwen3-1.7B and 4B on POLARIS-53K): Requires 2–4 nodes of 4×GH200 for approximately 1–1.5 weeks per full training run.
Scaling across two clusters brought practical challenges, and the Delta and DeltaAI support staff were immensely helpful throughout. The team used Open OnDemand for job submission, monitoring, and interactive debugging sessions on both systems, particularly useful when iterating on environment setup. The most significant technical issue was package compatibility: while building Flash-Attention and vLLM is straightforward on Delta, DeltaAI's aarch64 (ARM) CPUs often lack pre-built package wheels. The team worked around this by first compiling from source, then later using pre-compiled aarch64 wheels available as open-source modules (flash-attention-prebuild-wheels). Multi-node communication and checkpoint management were also non-trivial, since MaxRL's longer training runs often exceed a single job allocation. Globus handled data and checkpoint transfers between Delta and DeltaAI reliably across the two systems. The team's solutions, including proper use of the ray Python package for resource management and automatic checkpoint resumption, are open-sourced in the MaxRL code repository.
The team used XDMoD throughout the project to track compute usage across both clusters, monitoring CPU efficiency and idle cores in real time, complementing the job-level visibility they got through Open OnDemand. This proved directly useful: XDMoD revealed that vLLM, which is optimized for large models (1B+ parameters), introduced significant inefficiencies when applied to the small 3M-parameter Maze models. That insight led the team to build a custom inference system for the Maze experiments, recovering wasted compute hours that would otherwise have gone undetected.
MaxRL consistently matched or exceeded GRPO on pass@1 accuracy while dramatically improving pass@k, meaning the trained models find a correct answer far more reliably when given multiple attempts. Achieving the same coverage as MaxRL required GRPO to use 7.9× to 19.2× more samples at test time. MaxRL also showed less overfitting to the training distribution and maintained more solvable problems throughout training, enabling continued learning in later epochs.
For researchers looking to adopt MaxRL, an Explore or Discover tier allocation is a reasonable starting point. The team's recommended approach: find a smaller, representative setting that runs on a single node (8×H200 or 4×GH200 at most), validate the approach there, then scale up incrementally, the same ladder the paper itself follows.
