PaperlayerView exact revision on arXiv

Source-linked AI summary

The Best of N Worlds: Aligning Reinforcement Learning with Best-of-N Sampling via max@k Optimisation

Farid Bagirov, Mikhail Arkhipov, Ksenia Sycheva, Evgeniy Glukhov, Egor Bogomolov

arXiv:2510.23393v1cs.LG

TL;DR

RLVR can improve single-generation performance while reducing diversity and pass@k at larger sampling budgets, and existing pass@k estimators leave an off-policy gap. This paper derives on- and off-policy gradient estimates for continuous-reward max@k; off-policy BoN performs best on most evaluated datasets.

  • Problem

    RLVR can improve pass@1 while reducing diversity and pass@k, and existing pass@k estimators are on-policy despite modern RLVR methods also using off-policy updates.

    Source: §1 · pp. 1–2

  • Method

    The paper derives unbiased policy-gradient estimates for directly optimizing max@k, including off-policy updates.

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

  • Results

    3.7 p.p. was the highest max@128 gain, on MBPP; off-policy BoN performed best on all datasets except LiveBench, where it performed similarly to BoN mean.

    Source: §1 · p. 2, §4.5 · p. 8

  • Takeaways & Limitations

    The proposed objective effectively optimizes models for BoN inference on coding tasks.

    Source: §6 · p. 9

    BoN objectives require multiple completions per task, increasing inference wall-clock time for large N and reducing the number of prefixes that fit in a batch.

    Source: §5 · p. 9

Abstract

from arXiv · show

The application of Reinforcement Learning with Verifiable Rewards (RLVR) to mathematical and coding domains has demonstrated significant improvements in the reasoning and problem-solving abilities of Large Language Models. Despite its success in single generation problem solving, the reinforcement learning fine-tuning process may harm the model's exploration ability, as reflected in decreased diversity of generations and a resulting degradation of performance during Best-of-N sampling for large N values. In this work, we focus on optimizing the max@k metric, a continuous generalization of pass@k. We derive an unbiased on-policy gradient estimate for direct optimization of this metric. Furthermore, we extend our derivations to the off-policy updates, a common element in modern RLVR algorithms, that allows better sample efficiency. Empirically, we show that our objective effectively optimizes max@k metric in off-policy scenarios, aligning the model with the Best-of-N inference strategy.

1 INTRODUCTION

RLVR improves single-generation performance, but pass@1-focused training can reduce diversity and pass@k; existing pass@k estimators also leave an on/off-policy gap. This work derives unbiased max@k gradient estimates for both update settings and reports gains over baselines.

  • Best-of-N generates multiple solutions and selects the best using a verifier, making pass@k a natural evaluation metric; the study evaluates coding tasks with continuous test-based rewards.Evaluation uses CodeContests for training and evaluation, plus LiveCodeBench, LiveBench, CodeForces, and MBPP.Source: §1 · pp. 1–2
  • Pass@1-focused RLVR can reduce generation diversity and pass@k, while existing pass@k estimators are designed for on-policy RL despite modern methods also using off-policy updates.Many estimators use binary rewards, which are sparse and can limit stability and sample efficiency.Source: §1 · pp. 1–2
  • Up to +3.7 p.p. increase in max@k over baselines is reported, with the objective boosting higher k while also improving lower k.The highest gain is reported on MBPP.Source: §1 · p. 2
  • The paper derives unbiased gradient estimates for directly optimizing max@k with both on-policy and off-policy updates.max@k generalizes pass@k to continuous rewards.Source: §1 · p. 2

2 BACKGROUND

The background motivates max@k by describing how RL fine-tuning can trade high-k performance and diversity for low-k gains. It defines the continuous-reward objective and situates the work among RL, Best-of-N, and inference-aware fine-tuning approaches.

2.1 MOTIVATION

  • RLVR raises pass@k at low k but lowers it at high k, where some correct generations become nearly unattainable; entropy shifts toward zero as confidence rises and diversity falls.The paper attributes these observations to GRPO fine-tuning on CodeContests and describes the entropy distribution before and after RL fine-tuning.Source: §2.1 · p. 2
  • Binary-reward training degrades coding ability, while continuous rewards improve pass@1 but slightly reduce pass@128, motivating direct optimization of max@k.Binary reward requires passing all tests; continuous reward is proportional to the fraction passed.Source: §2.1 · p. 3

2.2 PROBLEM STATEMENT

  • For coding tasks, pass@k is the probability that at least one of k generations is correct, whereas max@k is the maximum reward among k generations.The reward may be binary or continuous, based on unit-test outcomes.Source: §2.2 · p. 3
  • The objective generalizes inference-time aggregation through a function selecting the best generation; choosing max yields expected max@k, and binary rewards yield expected pass@k.The paper seeks unbiased gradient estimates for this objective.Source: §2.2 · p. 3

2.3 RELATED WORK

  • Prior work reports that RLVR can improve pass@1 while narrowing reasoning capability or reducing pass@k, with proposed explanations including training-data diversity and updates on positive samples.The cited studies also report a trade-off between performance and policy entropy.Source: §2.3 · p. 4
  • Prior BoN-mimicking fine-tuning addresses computational cost but assumes a single inference completion; inference-aware pass@k objectives instead have focused on on-policy updates.The paper considers both on-policy and off-policy scenarios.Source: §2.3 · p. 4

3 PASS@K OPTIMIZATION

The section derives unbiased gradient estimates for optimizing the pass@k objective in both on-policy and off-policy settings. It uses GRPO as the primary optimization method and develops a max-reward estimator for on-policy updates alongside an approximation for off-policy updates.

3 PASS@K OPTIMIZATION

  • The section derives unbiased gradient estimates of the objective for both on-policy and off-policy updates.The derivations target the objective in equation 3.Source: §3 · p. 4

3.1 PRELIMINARIES

  • GRPO is the paper’s main optimization method; it is critic-free and computes advantages by z-scoring sampled completions’ rewards.The paper notes that methods with different variance-reduction techniques use their original advantage calculations instead.Source: §3.1 · p. 5

3.2 ON-POLICY CASE

  • Monte Carlo estimation replaces expectations with sample averages and averages over all sampled subsets of size k to estimate the on-policy objective without bias.The passage identifies the resulting objective as an unbiased estimate of equation 6.Source: §3.2 · pp. 5–6
  • For max reward, the on-policy estimator weights each completion using counts of size-k subsets in which another completion is the highest-scoring member.The weights count subsets containing the completion whose gradient is used and subsets where the other completion has the highest score.Source: §3.2 · p. 6

3.3 OFF-POLICY

  • Exact off-policy estimation for max reward is exponentially hard, so the method approximates the gradient to first order when recent-policy probability ratios are near 1.The probability ratio is the optimized policy’s probability divided by the sampling policy’s probability.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

The experiments compare inference-aware fine-tuning methods on coding benchmarks using Qwen2.5-7B-Coder-Instruct. The proposed off-policy BoN approach leads on max@128 across four datasets and performs similarly to BoN mean on LiveBench; max@1 results are more mixed.

4.1 DATASETS

  • CodeContests supplies 2,490 training and 276 test examples after filtering; LiveCodeBench, LiveBench, CodeForces, and MBPP are evaluation-only datasets.CodeContests tasks are filtered by complexity and length; LiveCodeBench uses version 6.Source: §4.1 · p. 7

4.2 MODELS

  • Qwen2.5-7B-Coder-Instruct is the base model, selected because the listed datasets target instructed code generation..Source: §4.2 · p. 7

4.3 BASELINES

  • Baselines include best-completion-only BoN-max with mean or second-best variance reduction, unbiased BoN LOO-1, and a BoN mean variant using GRPO-style standardization..Source: §4.3 · p. 8

4.4 HYPERPARAMETERS

  • All methods share hyperparameters: eight completions per sampling step, learning rate 5e-6, and one training epoch; off-policy updates use three PPO iterations.Sampling uses temperature 1.0, top-p 1.0, and a 256-token generation limit.Source: §4.4 · p. 8

4.5 RESULTS

  • 3.7 p.p. is the largest max@128 gain, on MBPP; off-policy BoN significantly outperforms baselines on four datasets and performs similarly to BoN mean on LiveBench.The highest p-value for rejecting the null against the second-best method is 0.039 on the four datasets; LiveBench has p = 0.08.Source: Table 2 · p. 8, §4.5 · p. 8
  • The proposed method has the best max@1 in two of five settings; three other objectives lead the remaining settings, while the proposed method ranks second on MBPP.The authors describe this alongside high pass@128 values as a potentially fruitful exploration/exploitation trade-off.Source: Table 3 · p. 9, §4.5 · p. 9

5 LIMITATIONS AND FUTURE WORK

The approach requires sampling multiple completions per problem, which can increase inference time and constrain batch size. Future work includes extending the method to math and investigating alternative training schedules and sampling schemes.

5 LIMITATIONS AND FUTURE WORK

  • Sampling multiple completions per problem can increase inference wall-clock time at large N and reduces how many prefixes fit in a training batch.Independent samples can be generated in parallel to reduce latency, and gradient accumulation can address the batch constraint.Source: §5 · p. 9
  • Future directions include math tasks with carefully constructed continuous rewards, dynamic transitions between max@k and max@1, and conditional sampling for iterative refinement.The authors also propose investigating how to close the max@k–max@1 gap by aligning single-generation and best-of-k distributions.Source: §5 · p. 9

6 CONCLUSION

The paper derives unbiased on-policy and off-policy gradient estimates for directly optimizing max@k and reports that the proposed estimate optimizes the BoN inference strategy on coding tasks. Its appendix describes baseline variants that differ in how they assign advantages and reduce variance.

6 CONCLUSION

  • The paper derives an unbiased Monte Carlo gradient estimate for max@k, extends it to off-policy updates, and reports effective optimization of BoN inference on coding tasks.The conclusion identifies continuous rewards as crucial to successful optimization.Source: §6 · p. 9

A BASELINES

  • BoN-max assigns advantage only to the highest-scoring completions, using either the mean reward or the second-best reward for variance reduction.These variants are identified as BoN-max mean and BoN-max second.Source: A BASELINES · p. 13
  • BoN LOO-1 modifies the on-policy BoN estimate to reduce variance, while BoN mean applies GRPO-style z-scoring to rewards calculated with the BoN objective..Source: A BASELINES · p. 13

B ADDITIONAL RLVR RESULTS

Extended pure RLVR training improves max@1 on both evaluated datasets, but its max@128 outcome depends on the dataset and training setup.

B ADDITIONAL RLVR RESULTS

  • Four epochs of pure RL training attain the best max@1 on both CodeContests and LiveCodeBench..Source: B ADDITIONAL RLVR RESULTS · p. 14
  • On CodeContests, max@128 degrades with additional pure RLVR training, even as max@1 improves..Source: B ADDITIONAL RLVR RESULTS · p. 14
  • On LiveCodeBench, three PPO steps attain the best max@128; the authors suggest diversity loss may be less pronounced on some datasets.They note that max@k degradation may still occur at larger k values.Source: B ADDITIONAL RLVR RESULTS · p. 14

C.1 PASS@K AND MAX@K EVALUATION

This appendix describes how pass@k and max@k are evaluated, then gives implementations of on-policy and off-policy BoN objectives.

C.1 PASS@K AND MAX@K EVALUATION

  • For pass@k, the appendix uses an unbiased estimator; for max@k, it averages the maximum reward across subsets of size k.Evaluation samples 256 completions per prompt.Source: §C.1 · p. 14

C.2 ON-POLICY BON

  • The on-policy implementation computes binomial subset weights, applies them to rewards, and returns the resulting BoN rewards and weights.The code constructs separate off-diagonal and diagonal scales.Source: §C.2 · p. 15

C.3 OFF-POLICY BON

  • The off-policy implementation clamps deltas and uses delta-based correction terms to produce policy-gradient weights and weights.The code includes cumulative deltas in its off-diagonal and diagonal corrections.Source: §C.3 · p. 16

D ON-POLICY DERIVATION

The on-policy derivation rewrites the policy-gradient sum by grouping subsets according to their maximum-scoring completion, then counts the subsets contributing to each weight.

D ON-POLICY DERIVATION

  • Sorting completions by score lets the derivation replace a sum over subsets with a sum over possible maximum values.A subset's maximum is the completion at index j when it contains j and its other indices are no greater than j.Source: D ON-POLICY DERIVATION · p. 17
  • The derivation defines w_ij as the count of eligible size-k subsets containing i and with j as the maximum index.These weights express the policy-gradient sum in terms of subset counts.Source: D ON-POLICY DERIVATION · p. 17
  • For j = i, w_ij counts subsets containing i by choosing k−1 other elements up to i; for j > i, it counts subsets containing both i and j by choosing k−2 others.The diagonal case is zero when i < k.Source: D ON-POLICY DERIVATION · p. 17

E OFF-POLICY DERIVATION

The off-policy derivation groups subsets by their maximum-scoring completion and incorporates policy deltas into the resulting weights.

E OFF-POLICY DERIVATION

  • For j > i, the off-policy correction groups deltas, counting the repeated 1 + δ_i + δ_j term and summing contributions from other deltas.The passage describes this as an approximation of an unbiased off-policy derivation.Source: E OFF-POLICY DERIVATION · pp. 17–18
  • As in the on-policy derivation, sorting completions by score allows the off-policy sum to be grouped by each subset's maximum index.The subset contains j, and its other indices are no greater than j.Source: E OFF-POLICY DERIVATION · p. 18
  • The derivation separates the j = i and j > i cases, counting subsets that contain i alone as the maximum or contain both i and j.For j = i, the sum is empty when i < k.Source: E OFF-POLICY DERIVATION · p. 18

F FULL EVALUATION RESULTS

The appendix lists full pass@k and max@k evaluation tables across five coding benchmarks, but the supplied passages do not report their cell values.

F FULL EVALUATION RESULTS

Loading 2510.23393v1…