Thank you for the thoughtful feedback and for highlighting the interesting insights that our work provides and its clarity. We treat your comments and questions below. If our response addresses your comments, we would greatly appreciate it if you would consider raising your score. Please let us know if you have any further questions; we will gladly respond.

W1: Scope of theoretical results. The core issue that we analyze – slow optimization due to low reward variance and its interplay with accuracy – does not depend on the neural network parameterization. It stems from the RLHF objective and applies to any softmax parameterized policy. That is, any policy that produces a distribution over outputs through the softmax function, which includes language models as a special case.

Our work includes three main theorems showing that: (i) the time it takes to increase the reward grows inversely with the initial reward variance (Theorem 1), (ii) more accurate reward models are not necessarily better teachers (Theorem 2), and (iii) for different language models different reward models can work better (Theorem 3). Indeed, while we prove Theorem 1 for general differentiable parameterizations, Theorems 2 and 3 are proved for tabular policies since they require establishing separations between the optimization time of reward models, which entails substantial technical difficulty (as further discussed below in response to other comments). However, given that the phenomenon that we analyze stems primarily from the objective rather than any particular neural network parameterization, we do not believe that this is a major limitation. This belief is strengthened by the experiments of Section 4, which empirically demonstrate that conclusions drawn from our theory carry over to practical settings.

W2 + Q3: Relation to the analysis in [65] and technical challenges. While our work is inspired by the observations made in [65], our theory (and experiments) go far beyond them. In particular, all three main theorems in our work are by no means corollaries of the results in [65], as can be seen from the proof appendix (Appendix E). Specifically, as discussed in the proof sketch of Theorem 1 (Section 3.2), [65] only provides a bound on the gradient norm at a given point. This alone does not imply a lower bound on optimization time. It is necessary to show that the gradient norm will not increase during training by proving that reward variance cannot increase rapidly along the trajectory of optimization. Moreover, Theorems 2 and 3 require proving separations in terms of optimization time between reward models. That is, aside from the lower bound on optimization time, we established upper bounds. This required proving a sufficient condition (Proposition 4 in Appendix E) for efficient optimization via a fine-grained analysis of the optimization trajectory when maximizing the non-concave RLHF objective. Lastly, aside from the technical novelty of our theoretical results, the focus of our work is different from that of [65], which did not consider the question of what makes a good reward model for RLHF or the interplay between accuracy and reward variance.

Regarding challenges in extending Theorems 2 and 3 to more general neural network architectures (considered in Theorem 1). Analyzing optimization for practical neural networks is still a major open problem even in simpler settings where the objective with respect to the neural network outputs is convex (e.g., supervised learning with the squared or cross entropy losses). While we were able to establish a lower bound, with the architecture contributing a dependence on its Jacobian norm, the challenges in establishing upper bounds on convergence rates are exacerbated due to the non-concavity of the RLHF objective with respect to the neural network outputs. Specifically, the architecture introduces complex interaction between the evolution of logits for different outputs, and as mentioned earlier, the understanding of these interactions is limited even in simpler settings. Hopefully, future technical advances may allow generalizing our theoretical results, as well as other optimization analyses of neural networks in the literature, to additional architectures.

Q1: Dependence on Jacobian norm in Theorem 1. Good question! Our interest lies in comparing, for a given language model, how a reward model affects RLHF optimization. Initially, the Jacobian of the language model does not depend on the reward model (rather just on the language model architecture and parameters) and we empirically found that its norm does not change much during policy gradient (see details below). Thus, for simplicity, we omitted the dependence on the Jacobian norm from the abridged Theorem 1, but of course included it in the full theorem statement (Theorem 4). The goal in doing so was to emphasize the parts which we found to be the most crucial and relevant to our study. Please let us know if you believe that it would be clearer to include this term in the abridged version as well and we will gladly accommodate.

Details: In a setting analogous to Section 4.1, with a Pythia-410m language model, we ran policy gradient using two different reward models (one trained on 100% on-policy data and another on 0% on-policy data). For efficiency, instead of the spectral norm, we tracked the average Frobenius norm of randomly selected rows in the Jacobian throughout training. In both runs the average Frobenius norm of rows hovered around 200 – 300, i.e., it was on the same order of magnitude throughout policy gradient and for both reward models. Note that since the Frobenius norm of the Jacobian upper bounds its spectral norm, this indicates that the spectral norm does not grow by much as well.

Q2: Applicability of our results to other RL settings. This is a great point! Thank you for raising it. While the positioning of our work focuses on RLHF, the theoretical results can be straightforwardly extended to support any RL setting with softmax parameterized policies (i.e., policies that produce a distribution over outputs via softmax). In accordance with your comment, we believe it can be a valuable direction to examine whether low reward variance empirically poses a problem in other RL settings and design potential-based reward modifications to increase reward variance. We will clarify this point in the manuscript.

Note 1: As discussed in Appendix A, we are not aware of existing reward shaping techniques that aim to increase reward variance or consider the effect of reward shaping on policy gradient optimization.

Note 2: Indeed, in bandit settings reward variance is equivalent to the expected squared advantage function, where the expectation is with respect to outputs sampled from the policy.

Q4: Effect of inexact gradients. The limitations paragraph in Section 3.1 and the first paragraph of the future work section (Appendix B) briefly discuss this point. In light of your question, we will expand the discussion on the effect of inexact gradients in the future work section along the lines of the answer below.

Our experiments consider standard policy gradient methods that compute gradient updates based on batches of sampled outputs, because computing the exact expected gradient is computationally infeasible. Thus, they empirically demonstrate that the implications of our theory carry over to settings with inexact gradients. Intuitively, the reason that low reward variance remains problematic in the inexact gradient setting is as follows. When the reward variance is low, the expected gradient vanishes. In this case, batch gradients may not necessarily be small (whether they will be small or not can depend on which policy gradient method you use). However, since the expected gradient is small, this implies that the batch gradient is often noisy, and so updates do not lead to substantial progress. Formally extending our analysis to the inexact gradient setting is highly non-trivial and we believe that it could be an interesting challenge for future work, which may further shed light on the relative advantages and disadvantages of different policy gradient methods.

W3: Experimental results on additional benchmarks. Our experiments are mostly based on the UltraFeedback dataset, which is a standard dataset for language model alignment. Since running the RLHF pipeline with large models is computationally expensive, we preferred to cover multiple language models, reward models, and ground truth rewards instead of additional datasets. The current experiments (in Section 4 and Appendix F) clearly demonstrate that the conclusions of our theory apply in practical RLHF settings. Thus, while we agree that additional datasets can give further evidence, we do not believe they are necessary to strengthen the claims made in the paper. Nonetheless, if there is any particular benchmark that you find is especially important, we would be happy to accommodate.

Thank you again for the fruitful and insightful discussion. In the revised manuscript, we will modify the notation in Theorem 4 and incorporate a discussion of the dependence on the Jacobian norm, following the line of our responses above.

Thank you for the thoughtful feedback and for highlighting its clarity, the theoretical results, and their comprehensive empirical validation. We treat your comments and questions below. If our response addresses your concerns, we would greatly appreciate it if you would consider raising your score. Please let us know if you have any further questions; we will gladly respond.

W1 and Q1: Relation of low reward variance to sparse rewards. When rewards are sparse (i.e., only a few outputs/trajectories receive a non-zero reward), often the reward variance is low. However, the reward variance can be low even if rewards are not sparse (i.e., dense). Importantly, in RLHF rewards are typically not sparse – reward models provide a reward for each output, which is usually non-zero. As our work highlights, despite rewards not being sparse, policy gradient still suffers from a flat objective landscape that leads to slow optimization when reward variance is low. This occurs when the reward model does not sufficiently distinguish between outputs that are probable under the policy (i.e., it gives them a similar, not necessarily near-zero, reward) and stems from the objective landscape induced by policies that produce a distribution over outputs via the softmax function.

As discussed in the related work section (Appendix A), the literature on reward shaping can indeed be seen as relevant to our work. Though, to the best of our knowledge, prior work did not characterize how reward shaping affects policy gradient optimization. In the related work section we further discuss the fact that, in theory, the poor optimization landscape of policy gradient can be solved via natural policy gradient (a preconditioning method). However, this approach is currently impractical in the context of large language models. Specifically, natural policy gradient requires inverting a $P \times P$ matrix, where $P$ is the number of language model parameters.

Lastly, while our work does not give a solution to the issue of low reward variance, it highlights limitations of existing reward model benchmarks and provides insight into what determines the effectiveness of a reward model, beyond accuracy. We hope that it can inspire future work on developing reward model training procedures that encourage a higher reward variance and policy gradient methods that can optimize more efficiently under low reward variance. Due to the extent of our theoretical and empirical results, we believe that this important and highly non-trivial endeavour falls outside the scope of the current paper.

W2 and Q2: Reward scale. You raise an excellent point! It is true that reward scale can impact performance. In particular, it can have an effect similar to modifying the learning rate (though not always exactly the same effect when applying KL regularization and using adaptive optimizers). For this reason, to ensure fair comparison across reward models, in our experiments we normalized all reward models so that they produce rewards of roughly the same scale. This normalization is also necessary for ensuring that reward variance is a meaningful measure when comparing different reward models.

Nonetheless, per your request we ran experiments using the same reward model, in which we scaled the reward by different constants. As expected, lower reward scales led to slower optimization while the larger reward scales that we tried led to unstable optimization that resulted in a lower ground truth reward.

Details: To obtain results in time for the response period, we considered a setting similar to Section 4.1 with Pythia-1B models and ran a single policy gradient run per scale factor (as shown in the paper, results were typically stable across random seeds). In all runs, we used the RM trained on 75% on-policy data.

As noted in the implementation details appendix, we observed similar results when tweaking the learning rate. Trends across reward models of the same learning rate remained intact.

W3: Definition of reward variance. Similarly to accuracy and many other evaluation measures, reward variance is defined for a given prompt/input. Such measures are typically aggregated over a set of inputs by taking the mean. In that regard, there is nothing special about reward variance and its definition. Indeed, the reward variance values we report in Table 1, as specified in its caption, are an average over the training prompts. In other words, the definition of reward variance does not ignore the existence of multiple examples in a dataset or batch, both in theory and in practice.

Q3: Prompt order during training. We did not fix the order of prompts in the experiments of Section 4. Instead, we used three different seeds for random data orders and reported the standard deviation of results across the seeds. As our results show (e.g., the almost indiscernible error bars in Figure 2), the data order did not have a significant impact on the results. Nonetheless, as detailed below, during the response period we ran additional experiments to further verify that our results are robust to fixing the data order.

Details: To obtain results in time for the response period, we considered a setting similar to Section 4.1 with Pythia-1B models and ran a single policy gradient per reward model (as shown in the paper, results were typically stable across random seeds). We include a table analogous to Table 2, showing the correlation between different reward model properties and the increase in reward used for training and ground truth reward. As was the case in our original experiments, reward variance is highly correlated with reward (used for training) increase, while accuracy was not, with a combination of both being most indicative of ground truth reward increase.

Note: We omitted off-policy accuracy for brevity. It had lower correlations than on-policy accuracy.

Thank you for the thoughtful feedback and for highlighting the importance of our results. We treat your comments and questions below. In particular, we emphasize that to a large extent the concerns raised in the review are already accounted for in the paper and provide additional experimental results, as requested in the review. If our response addresses your concerns, we would greatly appreciate it if you would consider raising your score. Please let us know if you have any further questions; we will gladly respond.

W1 and Q1: Experiments isolating confounding factors. We first want to emphasize that the paper already includes experiments that isolate the effects of reward variance from potential confounding factors (e.g., data distribution over which the reward model was trained). Specifically, in the experiments of Section 4.1 and Appendix F.1, the perfectly accurate reward model with low reward variance (e.g., the red reward model in Figure 2) is obtained from the ground truth reward by reducing in a controlled manner how well it separates different outputs while preserving their rankings. Thus, the ground truth reward (in yellow) and the perfectly accurate reward model with low reward variance (in red) were trained on the same data, have the exact same (perfect) accuracy, yet differ in their reward variance due to the controlled intervention. As our results show (e.g., Figure 2), the perfectly accurate reward model with low reward variance substantially underperforms the ground truth reward, in the sense that it leads to substantially slower optimization.

W1 and Q2: Further experiments comparing reward models trained on the same data. This is a good point! Thank you for raising it. Studying how reward variance and accuracy affect the outcome of RLHF requires reward models that vary in these properties. Our approach of training reward models on different data mixtures is one way of achieving this. Indeed, it is also possible to train reward models over the same data and vary instead different training hyperparameters.

We first note that, as discussed in the response to the previous comment, Section 4.1 and Appendix F.1 already include experiments that demonstrate the effect of reward variance for reward models trained on the same data distribution. To further address your concern, as requested in the review, during the response period we conducted additional experiments with reward models trained on the same data distribution (details and results are provided below). The results are analogous to the ones reported in Section 4.1, demonstrating that the connection between reward variance on optimization persists when reward models are trained on the same data.

Details: We considered a setting similar to Section 4.1 and trained, based on the Pythia-1B language model, 10 reward models over the same data (100% on-policy) using different training hyperparameters (e.g., learning rate and reward centering regularization coefficient). To obtain results in time for the response period, using each reward model we carried out a single policy gradient run as opposed to three (note that, as shown in the paper, results were typically stable across random seeds). Analogously to Section 4.1, we found that more accurate reward models are not necessarily better – e.g., a reward model with 0.67 accuracy but low reward variance led to an increase of 0.07 in ground truth reward while another reward model with 0.65 accuracy and higher reward variance led to an increase of 0.12 in ground truth reward. Furthermore, the table below presents the correlation between different reward model properties and increase in reward used for training and ground truth reward. As was the case in our original experiments, reward variance is highly correlated with increase in reward (used for training), while accuracy was not, with a combination of both being most indicative of increase in ground truth reward.

Note 1: We omitted off-policy accuracy for brevity. It had lower correlations than on-policy accuracy.

Note 2: The reward models were trained using learning rates 1e-7, 5e-7, and 1e-6, reward centering regularization coefficients 0, 0.1, 10, and scaling factors applied to the rewards of a fraction of the prompts (after scaling we still normalized reward models such that on average over prompts they have roughly the same scale).

W2 + Q3: Theory-practice validation gap. We respectfully disagree that the “paper fails to provide convincing positive validation of its core theoretical claims”. Each of our theoretical claims is backed by extensive empirical evidence in Sections 4.1 and 4.2 and Appendix F. Specifically, our theory establishes that: (i) low reward variance can lead to slow optimization, (ii) more accurate reward models are not necessarily better teachers, and (iii) for different language models different reward models can work better. Our experiments clearly demonstrate that these conclusions apply in practical RLHF settings.

Furthermore, the discussion in the future work section on why naively scaling the reward by a constant $c > 1$ is problematic does not contradict our theoretical analysis. Namely, by no means does our analysis suggest that naively scaling the reward by a large factor will improve performance. As explained in the future work section, when the reward variance is low the expected reward gradient vanishes. Thus, sample-based estimates of the gradient may be dominated by noise, which will be amplified by such reward scaling. We note that for reward variance to be a meaningful measure that allows comparing reward models, one must normalize reward models to have roughly the same scale, which is standard in practice and is already accounted for in our experiments and theoretical analysis. In light of your comment, we will clarify this point in the manuscript.

Regarding practical guidelines. The main purpose of our work is to advance the understanding of what makes a good reward model. While it does not give a solution to the issue of low reward variance, it highlights limitations of existing reward model benchmarks and provides insight into what determines the effectiveness of a reward model, beyond accuracy. We hope that our work inspires further research into reward model training procedures that encourage higher reward variance, as well as policy gradient methods that can optimize more efficiently under low reward variance. The future work section outlines a couple of possible directions going forward. However, due to the extent of our theoretical and empirical results, we believe that this important and highly non-trivial endeavour falls outside the scope of the current paper.

W3 + Q4: Ground truth validation. As discussed in Section 4, since evaluating policies via human labelers can be prohibitively costly, it has become common practice to use an existing large reward model as the ground truth reward (e.g., see references [27, 13, 74, 7, 10, 81, 73] in the paper). This is especially true for studies, such as ours, whose main goal is to further the understanding of different aspects of the RLHF pipeline, as opposed to training state-of-the-art models. The ArmoRM reward model in particular has already been used as the ground truth reward in numerous papers, and was validated in [77] to accurately accord with existing labeled preferences (better than GPT4 judges).

Moreover, using an existing reward model as the ground truth reward actually allows avoiding any potential “cascading errors” and inconsistencies due to noise in human preferences, since the exact same ground truth reward is applied for both labeling preferences and evaluating policies. Thus, in our experiments, the ground truth reward already plays the role of a verifiable reward, as you suggested we use in your review. Furthermore, note that we also verified our results are robust to the choice of ground truth reward — Appendix F provides experiments showing analogous results with a different ground truth reward.

Thank you for the swift response! We greatly appreciate your engagement and address your additional questions below. We believe that this should treat any remaining doubts.

W1 and Q1. As detailed in Section 4.1 and noted in your review, our experiments indeed contained reward models trained on different data mixtures. This is one possible approach to obtain reward models that differ in terms of their accuracy and reward variance. However, as specified in our original response above, the paper already contains experiments accounting for this potential confounding factor.

Furthermore, the previous response provided additional experiments, conducted per your request, with reward models that are trained on the exact same data, but with different training hyperparameters. The results demonstrate the validity of our conclusions regarding the interplay between accuracy, reward variance, and optimization with reward models that are trained on the same data.

W1 and Q2. We gladly provide in part 2 of the response the full implementation details of our additional experiments, including the hyperparameters used for training the reward models, the accuracy and reward variance of these models, and reward increase after policy gradient. These details were omitted from our first response due to space limitations.

As for reward hacking, we believe that there may be a misunderstanding with regards to what the ground truth reward stands for in our experiments. The ground truth reward stands for the true “downstream” performance. As discussed in the paper and our previous response, the approach of using an existing reward model as the ground truth reward is common in the literature. Since the ground truth reward is used for both labeling preferences and measuring ground truth performance, this methodology avoids any potential inconsistencies or reproducibility issues due to noise in human preferences, while circumventing the prohibitive cost of using human labelers. Note that reward hacking occurs when the reward used for training increases, but the ground truth reward decreases. This is why we report both rewards — it allows examining the effect on optimization speed, as manifested in the reward used for training increase, as well as how that translates to ground truth improvements.

W2 + Q3. Thank you for clarifying your comment. Our theory establishes that low reward variance leads to slow optimization. Yet, as discussed in Appendix B, this does not imply that any method for increasing reward variance necessarily leads to faster optimization. Specifically, our experiments demonstrate, using reward models that we trained as well as publicly available reward models, that reward variance is indicative of the rate at which the reward increases. However, artificially inflating reward variance by scaling rewards by a large factor $c > 1$ can lead to instabilities. We therefore view developing methods for increasing reward variance in a meaningful way (e.g., higher reward variance at a given reward scale) as a valuable direction for future work.

Note that scaling rewards by a factor $c > 1$ has an effect similar (or identical, depending on the particular policy gradient implementation) to scaling up the learning rate. Indeed, simply scaling them up does not give a “free lunch” of faster optimization. As mentioned in Appendix G.1, we experimented with different learning rates and found that analogous results are obtained when using reasonable values, though too large values led to instability.

Lastly, the explanation based on vanishing gradient does fall under our theoretical framework. In particular, the proof of Theorem 1 builds on the fact that the expected gradient and higher-order derivatives vanish when the reward variance is low. The arguments on why in this case scaling up rewards can lead to instabilities in practice due to noise in inexact batch gradients are made on an intuitive level. Formally extending our analysis to the inexact gradient setting is highly non-trivial and we believe that it could be an interesting challenge for future work, which may further shed light on the relative advantages and disadvantages of different policy gradient methods.

Conclusion. We would like to thank you again for the engagement and for taking the time to read our responses. Please let us know if you have any further questions.

Thank you for the support and for highlighting the significance and clarity of our work!

The future work section, which was deferred to Appendix B due to lack of space, discusses why simply scaling up the reward is insufficient and lays out potential avenues for increasing reward variance in a meaningful way. We plan to expand this discussion and incorporate it into the body of the paper given the additional content page allowed for accepted papers. Thank you for the suggestion!

If any questions arise during the discussion period, please let us know and we will gladly respond.

Thank you for the thoughtful feedback and for highlighting the clarity, significance, and originality of our work! Your review focuses on an important, yet nuanced, point regarding the relation between reward scale, “effective learning rate”, and reward variance. As detailed below, the effect of reward variance on optimization is not solely due to differences in reward scale/effective learning rate. In fact, for reward variance to be a meaningful measure that allows comparing reward models, rewards need to be normalized to have roughly the same scale. The paper already accounts for this and we further provide below additional empirical evidence to address your comments. Nonetheless, we recognize that the discussion of this point in the paper can be improved (currently, we mention it in the experimental section and future work section). In light of your review, we will incorporate a dedicated discussion of it after Definition 2. We hope that our response fully addresses your concern and would greatly appreciate it if you would consider raising your score. Please let us know if you have any further questions; we will gladly respond.

Effect of reward variance on optimization is not just due to scale

First, we would like to emphasize that reward variance captures more than just scale: two reward models can have roughly the same scale but different reward variances (note that reward variance can be low even when the reward scale is not small). For reward models of roughly the same scale, reward variance measures the degree of separation between outputs. A useful analogy is margin in linear prediction. Just as the linear predictor needs to be of normalized scale for margin to be a meaningful measure, the scale of rewards needs to be normalized for reward variance to faithfully measure degree of separation between outputs.

Indeed, as mentioned in the review, scaling up the rewards affects the effective learning rate and reward variance. Accordingly, to avoid differences in effective learning rate and allow comparing reward models based on reward variance, we followed common practice and normalized rewards to have roughly the same scale. In general, there does not exist a scale $\alpha$ such that for two reward models RM and RM’ it holds that RM’ = $\alpha$ * RM. So usually rewards are normalized such that if two reward models did satisfy RM’ = $\alpha$ * RM, then after normalization they would be exactly equal.

• As detailed in Appendix G, we normalized reward models based on outputs sampled from the initial policy over prompts in the training set. This ensured that the reward scale, i.e., average of absolute rewards, was roughly the same for all reward models. Interestingly, in the experiments of Figure 2, after normalization the reward models with lower reward variance (e.g., the perfectly accurate RM) had a slightly higher scale than reward models with a higher reward variance (e.g., RM 100% on-policy). Yet, the former still led to slower optimization. Thus, the results of Figure 2 already demonstrate that the faster optimization of reward models with a higher reward variance is not due to a larger reward scale.

  Details: Reward scale, i.e., average absolute reward over training prompts and outputs from initial policy, for the reward models in Figure 2 (after normalization).

  RM 0%	RM 25%	RM 50%	RM 75%	RM 100%	Perfect RM w/ low variance	GT reward
  0.8015	0.7783	0.7834	0.8042	0.7917	0.8411	0.8254

• The normalization above is computed based on the initial policy. It is also possible to normalize rewards online. We ran additional experiments, where in each batch rewards are normalized online by the standard deviation across the batch of prompts and outputs. As shown below, the results are analogous to ones reported in the paper. This further shows that differences in reward scale do not explain the relation between reward variance and optimization in our results.

  Details: To obtain results in time for the response period, we considered a setting similar to Section 4.1 with Pythia-1B models and ran a single policy gradient run per reward model as opposed to three (though, as shown in the paper, results were typically stable across random seeds). As was the case when normalizing rewards based on initial policy statistics, the perfectly accurate reward model led to a slower increase in ground truth reward than the reward models trained on 100%, 75%, and 50% on-policy data. Furthermore, reward variance was highly correlated with reward (used for training) increase, while accuracy was not, with a combination of both being most indicative of ground truth reward increase.

  	Reward Increase Pearson	Reward Increase Spearman	GT Reward Increase Pearson	GT Reward Increase Spearman
  Reward Variance	0.937	0.828	0.292	0.771
  On-Policy Acc.	-0.529	-0.085	0.416	0.657
  Reward Variance & On-Policy Acc.	0.781	0.714	0.880	0.942

  Note: We omitted off-policy accuracy for brevity. It had lower correlations than on-policy accuracy.

Challenges in making online reward variance normalization work

The above showed that the relation between reward variance (after normalization) and optimization cannot be explained by reward scale. Instead of normalizing rewards to have the same scale, an alternative option is to normalize rewards online so that each prompt in a batch has an empirical reward variance of 1. We believe that this is the experiment described in the review.

First, it is important to note that this approach is not captured merely by reward scale or an effective learning rate since it is an online per-prompt adaptive normalization scheme. Second, we experimented with such approaches during our work on the project. However, we found that such normalization can be problematic due to the difference between the empirical reward variance, which is estimated online over a few outputs, and the true reward variance (results given below). Our results suggest that the reason is similar to why scaling rewards by a large factor $c > 1$ does not help, which we discussed in the future work section (Appendix B). Namely, the expected gradient vanishes when the true reward variance is low. Thus, gradient estimates based on a few online samples may be dominated by noise, which will be amplified by scaling. We believe that a promising direction for future work can be to further investigate whether similar normalization techniques, along with substantially increasing the number of outputs sampled online per prompt, can lead to better performance within a given compute budget.

Details: In the setting of Section 4.1, we normalized online the rewards of each prompt separately to ensure its empirical reward variance is one. We ran experiments using two reward models – the ground truth reward and the 75% on-policy reward model – and different numbers of outputs sampled per prompt (2, 8, 16). The table below reports the increase in ground truth reward for each experiment. Notably, with 2 outputs normalization harms performance. Increasing the number of outputs sampled for each prompt led to higher ground truth reward. Yet, even with 16 outputs the imperfect reward model still worked better than using the (perfectly accurate) ground truth reward directly.

Theory

In our analysis, we account for reward models needing to be of a similar scale by assuming that rewards are bounded within some interval – a standard assumption in the reinforcement learning literature. This precludes circumventing the lower bound in Theorem 1 by artificially multiplying all rewards by a large factor.

Less relevant questions

1. As reported in Table 2, the Spearman correlation is 1 for the increase in reward used for training. For the ground truth reward it is 0.714.

2. Thank you for the suggestion! This can be useful intuition alongside the connection to the objective landscape. We note, however, that it does not fully capture the reason why reward variance affects optimization. The connection between reward variance and optimization depends on the fact that language models produce distributions through a softmax function (as mentioned in the proof sketch of Theorem 1). In theory, for other, non-practical, parameterizations (e.g., each token probability is its own trainable parameter), the objective landscape is not flat when the reward variance is low. Though, we are not aware of any practical replacement for softmax that avoids this issue, so exploring whether one can construct such replacements can be a valuable direction for future work.

Thank you very much for the extremely quick and thorough follow-up. I appreciate the time you took to explain your work to me, and hope that it helps you present it to other people in the future, though you already did a great job.

I am very confused about the results you present in the sense that they don't fit any model I can think of.

Perhaps somehow, the initial model variance also indicates how easy it is to increase variance by moving slightly off-policy, in a way in which re-scaling does not affect? Seems possible but still a bit unlikely to me. It would be interesting to see how reward variance evolves over training to see if that explains the velocity almost entirely.

Still, the experiment above is firmly in the square of Future Work and I think the authors have shown that the reward variance matters a lot, independent of just a scaling factor, to my satisfaction.

However, when considering the correlation between reward variance as measured with our original normalization and reward increase (after training using reward models normalized via the new normalization scheme), the results are analogous to our original experiments

Yeah, this is very strange, as detailed above. Nonetheless I believe you and just disbelieve my previous hypothesiszx`.

To give a concrete example, consider the non-practical parameterization mentioned in the previous response, in which each output probability is its own trainable parameter

Thank you for working through this example. I understand why I was wrong now: in this parameterization the gradient is nonzero whenever there are any reward differences anywhere, but the softmax parameterization eats this difference.

Regarding the final experiment, it's interesting that the correlations are lower; but they're still firmly positive and so aren't really evidence against the main theorem being meaningful.

I'm recommending this paper for acceptance. You probably want to put the experiments you did here (the one that controls for variance at initialization and still finds a correlation between variance before controlling, and final increase in GT reward) somewhere in the paper or appendix.