Tail-Likelihood Reinforcement Learning
Shrinivas Ramasubramanian, Daman Arora, Fahim Tajwar, Guanning Zeng, Qingyang Wu, Zhongzhu Zhou, Chenfeng Xu, Haiwen Feng, Yuda Song, Aarti Singh, Ruslan Salakhutdinov, J. Andrew Bagnell, Jeff Schneider, Andrea Zanette
cs.LG, stat.ML
2026-09-02
TailRL maximizes log-prob of beating a random reward threshold via an advantage swap. Code speedup hits 7.7× Best-of-1024; GRPO and RLOO collapse to copying.
Standard RL maximizes expected reward. For generative policies that average hides a real split: two policies can share a mean while assigning very different mass to rare, high-reward rollouts. Extra samples at train time and Best-of-k at inference only help if that upper tail still has coverage.
Recent work has seen the failure mode directly. Expected-reward training drops rare high-reward traces, and Best-of-k falls with them. Once those traces become too rare to resample, improvement stalls. The same coverage loss caps test-time scaling: a policy can look fine at k=1 and barely move as k grows. MaxRL already treats binary rewards by maximizing the log-probability of success, whose gradient is a harmonic mix of Pass@k gradients. Continuous rewards had no matching objective.
Every reward threshold defines a success event. The tail probability pθ(x,τ) is the chance a rollout beats τ. Expected reward is the area under that curve, an arithmetic mean. TailRL samples τ uniformly and maximizes log pθ(x,τ), a geometric mean. Harder thresholds are rarer, so the 1/p weight pushes the gradient onto the thin end of the distribution.
The gradient is a harmonic series of Best-of-k gradients. The 1/k weights come from the logarithm, so nobody has to pick an inference budget in advance. For binary rewards every nontrivial threshold is the same event, and TailRL reduces exactly to MaxRL.
A group of N rollouts defines an order-N truncation: N=1 recovers expected reward, N→∞ approaches the population objective. After sorting rewards, each threshold splits one unit of credit equally among rollouts that clear it; integrating those credits yields unbiased weights. In code this is an advantage swap. Center the weights inside the group and drop them into an existing GRPO or RLOO loop.
Four settings probe four failure modes.
ImageNet box localization uses IoU as a continuous reward. Population TailRL, with closed-form tail probabilities, beats L1+GIoU supervision on [email protected] and @0.75 and matches mean IoU. Finite-sample gradients move toward the population gradient as N grows, measured by cosine similarity. TailRL at N=16 already beats GRPO and RLOO trained at N=1024.
Text-maze navigation scores 1 only on a shortest path and gives partial credit for ending closer to the goal. When the initial shortest-path success rate drops from about 1% to 0.01%, GRPO and RLOO fail to learn reliably. TailRL still lifts Pass@1.
GUI grounding fine-tunes Qwen2.5-VL-3B and 7B and evaluates on ScreenSpot-Pro. Pass@1 is close to RLOO, but TailRL keeps climbing as k grows while RLOO plateaus earlier. At 3B, 8 samples match RLOO's Pass@1024; at 7B, 4 samples suffice. That is 128× and 256× fewer inference rollouts. A single TailRL sample already beats GRPO's Pass@1024.
Code runtime optimization is the sharpest stress test. The policy rewrites slow C++ programs from PIE; only correct-and-faster rewrites score above 1×. The base model is wrong 74.4% of the time, correct-but-not-faster 23.5% (a spike at 1×), and actually faster 2.1%. GRPO and RLOO push correctness above 98% while mean reward settles just below 1 and entropy drops one to two orders of magnitude: they copy the input. TailRL reaches mean reward 2.92 and 7.7× Best-of-1024 speedup on the test set, against 0.98× and 0.96×. Training is one epoch; TailRL is still climbing at step 300.
Swapping the advantage changes the objective without a new hyperparameter or a critic. Two regimes benefit: high-reward traces exist but are rare, and a safe moderate-reward shortcut would otherwise absorb exploration. If deployment relies on Best-of-k, training should keep mass on the tail rather than adding a diversity regularizer after the fact.
This is a clean incremental change: the objective moves, the pipeline barely does.
There is no dedicated limitations section. A few caveats sit in the experiments. GUI rewards live in [0, 2.5] while localization and maze use [0, 1]; the extension leans on an appendix treatment of bounded rewards that the main text does not spell out. The code run is matched compute, not a converged endpoint, so it is unknown whether longer GRPO training would escape the copy shortcut. PKPO optimizes a chosen kopt, which is a different contest from a harmonic mix over all k. Maze gains concentrate at very low initial success and shrink once the start policy is decent. All four tasks use verifiable rewards; there is no preference model or process reward.