PaperlayerView exact revision on arXiv

Source-linked AI summary

SERA: Scale-Equalized Rollout Allocation for Maximum Likelihood Reinforcement Learning

Zihao Chen, Fanxiang Xiong, Hongran Ren, Xuefeng Bai, Zhongxiang Dai, Kehai Chen, Zhiguo Zhang, Zhiyong Wang, Yu Cheng

arXiv:2609.36552v1cs.LGcs.AI

TL;DR

Finite-rollout MaxRL can scale low-success prompts’ likelihood gradients more strongly under uniform rollout allocation, distorting their relative contributions. SERA reallocates a fixed budget to equalize scaling and corrects weighting from unequal rollout counts. Experiments report closer gradient alignment on ImageNet and improved multi-sample coverage over MaxRL on maze and reasoning tasks.

  • Problem

    Under finite rollout budgets, MaxRL’s estimator attenuates prompt-wise likelihood gradients according to success probability and rollout count, so uniform allocation can disproportionately attenuate low-success prompts.

    Source: §1 · p. 1

  • Method

    SERA reallocates a fixed rollout budget to approximately equalize finite-rollout scaling coefficients, using a max–min waterline allocation and a multiplicity correction for unequal rollout counts.

    Source: §1 · p. 2, §3 · p. 3, §3.1 · p. 4

  • Results

    SERA aligns sampled updates more closely with exact likelihood gradients on ImageNet and improves multi-sample solution coverage over MaxRL on maze navigation and mathematical reasoning under matched rollout budgets.

    Source: §1 · p. 2, §4.2 · p. 7, §4.4 · p. 8

  • Takeaways & Limitations

    Across the evaluated settings, redistributing matched training rollout budgets improves high-budget solution coverage while more closely aligning ImageNet updates with exact likelihood gradients.

    Source: §1 · p. 2, §6 · p. 10

Abstract

from arXiv · show

Maximum Likelihood Reinforcement Learning (MaxRL) targets prompt-wise log-success and has shown strong performance on reasoning tasks. Under finite rollout budgets, however, the estimator used by MaxRL attenuates each prompt's likelihood gradient by a factor that depends on its success probability and rollout count. Under uniform rollout allocation, the common rollout count fails to compensate for success-dependent attenuation, leaving low-success prompts more strongly attenuated and distorting their relative contributions to the expected aggregate gradient. We introduce SERA (Scale-Equalized Rollout Allocation), which redistributes a fixed rollout budget to approximately equalize these finite-rollout scaling factors. Building on our theoretical analysis of how finite rollouts distort prompt-wise likelihood gradients, we formulate the allocation as a fixed-budget max--min problem, derive a waterline solution to its continuous relaxation, and introduce a multiplicity correction to remove the additional prompt weighting induced by heterogeneous rollout counts. Experiments show stronger alignment with exact likelihood gradients in a controlled ImageNet setting and improved multi-sample solution coverage over MaxRL on maze navigation and mathematical reasoning under matched training rollout budgets.

1 INTRODUCTION

Finite rollouts can distort MaxRL’s prompt-wise log-success gradient, despite its emphasis on low-success prompts. SERA reallocates a fixed budget to equalize scaling, and experiments report closer gradient alignment and better multi-sample coverage than MaxRL.

Motivation

  • MaxRL targets average prompt-wise log-success, emphasizing improvements on low-success prompts, while connecting its likelihood gradient to Pass@K gradients across sampling budgets.The introduction motivates multi-sample objectives because gains in Pass@1 need not translate to larger-budget coverage.Source: §1 · p. 1
  • Under uniform rollout allocation, low-success prompts are more strongly attenuated, reducing their relative gradient contributions and potentially shifting the expected batch-gradient direction.This finite-rollout mismatch motivates asking how to allocate a fixed training rollout budget across prompts.Source: §1 · pp. 1–2

SERA

  • SERA approximately equalizes finite-rollout scaling coefficients; its fixed-budget max–min allocation has a waterline solution in the continuous relaxation.A multiplicity correction removes the extra prompt weighting caused by unequal rollout counts under response-mean aggregation.Source: §1 · p. 2

Experimental validation

  • SERA improves multi-sample solution coverage over MaxRL on maze navigation and mathematical reasoning under matched rollout budgets, and aligns sampled updates more closely with exact likelihood gradients on ImageNet.The experiments evaluate SERA against MaxRL across the controlled ImageNet setting, maze navigation, and mathematical reasoning.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries define prompt success under binary verifier rewards and introduce the centered MaxRL estimator. Its finite-rollout fidelity varies with prompt success, so uniform counts can scale prompt gradients unequally.

Preliminaries

  • For binary verifier rewards, a prompt’s success probability is the probability that a response sampled from the policy receives reward one; MaxRL considers prompt-wise log-success.The setup draws N i.i.d. responses per prompt and defines their rewards and score gradients before presenting the centered estimator.Source: §2 · pp. 2–3
  • At a fixed rollout count, finite-rollout fidelity increases with success probability, so uniform allocation attenuates lower-success prompts more strongly and can change the aggregate gradient direction.Figure 1 contrasts fixed-count scaling, its possible batch-direction effect, and allocating more rollouts to lower-success prompts under Nmax = 64.Source: §2 · p. 3, Figure 1 · p. 3

3 METHOD

SERA allocates a fixed rollout budget to equalize centered-MaxRL scaling across retained prompts, with safeguards and weighting corrections for practical training. Its procedure combines bounded integer allocation with success estimates and response weights.

3 METHOD

  • SERA pairs scale-equalized allocation with multiplicity correction to prevent unequal rollout counts from adding prompt weighting under response-mean aggregation.The correction preserves uniform weighting of retained prompts while rollout counts control finite-rollout attenuation.Source: §3 · p. 3, §3.3 · p. 5

3.1 ROLLOUT ALLOCATION

  • The fixed-budget max–min allocation seeks to improve the least finite-rollout fidelity; its continuous optimum equalizes prompt coefficients, assigning more rollouts to lower-success prompts.Practical allocation bounds are incorporated into the waterline solution, then counts are rounded to integers while preserving budget and bounds.Source: §3.1 · p. 4

3.2 OBSERVABILITY SAFEGUARD

  • The observability safeguard selects capacity-feasible prompt prefixes using estimated success and checks that retained prompts meet a nonzero-signal threshold.The centered estimator requires mixed success and failure outcomes within a group for nonzero advantages.Source: §3.2 · p. 4

3.3 IMPLEMENTATION

  • Training uses historical reward observations for success estimates, freezes estimates within each epoch, and applies response weights before the policy update.If a candidate prompt lacks an available estimate, the batch uses uniform rollout counts; new outcomes enter at the next epoch boundary.Source: §3.3 · pp. 4–5

4 EXPERIMENTS

Experiments compare SERA with MaxRL under matched training rollout budgets, using an exact-gradient ImageNet control, maze navigation, and mathematical reasoning. Pass@K measures multi-sample solution coverage.

4.1 EXACT-LIKELIHOOD REFERENCE ON IMAGENET

  • At the larger ImageNet rollout budget, SERA beats MaxRL across all displayed final Pass@K budgets and narrows the gap to direct CE training.It also improves validation Pass@1 and Pass@128 during training; at the smaller budget, it leads over most of the displayed range.Source: §4.1 · p. 6, Figure 2 · p. 6
  • At N0 = 128 on ImageNet, SERA’s sampled-to-exact gradient cosine is 0.6266, compared with 0.4317 for MaxRL.The advantage persists into later training trajectories.Source: §4.1 · p. 6, Figure 2 · p. 6

4.2 MAZE NAVIGATION WITH A LIGHTWEIGHT TRANSFORMER

  • In mazes, SERA retains higher Pass@128 and Pass@1024 during later training and beats MaxRL across the final displayed Pass@K profile under Seqnorm.Under token-mean normalization, its advantage over MaxRL grows at larger evaluation budgets.Source: §4.2 · p. 7, Figure 3 · p. 7

4.3 MATHEMATICAL REASONING WITH SMOLLM2

  • Under both normalizations, SERA improves SmolLM2 Pass@32, Pass@128, and Pass@256 over MaxRL, with little change in Pass@1.SERA also exceeds GRPO at every reported K and beats DARS-HW at Pass@32, Pass@128, and Pass@256 while using fewer training responses.Source: §4.3 · p. 8, Table 1 · p. 8

4.4 MATHEMATICAL REASONING WITH QWEN

  • Across Qwen model scales, SERA beats GRPO at the largest displayed K on all eight model-benchmark pairs and improves high-budget coverage over MaxRL.For Qwen2.5-Math-1.5B, gains reach 7.00 percentage points on BeyondAIME and 6.67 on AIME 2025 at K = 512; Qwen3-4B gains 8.00 points on BeyondAIME and 4.78 on Minerva Math.Source: §4.4 · pp. 8–9, Figure 4 · p. 9

5 RELATED WORK

Related work studies reasoning objectives and adaptive rollout strategies. SERA retains MaxRL’s log-success target while controlling its estimator’s prompt-dependent finite-rollout scaling under a fixed budget.

RL objectives and reasoning coverage

  • Pass@K captures multi-sample solution coverage, which may not improve when Pass@1 improves; inference-aware and diversity-oriented objectives explicitly target these behaviors.The cited examples include PKPO, SetPO, and UCPO.Source: §5 · p. 9

Adaptive rollout allocation and selection

  • Adaptive methods allocate or filter rollouts using learning gain, variance, success probability, nonlinear objectives, or selection and continuation strategies.Examples include Knapsack RL, VIP, DynaMO, HORA, DARS-HW, Reinforce-Ada, DAPO, Pilot-Commit, and SARA.Source: §5 · p. 9
  • Unlike methods targeting other signals, SERA solves a bounded max–min problem over centered MaxRL’s finite-rollout scaling coefficients for a fixed retained set and total rollout budget.Historical success estimates supply allocation inputs, and multiplicity correction removes response-count weighting from unequal group sizes.Source: §5 · p. 9

6 CONCLUSION

The conclusion reports that SERA aligns sampled updates more closely with exact likelihood gradients on ImageNet and improves multi-sample coverage over MaxRL on maze navigation and mathematical reasoning. The accompanying analysis explains how finite rollouts alter prompt-gradient contributions and why equalizing their scaling matters.

6 CONCLUSION

  • SERA improves multi-sample coverage over MaxRL on maze navigation and mathematical reasoning under matched training rollout budgets, while mathematical-reasoning Pass@1 remains comparable or improves.On ImageNet, sampled updates also align more closely with exact likelihood gradients.Source: §6 · p. 10
  • Uniform rollout counts more strongly attenuate low-success prompts, while unequal scaling can shift the expected update direction away from the target likelihood gradient.Equalizing scales preserves relative likelihood-gradient contributions for a fixed retained set.Source: §6 · p. 10
  • SERA redistributes a fixed rollout budget to approximately equalize finite-rollout scaling and corrects the additional prompt weighting caused by unequal rollout counts.Its allocation is formulated as a fixed-budget max–min problem with a waterline solution.Source: §6 · p. 10, A EXTENDED RELATED WORK · p. 17

A EXTENDED RELATED WORK

  • DARS-HW and Reinforce-Ada share a likelihood-oriented formulation with MaxRL, but the comparison with SERA concerns which estimator is controlled and how allocation is used.DARS-HW’s idealized schedule targets maximum likelihood; Reinforce-Ada can realize the log objective through sampling, gradient weighting, or both.Source: A EXTENDED RELATED WORK · p. 17

B.1 FINITE-ROLLOUT FIDELITY OF THE CENTERED ESTIMATOR

  • The centered estimator’s finite-rollout expectation scales each prompt’s log-success gradient, and heterogeneous counts assign different truncation orders of the same log-success expansion.Unequal coefficients can change the relative contributions of linearly independent prompt gradients and thus the aggregate direction.Source: §B.1 · p. 17, §B.2 · pp. 18–19

B.3 DERIVATION OF THE FIXED-BUDGET ALLOCATION

The continuous fixed-budget allocation equalizes finite-rollout fidelity coefficients through a waterline rule, with practical bounds and observability safeguards shaping its use. The appendix also contrasts allocation with direct inverse-fidelity reweighting.

B.3 DERIVATION OF THE FIXED-BUDGET ALLOCATION

  • The continuous fixed-budget max–min allocation assigns counts so every prompt reaches the same fidelity, with the water level determined by the total budget.Lower-success prompts require more rollouts for a common fidelity; practical counts are bounded and integerized while preserving the budget.Source: §B.3 · p. 19

B.4 MIXED-GROUP OBSERVABILITY AND BUDGET CONCENTRATION

  • The coefficient κ measures gradient scaling, whereas U measures the probability of observing nonzero centered advantages; SERA allocates by κ and uses U as an observability safeguard.When success probability approaches zero, maintaining a given fidelity requires increasingly many rollouts, so such prompts can absorb a large budget share and lower the attainable common fidelity.Source: §B.4 · p. 20

B.5 COMPARISON WITH INVERSE-FIDELITY REWEIGHTING

  • Directly dividing by κ removes scaling in expectation but leaves zero-signal-group probability unchanged and amplifies estimator covariance when κ is small.Increasing rollout counts instead changes both finite-rollout scaling and mixed-group observability.Source: §B.5 · p. 20

C IMPLEMENTATION DETAILS

The implementation freezes historical success estimates for each epoch, then uses bounded allocation, observability-guided prompt selection, and response weighting to preserve the rollout budget and retained-prompt weighting.

C.1 HISTORICAL SUCCESS-PROBABILITY ESTIMATES

  • Historical success and failure statistics are discounted and updated at epoch boundaries, while estimates remain fixed within an epoch; batches with any unavailable estimate use uniform rollout counts.The prior alone does not make an unseen prompt’s estimate available.Source: §C.1 · pp. 20–21

C.2 BOUNDED ALLOCATION AND EXACT-BUDGET INTEGERIZATION

  • The bounded allocator solves for continuous counts under frozen success estimates, floors them, then assigns remaining rollouts by the smallest one-increment squared-fidelity errors.The completion preserves the total budget and count bounds.Source: §C.2 · pp. 21–22

C.3 ACTIVE-PREFIX SEARCH AND CAPACITY FALLBACK

  • Observability-guided selection ranks prompts by estimated mixed-group probability at the maximum count and tests capacity-feasible prefixes using the original budget and final integer counts.If no tested prefix passes, a capacity-preserving fallback allocates to the top-ranked minimum-size set, which may remain below the observability threshold.Source: §C.3 · p. 22

C.4 LOSS AGGREGATION AND MULTIPLICITY CORRECTION

  • Under response-mean aggregation, each response is weighted by N0/Nq so every retained prompt has the same total response weight despite heterogeneous rollout counts.The response weight is applied to the centered advantage before the policy loss is computed.Source: §C.4 · p. 22

D.1 EXPERIMENTAL SETUP

The experiments compare SERA with MaxRL under a fixed rollout budget across ImageNet, maze navigation, and mathematical reasoning, using task-specific training and evaluation settings. Success probabilities come from exact ImageNet probabilities or historical estimates for the sequence tasks.

D.1 EXPERIMENTAL SETUP

  • Across tasks, the full-batch budget is M = BN0; SERA varies active prompts’ rollout counts, while batches lacking historical estimates use uniform allocation.For sequence tasks, each active response is weighted by N0/Nq; historical estimates are updated at epoch boundaries and frozen within each epoch.Source: §D.1 · p. 23

D.1.1 IMAGENET

  • ImageNet treats each image as a prompt and each sampled class label as a one-step action rewarded for matching its ground-truth category.The exact class probability allows allocation to be studied without success-probability estimation error; training uses a randomly initialized ResNet-50 and evaluates Pass@K at step 100,000.Source: §D.1.1 · p. 24

D.1.2 MAZE

  • Maze experiments use 7,424 training mazes and 400 held-out 17 × 17 mazes, with both token-mean and fixed-length sequence-normalization comparisons evaluated at step 3,000.The policy generates a valid path from a start to a goal; Figure 7 illustrates row-major symbolic cell-token input and a valid output path.Source: §D.1.2 · pp. 24–25

D.1.3 SMOLLM2 ON GSM8K

  • SmolLM2-360M-Instruct is trained on GSM8K and evaluated at step 2,000 on 1,209 GSM8K-Platinum questions using 256 responses per question.Pass@K uses the combinatorial estimator, with the same response pool for all K values at a checkpoint and a separate pool for online validation.Source: §D.1.3 · p. 25

D.1.4 QWEN

  • The Qwen mathematical-reasoning experiments compare Qwen2.5-Math-1.5B and Qwen3-4B-Base using token-mean aggregation and model-specific training data and evaluation benchmarks.Final checkpoints are step 600 and step 1,000, respectively; both use 512 sampled responses per question, temperature 0.6, and top-p 0.95.Source: §D.1.4 · p. 26

D.2.1 IMAGENET

On ImageNet, SERA improves final Pass@K and sampled-gradient alignment relative to MaxRL, and historical success estimates track exact probabilities. On maze navigation, SERA improves coverage under both tested loss normalizations.

D.2.1 IMAGENET

  • At N0 = 128, SERA raises ImageNet Pass@1 from 36.59% to 54.80% and Pass@128 from 53.22% to 87.51% versus MaxRL.The gains substantially narrow the gap to direct cross-entropy training across evaluated sampling budgets; SERA also improves Pass@1 and Pass@128 at N0 = 16.Source: Table 2 · p. 28, §D.2.1 · p. 28
  • At N0 = 128, SERA’s mean cosine similarity is 0.6266 versus 0.4317 for MaxRL across 100 measurements, and 0.7960 versus 0.4447 at the final checkpoint.The measurements compare each run’s sampled training-loss gradient with the exact full-batch likelihood gradient along its observed training trajectory.Source: §D.2.1 · p. 28
  • Historical ImageNet estimates achieve MAE 0.0386 and mean Pearson correlation 0.8134 against exact success probabilities without current-batch pilot sampling.A pilot alternative with N0 = 16 and m = 8 commits half the budget uniformly before adaptive allocation; those samples are reused for training on active prompts.Source: Table 3 · p. 29, §D.2.1 · p. 29

D.2.2 MAZE

  • On held-out mazes, SERA’s Pass@128 gains over MaxRL are 4.84 percentage points with token-mean normalization and 3.28 points with Seqnorm.SERA beats MaxRL at every reported K ≥4 under token-mean normalization and improves the full reported Pass@K profile under Seqnorm.Source: §D.2.2 · p. 29

D.2.3 SMOLLM2 ON GSM8K

On GSM8K-Platinum, SERA finishes with stronger validation coverage than MaxRL and has the highest final Pass@K among the four compared methods at every displayed K ≥4. Paired-bootstrap intervals for its Seqnorm gains over MaxRL are above zero at all three reported budgets.

D.2.3 SMOLLM2 ON GSM8K

  • During later training, MaxRL’s validation coverage declines while SERA maintains higher Pass@32 and finishes with stronger coverage.The validation trajectory compares MaxRL, SERA, GRPO, and DARS-HW; the reported metric is bootstrap Best@32, labeled Pass@32 by the trainer.Source: §D.2.3 · p. 29
  • SERA achieves the highest final Pass@K among the four methods at every displayed K ≥4 and improves over MaxRL at every displayed K ≥2 while retaining nearly identical Pass@1.It also outperforms GRPO across the displayed evaluation-budget range under the reported comparison.Source: §D.2.3 · p. 30
  • The 95% paired-bootstrap confidence intervals for SERA’s Seqnorm gains over MaxRL are above zero at all three reported budgets.Intervals use 10,000 resamples of evaluation questions with checkpoints and response pools held fixed.Source: Table 5 · p. 30, §D.2.3 · p. 30

D.2.4 MATHEMATICAL REASONING WITH QWEN

The Qwen experiments report Pass@K profiles and paired-bootstrap uncertainty for two mathematical-reasoning models. SERA’s reported Pass@512 gains over MaxRL are positive on both Qwen3-4B-Base benchmarks, while Qwen2.5-Math-1.5B’s Pass@32 intervals are above zero on three benchmarks.

D.2.4 MATHEMATICAL REASONING WITH QWEN

  • Table 6 reports Pass@K profiles for MaxRL and SERA on Qwen2.5-Math-1.5B, with the pretrained base model as a reference; evaluations use step 600 and 512 sampled responses per question.Source: §D.2.4 · p. 30, Table 6 · p. 31
  • For Qwen2.5-Math-1.5B, the Pass@32 confidence intervals for SERA-minus-MaxRL differences are above zero on AIME 2025, MATH-500, and OlympiadBench.The intervals widen at larger evaluation budgets.Source: §D.2.4 · p. 30, Table 7 · p. 31
  • At K = 512, SERA improves over MaxRL by 8.00 percentage points on BeyondAIME and 4.78 points on Minerva Math for Qwen3-4B-Base.The results use step-1,000 checkpoints and 512 sampled responses per question.Source: Table 8 · p. 32, §D.2.4 · p. 32

D.2.5 TRAINING DYNAMICS AND ALLOCATION BEHAVIOR

In the analyzed SmolLM2 run, SERA achieves stronger multi-sample coverage than MaxRL while mixed-reward groups rise late in training despite a lower cumulative count. Across the reported task windows, all candidate prompts remain active, and SERA’s training diagnostics show higher late-training entropy than the compared methods.

D.2.5 TRAINING DYNAMICS AND ALLOCATION BEHAVIOR

  • SERA achieves stronger multi-sample coverage than MaxRL as its mixed-reward group fraction overtakes MaxRL late in training despite a lower cumulative count over steps 30–2,000.The comparison uses matched rollout budgets and Seqnorm; Figure 15 also includes token-mean GRPO as a reference.Source: §D.2.5 · p. 32, Figure 15 · p. 33
  • All candidate prompts remain active across the analyzed Maze, SmolLM2, and Qwen2.5-Math-1.5B windows while rollouts are redistributed across the full candidate batch.Prompt activation is guided by the observability criterion, while allocation targets equalizing gradient-scaling coefficients within the active set.Source: §D.2.5 · pp. 32–33, Figure 15 · p. 33
  • On SmolLM2, SERA maintains higher late-training token entropy than MaxRL under both normalizations and than token-mean GRPO, while response-length behavior depends on normalization.With token-mean normalization, MaxRL’s mean response length rises markedly late in training while GRPO and SERA produce shorter responses; under Seqnorm, MaxRL and SERA remain closer. Figure 16 also reports actor gradient ℓ2 norms.Source: §D.2.5 · p. 33, Figure 16 · p. 34
  • On Qwen2.5-Math-1.5B, SERA maintains higher late-training token entropy than GRPO and MaxRL, accompanied by longer generated responses over the 600-step window.All three methods use token-mean normalization and the same total rollout budget per step; Figure 17 also reports actor gradient ℓ2 norms.Source: §D.2.5 · p. 33, Figure 17 · p. 34
Loading 2609.36552v1…