Paper: Vanishing Gradients in Reinforcement Finetuning of Language Models (ICLR-2024)

Reviewer comment (Writing):
The writing has room for improvement. The authors tend to make heavy use of long sentences that are difficult to parse.

Author response:
We will make sure to further polish writing by rephrasing such sentences. We would greatly appreciate specific pointers to sentences you consider difficult to parse.

Reviewer comment (Writing):
The second paragraph on page 4, I am not sure how the sentence "the contribution of inputs... is still negligible..." is logically connected to the next sentence.

Author response:
The sentence you refer to explains why Adam is not expected to mitigate the vanishing gradients problem. The following sentence states that the experiments of Section 4 affirm this expectation. We have rephrased the second paragraph on page 4 to clarify this point (modified text is marked in red). Thank you for the helpful comment!

Reviewer comment (Experiments):
The experiment setup in Section 4.1 is worth closer scrutiny. It does not seem a fair comparison to me that RLFT optimizes for a task-specific metric while SFT is performed based on ground-truth completions, which means that RLFT *does not* have access to the ground-truth, whereas SFT *does*.

Author response:
We emphasize that the goal of Section 4.1 is not to compare the performance of RFT with SFT, in order to advocate for the use of one over the other, rather to demonstrate the prevalence and detrimental effects of vanishing gradients in RFT. In particular, we contrast RFT with SFT to showcase how RFT struggles to increase the reward of inputs with small reward standard deviation, compared to its ability to increase the reward of other inputs. For example, in Figure 1 and Table 1, SFT poses a baseline that allows quantifying whether the correlation between reward standard deviation and reward change during RFT is high or low. The comparison to SFT verifies that it is relatively high, as expected.

We hope that this addresses your concern. Please let us know otherwise and we will gladly elaborate.

Reviewer comment (Experiments):
It is worth noting that running RLFT with "task-specific" metrics, such as ROUGE, is no longer considered a good practice in the RLHF community. Specifically for the GRUE benchmarks, I believe that the community has reached a consensus that it is not a suitable playground for RLFT. I suppose that the authors would agree that the predominant approach in RLFT is to train a reward model using a preference (or ranking) dataset, and apply policy-gradient algorithms based on the reward model. I could imagine that the same issue identified in this manuscript is also present in such scenarios, but the authors did not provide sufficient evidence. (I wouldn't place too much weight on this weakness point, because, understandably, the authors might be short on compute resources.)

Author response:
The nature of our work is foundational, aiming to improve the understanding of existing methods. Accordingly, to allow for efficient experimentation, we mostly focused on the GRUE benchmark, which does not include reward functions learned from human feedback. We agree with the reviewer that a predominant approach in practice is to use reward functions learned from human preference data. Thus, as mentioned in the conclusion section, we consider studying the extent to which our empirical findings carry over to more complex finetuning pipelines, including human feedback, a valuable direction for future research.

Reviewer comment (Experiments):
Suppose that we are given a preference dataset: each prompt comes with two responses, one labelled as preferred. How should we run SFT on such dataset? Shall we only use the one preferred response for each prompt, or shall we also somehow leverage the other response as a negative sample? This is subject to debate in the LLM community, and I hope the authors can shed more light on this point.

Author response:
Our work focuses on a fundamental vanishing gradients problem in RFT. Among other contributions, we show that given access to a small number of ground truth labels (i.e. completions), an initial SFT phase can be effective in mitigating the problem. However, we do not treat the orthogonal question of how to best obtain ground truth labels for SFT. We believe investigating methods that use preference data to create SFT labels is an important subject, yet falls outside the scope of our findings.

Reviewer comment (Theory):
Can you discuss in more detail when small vanishing gradient becomes a problem? Assuming that an LLM is trained with RLHF to convergence, we would expect many of the completions to be of equal quality, i.e., similar rewards. Why would low standard deviation be a problem in this case? Similarly, at convergence, we expect gradient to be close to zero which also is not detrimental to training.

Author response:
Theorem 1 establishes that, if the reward standard deviation for an input is small, then its expected gradient vanishes even if the expected reward is far from optimal. While small reward standard deviation and vanishing gradients are expected when the model is near-optimal, they are problematic when the reward is suboptimal – they can hinder the ability of RFT to maximize the reward (as discussed both in the introduction and Section 3). Indeed, the experiments on the GRUE benchmark (Section 4.1) and controlled environments (Section 4.2) demonstrate that, in the presence of inputs with small reward standard deviation and low reward, RFT is unable to effectively maximize the reward.

As to when vanishing gradients, due to small reward standard deviation, occur in RFT. We expect inputs with small reward standard deviation (and low reward mean) to be common when the text distribution of the downstream task markedly differs from that of the pretraining corpora. In such a case, over some inputs, the pretrained model is likely to give non-negligible probability only to outputs of roughly the same suboptimal reward. The reward standard deviation for these inputs is small, implying that their expected gradients are near-zero. We note that the above was also discussed in the introduction and Section 3. Section 4.1 corroborates this expectation using the GRUE benchmark.

We hope this fully addresses your question. If not please let us know and we will happily expand!

Reviewer comment (Experiments):
While not guaranteed (also not always in practice), ADAM update equations suggest O(1) updates to the weights. For a batch size of 1, would this solve the vanishing gradient problem? What is the maximal batch size where this problem is negligible? Can you also plot some of the statistics from ADAM training, like $\mu$, $\nu$, and $u$ (update) norms? What would be the impact of $\epsilon$ when the significant portion of a batch has vanishing gradients, i.e., would $\epsilon$ dominate second order moments?

Author response:
Using a batch size of 1 is unlikely to effectively solve the problem since:

- The update in Adam is based on moving averages of gradients and squared gradients, which are computed across batches (when the batch size is 1, they are computed across individual samples). Hence, similar to the batch size > 1 case, the effect of inputs with vanishing expected gradients on these averages, and thus on the update, will be negligible compared to those of other inputs.

- Using a batch size of 1 is impractical since it does not allow efficient parallelism, and may lead to instability.

To affirm that a batch size of 1 does not mitigate the vanishing gradients problem in RFT, we ran experiments over the NarrativeQA dataset. Indeed, the reward achieved does not improve upon the default batch size used in our experiments.

**Results:**
- Dataset: NarrativeQA, Model: T5-base, Batch size: 64 (default value in our experiments) – train reward 0.107, test reward 0.114
- Dataset: NarrativeQA, Model: T5-base, Batch size: 1 – train reward 0.099, test reward 0.101

Furthermore, we include in the following (anonymized) link a figure, analogous to Figure 1. It compares the reward means and standard deviations of training samples under the pretrained, RFT, and RFT with batch size 1 models. The figure demonstrates that using a batch size of 1 does not improve the ability of RFT to maximize the reward in the presence of inputs with small reward standard deviation.

Link: https://ibb.co/F3JdLP9

Technical note: To obtain results in time for the discussion period, we reduced the overall number of optimization steps to 25% of the original number of steps. This is done for both the run with a batch size of 1 and the run with the default batch size of 64. Since reducing the batch size led to lower train and test rewards, it is likely that increasing the number of iterations will not result in a major improvement (if at all) over the default batch size.

> Can you also plot some of the statistics from ADAM training, like $\mu$, $\nu$, and $u$ (update) norms? 

Of course. In the additional experiments described above, we tracked the norms for the batch gradient (denoted $||\nabla||$), Adam update (denoted $||u||$), exponentially moving gradient average (denoted $||\mu||$), and the square root (taken entry-wise) of the exponentially moving squared gradients average (denoted $||\sqrt{\nu}||$). We include below (anonymized) links to plots reporting these quantities against the optimization step.

**Results:**
- Dataset: NarrativeQA, Model: T5-base, Batch size: 64 (default value in our experiments) – https://ibb.co/k3xRkmg
- Dataset: NarrativeQA, Model: T5-base, Batch size: 1 – https://ibb.co/nCbBLNF

As expected, the plots show that: (1) the Adam update and moving averages are not near-zero for both batch size 1 and 64; and (2) the gradient of individual samples, as can be observed in the leftmost plot for batch size 1, is zero for quite a few of the samples. Though, since the moving averages are not near-zero, the contribution of these samples (with near-zero gradients) to the update performed by Adam is negligible.

> What would be the impact of $\epsilon$ when the significant portion of a batch has vanishing gradients, i.e., would $\epsilon$ dominate second order moments?

If the second order moments average is practically zero (note that this requires multiple consecutive batches with a significant portion of vanishing gradients), then $\epsilon$ will indeed dominate the normalization term.

Reviewer comment (Theory):
To solve the suboptimal problem, it seems like designing a better reward function is more important. If there are many flat areas in the reward model, a different initialization might not necessarily help much.

Author response:
This is an excellent point. While modifying the initialization through an SFT phase is a promising solution, we certainly do not believe that it is the only viable one (nor do we claim so in the paper). We believe further investigating ways to overcome vanishing gradients in RFT, e.g., by designing better reward functions, is an important direction for future work.

Reviewer comment (Evaluation):
Is there a potential explanation of why the reward model generates small standard deviations on the GRUE benchmark?

Author response:
Certainly. As discussed in Sections 1 and 3, we expect inputs with small reward standard deviation (and low reward mean) to be common when the text distribution of the downstream task markedly differs from that of the pretraining corpora. In such a case, the pretrained model is likely to give non-negligible probability only to outputs of roughly the same suboptimal reward.

Section 4.1 corroborates this expectation – see text under the “Inputs with small reward standard deviation are prevalent” header. Representative examples for inputs with small and inputs with large reward standard deviation under the pretrained model are given in Appendix E.

Reviewer comment (Theory):
Could you list all the assumptions used in your proofs explicitly and discuss their realisticity?

Author response:
The gradient norm bound in Theorem 1, and the analogous results for PPO variants (Propositions 1 and 2), are general and entail no assumptions, beyond requiring that the model $f$ is differentiable with respect to the parameters $\theta$ and that the softmax operator is used to produce next-token distributions (as stated in the preliminaries section). Both requirements are upheld by standard language models.