Optimizing Language Models for Inference Time Objectives using Reinforcement Learning

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

Regarding the computational efficiency issue mentioned in weaknesses, can the authors provide another ablation study on the total number of samples? Namely, allow mean policy gradient to do k times of update than the proposed algorithms, and plot the figures.

In our manuscript the baseline (mean PG, and mean PPO), already uses the same number of samples as the other variants to enable a fair comparison in terms of compute cost.

We also conducted exp. using k samples for using the PPO loss but remove the value model. For mean objective, this corresponds to the GRPO loss. For the pass@k objective, note that the proposed biased pass@k objective already subtracts the mean out of k samples, which coincides with the Dr.GRPO loss recently proposed https://arxiv.org/abs/2503.20783 (only subtract the mean but not divide by the std. deviation).

For completeness, we also include:

1. the mean objective using Dr. GRPO loss
2. biased pass@k objective using GRPO loss, i.e., further divide by the std. Deviation

Due to time and resource constraints we are only able to finish the exp. using llama 3.1 8b model on the coding tasks. We will add the others in the revision.

Checkpoints are all trained for 8000 gradient steps using the same set of hyperparameters. We report the number by avg the perf of the last 5 ckpts, (8k, 7.8k, …, 7.2k) steps, to ensure a non cherry-picked results.

pass@1:

pass@10:

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

I think the gains in Pass@k come at a significant cost to Pass@1, particularly in coding tasks.

One major takeaway message is that there is a trade-off between pass@k and pass@1 in general. For coding tasks, pass@k is generally a more meaningful objective than pass@1 because we focus more on harder tasks, and obtaining one valid solution to a hard problem is enough [1].

[1] Li et al, Competitive level code generation with AlphaCode, 2022

Importantly, Mean PG outperforms other multi-sample objectives on the MATH test set, even at Pass@4.

We are being truthful about the train/test generalization gap here - the performance improvements of our algorithm are better assessed at the training set, since the algorithm does not account for train/test generalization. The generalization gap might depend on multiple factors out of our control, such as the choice of the base model and data used for pretraining the models.

It is true that the mean PG outperforms multi-sample objectives with MATH at pass@4, but the outperformance is marginal and limited to pass@4. The overall conclusions about the performance improvement from multi-sample gradient estimate still stand, but we are truthful about the fact that it does not necessarily generalize in all cases.

The approach depends on having a strong verifier, raising doubts about practical improvements from these objectives.

This is not true. We do not require verifiers at all during training, unlike Chow et al. Can the reviewer clarify on the claim that the approach depends on a strong verifier and where it is used?

For both training and evaluations, we use the standard numpy package to compute the rewards for training. This is common practice in the reasoning LLM literature.

The paper lacks a comparison with standard PG + entropy bonus. The entropy bonus in PG algorithms is complementary to the approach we take. In other words, in principle, our method can be combined with an entropy bonus as-is.

As a result, we do not think adding entropy bonus comparison offers new insights into the current algorithmic designs. Entropy bonus might lead to performance improvements as a result of avoiding premature collapse, but it also introduces a tunable hyper-parameter, whereas our approach is hyper-parameter free compared to the baseline.

The majority of the related work section is relegated to the appendix, omitting key citations from the main text. Important prior work, such as Variational Best-of-N Alignment, is not discussed.

We apologize for the inconvenience of moving most related work discussion to the Appendix due to ICML page limit. When the limit is relieved (camera ready), we can move much of the discussion into the main paper.

Variational BoN: this is a related work and we will discuss it. Variational BoN does not optimize for the original objective as-is, but rather requires a variational approximation. This leads to an algorithm akin to the one introduced in Balashankar et al, and in a pairwise comparison setting. In comparison, we focus on the point-wise reward case and require no approximation to the original objective of interest.

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

On the other hand, for general RL algorithms, I suspect the performance of algorithms like GRPO, that utilize the same number of samples, will even outperform the objective proposed in this paper.

We strongly disagree with the reviewer.

In our manuscript the baseline (mean PG, and mean PPO), as what the reviewer refers to as “general RL algorithm”, already uses the same number of samples as the other variants to enable a fair comparison. And we already provide empirical evidence in our paper that they do not outperform our proposed variants that optimize for pass@k.

In other words, general RL algorithms such as GRPO only optimize for the mean performance, and would similarly benefit from algorithmic improvements proposed in this work. GRPO is not designed to optimize for pass@k on its own.

Baselines like GRPO and other single-objective functions with the same sample complexity are necessary for comparison.

Please refer to our response to Reviewer bYLj for exp. on GRPO; we also include Dr. GRPO as a variant. Note that, our method could be integrated with various policy gradient variants, e.g., the PPO or vanilla Policy Gradient in the manuscript. Please refer to our response to Reviewer 61s6 for results of SFT and Rejection Sampling Finetuning (RFT) for additional comparison.

To be clear, GRPO [1] doesn't require any prior assumption on the test-time algorithm, but we can still evaluate it on both approaches.

GRPO (or any policy gradient method derived from the objective $\max E(f(y)))$ does have prior assumptions on the test-time algorithm: the one that uses the exact protocol as the $E(f(y))$, i.e., pass@1. The reviewer refers to [2] to argue that a model trained using $\max E(f(y))$ objective works well on pass@k metrics, i.e., $\max E(f(y_1, … y_k))$ objective. This could be also true on the other way round.

We motivate this from the first principle that such an objective is to maximize the k-sample objective $E(f(y_1, … y_k))$, contrary to the common formulation of maximizing cumulative return $E(f(y))$, where the return is a function for individual sample. Our proposed scheme works for settings that fit into this objective $E(f(y_1, … y_k))$. Specifically, we show that in the case of $f = r \circ \text{maj}$ and $f = \max$, this corresponds to the maj@k and pass@k metrics in math reasoning and code generation and gives empirical evidence.

That said, our method also can be evaluated in multiple sampling budgets in test time (we do not set a constraint that training on k = 8 only works for k = 8 in test time, see our section of generalization to k = 100).

Also, we share exactly the same hyperparameter as GRPO: the number of samples per prompt.

As shown in Figure 18 and Figure 19 from [2], the R1-Distill-Qwen-7B model, a single model, outperforms both maj@k and pass@k separately-trained 8B models reported in this paper on MATH. This is not a fair comparison since R1-Distill-Qwen-7B is a stronger base model, but my argument is mainly to question the necessity of separate objectives.

The reviewer suggests that the access to a stronger model R1-Distill-Qwen-7B questions the necessity of separate training objectives.

We object to this argument because the R1-Distill-Qwen-7B model is distilled from a capable mode R1, which is specifically trained via RL methods, which is the focus of our paper.

This does not undermine the necessity of separate objectives but rather reinforces the validity and potential benefits of explicitly targeting inference-time objectives during training, which results in a better model in terms of pass@k performance and could bring the pass@k superiority to other models by following the same distillation process. Also, Figure 18 and Figure 19 from [2] is reported on MATH500, while our numbers are on a full MATH test set (a total 2500 problems).

There are no implementation details in the paper. How are the hyperparameters selected for different settings in the paper?

Please find implementation details in the Appendix A (in particular A. and A.4).

For sampling, we use top-p=1 and temperature $\tau=1$ sampling with standard parameters, as they are compatible with the RL training, though evaluations can be executed with various sampling scheme.

For training, the experiments are rather robust to the choice of learning rate and other training parameters: Therefore, we fixed these hyperparameters among different settings (mean/pass@k objective) within math/code experiments respectively, to ensure a fair comparison.

Other hyper-parameters such as $k$ are ablated in the experiments. We bump the $k$ for code experiments since CodeContests is a challenging benchmark.

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

In the proof of Lemma 1, why is the expectation of two independent variables zero? I didn’t see any assumption made on either variable being zero.

This is because the expectation of the score function is zero $E[ \log \pi(y_i|x)]=0$. More precisely,

$E[f_{-i}\log \pi(y_i|x)] = E[f_{-i}] \cdot E[\log \pi(y_i|x)]=0$

Where the equality makes use of the statistical independence between $i$ and other samples $-i$. We will make this more clear in the revision.

The proof for Lemma 2 is hard to follow. Not sure which expression the authors refer to. I suggest the authors add a complete proof in the appendix for better understanding of the claims.

Thank you for the suggestion, we will add a more detailed explanation in the Appendix.

The goal of Lemma 2 is mainly to derive the analytic form of the biased gradient estimate, which turns out to optimize a biased objective with an interesting analytic form - it is the average of leave-one-out objectives from $k$ samples.

The variance reduction is not shown both theoretically and empirically.

The understanding that the control variate improves variance, either in the form of the leave-one-out control variate (for the unbiased estimator), or subtracting a mean baseline (for the biased estimator), was based on common consensus in the literature on this topic [1]. In fact, we cannot show theoretically that the variance is reduced in all cases, due to the multiplication with the score functions; but we can show that the advantage estimation has lower variance. We can make the theoretical argument more precise in the revision.

For empirical validation: we never implemented an estimator without control variate for variance reduction, as such estimators are plain REINFORCE gradient estimators and do not work well in practice. We can add such an ablation in the revision in case the reviewer finds it useful.

[1] Kool et al, Buy 4 REINFORCE Samples, Get a Baseline for Free!, 2023

why does pass@k policy gradient become much worse than the baseline when PPO is used?

We have some speculations regarding this observation: the pass@k policy gradient, compared to the mean policy gradient and biased pass@k policy gradient, has the sparsest gradient, i.e., the weight before the policy gradient is non-zero only for positive samples where it is the only 1 positive sample among k samples, while for the other samples the weight is 0. Biased pass@k policy gradient, by subtracting the mean among k samples, makes the policy gradient no-sparse in such cases. When incorporating PPO, adding a learnt value breaks this sparsity due to the imperfectness of the value model.

What is the performance comparison between the proposed approach (e.g., pass@k policy gradient) and purely SFT trained or pass@1 mean policy gradient method with pass@k test-time strategy? This would be meaningful as it shows how much benefit the proposed method can bring.

We’ve observed a performance drop when doing SFT on CodeContests training set. Given that Llama 3.1 70B has already been heavily tuned in the post-training, some code solutions in CodeContests training set are of less quality than the data presented in its post-training phase. For example, some imports in the Python codes are outdated (e.g., from fractions import gcd will throw an ImportError since Python 3.9).

Therefore, we included a baseline for Rejection Sampling Finetuning (RFT) instead, where the model generates 200 rollouts for each problem instance in the training dataset and collect the correct solutions, we downsample them to maximum 50 solutions per problem and conduct SFT on this dataset.

We show the comparison below, the best performance is bold and the second best is underscored. With these comparisons, the both pass@k (k=8) variants achieves the best performance in terms of the closest metrics pass@10.

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

The paper lacks key comparisons with InfAlign (Balashankar et. al.), and Chow et. al., that both propose RL training objectives that directly optimize pass@n or BoN performance.

We have discussed the technical differences between our work and Balashankar et. al, Chow et al, in Section 3.3, with more extended discussion in Appendix C.

As pointed out in the paper, Balashankar et. al and Chow et al consider structurally different problem formulation as we have in this work. Balashankar et. al considers a pairwise preference setting (vs. point-wise reward setting in our case) while Chow et al considers the case with auxiliary models such as verifiers or scorers (vs. no requirement for auxiliary models in our case). As a result, it is challenging to set up an apple-to-apple comparison against such concurrent work since they tackle very different settings. Further, we also feel that it’s fair not to make direct comparison to unpublished concurrent work.

The performance on MATH is underwhelming, especially Maj@k in Figure 15, and even pass@4 in Figure 2, is only better in a narrow KL regime, with the proposed policy gradient.

We acknowledge the fact that the performance improvement on MATH is not huge, though they are still statistically significant. The upper bound of the performance improvement also depends heavily on the base model we started with (which we do not have full control over). MATH as a benchmark task is also relatively easy since the pass@16 has also reached ~86%, at a competitive level as the frontier models such as O1 and Deepseek R1. This also puts a limit on the amount of improvement.

The narrow KL divergence regime: Fig 2 shows improvement across all KL values we observe during training. For Fig 15, since maj@k signals are much sparser, this makes it more difficult for the model to reach a larger KL divergence deviation during training. However, for the KL divergence regime it reaches, the performance improvements are clear. We have highlighted this point as an algorithmic trade-off in the paper.

Additionally, depending on practical applications, the “narrow KL” regime might suffice since for general post-training, ideally the RLHF stage does not lead to too much deviation from the reference policy. We leave this to the judgement of practitioners.

Do the authors think that there is a fundamental tradeoff between optimizing for pass@1 and pass@k performance?

In principle, pass@1 and pass@k share common goals but there is a conflict between the two objectives, in that pass@k always encourages higher diversity and coverage over the solution space, while pass@1 encourages more determinism.

Why is it that we cannot learn a single model with good pass@1 performance on easy problems and good pass@k on hard problems? Is this lack of representation capabilities, pre-training biases, or issues with training objectives/algorithms?

It should be possible in principle: this might require a hybrid objective, and switching between pass@1 vs. pass@k training objectives depending on the problem difficulty. We might even want to devise a novel algorithm that adjusts the objective based on the online sampled pass rate. This definitely is a promising path for future work that integrates the two objectives into a unified algorithm.

The pretraining biases and representation capabilities certainly plays a key role as well, as suggested by recent investigations to reproduce Deepseek R1 (e.g., Qwen models can reproduce R1 behavior more consistently, and model sizes matter too). That said, training algorithms / objectives should make a difference, as we have demonstrated in the paper, especially because RL has become increasingly important in post-training.

We thank the reviewer for the time and efforts in providing us with the valuable reviews. We will incorporate your feedback into our later revision. We address your comments below.

1. Generalization of improvements to evaluation time is not always consistent.

The generalization of improvements to evaluation is indeed less straightforward, as we have explicitly discussed in the paper. The generalization gap depends in a complex way on the training data and model sizes. The cleanest way to assess our algorithmic progress is via the training performance, though we have also included evaluation for completeness. We still see evaluation improvements overall, though they are not as strong as training rewards improvements.

2. The method may be less effective for problems where models already achieve high performance (e.g., MATH with LLaMA 70B). This is true and we have elaborated on this observation in the paper. The k-sample objectives are generally less effective for easy problems to the model, since the model does not have strong signals to make improvements. However, for more difficult problems (HARP for Llama 70B), we see more significant improvements, which might be a more useful case for many practitioners as well. Consider emphasizing earlier that these methods do not require auxiliary models (unlike best-of-k with reward models)

We will make sure to emphasize this point earlier on in this work. We have introduced the inference objectives of interest in Section 2.

Including a discussion on compute efficiency (vs benefit) can help practitioners decide when to use these techniques.

We will elaborate on such points more extensively in the revision. In general, we find that the newly proposed algorithms automatically improve inference time objectives with a similar computational budget as the baseline algorithms (which maximizes the mean objectives). The algorithmic change is also almost a few-line code change, which is quite convenient for practitioners.

Can you elaborate on why training-time improvements on pass@k/majority voting objectives do not always translate to better test-time performance? Could this relate to diversity collapse or overfitting to training prompts?

We conjecture that this might be more related to the generalization gap between training and test set, which depends in a rather complex way on factors such as the model size, and training data that went into the model pre-training and SFT (the released Llama 3.1 model might have SFT’ed on similar dataset to our RL prompt set, which might reduce the level of generalization we expect). As a result, maybe training set diversity and performance improvements do not transfer 100% to the evaluation set.

When should practitioners prefer biased vs unbiased gradient estimators? Do you recommend any heuristics for setting k?

We feel that a good approach is to decide on the estimator adapting to the specific use case. Overall, we find the two estimators to be both quite competitive, and since the unbiased gradient estimator optimizes for the exact objective, it might be the first thing to try.

For setting k: this depends on the application of interest. We can choose a specific k that will be used at inference time. We can also choose a generic value of k to maximize model performance while diversifying the samples as much as possible - in general, a large value of k indicates more diversity. We find that values of {4,8,16} are good starting points, depending on the application.