Tapered Off-Policy REINFORCE - Stable and efficient reinforcement learning for large language models

Thank you for your review and many questions.

1/ Comparison between TOPR and TIS for various proportions of positive examples

While we did not perform the experiment for TIS, Figure 8 in the appendix gives us insight on the relative impact of downsampling negative examples compared to downweighing them for TOPR. We can see that using 10% positive examples in the dataset performs comparably, albeit slightly better, to using 50% positive examples, regardless of the effective percentage of positive examples. Difference is on the order of a couple percentage points, which is less than the difference between TIS and TOPR shown in Figure 4. Should you place this as high priority (because of the incurred cost), we could run the equivalent of Figure 8 for TIS in the final version of the paper.

Figure 8 should also give an idea of what Figure 4 would look like with the actual proportion of positive examples on the x-axis.

Another big difference between TIS and TOPR comes when the dataset is not composed of trajectories generated from the base model but either from a more powerful model (distillation) or with expert trajectories. In the former case, the importance ratio can be close to 0 if the teacher model is very different from the student model and TIS will assign a very small gradient to these positive examples. In the latter case, one does not even have access to the sampling distribution and cannot use the importance ratio. We believe this is why, even though TIS performs well in our setup, SFT is the predominant method for fine-tuning LLMs.

As your comment was echoed by all other reviewers, we will strongly emphasize this difference in the final version should our submission be accepted.

2/ Claims of Table 1

Before discussing evidence, we want to highlight the fact that there would only be a significant difference between TOPR and TIS in the presence of desirable trajectories with low probability under $\pi$. We identify two cases where this can happen:

• First, when most of the datasets is comprised of negative examples, the optimization will decrease $\pi(\tau)$ for sampled trajectories $\tau$, which are generally those with larger $\mu(\tau)$ with $\mu$ the sampling policy. Hence, during the stochastic optimization, we might select positive trajectories $\tau$ for which $\pi(\tau)$ is much smaller than $\mu(\tau)$. The gradient for these trajectories will be much smaller for TIS than for TOPR, leading to poorer optimization. The evidence for this is in Figure 4, where the performance of TIS drops much more than that of TOPR for low effective percentage of positive examples.
• Second, when the dataset is composed of examples not drawn from the base model, for instance in the case of behaviour cloning. In that instance, even when there are only positive examples in the dataset, it might happen that $\pi(\tau)$ is much smaller than $\mu(\tau)$ and the model will not learn them. The empirical evidence for the superiority of TOPR over TIS in that context is that behaviour cloning uses an SFT-like objective, also used by TOPR for positive examples, and not the REINFORCE based one used by TIS.

However, in cases where a/ the dataset is generated from the base model and b/ there is a signification proportion of positive examples in the dataset, we can expect $\pi(\tau)$ to never be much smaller than $\mu(\tau)$ and for TIS and TOPR to perform similarly, as observed on the right-hand side of the two plots in Figure 4.

3/ Gradient accumulation

That is an interesting question. As experiments ran on 8 GPUs, we accumulated 1 gradient per GPU before applying the update (see section C.2 in the appendix). We did not try other values of gradient accumulation for this submission. However, we are in the process of open-sourcing a version of the code and we tested it with various batch sizes without observing any significant difference so far. Please note however that our experiments with that new codebase are not as thorough as the ones we have done for the submission.

Could you maybe explain why you think different methods might be affected differently by the percentage of negative examples? We have additional theoretical results, which we decided not to include in the submission for conciseness, showing that TOPR is likely to have lower variance than SFT as it minimizes the impact of a gradient step on trajectories not in the batch. We can definitely consider including these results in the Appendix should you think this is important.

4/ Error bars

All figures except Figure 6 and Figure 9 (in the Appendix) contain error bars corresponding to two standard errors (see section C.2). Error bars are indeed extremely small due to the number of solutions generated. Note that our error bars capture uncertainty in the generation of the test examples, not in the training process nor the creation of the training set as this would significantly increase the total computational cost of the experiments.

5/ Cost of discarding negative examples

The sole difference between SFT and TOPR is the inclusion of negative examples. Hence, the cost of this waste cannot be isolated from "differences in the behavior of different methods" as the different behavior between SFT and TOPR is precisely this waste. This is in contrast with, e.g., the difference between TOPR and Naive REINFORCE where all examples are included but with a different update. Thus, Figure 3 highlights this waste.

Thank you for the kind review. Regarding your two comments:

1/ Comparison to other methods

• We do compare to PPO: theoretically in Figure 2 and empirically in Figure 1.
• While GRPO uses advantages, it does not address the main limitation of PPO which is that the gradient does not flow as soon as the policy deviates significantly from the base policy, making it unsuitable for truly off-policy optimization. We can thus expect to observe the same suboptimal plateau as PPO.
• Although TRPO can be faster than PPO in terms of updates, making sure the constraint is satisfied can be computationally expensive.

2/ Larger networks

We appreciate the desire for more thorough experiments, a desire echoed by other reviewers. Sadly, the current set of experiments was already done at significant expense, especially given their thoroughness, for instance to evaluate the influence of the baseline. Larger networks also achieve high accuracy on MATH and GSM8K, leading to most examples being positive. Hence, a meaningful experiment using larger networks should involve harder datasets, leading to even larger compute requirements.

However, given the strong demand for the reviewers for such additional tasks, we would love if you collaboratively could have a prioritized list of tasks that should be added. This would allow us to strengthen the submission while keeping computational costs manageable.

3/ Figure 5

If you have the time, we would love to know how we can improve Figure 5. Are you talking about the right plot, where the legend crosses the lines? In any case, we will update the Figure to improve its readability.

Thank you for the review.

1/ More tasks

We acknowledge your desire, echoed by other reviewers, to see TOPR tested on more standard benchmarks. Unfortunately, our GPU budget (in dollars) required us to make the choice between shallow evaluation across many datasets, or focusing on a smaller set and going in depth. The current set of experiments was already done at significant expense given their thoroughness, and every additional task requires careful tuning of the baselines to make sure our results are trustworthy.

However, while our submission focused on math reasoning datasets, we also did preliminary evaluations of TOPR on query expansion. In these experiments, the website clinicaltrials.gov was used as a ground truth for expanded query terms. For instance, a researcher may search "IV" but expect "intravenous" to be included in their search results. We fine tuned a Llama 3 8b model to produce a domain-specific set of query expansions after having collected and verified a corpus of positive and negative samples.

TOPR was able to shift the increase the performance of the LLM, measured as the intersection over union of the expanded set of queries generated by the LLM compared to the ground truth. This was done without any domain knowledge. The results remain preliminary but could be included if desired, although they do not correspond to a standard benchmark.

2/ Compute budget

All methods that we evaluated require the same compute budget for generating samples.

For training, TOPR, PPO and REINFORCE all require about the same time to perform one gradient step (on a H100 node in our setup, about 1.3s per batch). PPO and TOPR compute the ratio of probabilities from the reference model, which requires an additional forward pass (the reference). We found this to add a small overhead (5-10% of per-step cost), with computing the gradients and the backward pass the dominant cost. DPO is slower because it performs the backward pass on two samples and computes two reference probabilities.

Overall, every figure that uses "Gradient steps" as the x-axis can be equivalently thought of as approximately using runtime on that axis.

We thank the reviewer for their comments and questions.

While we agree that TOPR can be viewed as a mixture of SFT (for the positive examples) and TIS (for the negative examples) on self-generated examples, we are a bit surprised to see this listed as a weakness rather than a strength. Had such a mix been straightforward, one would expect this method to have been developed earlier given its strong results. Most implementations that only rely on positive rewards, such as STaR, ReST, etc. emphasize complex prompting techniques (such as rationalization) and dataset curation, while here we deal with the negative samples directly.

Additionally, the derivation of TOPR is not a random idea. It is the result of an understanding of the behavior of several well-established methods and their own strengths and weaknesses. We are glad to have made our presentation clear enough that the idea of TOPR seems obvious but we dispute the claim that it is not a new contribution.

On the other hand, we agree with the statement that we provide no way to easily set the baseline parameter. As we stated in our submission, we found the best effective percentage to be around 10-20% for both GSM8K and MATH. Of note, our result shows that the GRPO baseline is not necessarily optimal. A deeper understanding of the baseline parameter would indeed be most welcome.

Moving on to your questions:

1/ Construction of the dataset for DPO

For each question, we generate 16 answers, leading to two sets $T_+$ and $T_-$. We then sample with replacement 16 contrastive pairs for each question. When one of the two sets is empty, i.e. all answers are incorrect or all answers are correct, then we cannot generate contrastive pairs and the question is discarded. We do this to compare things fairly across algorithms (same number of generated responses). Although we could attempt a different scheme (e.g., sample 8 distinct pairs for each question), this would be expensive for questions for which the model rarely generates the correct answer.

2/ Setting of the baseline

There is an explicit formula for the setting of the baseline parameter to achieve a particular effective positive ratio, shown in Eq. 7 in Appendix B. If the proportion of positive examples in the dataset is $p$ and the baseline is $c$, then the effective proportion of positive examples is $\tilde{p} =\frac{p(1-c)}{1 + c(1 - 2p)}$. Solving for $c$ readily gives the baseline parameter to use. Not that this is for rewards of +1 and -1 but the formula can be adjusted for any rewards.

3/ TIS vs. TOPR

The importance of the SFT update for the positive examples, as opposed to simply using TIS for all examples, is of prime importance. This concern was echoed by other reviewers and we thus intend to greatly clarify the difference in the revised submission.

Before answering your question, we want to clarify what might have been a misunderstanding: the method we call TIS in our submission uses truncated importance sampling for both positive and negative examples, as you suggest.

We were also surprised by the strong performance of TIS given the predominance of SFT-like methods to finetune LLMs. This said, TIS does struggle when there are few positive examples, because in this case their importance ratio goes to zero faster than the model can positively reinforce those examples.

TIS requires a reference or sampling distribution for the positive examples, allowing it to leverage expert data for instance. As such expert trajectories are often used to do behaviour cloning, TIS is hardly used despite its strong performance when the dataset is entirely generated from the base model. Further, even if we knew the sampling distribution for expert trajectories, some or many importance ratios for these trajectories would be close to 0 if the base model is far from an expert, leading to a small gradient for these trajectories. We would then likely be in a similar situation as that when there are few positive examples, i.e. TOPR would outperform TIS.

In addition to this property, we did not find a setting where TIS consistently outperformed TOPR. As shown in Figure 4, there are cases where TIS is superior to TOPR, but they require the baseline to be very carefully selected, and thus a potentially expensive hyperparameter sweep.

4/ Non-binary rewards

We decided to privilege a thorough experimental analysis on a smaller set of tasks rather than provide less conclusive answers on more tasks. We acknowledge the compromises of this decision and we do not have results for non-binary reward signals. Should the reviewer have suggestion on which such dataset, with manageable computational complexity, to include, we will give it strong consideration.