Vanishing Gradients in Reinforcement Finetuning of Language Models
Authors: Noam Razin
Hattie Zhou
Omid Saremi
Vimal Thilak
Arwen Bradley
Preetum Nakkiran
Joshua M. Susskind
Etai Littwin
Venue: ICLR 2024 poster (2024)
Keywords: Vanishing Gradients
Reinforcement Finetuning
Supervised Finetuning
Language Models
Primary area: general machine learning (i.e., none of the above)
TL;DR: We uncover a fundamental vanishing gradients problem in reinforcement finetuning of language models, demonstrate its prevalence and detrimental effects, and explore possible solutions.

Abstract: Pretrained language models are commonly aligned with human preferences and downstream tasks via reinforcement finetuning (RFT), which refers to maximizing a (possibly learned) reward function using policy gradient algorithms. This work identifies a fundamental optimization obstacle in RFT: we prove that the expected gradient for an input vanishes when its reward standard deviation under the model is small, even if the expected reward is far from optimal. Through experiments on an RFT benchmark and controlled environments, as well as a theoretical analysis, we then demonstrate that vanishing gradients due to small reward standard deviation are prevalent and detrimental, leading to extremely slow reward maximization. Lastly, we explore ways to overcome vanishing gradients in RFT. We find the common practice of an initial supervised finetuning (SFT) phase to be the most promising candidate, which sheds light on its importance in an RFT pipeline. Moreover, we show that a relatively small number of SFT optimization steps on as few as 1% of the input samples can suffice, indicating that the initial SFT phase need not be expensive in terms of compute and data labeling efforts. Overall, our results emphasize that being mindful for inputs whose expected gradient vanishes, as measured by the reward standard deviation, is crucial for successful execution of RFT.

Published as a conference paper at ICLR 2024
VANISHING GRADIENTS IN REINFORCEMENT FINETUNING OF LANGUAGE MODELS
Noam Razin"i, Hattie Zhou Omid Saremi', Vimal Thilak,, Preetum Nakkiran', Joshua Susskind", Etai Littwin
T Arwen Bradley',
i Apple Tel Aviv University $ Milla, Université de Montréal
ABSTRACT
Pretrained language models are commonly aligned with human preferences and downstream tasks via reinforcement finetuning (RFT), which refers to maximizing a (possibly learned) reward function using policy gradient algorithms. This work identifies fundamental optimization obstacle in RFT: we prove that the expected gradient for an input vanishes when its reward standard deviation under the model is small, even if the expected reward is far from optimal. Through experiments on an RFT benchmark and controlled environments, as well as a thcoretical anal- ysis, we then demonstrate that vanishing gradients due to small reward standard deviation are prevalent and detrimental, leading to extremely slow reward maxi- mization. Lastly, we explore ways to overcome vanishing gradients in RFT. We find the common practice of an initial supervised finetuning (SFT) phase to be the most promising candidate, which sheds light on its importance in an RFT pipeline. Moreover, we show that a relatively small number of SFT optimization steps on few as 1% of the input samples can suffice, indicating that the initial SFT phase need not be expensive in terms of compute and data labeling efforts. Overall, our results emphasize that being mindful for inputs whose expected gradient vanishes, as measured by the reward standard deviation, is crucial for successful execution of RFT.
INTRODUCTION
The dominant machine learning paradigm for textual data relies on pretraining a language model on vast corpora, and subsequently aligning it with human preferences and downstream tasks via finetuning (see, e.g., Yang et al. (2019); Raffel et al. (2020); Wei et al. (2022); Thoppilan et al. (2022); OpenAl (2023); Touvron et al. (2023); Zhao et al. (2023)). Supervised finetuning (SFT) provides straightforward procedure for adapting the model, which the cross entropy loss is minimized over pairs of textual inputs and desirable outputs. Despite the simplicity of SFT, the cost of obtaining high-quality labeled data, as well as the difficulty expressing human preferences through ground truth labels, has led to wide adoption of reinforcement learning based approach, herein referred to as reinforcement finetuning (RFT) (Wu and Hu, 2018; Ziegler ct al., 2019; Stiennon et al., 2020; Ouyang et al., 2022; Bai et al., 2022; OpenAI, 2023; Touvron et al., 2023).
In RFT, instead of minimizing a loss over labeled inputs, policy gradient algorithms are used for maximizing reward function. The reward function can be learned from human preferences, in which case RFT known as reinforcement learning from human feedback (RLHF), or tailored to a downstream task. Specifically, denote by V(x; 0) the expected reward that model parameterized by 0 achieves, given as input a sequence of tokens x. For a distribution over inputs, RFT maximizes Ex [V(V 0)] with respect to 0 by following estimates of its gradient. Proximal Policy Optimization (PPO) (Schulman et al., 2017)- - policy gradient algorithm widely used for RFT - adheres to this scheme, up replacing the expected reward with a surrogate objective.
In this work, we highlight a fundamental optimization obstacle in RFT: we prove that the expected gradient for an input x, i.e. (x; vanishes when the reward standard deviation of x under
Work done while interning at Apple.
Cooee for reproducing our experiments is available at
-rlgadd

---

Published as a conference paper at ICLR 2024
the model is small, even if the expected reward ((x; 0) is far from optimal (Section 3). The same holds for PPO, and stems from a combination of the reward maximization objective and the softmax operator used for producing distributions over tokens. In contrast, under SFT the expected gradient does not necessarily vanish in analogous circumstances.
Vanishing gradients in RFT can arise when attempting to reverse existing behavior of the model or when the text distribution of the downstream task differs from that of the pretraining corpora. In both of these common use cases, 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 such inputs is small, implying that their expected gradients are near-zero. Thus, RFT is anticipated to be ineffective for achieving a higher reward over them.
Using the GRUE benchmark (Ramamurthy et al., 2023) for RFT of language models, we empirically affirm the prevalence of inputs with small reward standard deviation, for which the expected gradient vanishes, and showcase its harmful effects (Section 4.1). Specifically, we find that three of seven GRUE datasets contain a considerable amount of inputs with near-zero reward standard deviation under pretrained models, while their expected reward is suboptimal. As anticipated, RFT has limited impact on the rewards of such inputs. Moreover, our experiments show that, in datasets where inputs with small reward standard deviation are more common, the reward that RFT achieves compared to SFT is worse. This indicates that vanishing gradients in RFT hinder the ability to maximize the reward function. Controlled experiments that remove possible confounding factors (Section 4.2), such as insufficient exploration, and a theoretical analysis for a simplified setting (Section 4.3), corroborate that vanishing gradients in RFT can lead to extremely slow optimization.
Lastly, we explore ways to overcome vanishing gradients RFT of language models (Section 5)- Conventional heuristics, including increasing the learning rate, using a temperature hyperparameter, and entropy regularization, are shown to be inadequate (Section 5.1). On the other hand, in agree- ment with prior evidence from Ramamurthy et al. (2023), we find the common practice of initial SFT phase (e.g., used in Ziegler et al. (2019); Stiennon et al. (2020); Ouyang et al. (2022); Dubois et al. (2023); Touvron et al. (2023)) to be beneficial, both in terms of the resulting reward and in sig- nificantly reducing the number of input samples with small reward standard deviation. This sheds light on the importance of SFT in an RFT pipeline- - it can help alleviate vanishing gradients.
If the benefits of an initial SFT phase stem, at least partially, from alleviating vanishing gradients for the subsequent RFT phase, our results suggest, one would expect that relatively few SFT steps on small number of labeled inputs can suffice. Meaning, we need only perform SFT to decrease the number of input samples with near-zero reward standard deviation, after which the potency of RFT will substantially improve. We show that this indeed the case (Section 5.2). Remarkably, for some datasets, SFT over as few as 1% of the input samples allows RFT to reach 96% of the reward achieved when SFT is performed over all input samples. This demonstrates that the initial SFT phase need not be expensive in terms of compute and data labeling efforts, highlighting its effectiveness in addressing vanishing gradients in RFT.
Overall, we identify fundamental optimization challenge in RFT, demonstrate its prevalence and detrimental effects, and explore possible solutions. Our results emphasize that being mindful for inputs whose expected gradient vanishes, as can be measured by the reward standard deviation, is crucial for successful execution of RFT. We believe further investigating ways overcome vanish- ing gradients in RFT, e.g., by modifying the objective (as recently done in Lu et al. (2022); Rafailov al. (2023); Hu al. (2023); Dong et al. (2023); Gulcehre et al. (2023)), is valuable direction for future research.
Related work: We discuss related work throughout and defer a concentrated account with additional references to Appendix A.
2 PRELIMINARIES: SUPERVISED AND REINFORCEMENT FINETUNING
Contemporary language models, such as GPT-4 (OpenAI, 2023) and Llama2 (Touvron et al., 2023), are trained on large corpora to produce probability distribution over sequences of tokens. The distribution is modeled in an autoregressive manner, often conditioned on an input sequence, where the output of the model is converted into distribution over tokens through the sofimax operator softmax(z) = exp(**)/ E exp(zer) for E K and k E (1,..., K).
2

---

Published as a conference paper at ICLR 2024
Formally, let X' be a finite vocabulary of tokens. We consider model f parameterized by 0 R that, conditioned on an input x = (Z1....,ILI... E X of length Lin, defines a distribution over outputs y = (y1,.... Vhowe) ) E Xoout of length Toout as follows:
(o(y|x)= II 1-1 po po(u x. = II softmax(f(x.yss-:1)))
where yst-1 == (y1...., and (x, 0) RR In particular, f can stand for a neural network that intakes a sequence of tokens and produces logits for the distribution of the next token.
Supervised finetuning (SFT): SFT is straightforward approach for adapting a pretrained language model to comply with human intent and downstream tasks. Given pairs of textual inputs and ouputs, e.g. instructions and desirable completions (Wei al., 2022: Chung et al., 2022; Wang al., 2022; Longpre et al., 2023), minimize the cross entropy loss through gradient-based methods. The SFT objective (OF is defined by:
CSFT(O) :- Ex-D (X;
CSFT(X; 0) :- Ey-DD()x) In [o)(xxx))
(I)
where a distribution over inputs, D(-|x) ) is ground truth conditional distribution over outputs, and LsFT(x: (x; 0) is the loss over single input x. Although the simplicity of SFT is appealing, it
comes with a few drawbacks. Namely: (i) human preferences can be difficult to express through ground truth labels, e.g., reducing the toxieity and bias of the model's outputs; and (ii) obtaining high-quality labeled data can be expensive.
Reinforcement finetuning (RFT): Given reward function r X x & + [-1, 1], instead of minimizing a loss over labeled inputs, the goal of RFT is to maximize the expected reward V(0) with respect to the model:
V(0) == EX~ [V< 0)
V(x; == Ey-po(-lxx) [r(x,
In reinforcement learning terminology, RFT forms a horizon-one (bandit) environment, where each input is state, each output an action that the model (i.e. policy) can take, and D a distribution over initial states. The reward function r can be learned from human preferences, e.g. based on rankings of the model's outputs, which are often easier obtain than ground truth labels (Ziegler et al., 2019; Stiennon et al., 2020; Ouyang et al., 2022). Alternatively, it may be task specific metric, such ROUGE (Lin, 2004), METEOR (Banerjee and Lavie, 2005), or a learned sentiment classifier (Ramamurthy et al., 2023). Policy gradient algorithms are used to maximize V(f), by updating the parameters 0 according to sample-based estimates of Vo V().. PPO (Schulman al., 2017) - policy gradient algorithm widely used for RFT adheres to this scheme, up to replacing (xx; 0) with surrogate (clipped) objective.
3 SMALL REWARD STANDARD DEVIATION IMPLIES VANISHING GRADIENTS IN REINFORCEMENT FINETUNING
We begin by introducing a, yet unnoticed, cause for vanishing gradients in RFT: the expected gradi- ent for an input is near-zero if its reward standard deviation is small.
Theorem 1. For parameters 0 E IR and input x E XLin, denote the reward standard deviation of under the model by yTDyyppol (-2g [rxx, y)] := / Exy-pot-px)ll (r(x,))))))))))))))))))))))))[ V(x; where V 0) is the expected reward defined in Equation (2). Then, it holds that:
(x; 0)1 < (x; -TD-STDDSS STDDg-no((s
where 7xx;0) 0) == maxie (1. [I...., booaayssa ex f(x 197112 with Js(xysI-199 standing for the Jacobian of (xx, 0) with respect to 6, and |1-11 and 1l.112 denote the Euclidean and operator norms, respectively.
The gradient VaV (x; 0) vanishes, when the reward standard deviation of x is small, due to a combi- nation of the reward maximization objective and the use of softmax for producing distributions over output tokens (see proof of Theorem in Appendix D.I). Notably, this occurs regardless of what the expected reward V V(x: 0) is and, in particular, whether is nearly optimal or not.
Vanishing gradients in RFT can be problematic when attempting to reverse an existing behavior of the model when the text distribution of the downstream task differs from that of the pretraining
3

---

Published as a conference paper at ICLR 2024
corpora. In both of these common use cases, over some inputs, the pretrained model is likely to give non-negligible probability only to outputs of roughly the same suboptimal reward. The reward stan- dard deviation for such inputs is small, implying that their expected gradients are near-zero. Thus, RFT is anticipated to be ineffective for achieving higher reward over them. We empirically affirm this expectation in Section 4 through experiments an RFT benchmark and controlled environ- ments.
We make two important distinctions. First, Theorem considers the expected gradient for an indi- vidual input x, as opposed to the gradient across a batch or the population of inputs. Thus, although adaptive optimizers such Adam (Kingma and Ba, 2015) the standard optimizer choice for RFT - normalize the batch gradient, the contribution of inputs whose reward standard deviation is small is still negligible compared to the contribution of other inputs. Indeed, the experiments of Section 4 verify that adaptive (batch) gradient normalization does not solve the harmful effects of vanishing gradicnts in RFT, since they are carried out using the Adam optimizer. Second, in contrast to RFT, under SFT the expected gradient VgLSFT(X) (x; 0) need not vanish unless LSFT (x: 0) is nearly minimal (and usually does not, e.g., demonstrated by the experiments of Section 4.2).
Vanishing gradients in PPO due to small reward standard deviation: A common practice for maximizing the reward in RFT is to use PPO, which replaces the expected reward with a surrogate objective. In Appendix B we establish that, if the reward standard deviation for an input small, then its expected gradient under PPO vanishes, just as its expected reward gradient vanishes.
Known causes of vanishing gradients: Focusing mostly on tabular Markov Decision Processes, prior works identified other conditions that lead to vanishing gradients in policy gradient algorithms. Agarwal al. (2021) showed that the gradient norm can decay exponentially with the horizon when the rewards are sparse. In RFT, however, the horizon is equal to one and the rewards are not sparse. More relevant is the fact that the expected gradient for a softmax parameterized model vanishes if outputs near-deterministic distributions (Ahmed et al., 2019; Hennes et al., 2019; Schaul et al., 2019; Mei et al., 2020a;b; Agarwal et al., 2021; Garg et al., 2022) - a special case of the reward standard deviation being small, i.e. of Theorem 1. Though, this is unlikely to be an issue for RFT of language models since, barring extreme cases, the output distributions produced by language models are not near-deterministic..
PREVALENCE AND DETRIMENTAL EFFECTS OF VANISHING GRADIENTS IN REINFORCEMENT FINETUNING
In this section, first show that the vanishing gradients phenomenon from Section 3 is prevalent in the GRUE benchmark, using the GPT-2 (Radford et al., 2019) and T5-base (Raffel et al., 2020) pretrained language models considered in Ramamurthy et al. (2023) (Section 4.1). Specifically, among its seven datasets, three contain a substantial amount of inputs with small reward standard deviation under the corresponding pretrained model, i.e. their expected gradient vanishes, although their expected reward is suboptimal. As anticipated, RFT has limited impact on the rewards of these inputs compared to its impact on the rewards of other inputs. Moreover, in datasets where inputs with small reward standard deviation are more common, the reward that RFT achieves relative to SFT is worse, indicating that vanishing gradients in RFT hinder the ability maximize the reward function. Controlled experiments that remove possible confounding factors (Section 4.2), such as insufficient exploration, and a theoretical analysis for simplified setting with linear models (Section 4.3), further establish that vanishing gradients in RFT can lead to extremely slow reward maximization.
For brevity, we often refer the reward standard deviation of input under pretrained model as the input's "pretrain reward standard deviation."
4.1 VANISHING GRADIENTS IN THE GRUE BENCIMARK
We followed the experimental setup of Ramamurthy et al. (2023), up to slight adjustments (spec- ified in Appendix F.2.1) for fairer comparison between RFT and SFT. In particular, we adopted their default hyperparameters and considered the following text generation datasets from GRUE: N (Kočiský et al., 2018). ToTTo (Parikh et al., 2020), CommonGen (Lin et al., 2020), IWSLT 2017 (Cettolo et al., 2017), CNN/Daily Mail (Hermann et al., 2015), DailyDialog (Li et al., 2017), and IMDB (Maas et al., 2011) see Table in Ramamurthy et al. (2023) for complete de-
4

---

Published as a conference paper at ICLR 2024
Pretrained Model
RFT Model
SFT Model
Dataset: Narrative
Reward kots
0.33
Rewe STD
0.23
2.4
1093
0.0
0:59
Sample Order (by protrain
Sample Order Iby pretrain
Sample Order (by pretrain
Dataset: ToTTo
Reward Mosn
-
Reword
6.23
30
1000
2003
3000
1000
5000
3.0
1000
3340
5000
0.0
1000
2000
3000
5003
0.03
Sample Ordar (by protrain reward std)
Sample Order (by pretrain reward sto)
Sample Order (by pretrain reward std)
Dataset: IMDS
Rewerd Moan
3-8-
38-
0.5
STD
0.33
2.4
2A-
22-
1000 2003
3000
.000
5000
0
1000
2000
5000
4003
5000
1000
2000
5000
4000
5003
Sample Order (by pretrein reward std)
Sample Order Iby pretrain reward stol
Sample Order (by pretrain reward std)
Figure 1: Inputs with small reward standard deviation under the pretrained model, Le. with vanishing expected gradient, are prevalent in the GRUE benchmark. For randomly chosen subsets of 5000 train samples from the NarrativeQA, ToTTo, and IMDB datasets, presented are the reward means and standard deviations (estimated based samples are ordered according to their pretrain reward standard deviation, with each marker representing reward mean and standard deviation (depicted by color) of an individual sample. Notice that a significant number of samples from NarrativeQA and ToTTo have small pretrain reward standard deviation, while their reward mean is low, Accordingly, RFT struggles improve the reward of these inputs, especially compared to SFT. In contrast, IMDB does not suffer from this issue, and the effect of RFT and SFT over it is more similar. Figure in Appendix for identical experiments
generations per input) under the pretrained, RFT, and SFT models. The
remaining GRUE datasets.
scription of the datasets. We ran PPO using the Adam optimizer for RFT, with the reward function in each dataset being either a task specific metric or a learned reward model (as specified in Ap- pendix F.2.1). In contrast, SFT was performed based on ground truth completions (i.e. labels), also using Adam. Similarly to Ramamurthy et al. (2023), we used GPT-2 as the pretrained language model for DailyDialog and IMDB, and T5-base for the remaining datasets.
Our findings are detailed below. For conciseness, we defer some experiments, egg., using NLPO (Ra- mamurthy al., 2023) instead of PPO for RFT, and implementation details Appendix F.
Inputs with small reward standard deviation are prevalent: As evident from Figure 1, Narra- tiveQA and ToTTo contain considerable amount of input samples whose pretrain reward standard deviation is near-zero, i.e. their expected gradient vanishes, while their reward mean is low. On the other hand, in IMDB the vast majority of samples have relatively large pretrain reward standard de- viation or a high reward mean. Figure 8 in Appendix F.1 further shows that the remaining datasets lie on spectrum between these two extremes, with CommonGen also exhibiting a noticeable amount of input samples with small pretrain reward standard deviation.
As expected (in Section 3), inputs with small pretrain reward standard deviation mon in datasets whose text distribution differs markedly from the pretraining corpora. Specifically, NarrativeQA consists of questions concatenated form answers based on the texts, while in ToTTo each input is a table and the task is to produce one-sentence description of it. By manual inspection, unsurprisingly, the pretrained model often generates continuations for the input text instead of obliging to the task at hand. When the desired output does not overlap with "natural" (according the pretraining corpora) continuation of an input, the model assigns non non-neglinee negligill probability only to outputs of roughly the same low reward, leading small reward standard deviation and vanishing expected gradient. We provide repre- sentative examples of inputs with small and inputs with large pretrain reward standard deviation, including outputs produced by the pretrained, RFT, and SFT models, in Appendix E.
com-
text segments and the task is output free-
5

---

Published as a conference paper at ICLR 2024
RFT Reward
0.0
NarrativeQA
ToTTo
IMDB
-0.1
RFT
0,48
0,46
0.72
-0.2
Train Test
IWSLT 2017
SFT
0.05
0.16
0.72
CNN/Daily Mail
-0.3
NarrativeQA
DailyDialog
IMDB
Table 1: RFT affects inputs with small reward standard deviation under the pretrained model less than it affects other inputs. Per dataset, reported is the Pearson correlation between the pretrain reward standard deviation and the abso- lute reward mean change due to finetuning, across the train samples. The correlations are computed based on the subsets train samples from Fig- ure 1. As anticipated, the correlation is consid- crably higher for RFT compared to SFT on Nar- rativeQA and ToTTo, which contain substan- tial amount of samples with small pretrain reward standard deviation. In contrast, relation is roughly the same since most samples with a small pretrain reward standard deviation al- ready have high reward mean, thus, their reward mean does not change much under both RFT and SFT. See Table in Appendix F.1 for the correla- tions on the remaining GRUE datasets.
-0.4
CommonGen
0.00
0.02
0.04
0.06
0.08
0.10
0:12
10'th Pretrain Reward STD Percentile
Figure 2: RFT performance (relative to SFT) is worse when inputs with small reward standard deviation are prevalent. Per dataset, the difference between the mean reward achieved by RFT and SFT is plotted against the 10'th percentile of the reward standard deviation un- der the pretrained model. Mcans are taken over five runs and error hars (indiscemiblec) mark standard deviations. We exclude inputs with near-optimal reward mean under the pretrained model (higher than 0.9) when computing the percentiles, since small reward standard deviation is only problematic if the reward mean is not high to begin with. Observe that, the lower the pretrain reward stan- dard deviation percentile is, i.e. the more train samples have small pretrain reward standard deviation, the worse the reward that RFT achieves relative to SFT.
IMDB the cor-
RFT affects inputs with small reward standard deviation less than other inputs: Figure I qual- itatively demonstrates that the rewards of train samples with small reward standard deviation change less due to RFT than the rewards of other train samples, whereas under SFT the change is more uni- form across the samples. To quantify this observation, for each dataset, Table 1 reports the Pearson correlation between the pretrain reward standard deviation and the absolute reward mean change due to finetuning, across the train samples. Indeed, the correlations for RFT are positive and significantly higher than those for SFT.
RFT performance worse when inputs with small reward standard deviation are prevalent: Lastly, we examine whether there is connection between the prevalence of inputs with small pre- train reward standard deviation, í.e. with vanishing expected gradient, and the reward that RFT achieves. A simple way to measure the prevalence of such inputs in dataset is examine some percentile the pretrain reward standard deviation, over the train samples. Remarkably, Figure shows that the lower the 10' th percentile of the pretrain reward standard deviation is, the worse the reward that RFT achieves relative to SFT. This observation is not sensitive to the choice of per- centile, as long it is not too large- Figure 9 in Appendix F.1 demonstrates similar trend for varying percentile choices.?
Our investigation reveals that the extent of vanishing expected gradients in a dataset, as measured by the reward standard deviation, is indicative of how well RFT performs. Yet one cannot infer a causal relation due to possible confounding factors. Mainly, the large output space in language generation introduces the well-known challenge of exploration (Ranzato et al., 2016; Nguyen et al., 2017; Choshen et al., 2020), which in the context of RFT translates to the challenge of obtaining accurate estimates for expected gradients (Greensmith et al., 2004; Schulman et al., 2015; Tucker et al., 2018). A natural question is whether the difficulty of RFT to maximize the rewards of inputs with small reward standard deviation indeed stems from their expected gradients vanishing. or rather solely from insufficient exploration. We address this question via controlled experiments (Section 4.2) and a theoretical analysis (Section 4.3). They show how vanishing gradients in RFT, due to small reward standard deviation, can lead to extremely slow reward maximization even under perfect exploration, iee. even when we have access to expected gradients.
4.2 CONTROLLED EXPERIMENTS
To account for the possible confounding effect of insufficient exploration, i.e. of inaccurate gradient estimates, we consider environments where exploration is not an obstacle. Namely, environments
For completeness, the train and test rewards that the pretrained, RFT, and SFT models obtain over all datasets are provided in Table of Appendix F.I.

---

Published as a conference paper at ICLR 2024
Model: MLP Dataset: MNIST
Model: ResNet18
Dataset: CIFAR10
Model: BERT-mini Dataset: STS-B
104
0.-
6.5
Reward
SFT
0.3-
small pretrain reward std
0.0-
0.0
-05
10-
1.0
-1.0
1000
2000
Epoch 4200
3000
5000
0
2020
2000
1000
4000
1000
2020
2030
3000
4000
2000
Epoch
Epoch
Figure 3: RFT struggles to maximize the reward over inputs with small reward standard deviation under the pretrained model, L.e. inputs with vanishing expected gradient, even with perfect exploration. On the contrary, SFT easily leads to maximal reward. For the controlled environments described in Section 4.2, in which RFT has access to expected gradients, displayed is the train reward that RFT and SFT attain throughout optimization, separately for train samples with small and train samples with large pretrain reward standard deviation. In Appendix F.I, Figure 12 provides plots showing the reward standard deviation and gradient norm throughout optimization, and Figure 13 presents identical experiments with stochastic gradient descent instead of Adam.
with a relatively small number of outputs in which one can compute the expected gradient for an input instead of estimating it. We observe that RFT is still unable to achieve high reward over inputs whose pretrain reward standard deviation is small, in stark contrast to SFT which easily does so. This demonstrates that, in the presence of inputs with vanishing expected gradient, RFT can suffer from extremely slow optimization despite perfect exploration. For conciseness, we defer some experiments and implementation details to Appendix F.
We seek to create finetuning environments in which: (i) it is possible to compute the expected reward gradient VoV(x; 0) for an input x; and (ii) there are inputs whose pretrain reward standard deviation is small. To that end, we take classification dataset with a modest number of labels, e.g. STS-B (Cer et al., 2017).² This ensures (i): we can compute the expected gradient given an input by simply querying the reward for all labels.
For pretraining, we assign to each train sample multiple ground truth labels by randomly choosing additional labels, and minimize the cross entropy loss starting from randomly initialized model, egg. BERT model (Devlin al., 2018). At the end of this process, given a train sample, the model roughly outputs a uniform distribution over the sample's 's pretraining labels. During finetuning, every train sample is assigned a (possibly different) single ground truth label, whose reward for RFT is set to 1, whereas the reward of an incorrect label is set to -1 For subset of the samples, we set their finetuning label such that it was not included (īi): the pretrain reward standard deviation of these train samples is small since the model (roughly) produces uniform distribution over labels with the same suboptimal reward. On the other hand, the reward standard deviation is large for the remaining train samples, whose finetuning label existed in their pretraining labels.
their pretraining labels. This guarantees
We followed the blueprint above for three different model and dataset configurations: multilayer perceptron (MLP) on MNIST (LeCun, 1998), ResNet18 (He et al., 2016) on CIFARIO (Krizhevsky et al., 2009), and BERT-mini (Ture et al., 2019) on STS-B (Cer et al., 2017). RFT was performed by maximizing the expected reward (Equation (2)) via the Adam optimizer, and SFT was performed by minimizing the cross entropy loss, also via the Adam optimizer. Figure 3 tracks the train reward that RFT and SFT attain throughout optimization, separately for train samples with small and train samples with large pretrain reward standard deviation. The results confirm that, despite having access to expected gradients and the small output space, RFT unable to achieve high reward in a reasonable time over inputs with small pretrain reward standard deviation. On the other hand, it does not encounter difficulty for inputs with large pretrain reward standard deviation. The behavior of SFT strikingly differs- - quickly reaches the maximal reward for both types of inputs.
The pretrain expected reward (i.e. expected re- ward under the pretrained model) of inputs with small pretrain reward standard deviation is often close to minimal, as shown on the GRUE benchmark in Section 4.1. This is also the case in the controlled experiments of Figure 3. However, we highlight that the vanishing gradient phenomenon established in Section 3 is independent of the expected reward- - the expected gradient for an input vanishes if its reward standard deviation is small, regardless of what the expected reward is. We empirically demonstrate this independence through experiments analogous to those of Figure 3, in
Insignificance of the pretrain expected reward:
We quantize the continuous labels of STS-B, which reside in [0, 5), to six values via rounding.

---

Published as a conference paper at ICLR 2024
which, although the initial expected reward of inputs with small pretrain reward standard deviation is relatively high, RFT is still unable to maximize their reward - see Figure 14 in Appendix F.I.
4.3 THEORETICAL ANALYSIS: SLOW OPTIMIZATION DUE TO VANISHING GRADIENTS
Through an analysis for a simplified setting, we theoretically demonstrate that vanishing gradients in RFT, due to small reward reward standard deviation, can lead to impractically slow opumization. Specifically, we prove an exponential optimization time separation between RFT and SFT, with re- spect to the pretrain reward standard deviation of an input x: the time it takes RFT to assign maximal probability for the correct label of x is (2(1 STD STDD prrp (-1x) [r(x, (-)x))((x,
)) where ((0) is the pretrained model parameters, while the time it takes SFT to do is only (In(11 STDy~p,00 (1x)))o))//)/r (r(rrx y)))).
Our analysis considers linear models trained over orthonormal inputs using small learning rate. Although the setting is rather simplistic, it gives insight into the detrimental effects of vanishing gra- dients due to small reward standard deviation. The details of our analysis are given in Appendix C; we regard extending it to more complex settings as an interesting avenue for future work.
5 OVERCOMING VANISHING GRADIENTS IN REINFORCEMENT FINETUNING
Section established that small reward standard deviation leads to vanishing gradients in RFT, which Section 4 demonstrated to be detrimental in a world benchmark and controlled environments. In this section, we explore ways to overcome vanishing gradients in RFT. We show that the con- ventional heuristics of increasing the learning rate, applying temperature to the logits, and entropy regularization are inadequate (Section 5.1). On the other hand, the common practice of an initial SFT phase (e.g.. used in Ziegler et al. (2019); Stiennon et al. (2020); Ouyang et al. (2022); Dubois al. (2023); Touvron et al. (2023)) is shown to be highly effective, both in terms of the resulting reward and in reducing the number of inputs with small reward standard deviation. This sheds light on the importance of SFT in an RFT pipeline it can help alleviate vanishing gradients. Moreover, our experiments reveal that a relatively small number of SFT optimization steps on as few as 1% of the input samples can suffice, indicating that the initial SFT phase need not be expensive in terms of compute and data labeling efforts (Section 5.2).
We focus on three datasets from the GRUE benchmark, which were found in Section 4.1 suffer the most from vanishing gradients: NarrativeQA, ToTTo, and CommonGen. For conciseness, we defer to Appendix F the results for ToTTo and CommonGen, as well as some implementation details.
5.1 INADEQUACY OF CONVENTIONAL HRURISTICS
We validate the resilience of vanishing gradients in RFT to common heuristics. Three sensible meth- ods that may seem suitable for addressing the problem are: increasing the learning rate, applying temperature to the logits, and entropy regularization (which is known theoretically to improve con- vergence rates in simple settings (Mei et al., 2020b)). However, we do not expect modifying the learning rate help since it increases the effective gradient norm over batch of inputs. Thus, the contribution of inputs whose reward standard deviation is small will continue to be negligible compared to the contribution of other inputs. Furthermore, while applying temperature the logits or entropy regularization may yield non-zero expected gradients, the gradients will likely not be in a direction that aids maximizing the reward due to the large output space of language models.
Based on the default hyperparameters from Ramamurthy et al. (2023), we carried out RFT runs with: (i) several values of larger learning rates, temperature, and entropy regularization coefficient; and, for comparison, (ii) an initial SFT phase, whereby the pretrained model is first finetuned in supervised manner before applying RFT. Confirming our expectation, none of the RFT config- urations led to a higher train reward compared to the default one. By contrast, in line with prior evidence (Ramamurthy al., 2023), an initial SFT phase greatly increases the reward that RFT at- tains- - see Table 6 in Appendix F.I. Furthermore, to gauge the effect of the considered heuristics on inputs with small reward standard deviation, Figure 15 in Appendix F.I presents the reward mean and standard deviation of individual train samples after applying RFT with each of the heuristics. In comparison to the effect of the default RFT configuration and SFT (shown in Figure 1), it evi- dent that the considered heuristics do not reduce the number of samples with small reward standard deviation substantially improve their rewards, whereas SFT is highly effective in doing so.
8

---

Published as a conference paper at ICLR 2024
RFT After Partia SFT Reward
RFT After Partial SFT Reward - Partia SFT Reward
RFT After Fu I SFT Reward - 0.538)
RFT Reward (- B 092)
Pretrain Reward [- 0.110Ij
Optimiization Steps
100
0.96
1.01
1:01
1.01
1,02
1.00
100
5.6
5.6
5.7
5.7
68
0.95
1,00
1:00
0.99
0.39
1.01
68
7.2
6.2
6.4
5.9
6.0
6.6
0.96
0.99
0.99
0.98
9.90
0.99
66
9.8
0.4
8.0
0.0
7.7
0.2
40
0.96
0,96
0.96
0.95
0.96
40
13.2
13.0
11.8
11.7
12
12.0
% SFT
20
0.91
0.86
9.91
20
18.6
17.1
18.6
17.4
10
0.60
59
10
11:1
12.2
11
11.4
10
%
% of SFT Samples
% of SFT Samples
Figure 4: On datasets in which RFT suffers from vanishing gradients, few initial SFT optimization steps on small number of labeled inputs substantially boost the efficacy RFT. For the NarrativeQA dataset, reported metrics are based on the mean train reward achieved when performing RFT after initial SFT phase with various percentages of optimization steps and labeled inputs (over three random seeds). A "full" SFT phase refers to 100% of the steps and labeled inputs used by Ramamurthy et al. (2023). Observe that the number of optimization steps and labeled inputs can he greatly reduced without causing a significant degradation in reward (left). Furthermore, RFT becomes roughly 5 to 18 times more potent after the initial SFT phase (right). We refer Appendix F.1 for analogous plots reporting metries hased on the mean reward (Figure 16), as well as identical experiments on the ToTTo and CommonGen datasets (Figures 18 and 19, respectively).
5.2 A FEW SUPERVISED FINETUNING STEPS ON A SMALL NUMBER OF INPUTS SUFFICE
Although an initial SFT phase often yields improved performance, it comes with a major drawback: it requires additional compute and, more importantly, relies on labeled data, which can be especially expensive to obtain for language generation. The results of Section 5.1 suggest that the benefits of an initial SFT phase stem, least partially, from mitigating vanishing gradients for the subsequent RFT phase. This gives rise to the possibility that relatively few SFT steps on small number of labeled inputs suffice. Meaning, we need only perform SFT decrease the number of input samples with near-zero reward standard deviation, after which the potency of RFT will substantially improve.
We investigate this prospect by decreasing the percentage of optimization steps and labeled inputs used for the initial SFT phase, with 100% corresponding to the default setup from Ramamurthy et al. (2023). As Figure 4 shows, the number of SFT optimization steps and labeled inputs can be greatly reduced without significantly harming the reward achieved by RFT. In particular, using 40% of the optimization steps and as few as 1% of the input samples allows RFT to reach 96% of the reward achieved when performing a "full" initial SFT phase. Moreover, the efficacy of RFT, as quantified by the difference between the reward after and before RFT, exhibits substantial increase as a result of the initial SFT phase. To understand the cause for this increase, we examine the impact of a partial SFT phase on individual train samples (see Figure 17 in Appendix F.1). We observe that the number of samples with small reward standard deviation, i.e. with vanishing expected gradient, is considerably reduced. This supports the perspective that the initial SFT phase beneficial due to mitigating vanishing gradients for the following RFT phase, and showcases that the SFT phase need not be expensive in terms of compute and data labeling efforts.
CONCLUSION
In this work, we identified a fundamental optimization challenge in RFT: the expected gradient for an input vanishes when its reward standard deviation is small. We then demonstrated the prevalence and detrimental effects of this phenomenon, as well as explored possible solutions. Overall, our results emphasize that being mindful for inputs whose expected gradient vanishes, as can be measured by the reward standard deviation, is crucial for successful execution of RFT.
L Limitations and future work:
To facilitate efficient experimentation, we considered language mod- els of small to moderate size and did not incorporate reward functions learned from human feedback. An exciting next step is to study the extent to which our empirical findings carry over more complex finetuning pipelines, including larger models and iterative human feedback (ef. Bai et al. (2022); OpenAl (2023); Touvron al. (2023)). Additionally, Section 5 demonstrated that three common heuristics, which may seem suitable for overcoming vanishing gradients in RFT, are in- adequate, whereas performing few initial SFT steps over small number of labeled samples is promising solution. Our investigation of possible solutions is not exhaustive. We believe further ex- ploring ways overcome vanishing gradients in RFT, e.g., by modifying the objective (as recently done in Lu et al. (2022); Rafailov et al. (2023); Hu et al. (2023); Dong et al. (2023); Gulcehre et al. (2023)), is valuable direction for future research.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
We thank Eshbal Hezroni for aid in preparing illustrative figures. NR is supported by the Apple Scholars AI/L PhD fellowship.
REFERENCES
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. theory of policy gradient methods: Optimality, approximation, and distribution shift. The Journal of Machine Learning Research, 22( 1))34111-- 4506, 2021.
7afarali Ahmed, Nicolas I.e Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy on policy optimization. In International conference on machine learning, pages 151-160. PMLR, 2019.
Dzmitry Bahdanau, Philemon Brakel, Kelvin Xu, Anirudh Goyal, Ryan Lowe, Joelle Pineau, Aaron Courville, and Yoshua Bengio. An actor-critic algorithm for sequence prediction. In International Conference on Learning Representations. 2017.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.0885 05862, 2022.
Satanjeev Banerjee and Alon Lavie. Meteor: An automatic metric for mt evaluation with improved correlation with human judgments. In Proceedings of the acl workshop on intrinsic and extrinsic evaluation measures .or machine trunslation und/or summarization, pages 65-72, 2005.
Daniel Cer, Mona Diab, Eneko Agirre, lñigo Lopez-Gazpio and Lucia Specia. SemEval-200 2017 task I: Semantic textual similarity multilingual and crosslingual focused evaluation. In Proceedings of the 1 11th International Workshop on Semantic Evaluation (SemExal-2017) Association for Computational Linguistics, 2017
Mauro Cettllo, Marcello Federico, Luisa Bentivogli, Niehues Jan, Stüker Sebastian, Sudoh Katsuitho, Yoshino Koichiro, and Federmann Christian. Overview of the iwslt 2017 evaluation campaign. In Proceedings of the 14th International Workshop on Spoken Language Translation, pages 2-14, 2017.
Leshem Choshen. Lior Fox, Zohar Aizenbud, and Omri Abend. the weaknesses of reinforcement learning for neural machine translation. In International Conference on Learning Reprexentations, 2020.
Hyung Won Chung, Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al. Scaling instruction-finetune language models. arXiv preprint arXiv:2210.1 1416, 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lec, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint arXiv:1810.04805, 2018.
Hanze Dong, Wei Xiong, Deepanshu Goyal, Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint arXiv:2304.0611 006677 2023.
Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes parameterized neural networks. 2019.
Yann Dubois, Xuechen Li, Rohan Taori, Tianyi Zhang, Ishaan Guirajani, Jimmy Ba, Carlos Guestrin, Percy Liang, and Tatsunori B Hashimoto. Alpacafarm: A simulation framework for methods that learn from human feedback. arXiv preprint arXiv:2305. 14387, 2023.
Omer Elkabetz and Nadav Cohen. Continuous vs. discrete optimization of deep neural networks. Advances in Neural Information Processing Syxtems, 34:4947-4960, 2021.
Spencer Frei, Gal Vardi, Peter Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. In The Thirty Sixth Annual Conference on Learning Theory, pages 3173-3228. PMLR, 2023.
Shivam Garg, Samuele Tosatto, Yangchen Pan, Martha White, and Rupam Mahmood. An alternate policy gradient estimator for softmax policies. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, pages 6630-6689, 2022.
10

---

Published as a conference paper at ICLR 2024
Evan Greensmith, Peter L Bartlett, and Jonathan Baxter. Variance reduction techniques for gradient estimates reinforcement learning. Journal of Machine Learning Research, 5(9), 2004.
Caglar Gulcehre, Tom Le Paine, Srivatsan Srinivasan, Ksenia Konyushkova, Lotte Weerts, Abhishek Sharma, Aditya Siddhant, Alex Ahem, Miaosen Wang. Chenjie Gu, et al. Reinforced self-training (rest) for language modeling. arXiv preprint arXiv:2308.08993 2023.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
Daniel Hemnes, Dustin Morrill, Shayegan Omidshafiei, Remi Munos, Julien Perolat, Marc Lanctot, Audrunas Cruslys, Jean- Baptiste Lespiau, Paavo Parmas, Edgar Quenez-Guzman, al. Neural replicator dynamics. arXiv preprint arXiv:1900 1000000990000 2019.
Karl Moritz Hermann, Tomas Kocisky, Edward Grefenstette, Lasse Espeholt, Will Kay, Mustafa Suleyman, and Phil Blunsom. Teaching machines read and comprehend. Advances neural information processing systems, 28, 2015.
Jian Hu, Li Tao, June Yang. and Chandler Zhou, Aligning language models with offline reinforcement learning from human feedback. arXiv preprint arXiv:2308. 12050, 2023.
Natasha Jaques, Asma Ghandeharioun, Judy Hanwen Shen, Craig Ferguson, Agata Lapedriza. Noah Jones, Shixiang Gu, and Rosalind Picard. Way off-policy batch deep reinforcement learning of implicit human preferences in dialog. arXiv preprint arXiv: / /907 997700056, 2019.
Samuel Kiegeland and Julia Kreutzer. Revisiting the weaknesses of reinforcement learning for neural machine translation. In Proceedings of the 2021 Conference the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Association for Computational Linguisties, 2021.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations, 2015.
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Ed- ward Grefenstette. The narrativeqa reading comprehension challenge. Transactions of the Association for Computational Linguistics, 6:317-328, 2018.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Yann LeCun. The mnist database of handwritten digits. http://yann. lecun. com/exdb/nnist/, 1998.
Gen Li, Yuting Wei. Yuejie Chi, Yuantao Gu, and Yuxin Chen. Softmax policy gradient methods can take exponential time to converge. In Conference Learning Theory, pages 3107-3110. PMLR, 2021.
Jiwei Li, Will Monroe, Alan Ritter, Michel Galley, Jianfeng Gao, and Dan Jurafsky. Deep reinforcement learning for dialogue generation. arXiv preprint 1606::1606.0154), 2016.
Yanran Li, Hui Su, Xiaoyu Shen. Wenjie Li, Ziqiang Cao, and Shuzi Niu. DailyDialog: A manually la- belled multi-tum dialogue dataset. In Proceedings of the Eighth International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pages 986 16-9995. Asian Federation of Natural Language Processing, 2017. URL https:/ http///aaaaahhhhcccc org/ 11t-/117719.
Bill Yuchen Lin, Wangchunshu Zhou, Ming Shen, Pei Zhou, Chandra Bhagavatula, Yejin Choi, and Xi- ang Ren. CommonGen: A constrained text generation challenge for generative commonsense reason- ing. In Findings of the Association for Computational Linguisties: EMNLP 2020, pages 1823-1840. Association for Computational Linguistics, 2020. doi: 10. 0005:////222..nnnnnn-oooo
URL
hittps://aat https////aaaaathogc.... oog. 2020.002. findingg-emnnn emnlp 165.
Chin-Yew Lin. Rouge: package for automatic evaluation summaries. Text summarization branches out, pages 74 81, 2004.
Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Tay, Denny Zhou, Quoe V Le, Barret Zoph, Jason Wei, et al. The flan collection: Designing data and methods for effective instruction tuning. arXiv preprint arXiv:2301. 13688, 2023.
Ximing Lu, Sean Welleck, Jack Hessel, Liwei Jiang, Lianhui Qin, Peter West, Prithviraj Ammanabrolu, and Yejin Choi. Quark: Controllable text generation with reinforced unlearning. Advances in neural information processing systems, 35:27591-27609, 2022.
11

---

Published as a conference paper at ICLR 2024
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In Inter- national Conference on Learning Representations, 2020.
Andrew Maas, Raymond E. Daly, Peter T. Pham. Dan Huang, Andrew Y. Ng. and Christopher Potts. Learn- ing word vectors for sentiment analysis. In Proceedings of the 49th Annual Meeting of the Association for Computational Linguisties: Human Language Technologies, pages 142-150. Association for Computational Linguistics, 2011. URL https:/ aalanttooooooooooooooooooooooola org/P11-1011
Jincheng Mei, Chenjun Xiao, Bo Dai, Lihong Li, Csaba Szepesvári, and Dale Schuurmans. Escaping the gravitational pull of softmax. Advances in Neural Information Processing Systems, 33:21 100211440, 2020a.
Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global convergence rates of softmax policy gradient methods. In International Conference on Machine Learning, pages 6820-6829. PMLR, 2020b.
Khanh Nguyen, Hal Daumé III. and Jordan Boyd-Grabee Reinforcement learning for bandit neural machine translation with simulated human feedback. In Proceedings of the 2017 Conference Empirical Methods in Natural Language Processing. Association for Computational Linguistics, 2017.
OpenAl. Gpt-4 technical report. arXiv preprint arXiv:2303. 08,7440,,8777
Long Ouyang, Jeffrey Wu, Xu Jiang. Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray. et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 35:27730- 27744, 2022.
Ankur Parikh, Xuezhi Wang. Sebastian Gehrmann, Manaal Faruqui. Bhuwan Dhingra, Diyi Yang. and Dipan- jan Das. ToTTo: A controlled table-to-text generation dataset. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 1173-1186, Association for Compu- tational Linguisties, 2020. doi: 10. 12118665/112222.mmmll emw.:p00e.mmlmmmmmmmmm emnlp-main.89.
URL https: aal.atoonnnoooy.....oo..
org/2020.. mommlp-mainn.. 89.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytonch. In NIPS- W. 2017.
Romain Paulus, Caiming Xiong, and Richard Socher. deep reinforced model for abstractive summarization. arXiv preprint arXiv:1705.044 7030444 2017.
Matt Post. call for clarity in reporting bleu scores. arXiv preprint axXiv:1804.000 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. arXiv preprint arXiv:2305. 182900 2023.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. The Journal of Machine Learning Research, 21(1):5485-555 2020.
Rajkumar Ramamurthy, Prithviraj Ammanabrolu, Kianté Brantley, Jack Hessel, Rafet Sifa, Christian Bauck- hage. Hannaneh Hajishirzi, and Yejin Choi. Is reinforcement learning (not) for natural language processing?: Benchmarks, baselines, and building blocks for natural language policy optimization. In International Con- ference on Learning Representations, 2023.
Mare' Aurelio Ranzato, Sumit Chopra, Michael Auli, and Wojeiech Zaremba. Sequence level training with recurrent neural networks. In International Conference Learning Representations, 2016.
Noam Razin and Nadav Cohen. Implicit regularization in deep learning may not be explainable by norms. Advances neural information processing systems, 33:21 174-211877 2020.
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions the nonlinear dynamics of
learning in deep linear neural networks. In International Conference on Learning Representations, 2014.
Tom Schaul, Diana Borsa, Joseph Modayil, and Razvan Pascanu. Ray interference: source of platcaus in deep reinforcement learning. arXiv preprint arXiv:1904. 1455, 2019.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional contin- control using generalized advantage estimation. arXiv preprint arXiv: 1506.02438, 2015.
12

---

Published as a conference paper at ICLR 2024
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov, Proximal policy optimization algorithms. arXiv preprint arXiv:1707.0634, 2017.
Nisan Stiennon, Long Ouyang. Jeffney Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Lccrning summarize with human feedback. Advances in Neural Infor- mation Processing Systems. 33:3008-3021, 2020.
Romal Thoppilan, Daniel De Freitas, Jamie Hall, Noam Shazeer, Apoorv Kulshreshtha, Heng-Tze Cheng. Alicia Jin. Taylor Bos, Leslie Baker, Du, et al. Lamda: Language models dialog applications. arXiv preprint arXiv:2201 08239, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bash- lykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.092880 2023.
George Tucker. Surya Bhupatiraju, Shixiang Gu, Richard Turner, Zoubin Ghahramani, and Sergey Levine. The mirage of action-dependent baselines in reinforcement learning. In International conference on machine learning, pages 5015-5024, PMLR, 2018.
Iulia Turc, Ming-Wii Wei Chang. Kenton Lee, and Kristina Toutanova. Well-read students learn better: On the importance of pre-training compact models. arXiv preprint arXiv: 1901199988,,,2,, 2019.
Yizhong Wang, Swaroop Mishra, Pegah Alipoormolabashi, Yeganch Kordi, Amirreza Mirzaci, Anjana Arunk- umar, Arjun Ashok, Arut Selvan Dhanasekaran, Atharva Naik, David Stap, et al. Super-naturalinstructiin Generalization via declarative instructions on 1600+ nlp tasks. arXiv preprint arXiv:2204.077000 077000 2022.
Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoe V Lc. Finetuned language models are zero-shot learners. In International Conference on Learning Representations, 2022.
Thomas Wolf, Lysandre Dehut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cis- tac, Tim Rault, Rémi Louf, Morgan Funtowiez, al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint 1rX.0:1910.0377,, 2019.
Yuxiang Wu and Baotian Hu. Learning to extract coherent summary via deep reinforcement learning. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018.
Zhilin Yang. Zihang Dai, Yiming Yang. Jaime Carbonell, Russ R Salakhutdinov, and Quoc Le. Xlnet: Gen- cralized autoregressive pretraining for language understanding. Advances in neural information processing systems, 32, 2019.
Wayne Xin Zhao, Kun Zhou, Junyi Li, Tianyi Tang, Xiaolci Wang, Yupeng Hou, Yingqian Min, Beichen Zhang, Junjie Zhang. Zican Dong, et al. A survey of large language models. arXiv preprint arXiv:2303. 20t.:2223..8223.... 2023.
Li Zhou, Kevin Small, Oleg Rokhlenko, and Charles Elkan. End-to-end offline goal-oriented dialog policy learning via policy gradient. arXiv preprint 1rXiv:1712.0083
2017.
Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Chris- tiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv:1909.00 08593, 2019.
13

---

Published as a conference paper at ICLR 2024
A RELATED WORK
Reinforcement learning for language generation: Our work highlights an optimization challenge in policy gradient-based RFT, motivated by its widespread use for adapting pretrained language models (Ziegler et al., 2019; Stiennon et al., 2020; Ouyang et al., 2022: Bai et al., 2022; Dubois al., 2023; OpenAI, 2023; Touvron al., 2023). Reinforcement learning has been applied more broadly in task specific settings, for example: translation (Ranzato et al., 2016; Bahdanau et al., 2017: Nguyen et al., 2017: Kiegeland and Kreutzer, 2021), summarization (Ranzato et al., 2016; Paulus et al., 2017: Wu and Hu, 2018), and dialog (Li et al., 2016; Zhou et al., 2017: Jaques et al., 2019). Policy gradient algorithms are most commonly the method of choice. Though, value-basee methods have also been used (as in Jaques et al. (2019)), which notably do not suffer from the vanishing gradients phenomenon identified in this work.
Optimization challenges in policy gradient: The objective landscape of policy gradient is non- convex even in simple settings (e.g. with linear models). Thus, deriving conditions under which efficient optimization occurs has proven difficult (Mei et al., 2020a;b; Agarwal et al., 2021; Li et al., 2021). In environments with a large output (i.e. action) space, for which language generation is prominent example, another well-known challenge is that of exploration (see, e.g., Ranzato et al. (2016); Nguyen al. (2017); Choshen et al. (2020),, which in the context of RFT boils down to obtaining accurate expected gradient estimates (Greensmith et al., 2004; Schulman et al., 2015; Tucker al., 2018). Interestingly, our results highlight that, for inputs with small reward standard deviation, obtaining more accurate estimates is unlikely to be useful since their expected gradient vanishes.
As discussed in Section 3, the expected gradient for a softmax parameterized model is known to vanish when the model outputs near-deterministic distributions (Mei al., 2020a;b; Agarwal al., 2021) -a a special case of the reward standard deviation being small, i.e. of Theorem 1. Under this special case, was empirically demonstrated that policy gradient algorithms struggle maximize the reward, even with perfect exploration, i.e. even with access to expected gradients (Ahmed et al., 2019; Hennes al., 2019; Schaul et al., 2019; Mei et al., 2020a;b; Garg al., 2022). However, the relevance of near- -etermmteriiiiiicc output distributions to RFT is limited since language models rarely produce such distributions. In contrast, Section establishes more general condition leading to vanishing gradients, relevant to RFT, that is demonstrated in Section 4 to result in optimization difficulties.
Overcoming optimization challenges in policy gradient: To address the poor optimization land- scape of softmax parameterized policy gradient, prior work proposed replace softmax with an alternative transformation (Mei et al., 2020a) and modify the gradient update step (Hennes et al., 2019; Garg et al., 2022). Examining the utility of these methods for RFT may be an interesting topic for future work.
B SMALL REWARD STANDARD DEVIATION IMPLIES VANISHING GRADIENTS IN PPO
Theorem 1 (in Section 3) established that the expected reward gradient VoV(x; 0). for an input x, vanishes if the reward standard deviation of x is small. In this appendix, we prove that the expected gradient under PPO - policy gradient algorithm widely used for RFT vanishes as well in that case.
PPO most often replaces the expected reward V(x; 0) with a surrogate (clipped) objective (Ouyang al., 2022; Bai et al., 2022; Ramamurthy et al., 2023; Touvron et al., 2023). That is, it maximizes (0) := Ex-D (VPPO 0)) with (cf. Schulman et al. (2017)):
Vxpo(x;0) == --y-yssp11 min [o(y))) ()(() Ag y), clip
(Po(y)x) po(y!x)
5, 6,1+ o,))+)))
(3)
where Ō PR stands for fixed reference parameters, which updated to the current parameter values at the start of each PPO iteration, Ag(x, = r(x, V(x; 0) is the advantage function
with respect to 0, and clip(2, 8, 8) projects z R onto the interval [1 5, 5] for
14

---

Published as a conference paper at ICLR 2024
5 E (0, Another, less common, variant of PPO includes the Bullback-Leibler (KL) divergence between pe(- (x) and pa(-1x) ) as regularization term instead of clipping (ef. Schulman et al. (2017)):
( x) Ax(x. xY) -
(xx:- == Ey Ex-pp(-1xx)
(4)
p)(y
for X20. 0.
Propositions and 2 below establish that, the reward standard deviation for an input small, then its expected gradient under both PPO variants vanishes, just as its expected reward gradient vanishes. Specifically, we show that the distance between (xV(x; 0) and Vo Vpro (x; 0) depends on the total variation distance between po(+ x) and p x), where 8 denotes the PPO reference parameters. The same is true for the distance between VoV (x; 8) and (yVPOO-K((S 8). At the beginning of each PPO iteration 0 is updated to the current parameter values. Hence, the distance between po po((xx) x) and p (x) typically remains small throughout optimization, and is exactly zero the beginning of each PPO iteration. The proof' of Propositions and 2 are deferred to Appendices D.2 and D.3, respectively.
_ be the reference parameters some iteration of PPO (Equation (3)). For an input x E Xlinn, we denote the total variation distance between po(-Ix) and po(-1x) by (V(po(- (((po(((xx)(lp(((x)))) it holds that:
Proposition 1. Under the notation of Theorem 1, let 0 R.
Eyex (po (ylx) po(y|x)]. Then,
VoVPO 0) VoV(x;0)1 (x; TV (xxx(<((x))lpp((xx))
8
where 8 (0,1) the PPO clipping hyperparameter, Consequently, by Theorem 1:
2 - (T(pp(()x)lllaa())))
V (Pxo(x)( 0)1 < (x; (<)))>p(((()))(
where the beginning of each PPO iteration 0, meaning TV(po(->x)llpatt =
Proposition 2. Under the notation of Theorem let 0 TRP be the reference parameters at some teration of the KI. regularized PPO variant (Equation (4)). For an input x E XLis, we denote the total variation distance between po(-Ix) and pa(-Ix) by TV(Do(>>x)llattH Eyextoou (yo(y|x) po(y|x))- Then, it holds that:
0) VxV(x; 0) S 4T 0); (V(po((ss)llaa(t))))
where X 2 0 is the KL regularizationn coefficient. Consequently, by Theorem 1:
[XVV 0)1l S 6 (x;: 0) STD y 2/3 + (V(po(()x))lpa((xxx) 3
2A
where at the beginning of each PPO-KL iteration 0 0, meaning TV TV(pp( 0.
C DETAILS OF THEORETICAL ANALYSIS: SLOW OPTIMIZATION DUE TO VANISHING GRADIENTS
In this appendix, we deliver our formal analysis of slow optimization due to vanishing gradients in RFT, referred to in Section 4.3.
Consider a linear multiclass classification problem with D-dimensional inputs and K 2 labels, where the model f : R x I K + R² is parameterized by matrix f E RR K Conditioned on an input E R., the model produces distribution over labels y E {1,..., K} through the softmax operator as follows:
p = softmax(f(x;)))) 0)), = softmax(Ox)..
Suppose that we are given parameters (10) kx obtained through some pretraining process, and a finetuning training set (xI, NNN E R² x (1...., K} consisting of N orthonormal
We consider horizon-one formulation of PPO. One can also treat the environment having a horizon of Tout, with all intermediate rewards heing zero and discount factor of one.
15

---

Published as a conference paper at ICLR 2024
inputs (meaning, NN Σ D) and corresponding ground truth labels. A natural way define a reward function in this case is, given a train sample Xn, set the reward of the ground truth label to 1, f.e Vn) 1, and the reward of all other labels to - 1, Le. r) - for all y # yn-
We analyze the optimization dynamics of RFT when carried out via gradient descent with the re- wards specified above over the training set, i.e. of gradient descent over (0) (Equation (2)) with D being the uniform distribution over XI.... :NN. The dynamics of gradient descent are commonly studied via the infinitesimal limit of small learning rate, known as gradient flow (cf. Saxe et al. (2014); Du et al. (2019); Lyu and Li (2020); Razin and Cohen (2020); Elkabetz and Cohen (2021); Frei et al. (2023)), which in the context of RFT boils down to assuming the parameters are governed by the following differential equation:
(RFT(O) = (
(t) == Va (OR((())))). , di
where RRFT (t) denotes the parameters at time t 2 0 of optimization. Analogously, in SFT, gradient flow for minimizing the cross entropy loss (SFT(0) (Equation (1)) over the training set amounts to:
d
(SFT(0) 0(0)
0SPT (1) == VoLSPr(OsFT(()) .20
dt
with the distribution DD in the definition of (fF being uniform over the training set, and 6SFT (t) denoting the parameters at time 2 0.
Theorem 2 establishes an exponential optimization time separation between RFT and SFT with respect to the reward standard deviation of input under g()) Resonating with the controlled experiments of Section 4.2, we assume that there exists train sample misclassified by g(0), for which (10) assigns the same probability to all incorrect labels.
Theorem 2. Consider the setting described above (in Appendix C), and let Xn be train sample that ( misclassifies, i.e. Potor (yn(xx) & yrgnaxye(11...K] Pot: (y(xx), ), for which Φ(0) assigns the same probability to all incorrect labels, i.e. Poços (y xn) Poços (y'[xn) for all y.y' + y Denote by IRFT, 1SFT 2 the initial times at which RFT and SFT classify Xn correctly, respectively. Formally:
RRFT := min t 2 0 (ORFT(() (yn( E argmaxy< (I..,,, PÖRFT (4) (y (y/xx))
tsFT:
min 2 0: PearT (V argmaxye(1....t ((ter(t) (y|xx))
Then, holds that:
1
tsFT = =0 (mm( STD y-Pooo (1x, (1x)))((x Ir(xx. xx] x))))
tRET = =Ω
STD y-DDy-aos
(xx))) ( (xx)))())))))
That is, there exist constants e,, d E Rsso and c2 E (ie e,, G 02, C22 do not depend on the initial parameters )(0) such that tRET 2 ci . 10..11/517 (1/STDy~p, y~ (1x)
(1x))r(x [r(x, (xx,,)) + cz whereas
tsFT
SIDDyppoos (-b)) (xx,((xx)))) + c2.
Proof sketch (full proof in Appendix D.4). Since X1,..., XN are orthogonal, the logits that the model produces for X,, i.e. f(x;; (t)) in RFT and f(x,; (trr(t) in SFT, evolve indepen- dently of the remaining train samples. Furthermore, the assumption that g(0) assigns equal proba- bilities to all incorrect labels of x,, implies that they remain equal throughout optimization, and so do the respective logits. Let put) be the difference between the logit of yn and of some y * yn time I. given X,, and note that the initial time at which x, is classified correctly is the initial time at which ju(t) = 0. The above facilitate deriving the exact dependence of /RFT and (SFT on (x,; g(")), i.e. on the initial logits produced by the model given x,. from the differential equation governing (t).. The proof concludes by relating f(x,; 8(0)) to the reward standard deviation of x under ((0).
We note that our lower bound on TRET is similar in spirit to Theorem of Mei et al. (2020a). For multi-armed bandit problem, Mei et al. (2020a) lower bounds the time until reaching near-optimal reward, in terms of the initial optimal action probability. The main difference between the results, besides the exact technical conditions, is that we derive the dependence of optimization time on the initial reward standard deviation and contrast it with the optimization time under SFT.
16

---

Published as a conference paper at ICLR 2024
D DEFERRED PROOFS
D.1 PROOF OF THEOREM 1
Differentiating (x;; 0) with respect to using the log-derivative trick we get:
Vo V (x; 0) Ey-po(-lx)
r(x, y) Vo po()
= Ey~po (x)
r(x,y) 1-1
Vo (olpo(unix.yss-1))
Louut
r(x,,) y)
Ve In softmaax (xx (
= -E, Lyxtoon Po (x(x(x))(xx,)))
Enn1 (0) p-x-
Lout
where Cy. E R denotes the yr 'tth standard basis vector, i.e. the 'th entry of By is and its other entries are zero, and po ys1-1) stands for the vector holding the next-token probabilities, ie. pa ysi-i) = softmax (x Vzi-1:0) E x
For 0 to be determined later, denote by 3i the set of outputs whose rewards deviate by more than c from the expected reward, i.e.:
V = {y E Xhouu |r(x,)) V(x; 0)1 c).
Defining modified reward function F X X -> 1-1, 1] by:
f(x,y) (r(x,y)
y $ Ve
(xx;0) . E Ve
we may write Vo (xV((; as follows:
Vo (x:0) Eyextome po (o(y(x)F(x,)) E 1=1
Loue
J/xxys1110)) -o(((x,)21-1))
Σ Eyev. po(ylx)(((x,y) f(x, y)1
Σ Ji ys--)))
Loue
+
1=1
(IT)
We upper bound the Euclidean norms of (7) and (IT) separately. Starting with (IT), by Chebyshev's inequality we know that:
(o())()== < STD (xD
6
Notice that j,e, - pa(-Ix -))l < lle,,
1)1 S for all yst E Xt, where II-I1
denotes the E norm. Thus, for any y E V. and / E {l,..., Loue):
| (1 pe(-Ix,ys1-)))
< || | [(e, $ ():
-o(-(x,,
Since rewards lie in [-1, 1), by the triangle inequality this implies that:
II(TT))I < 4Lom (x: 0): (o(())))= (1;m)((x00) STD (TDyypa((s)/rxx,.) As for (I), denoting for y51- E x111
(5)
(Vsi 1) == Eysexxtou-ty tn1 x X
where y21 = (ye,.... Wlove ) we have that:
(T) Σ 1=1 Eye-ex p
Lout
17

---

Published as a conference paper at ICLR 2024
The y 'th entry of a asss 1) is given by:
(x)1-1) = - Po
-po(vlx,ys--)) B < wtl.://wwss.....................
(
Since by the definition of f it holds that [ñ(x, y) ((x; [)] S for all y Xhout, bound the difference of expectations the equation above through adding and subtracting V(x; 0) and the triangle inequality:
$ |e -((x;) + [x( 0) Byg [Fx
s 2c.
Thus:
||a\\)-1)]| < (a S > yEX Pe(1/xx,51-11 2c 2c.
Recalling that:
(1) (1)= > 1 E -x Po po be ((--))))
taking the norm of both sides and applying the triangle inequality yields:
(I(T)II
Enn
i-l Ey 1EX1-1 poly (ol I
Lout
Eien |=l 27(x: 0)
Loui
Ey Po
(6)
(x; -c.
Combining the bounds on the norms of (T) and (JT) (Equations (5) and (6)) then leads to:
VoV (x; < 4Lout (x; 0) STD (1) 2Lo 0) c.
Lastly, choosing c = STD, (IDDy-pu(-l)))) [x 2)
concludes the proof:
[VoV(x; S 6I (x;;) STD 1x) [x),
D.2 PROOF OF PROPOSITION 1
Define:
Poly x)
R$ = y VExloN!
> 11+8, 18, As(x,y) > >0}
po(y x)
R3 = { {y E Ex
Poly x)
< - 8, As(x,y) < <0
Poly x)
Rs = { y {yEx
Po Pa (x)
og.
i.e. R2 is the set of outputs for which PPO will clip the objective when the probability ratio is strictly greater than Rg is the set of outputs for which clipping occurs when the probability ratio is
18

---

Published as a conference paper at ICLR 2024
strictly less than 1 - 6, and R& is the set of outputs for which the probability ratio is either strictly
greater than + 6 or strictly less than - 8. Note that Rg U UR, - Rs..
We can decompose Vppo(x: 0) according to R2z Ry and the remaining outputs as follows:
Vxro(x; 0) = Σ yE \RR) URZ po(ylxx)4gxxxy) 10t
Σ YER (1+
EYER, -
Since the second and third sums in the equation above do not depend on 0 we have that:
V Vx)ppo((:0) ()
E (x ((FURR)) Vepelyl))
Adding and subtracting Σ EYeRjUR y) yields:
VoVppo 0) = E yexhuu! ( Ey Ag(x, y) Vopoly/x. (7)
Recall that Az(x.y) = r(x, V(x; 0). Because V(x: 8) does not depend on we get that:
Eycxeme V)peyy)) = EyExtome r(x, y) (y - v(x;)) Σ Eyextour (yps(y|x)
= VoV 0) V(x; V Σ yEX Po(y/x)
(xV(x; V(x; 0 Vol
= ();
Plugging this into Equation (7), rearranging the equality, and taking the Euclidean norm of both sides, it follows that:
VoVpro(x;0) V (xV(x; < EyERRJRR; Ag(x, V
Rewards lie in the interval [-1, 1], hence, |Ag(x, I 2. So, applying the triangle inequality we obtain:
VoVrro(x;0)) VoV(x; $ Σ YERO UR; Vopo(y|x)]]
= 2 =2E > EVERTUR Vo
= 2 Σ YE po(y)x)
Σ Volnpo(w/x.vssi)|
Laue
1=1
where the first equality is by the log-derivative trick. Since
(o( x, y s)) softmax ((x.
for any / E (I,..., Louth and E X hous we have:
Va In (x x, ys--) = (iixss-roo(e,.) -xo
where Gg. R|2ll denotes the yı 'th standard basis vector, iee. the yi 'tt entry of cy one and its other entries are zero, and po(-|x, y sI-1) = softmax(/ (x, y<(-1; E [RIX Noticing that lleg pa( xx--)) < ||egr po x, vel-)) < 2, where II-I[, denotes the & norm, we have: | In (mn (x, sll-1) < 2x(xx))
By the triangle inequality this implies that:
V VXVPoo 0) voV(x;0) < Σ
SVERTRR po(y(x)
Eien bout i-
< 4T T 0) Σ URR po(y|x)
(8)
4hout (x: 0) Pe (R UR, |x)
(x: 0) (po (R3 (x) + po (R5 (x))
19

---

Published as a conference paper at ICLR 2024
where the last transition is due to R3 and R2 being disjoint. The proof concludes by showing that both po xx) 25- (Vo (-)))()))llxt (|)]((((())))))))))) (x)) and (o |x) - 8-1 .VV (P))
Focusing po (R3 xx), by the definition of the total variation distance:
po (Rz [x) - pā ( [x] - (V(po
from which it follows that:
pø R J (x) po R2 x)
po (R xx)
where note that (R3 x) jf due to the softmax parameterization, under which the probability of any output cannot be equal to zero. For each y R3 we know that pa(y xx)/payl/x 1 + 6. Thus, po (R (R) (x) > 1 + 5 and Pa(R xx po (R) xx) 146 Plugging this into the inequality above, we have that:
Py (R5 $ [x) po (R (x)
Pe (R
< Po (R
1+5
and so:
+ 8
po [ x) I
(V(po(->x))lpa( - TV(po(1b))lla(ttx)
(9)
Turning our attention to (R, xx). by an analogous derivation we get:
Po (R5 (x)
po (R5 (x)
S (V(6o(11x)llla((
Po R xx
For each Rg we know that p < 1 - 5, implying that po (R, xx) (Rz xx)
1 - 8 and pa (R5 (x) (R) xx) > 1 + I Consequently:
Pa (R5 (x)
po (x)
1-6 < po (R5 Po (x)
< (V(p((()))lloa((bb)))
po (Rg (x)
from which it follows that:
1
pa R [x) TV(po(->x))ppa(())) - (V(po((xx)llaa(xx))
(10)
Finally, combining Equations (8), (9), and (10) establishes the sought-after result:
VVVPpo((x - V (x: 0) < A 0) (po (R₂ (x) (R (x))
12 (xx 0) (V
D.3 PROOF OF PROPOSITION 2
We may write the PPO-KL objective (Equation (4)) as:
Pe(yy'x x)
VPPO-KL(x;0) 0) = Ey-paa (x)
Ag(x,)) y) A (pxx)xx-)))
( x)
= Ey-po (- [Ag y)] A DK (x)eee
Since Ag(x,y) y(x,y) V(x; 0), it holds that Ey~ (Aa(x, y) V(x; Vxx 0), where
note that V(x; 8) does not depend on 6. Differentiating with respect to f thus yields:
VOVPPO-KI(X
VoV(x: 0) A VolDeer (pa( (x)()) (x)olplll(x)))
20

---

Published as a conference paper at ICLR 2024
and so:
(oVPPo-KL(x:0) VoV(x: - (Ps(
We therefore need only show that V (x)xx)))) is upper bounded by 4Lout 0)
Writing the KL divergence term explicitly:
yEX Po(y In pe (y(x)
(Px
Vo
Σ
Po (y(x))
Ve yEX pe (y In pe (y(x)
+ V yEAthut (o(y In (x()x
Σ yEx pa(y(x)
Ve Volnpolw(x.sss-)) 1=1
Lous
For any E (1,..., Loud] and Xoue we have that:
V (n Y51-1) = J J/xvaa- - po ((1-)))
where cy. E I R' denotes the y 'th standard basis vector, iee. the y 'th entry of Cy is one and its other entries are zero, and po yo(-)))
R! Thus, a straightforward
computation shows:
Voll (po()x/ooo-)))llee
--E Trout
Eycreex Po(yst-ilx)) Jirver 100) (xxx -xx
Then, by the triangle inequality and the fact that:
Jiins 100) (po(-1x, v5(-1) po ((-x))))))))))
1 (K(pa((x. po(-
() -
where I|-II, denotes the L norm, we obtain:
S 7x(x; 0)
Ein
Toou
Eye -EX-11 Polys1-11x) (ed (II)
Notice that for any / E {l,..., Louth and y E 2111;
pe(ys1-1lx)) (-11,,,, (x
= -Pelysi-1x1 Σ
ex Ips Ip | polol
Σ M
- 1x
WEE
Σ mEx [paly ( (y + po
po (y -x(()==-
- + 1o po(yst)x))
where the penultimate transition is by adding and subtracting poly (x1, in each summand and the triangle inequality. Summing over yst-1 E X we thus get:
P Polys1-1(x) (1po(-Ix.ysı-1) - po (x,
E (o -pgys(-xx)) E
Epsisxlallysll) E (pa(ys())) po(ys(/x))
21

---

Published as a conference paper at ICLR 2024
Both sums in right hand side of the inequality above equal to twice the total variation distance between marginal distributions of (-(xx) and pa(-(x). Hence, each one of them is upper bounded by 2TV (1)o0-/))lloo((px))) Le::
Eys 1621-1 (5 (y s-- (epxxx)))lllxxxxxxx (1xx(((x.
Going back to Equation (11):
VoDKL (Pol & (xx TV(90()
from which it readily follows that
(x voV(x;0) & A- ())
I o (x)p(()
D.4 PROOF OF THEOREM 2
The proof is sectioned into two parts. We first derive the lower bound on the time until correct classification of Xn, for RFT, after which prove the upper bound for SFT. For simplicity, assume without loss of generality that X₁ = ej...., XN = eN, with e eN being the first N standard basis vectors of igD The analysis for arbitrary orthonormal inputs .I.....X; Xy follows from analogous arguments, by tracking the evolution of inner products between rows of and x X instead of individual entries of 0.
D.4.1 OPTIMIZATION TIME LOWER BOUND FOR RFT
Denote by Wy.alt) E R the (y, d)'th 'hh entry of (RFT(t), for y E (1....,K), K),d E D). Under this notation, since Xn'
en' we have that f(xx;; ORF((() for n' E NN,...,N.. Now, recalling that r(xar, = and r(xn',) - for all y + Yn', the expected reward at time 2 0 can be written as:
ÓRFT(1)Xx'
(t), WK,
(t)) WK,,...,
N Σ N
K
V(ORFT (1))
y=l PoRPT (t) (y Xte (x(x),)))
n'=1
Σ N
P (OaPT(i)
(1 -PORET (N [)]
N
Σ N
exp(wy, nn (1))
E exp(uyyn (t)) )
exp(uy, ,n (1))
N
y= exp( Wy (t))
2 eexp(uy (t))
N
m==1
ンý. exp(py n² (t))
By introducing the shorthand zn/ zn) == Exn exp( Wy (t), we have that:
1 .N' Σ n'=1 N
2 2exp(wy, m on (t))
V(ORFT (t))
Zn
Focusing X,, ea. under gradient flow the parameters wi, (t),..,, WK,n (t), which determine the output distribution given Xn, obey the following differential equations:
Texp Sxpluynnnn (()) on (1)
exp( Mxx (t))
d My uy V Vuy.n ((ORFT(t))
Su (t)
y yn
(12)
expli (t)) (ND(
y * yyn
(0)
for all y (1...., K}. As can be seen from Equation (12), the evolution of (t,..t)..... WKKn (t) independent of the remaining parameters (this is a result of the orthogonality assumption on Xi...., XN). Furthermore, the differential equation governing each parameter Wy,n(() (1) correspond- ing to an incorrect label y * y. is the same. Since, by assumption, the softmax probabilities of the incorrect labels for x,, are equal under (R), know that Wy. (0) wy (0) for all
22

---

Published as a conference paper at ICLR 2024
y.y' * Yn- Hence, the uniqueness of the solution to Equation (12) for an initial condition implies that Wy,n (t) = Wy', (t) for all 2 0 and y. * Yn- Meaning, the parameters corresponding to incorrect labels of x,, are equal throughout optimization.
Given this observation, we may denote e(t)) == wy,((t) for any y + Yn- The dynamics for the n 'th column of ÓRET (t) thus boil down to:
2(K 1) expl exp(piy......e (t) e()))
dt W yyn,n(t) Mexpp(Wy..nl + (K exp(v
(13)
expp(1g Wy
+ v(l))
dt v(t) = M Wy (t)) + (K 1) exp(v(t)))*
where we used the fact that z(1) = exp("yy... (1)) + (K - 1) exp(v(1)).
Let su(t) == "y (t) v(t). Dividing the nominator and denominator of the terms on the right hand side in Equation (13) by exp(2e(t)), we may describe the dynamics of puf) through a single
differential equation:
= d
d
2K exp(p(t))
(t)
dt y
(t) vt))
dt N(exp(n(t)) + +K -K-11 1]
Equivalently, this can be written as:
d
2K
dt (() M(exp(u(t) + 2 1) (K - 1)2 exp((x(()))
(14)
It can be verified by differentiation that the solution to the above differential equation satisfies:
2K N + exp()(0)) + 2(K 1))(0) (K 1)2 exp(-x(0))
exp(p(t)) + 2(K (t( - (K - 1)² exp(-x(t)). d Notice that by Equation (14) we have de (1()) > 0, and so jut) is monotonically increasing. Since ((0) misclassifies X.. i.e. p()) this implies that IRTT the initial time which RFT classifies x, correctly- is the unique time which yut) 0. Such time must exist since de ((t) is bounded
away from zero when ((() E [n(0), OJ. It therefore follows that:
fRFT N(K 1)2
exp(-x(0)) N(K 1)a
N
N(K 11)
+
exp(p(0))
p(0)
2K
2K
2K
K
(15)
N(K- 1)²
- -1),
2K
where the inequality is due to v()) < 0 and 1 - exp(n(0)) >
Now, recalling that exp(ur. (1) z, (t) the probability for outputing the correct class, which receives reward of 1, while incorrect labels receive reward of -1, the standard deviation of the reward for x,, at time is given by:
exp(uyy... (t)) In (1)
( (1 (1_xxp(uy... exp(uyy.. (1))
STD, Perrre (ylx.))))))))))))))))))))) (x y)1
zn((1
4(K 1) exp( (Wy.. (t) t(t))
Zn (t)²
4(K -
exp(u(t)) + 2(K 1) (K 1)2 exp(-u(t))
where the first transition is due to fact that the standard deviation of random variable that is with probability q and -1 otherwise is /I- (q (I - q)}², and the last transition by dividing the nominator and denominator of the fraction by exp(wy.,.. (t) + ut)) and plugging in p(t)
Wy Wy... v(t). Specifically, at time & 0:
exp((((0) + 2(K - 1) + (K - 1)2 4(K -1)
(16)
STD yrupo (yl (y)x.)) (r V)]
23

---

Published as a conference paper at ICLR 2024
Applying exp()(0)) <1 < 2(K - 1) and squaring both sides of the inequality leads to:
<1+ (K - 1) exp(-<(0)),
STD you Potos (vix.)
which implies:
4
"STTD Potto (STDy-paoo x.,))) [x(x..))]] y)l
-1
exp(-M(0)) >
K
Combining the inequality above with Equation (15), we may conclude:
2N(K - 11
K K-1
URET 2
K
STD y-ppat (-bx,.) |r (Xn, v))
4
= Ω STD STDY-P y- ( (x,.,) (xx,)) (xxxx. v))
D.4.2 OPTIMIZATION TIME UPPER BOUND FOR SFT
Similarly to the analysis for RFT, we denote by (y.a(1) IR the (y, dy,t)th entry of OSFT(!), for E fl...., K),d d f1..... D}. and let (t) = Ey=1 exp(wy we (t)) for n' E [I.....N} Under this notation, the SFT loss is:
CSrT(OsFri
Σ
N
In (N
n'=1
exp eep((Wyeend (uy..
In
m==1
Zn (t)
Focusing on X,, = en, the parameters W (t),.))..... ..., WK (t), which determine the output distribution given X,, obey the following differential equations:
(1
exp exp(uymn. (()) (n1(1)
d dt "y -V wy. CSFT (Os(T(((()))
y
(17)
exp(uy...(1) (t))
y x Yn
N
for all y f1.....,AA K}. As in RFT, the evolution of Wi,n (t), ..... WK, (t) is independent of the maining parameters due to the orthogonality assumption on X corresponding to incorrect label y A Yn is governed by the same differential equation, and, by assumption, initially Wy.n (0) = Wy', o (0) for all y, * yn. As a result, the uniqueness of the solution to Equation (17) for an initial condition implies that (y (1) (a) for all / > 0 and y,y' * Yn-
XN- Furthermore, every Wy,n (t)
Denoting vtt) (t) for some y Yas the dynamics for the n 'tth column of 9rrt (f) therefore given by:
d
(K 1) exp(v(() Wy (1))
16 yn. n (1)
dt
N + (K - 1) (ut) Wy.. (t))]
(18)
exp(u(t) - (y..((t)
v(t) = N[1 + (K - 1) exp(v(t) - Wy (t))
dt
Let ((t) = Wy.. (f) - v(t). We can express the dynamics of p(t) through:
d
K
dt (()
dt Wy... (t)
dt v(t) = N((xp(u(()) + (K (K-1)) 1)]
for which it can be verified by differentiation that the unique solution satisfies:
K + (K 1)p(0) + exp(x(0)) = (K - 1)y(() + exp(u(t))
(19)
N
24

---

Published as a conference paper at ICLR 2024
We now in position to extract from Equation (19) upper bound on tsFT the initial time which SFT classifies Xn correctly. Notice that d (²((() > O, and pu(() is monotonically increasing.
Since ((0) misclassifies Xx, f.e. (0) O, this implies that 1SFT the unique time at which p(() = 0. Note that such time must exist since ă (u(1) is bounded away from when pu()) [n(0)) 0j.
Thus, from Equation (19) it follows that:
MI1 (K 1)n(0) - exp(p(O))
1SFT
K
(20)
N
N(K 1)
(d)).
K
K
where the inequality due to 1 exx((u(0)) < 1. Recall from Equation (16) that:
1
(p((4(0))+2(4 -1) + (K 1)2 exp(-m(0))
STD y (v/x,,
(o)x,.(((xx)) ((xx:)))
4(K -1)
KK 1) exp((x((0))
4
Taking the natural logarithm of both sides and rearranging the inequality yields:
(m( s In
(K (K-1) SDDDD p (r|x..) (rxx,)((xx, (r(x.. y)}²
Hence, combined with Equation (20), we arrive the desired result:
N N(K
- 1). In
4
1SFT
K
MTDDDDD>ppoo (xxx)) (t
(K- - STD ym (y x,, (x.x..)1((x.)))) xx)
=1( STD STDDyppaios
(r y)j
E EXAMPLES OF INPUTS WITH SMALL AND INPUTS WITH LARGE REWARD STANDARD DEVIATION UNDER THE PRETRAINED LANGUAGE MODEL
Figures 5, 6, and 7 give examples from the NarrativeQA, ToTTo, and CommonGen datasets, spectively, of inputs with small and inputs with large pretrain reward standard deviation. We also include outputs produced by the pretrained, RFT, and SFT models along with their rewards. The examples showease how: (i) the pretrained model often generates continuations for the input text instead of obliging to the task at hand; and (ūi RFT struggles to maximize the reward of inputs with a small pretrain reward standard deviation compared to those with a large pretrain reward standard deviation.
F FURTHER EXPERIMENTS AND IMPLEMENTATION DETAILS
F.1 FURTHER EXPERIMENTS
Listed below are additional experiments, omitted from Sections 4 and 5.
Vanishing gradients in the GRUE benchmark (Section 4.1):
Figure 8 supplements Figure 1, by plotting the reward mean and standard deviation of train samples, under the pretrained, RFT, and SFT models, for the remaining GRUE datasets: Com- monGen, IWSLT 2017, CNN/Daily Mail, and DailyDialog.
Table 3 supplements Table by reporting, for the remaining GRUE datasets (see above), the Pearson correlation between the reward standard deviation under the pretrained model and the absolute reward mean change due finetuning, across the train samples.
25

---

Published as a conference paper at ICLR 2024
Input
Pretrained Outputs
Outputs
SFT Outputs
Where co john and Jenny travel for be atedl haneymoon?
John has many oppo nunities for his work and a seeikty column he nevoppoorss
John and andjenny have their own bs siness and insaluible experience oditars
Rewoent
Sanwed
Grogan escape wedeing, and Jenny Michigan winters and refocate to cotage South Horids, where they are hired as repor ters for competing newsaperrs... She and John trisel treland for belated honeymmonn, young womon.. family pay respects their beloved per they bury him beneath a tree in their yoont
Ireland
They wark three newspapers. John hired for several months in paln beach
They curselves warking at Palm Beach Post and Sun- Sent und they are happlly
Ireland
He stars learn basic Fnngishh summers chilly Hlriica
They have children and begin wedky column articies competing newspaper reports
Ireland
John takes his salory and resigns as resuor ten for meveggaper
in
Ireland
johni fised colummiis for rival paper South Florida
Reference
Ireland
Deviation:
Input
e Pretrained Model Outputs
RFT Model Outputs
SFT Model Outputs
whom does
Ratties ambitious well- spoken phiannhroppe
Rafttes How Meintyee
Raffles Haw, comees to reside n Tamriill... he seeks the counsel of his fiance Laura Mcl -iyre, the one person he supposes
mysterious
Mclroyre
Lauren is marnied to Hector Sparing the son the vicar
Raffles How is engages Laura Mcnnyre
0.17
0.56
true to him. Unbekr ownstt him, Leurs has eccested his
becomnes engaged Spurling: both opposite
Raflles
marriaga previous engage ment Hector Spurling..
ending
Mcnntyra
Laura engagedd to Hector Spurting, the son of the vicar
References
The parents the Mctnoyyr
0.31
Hector Souriigg
famiy How due
He becomes engaged to Laura
Standard Deviation: 0.16
Laura Mcnntyre
Figure 5: Representative example of an input with small (top) and an input with large (bottom) reward standard deviation under the pretrained model from the Narrative A dataset.
Figure supplements Figure 2, by showing that the correlation between the pretrain reward standard deviation percentile and the reward difference between RFT and SFT is not sensitive to the choice of percentile.
Table reports the train and test rewards that the pretrained, RFT, and SFT models obtain over each of the considered GRUE datasets.
Figure 10, Table 5, and Figure 11 extend Figure 1, Table 1, and Figure 2, respectively, by considering RFT runs carried out using NLPO (Ramamurthy et al., 2023), adaptation of PPO for language generation tasks, instead of PPO.
Controlled experiments (Section 4.2):
Figure 12 supplements Figure 3 by reporting the reward standard deviation and gradient norm throughout optimization, in addition the reward.
Figure 13 presents experiments identical Figure 12, except that RFT and SFT were optimized using stochastic gradient descent (SGD) as opposed to Adam.
Figure 14 is identical to Figure 12, except that the pretrain expected reward of inputs with small pretrain reward standard deviation is 0.5 instead of -1, i.e. it is relatively high at the start of finetuning.
O Overcoming vanishing gradients in RFT (Section 5):
Table 6 compares the train and test rewards obtained by RFT with increased learning rate, temperature, and entropy regularization coefficient, to the rewards obtained by the default RFT configuration and RFT after an initial SFT phase, on the NarrativeQA dataset.
Figure 15 presents reward means and standard deviations of individual train samples from NarrativeQA (analogously to Figure 1) after applying RFT with an increased learning rate, temperature, or entropy regularization coefficient.
Table 7 supplements Table by including results for identical experiments on the ToTTo and CommonGen datasets.
26

---

Published as a conference paper at ICLR 2024
Input
Pretrained Model Ouspucs
RIF Model Guaputs
Model Outputs
Rowerd
Reward
Snverd
suaito heeeee
Lacakh pegee
Ihlan's highani
staide--page, Younis sipage ottee epection mis one Dry interational camuries stacatio, jires <table- scotl 144 scol Iicadin- Score-i of col, heacer- <ieels <hable>
cricket
0.422
0.188
of omail crickt
Knon tite> seation tileo
ndla ihan
Corporation, 18th 1897-1596..
Vounm Grsmir Khen OOI centuryy (144) is the tighest score in ODI cridkct.
of all memerablee Test centur les was hile in 144 One Day Cricker conourtes Yourt mus than.
coll 144 oot hoader Score tol bouheader ite losll rab eo page, otler Lie of mereationall cridks centurnes by
011
KP an highest 0.31 score in 146e rurs
0.099
Reference
one day for centuries is cricket. Yourss when last appeared rocmestion 1997.
Younts thanis highest score was on ODI match.
tbb cell not headers Seoee fra heaeter tral rutris pape Oiter Lit intornatoonal contu fiss page
0.027
Yourss shan's highest soove rurs.
0
score
four days AICE crisee the rournd Towetyy Twverejc
0.017
pilee International be
1966 1983.
Knen spage thte section
introwiing international cricket inrings Youn
Standard Devlation: 0.000G
Input
Pretrained Model Ouspues
REF Model Ousputs
d SFT Model Outpuis
<tadds> scallo numertis Hapertzane
Aubiionallo Seng of the year, Operaa icol header English otle [Chnese thee) lool hear er leels redl Lew Vim-
Reward
uewir Matinnar grrudacer Another Lee Yim g (stainer)
Reward
Horanttcc Rapartionn is 2015 romantic co modeyy telexsion drama pr uced by Lee "im- foeg
<ool headers English sh tite (Chinese idet ot
0.81
0.16
0.14
col hear sloets cceli fiin-trng (pridtaeerg Moses Chan, Jnyoe koy tnuhue Los. Evergy
Beergrenn Mikk, Katy ki ng jacon Chan thisen Bomantte commedy. Opera
Katy Kung,
Chi san header- Cast and crewosf col
reell-
fong torouueent Mooes Joyce
Romantic Repomoire (Chinese: Hs fong. starning
Romantic
Phillgpe Lee Yim-tong reciort
Konean
corredyy Opera stable--gggo.tee Tive
Evergneen (producent
0.42
0.43
cipage pile>
Chan Romantic cometty Opera Seveson Drama
Joyce performing
settoon tites Second
section Jillo stabler readi Bomentic Rxperto sool, header Enghsinotie tee (Chinese ritel sico header- -los ocall- Lee Wim-tong (p recuoett Mones Chan, loyon Kol, Lou se ues, Evergreen Nalk Kity Kung, jason Char Chi san scol dhoodor- Cast and crew sital hender- sicell-cee <eel- Romande corredt. Opera reol beaders Gerre ot ool heacerr studll stabbee
Alluan n 201 soction Liteo tacler odlo-
Bomsntti Sepertoire is 2015 Hong Kong remante comedy ''V drama prod uced by Lee "Imm fong teaturing Mores Chan, Joyne Kol, Louise Les, but Les, Evergroen Malk Koty Kung and
Romonic Repertoone re ool headers Frigish tite (Ch nese tht et
0:82
Les rn-fonggio (co produding: Moses Chan Joyon Koj, Lostion Lee, evergreon Mak, katy Kung Jason Chan Chi san Romantic colredt. Opera.
0.65
0.44
tab coll Romoo 2012: Ohinose
drama lyter cramm only in the U.S.A. and atrowd jptt tahlee namanti
Remantic Repertoine is 2015 Hong Kong numantis surmel telovision drama produced fong starting joyne Hartian Eveng
Vimfooy (prodacert Moces Jhan Louise Lee. Bvergreen Makk, Rung jacon chan
Reference
0.44
Remandic Repertorre a Hong Kong romandic comiccy saleviaion druma produced TVB, staring Moses Chan, Joyoe Kol. Louise Lee. Evergreen Malk, Keyy Kung and Jance Cleir.
camoddy Opera Bomanni Rep
Standard Devistion: 0.175
Figure 6: Representative example of input with small (top) and input with large (bottom) reward standard deviation under the pretrained model from the ToTTo dataset.
Figure 16 supplements Figure by reporting the same metrics based on the mean test reward, in addition to the metrics based on the mean train reward, for the partial SFT experiments on NarrativeQA.
Figure 17 compares the reward means and standard deviations of individual train samples (anal- ogously to Figure 1), for the NarrativeQA dataset, after performing a full SFT phase and after partial SFT phase with 40% of the optimization steps and 1% of the labeled samples.
Figures 18 and 19 supplement Figure 4 by including results for identical experiments on the ToTTo and CommonGen datasets, respectively.
F.2 FURTHER IMPLEMENTATION DETAILSS
We provide implementation details omitted from our experimental reports (Sections 4 and and Ap- pendix F.I). Code for reproducing results, based on the PyTorch (Paszke et al., 2017), Hug- gingFace (Wolf et al., 2019), and RLALMS (Ramamurthy et al., 2023) libraries, can be found at https: / /i thub. cituub.com/ae cwm/ apple h/mmmmmlrllaarr Each finetuning experiment was run on four Nvidia A100 80GB GPUs, and each experiment from Section 4.2 was run on single Nvidia V100 32GB GPU, except for those with MLPs for which we used a standard laptop.
We note that, since the test sets of ToTTo and CommonGen are not publicly available, we report results over their validation sets instead. This does not pose an issue our interest lies in opti- mization, iee. the ability to achieve a higher reward over the train set, and we do not conduct any hyperparameter tuning beyond taking the default values from Ramamurthy et al. (2023).
27

---

Published as a conference paper at ICLR 2024
Inport
Dusputs
Model Gucputs
Outguts
Roverd
Roward
Rowir
jeans
your snauyt juan
0.11
thei young,
joarts
Reference
Cenerate senteecee
Gene
Pains of jeans our and sewn together to make sairt
Deer asyment kcal hemn out and sown
single
0.25
Gencrons sentence with out Jean sew:
Gone ate sentence with the sentencer out jean new out
0.10
A baby a wearingg jeans and newn a in chafion out into markers.
Woran jennyy mace that fiting dress sewn.
Gene att sentence with the co cape of ourjjnn sow.
0:11
0.14
man sewing new jean cut
Standard Deviation: 0.0002
Inpurt
Ourpues
Ourputs
Outputs
Reward
ste sentence
"used siving
Gene sentence
lowing working guager
0.23
& rrowing on grees light
0.16
Reference
combine
0.24
mblie harvest gener reverss chrondaogte Coeare genee sentence with porticulor styword
Hanvesting whoat with combine
Combane harded schwet.
0:13
0.21
Wheer surflower and combina dlenn horrard
0.45
Combin ind whoae
0:13
Trachianal cante combline wheat with Goche and rabes durinn trophry harvent in winter.
0.29
Gener sit sentence with the dea combéne suther wheat and gene ate harvest whrect
0.23
Combining poonnces with chickpess for the harvest
Standard Deviation: 037
Cambere hurvest wur genee sentamue sharing cropscrraps
Figure 7: Representative example an input with small (top) an input with a relatively large (bottom) reward standard deviation under the pretrained model from the CommonGen dataset.
Table 2: The reward functions used in our experiments for datasets from the GRUE benchmark.
Dataset
Reward
NarrativeQA
ROUGE-L-Max (Lin, 2004)
ToTTo
SacreBLEU (Post, 2018)
CommonGen
METEOR (Banerjee and Lavie, 2005)
IWSIT 2017
SacreBI.EU
CNN/Daily Mail
METEOR
DailyDialog
METEOR
IMDB
Learned sentiment classifier (Ramamurthy et al., 2023)
F.2.1 EXPERIMENTS ON THE GRUE BENCIIMARK (SECTIONS 4.1 AND 5)
Modifications to the default setup from Ramamurthy et al. (2023): In our experiments with the GRUE benchmark we followed the experimental setup of Ramamurthy et al. (2023), including default hyperparameters for RFT and SFT, up to the following adjustments.
For the NarrativeQA dataset, Ramamurthy et al. (2023) use an SFT model that was finetuned multiple question answering datasets. In contrast, finetune only NarrativeQA for fairer comparison with RFT, which does not have access to the additional datasets.
For the IMDB dataset, Ramamurthy et al. (2023) perform SFT using positively labeled train samples only, whereas RFT performed over the whole training set. For a fairer comparison, we use for RFT the same train samples used for SFT.
For the GommonGen dataset, we do not apply repeat penalty to the reward function.
Reward functions: The GRUE benchmark suggests a few reward functions for each dataset, from which we used those specified in Table 2.
28

---

Published as a conference paper at ICLR 2024
Pretrained Model
RFT Model
SFT Model
Dataset: CommonGen
Reward Mear
2.0
0.6
2.13
0-3
36-
06-
22
02
0.03
80
1993
2000
2009
1003
5350
20
1000
2003
30.00
4000
5000
00
1000
2000
1050
4000
5009
Sample Order (by pretrain reward
Sample Order (by pretrain reward std)
Sample Order protrain reward std)
Dataset: 2017
Reward Mean
6-3
2.0
0.0
9.40
Newword 910
12
02
1993
2006
2000
1003
5050
1000
2003
3000
1000
5003
1000
2000
3000
1000
5000
Sample Order (by pretrain roward sta)
Sample Ordor (by pretrain reward std)
Sample Order (yy protrain reward std]
Dataset: CNN /Dailtyyyy Maii
1.
16
Roward Mean
Moooon $10
14-
04
2.10
12
5000
2003
3000
4000
Sample Order (by pretrain reword
Sample Order (by protra reward std)
Sample protroin reward std)
Dataset: DailyDialog
3.20
Reward Mean
9.15
Reward S1D
0--
38-
0.-
30
00
3000
2.00
Sample Order (by
Sample Ordor
std)
Figure 8: Inputs with small reward standard deviation under the pretrained model, ie with vanishing expected gradient, prevalent in the GRUE benchmark. This figure supplements Figure 1, by including identical plots for the remaining GRUE datasets. further details see caption of Figure
F.2.2 CONTROLLED EXPERIMENTS (SECTION 4.2)
Models: We used the default ResNet18 and BERT-mini implementations from the PyTorch and HuggingFace libraries, respectively. The MLP had two hidden layers, of sizes 250 and 100, with ReLU non-linearity.
Data: To facilitate more efficient experimentation, we subsampled the MNIST and CIFARIO train- ing sets, choosing uniformly at random 1000 samples from MNIST and 10,000 from CIFARIO 10. For pretraining, we assigned 4 additional labels for each sample in MNIST and CIFARIO, i.e. each sample had a total of 5 pretraining labels, while each sample in STS-B B assigned 2 additional labels, i.e. a total of pretraining labels. During finetuning, every train sample was assigned a sin- gle, randomly chosen, ground truth label. For 10% of the samples the label was chosen such that was not included in their pretraining labels. As a result, their pretrain reward standard deviation was small. The rest of the samples were assigned a label such that it was one of their pretraining labels. Meaning, their pretrain reward standard deviation was relatively large.
In the experiments of Figures 3, 12, and 14, the cross-entropy loss was minimized via the Adam optimizer with learning rate 0.0001, default 31, 31,92 coefficients, and a batch size of 512. Pretraining and finetuning were carried out for 1000 and 5000 epochs, respectively. For Figure 13 we used an identical setup, except that Adam was replaced by SGD with learning rate 0.01 for the MLP and ResNet18 models. To improve optimization stability, for BERT-mini learning rate to 0.001. Varying the learning rate (for both experiments with Adam and SGD) did not yield noticeable improvement in the ability of RFT to maximize the reward of train samples with small pretrain reward standard deviation.
Optimization:
reduced the
29

---

Published as a conference paper at ICLR 2024
Table 3: RFT affects inputs with small reward standard deviation under the pretrained model less than it affects other inputs. This table supplements Table by reporting, for the remaining GRUE datasets, the Pearson corre- lation between the pretrain reward standard deviation and the absolute reward mean change due to finetuning, across the train samples. As anticipated, the correlation is considerably higher for RFT compared SFT. For further details see caption Table
CommonGien
IWSLT 2017
CNN/Daily Mail
DailyDialog
RFT
0.18
0.51
0.33
0.33
SFT
0.07
0.37
0.23
0.21
RFT SFT Roward
SFT Rowera
-01
Than
IWST 2017 IMDS
Tain
WSLT 2017
Train
INSLT 2017
That
Text
ChAlDaly Mall
Text
CNNuaall Mall
Com
CHilitally Mail
DailyClebog
NerrativaQy
TeT'
DetyCheicg
MDD
RFF
Narrali wage
T:
DaltyCii og
IND0
To
CommonGeen
-84
Commontion
0.00
d
155 Pretrain Reward STD Percencie
5thh Pritrain Reward STD Parcencil
th Pretram Reward STD Percentile
Figure 9: RFT performance (relative to SFT) is worse when inputs with small reward standard deviation are prevalent. This figure supplements Figure 2, by including identical plots with different choices for the mea- sured pretrain reward standard deviation percentile: 1%, 5%, and 15% (instead of 10%). As be seen, the correlation between the pretrain reward standard deviation percentile and the reward difference between RFT SFT is not sensitive to reasonable choices percentile. For further details see caption Figure
Table 4: Mean reward attained by pretrained, RFT. SFT models over datasets. Means and standard deviations are taken over five different random seeds per model dataset. Since the test sets of ToTTo and CommonGen are not publicly available, reported are the rewards over their validation sets instead. This does not pose an issue as interest lies in optimization. ie. the ability to achieve a higher reward over the train set, and we do not conduct any hyperparameter tuning beyond taking the default values from Ramamurthy et al. (2023).
considered GRUE
Train Reward
Test Reward
Pretrained Model
RFT Model
SET Model
Pretrained Model
RET Model
SFT Model
NarratieQA
0.110 + 0.000
0.102
0.012
0.493
+ 0.003
0.117 + 0.000
0.116
0.002
0.537
0.001
ToTTo
0.024 # 0.000
0.098 + 0.014
0.305
+ 0.005
0.036 + 0.000
u.144
0.019
0.443 # 0.003
Commoncien
0.130 + 0, 0.000
0.163
0.007
0:224
1: 0.000
0.179
0.000
0. 206
0.005
0.308
0.001
2017
0.239
0.005
0.307
0.337
CNN/aail Mail
= 0.000
0.260
0.011
+ 0.000
0.255
0.000
0.010
0.308
0.000
DailyDialog
0.191 i 0.000
0.224
0.003
0.189 1: 0.005
0.190 1 0.001
0.222 0.004
0.180 1: 0.065
IMDB
0.831 +: 0.001
0.890 1 0.005
0.860 1 0.001
0.198 1: 0.002
0.565 0.009
0.547 1: 0.002
Table 5: RFT (using NLPO) affects inputs with small reward standard deviation under the pretrained model less than it affects other inputs. This table supplements Tables and 3, by including the Pearson correlations for RFT performed via NLPO (Ramamurthy al., 2023) instead of PPO. Similarly RFT with PPO, the correlation between pretrain reward standard deviation of train sample its absolute reward mean change RFT with NLPO considerably higher than SFT. For further details see caption Table
NarrativeQA
ToTTo
IMDB
CommonGen
IWSLT 2017
CNN/Daily Mail
DailyDialgg
RFT (NLPO)
0.42
0.48
0.71
0.19
0.51
0.34
0.29
RFT
0.48
0.46
0.72
0.18
0.51
0.33
0.33
SFT
0.05
0.72
0.07
0.37
0.23
0.21
30

---

Published as a conference paper at ICLR 2024
Pretrained Model
RFT (NLPO) Model
SFT Model
Datoset: NarrativeQA
Rewenrt Meon
2.0
0.00
1090
2000
3006
4000
5000
2909
2006
s0ee
4030
300
3009
2000
3000
4093
5000
Sample Order (by pretrain reward std)
Sample pretrain rewerd
Sample Order (by pretrain reward std)
Detaset: ToT
Dawarr Mean
0.29
53-
62-
4230
1000
2000
dry
Sample
Sample (by
Dataset: IMD8
Becord
std)
Sample pretrain
(by std)
Dataset: CommonGen
815
Daward Meain
55-
2.-
0.00
3003
88
4530
2000
3000
8.00
Sample Order Iby pretrain reward shd)
Sample Order pritrain reward std)
Sample Order (ay pretrain reward std)
Dataset: IWSLT 2017
1.0
Rorourd Mean
0.40
250
1802
00
3090
sces
0.00
Sample Order Iby pretrain reward sta)
Sample Order Iby pretrain reward std)
Sample Order (by pretrain reward stdl
Dataset: CNN/DDally Mail
Roward Moan
0.20
0-19
2000
2006
1993
a5no
sees
2000
2000
3008
4250
5000
2009
3069
4300
3000
5.00
Sample Order Iby pretrain reward stol
Sample Order (by pritrain reward st.)
Sample Order (by pretrain reward stal
Dataset: DallyDlalog
620
Deoware Mean
34
5.00
2000
2008
1003
2000
2204
2000
42.0
300
2000
2000
4200
2000
Sample Order Iby pretrain reward sto)
Sample Order fby pretrain reward stdl
Sample Order Dby pretrain reward stdl
Figure 10: Inputs with small reward standard deviation under the pretrained model, f.e. with vanishing pected gradient, are prevalent the GRUE benchmark. This figure is identical to Figures that NLPO (Ramamurthy et al., 2023) is used instead of PPO for RFT. The pretrained model (left) and SFT model (right) columns are repeated here for case of comparison. Notice that, the effect of RFT with NLPO on the rewards of train samples is qualitatively similar to that of RFT with PPO. For further details see caption of Figure 1.
except
31

---

Published as a conference paper at ICLR 2024
RFT (NLPO) SFT Reward
0.0
-0.1
Train
IWSLT 2017
CNN/Daily
Narrative@A
DailyDiateg
ToTTo
IMDB
CommonGen
0.00
10'th Pretrain Reward STD Percentile
Figure RFT (using NLPO) performance (relative SFT) is worse when inputs with small reward standard deviation are prevalent. This figure is identical to Figure 2, except that NLPO (Ramamurthy et al., 2023) is used instead of PPO for RFT. For further details see caption of Figure 2.
Model: MLP Dataset: MNIST
Model: ResNet18
Dataset: CIFAR10
Model: BERT-mini
Dataset: STS-B
10
Reward
small pretrain reward
0.0
0.0
-0.5
ange protram reward sta
-05-
-6.5
2000
3020
4000
5000
1000
3000
5000
1000
2000
4003
Epech
Epoch
Spach
RFT
Reward STD
0.8
arge pretzein reward
pretrain
5.6
0.6
0.4-
2.4
0.2
2000
3000
5000
1000
6000
2000
3000
Epech
Spoch
RF1
Gradient Norm
1.-
protrain reward arge pretrain reware std
100
75
1.000
2000
3200
3000
10000
2000
30000
4000
5000
2000
3000
5000
Spach
Epoch
Figure 12: RFT struggles to maximize the reward over inputs with small reward standard deviation under the pretrained model, inputs with vanishing expected gradient, even with perfect exploration. On the contrary, SFT easily leads to maximal reward. This figure supplements Figure by reporting ation and gradient norm, addition to the reward (separately for train samples with small and train samples with large pretrain reward standard deviation). The top row is identical to Figure 3, and is repeated here for convenience. Notice that, in accordance with Theorem the gradient norm under RFT vanishes when the reward standard deviation is small, whereas under SFT the gradient near maximal. For further details see caption Figure 3.
reward standard devi-
vanishes only when the reward
32

---

Published as a conference paper at ICLR 2024
Model: MLP Dataset: MNIST
Model: ResNet18
Dataset: CIFAR10
Model: BERT-mini Dataset: STS-B
10-
1.0
Reward
0.0-
Small pretrein reward std
SFT
0.0
0.0-
-0.5
... large pretrain reward Mid
-05
-0.5
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
1000
2000
3000
4000
5000
spoch
Epoch
Epocch
1.0
RFTT SFT
0.8
1.0
Rewerd STD
0.6
small pretrain reward std pretrain
0.6
0.6
0.4
0.2
0.2
1020
2000
3000
4000
5000
1000
2000
2000
40:00
5000
2000
2003
5000
4000
5000
Epoch
Esoch
Epoch
12-
RFT SFT
100
Gradient Norm
30
pretzein rewant std
large protrain reward std
Hel
40
20-
1000
2000
3000
60C0
5000
0
1000 2000 3000
4000
5000
10200
2000
1000
4000
5000
Epach
Epoch
Epoch
Figure 13: RFT struggles maximize the reward over inputs with small reward standard deviation under the pretrained model, Le. inputs with vanishing expected gradient, even with perfect exploration (using stochatic gradient descent). On the contrary. SFT easily leads to maximal reward. This ligure is identical to Figure 12, except that RFT and SFT were optimized using stochastic gradient descent (SGD) as opposed to Adam (note that in RFT we use SGD to minimize the negative expected reward). As can be seen, the inability of RFT to maximize the rewards of inputs with small reward standard deviation instead of Adam. For further details see caption of Figure 12.
more when using SGD
Table 6: Conventional heuristics are inadequate for overcoming vanishing gradients in RFT. For the Narra- tiveQA dataset, which was shown Section 4.1 to suffer from vanishing gradients. reported are means and standard deviations of train and test rewards, taken over three random seeds, with best result in bold. While the common heuristics of increasing the learning rate, temperature, and entropy regularization coefficient are unable to improve upon the default RFT configuration, an initial SFT phase is highly effective. See Table 7 in Appendix F.1 for analogous experiments the ToTTo and CommonGen datasets.
Dataset: NarrativeQA
Train Reward
Test Reward
RFT' SIT RFT
0.101 + 0.009
0.116 + 0.000
0.537 0.005
0.544 1 0.003
RFT with learning rate 2 10-5
0.012 + 0.010
0.020 + 0.017
RFT with learning rale 10-4
0.053 0.010
0.048 1 0.016
RFT with leamning rate 2 10-3
0.039 + 0.018
0.012 + 0.020
RFT with temperature 1.5
0.077 + 0.021
0.118 + 0.002
RFT with temperature 2
0.060 + 0.018
0.104 + 0.009
RFT with temperature 2.5
0.044 0.004
0.088 1 0.009
RFT with entropy regularization 0.01
0.080 0.006
0.113 1 0.002
RFT with entropy regularization 0.1
0.024 0.005
0.019 1 0.007
RFT with entropy regularization
0.011 + 0.000
0.013 + 0.010
'itt default hyperparameters: learning rate 2 10 10-6, temperature 1, cntropy regularization 0
33

---

Published as a conference paper at ICLR 2024
Model: Dataset: MNIST
Model: ResNet18 Dataset: CIFAR10
Model: BERT-mini
1.00
1.0
0.8
0.6
0:55
RFT
0.4
0.00
reward
-0.59
large pretrain reward std
1003
Epech
IFT
0.5
SFT
0.8
0.8
Reward
0.5-
pretraln reward
0.6
0.4
0.4
0.0
0.0
0.0
2020
1000
3000
Epoch
RFT
30-
SPT
std
40-
Gradient
lange
30
20
Epoch
RFT struggles to maximize the reward over inputs with small reward standard deviation under the pretrained model, i.e. inputs with vanishing expected gradient, even with perfect exploration and when their expected reward under the pretrained model is relatively high. On the contrary, SFT casily leads maximal reward. This figure is identical Figure 12, except that inputs with small pretrain reward standard deviation (i.e. inputs for which the finctuning label was not included in their pretraining labels), the reward of an incorrect label is set to 0.5 as opposed to -1. Consequently, the expected reward of such inputs at the beginning of finetuning is 0.5, while it is -1 in the experiments of Figure 12. Nonetheless, RFT struggles to increase their reward, with the reward dynamics mimicking those from Figure 12. For further details see caption of Figure 12.
Figure 14:
With Increased Leaming Rate
Temperature
Entropy Regularization
Reporrt STD
Sample Order Dty pretriais roward sUl)
Sample Driar Iby pewrra..reeadd st2)
Samppl Order dby protrain roward skll
Sample Order prosain roward st29
Figure 15: Conventional heuristics are inadequate for overcoming vanishing gradients in RFT. For the same randomly chosen NarrativeQA train samples from Figure 1. presented are the reward means and standard devia- tions (estimated based on ten generations per input) under the pretrained model and RFT models optimized with increased learning rate, temperature, or entropy regularization coefficient. The learning rate, temperature, and entropy regularization values with which the plots were created those achieving the highest train reward (excluding the default values), as reported in Table Notice that the effect of these RFT runs is qualitatively similar that of the default RFT configuration (shown in Figure 1). Meaning, the heuristics do not alleviate the difficulty of RFT to improve the reward of inputs with small pretrain reward standard deviation.
34

---

Published as a conference paper at ICLR 2024
Table 7: Conventional heuristics inadequate for overcoming vanishing gradients in RFT. This table supple- ments Table 6 by including results for the ToTTo and CommonGen datasets. Reported means and standard deviations train and test rewards, taken over three different random seeds, with best result in bold. While the common heuristies of increasing the lcarning rate, temperature, and entropy regularization coefficient are unable to improve upon the default RFT configuration, an initial SFT phase is highly effective.
Dataset: ToTTo
Train Reward
Validation Reward
RFT" SFT RFT
0.093
0.026
0.137
0.037
0.342 0.001
0.439 1 0.001
RFT with learning rate 10-6
0.063
0.005
0.131 1 0.003
RFT with learning rate 2 10
0.000 + 0.000
0.000 + 0.000
RFT with learning
10-3
0.000 0.000
0.000 0.000
RFT with temperature 1.5
0.059 0.007
0.096 # 0.022
RFT with temperature 2
0.043 0.004
0.073 1 0.005
RFT with temperature 2.5
0.019 + 0.001
0.038 + 0.002
RFT with entropy regularization 0.01
0.061 1: 0.013
0.139 1 0.014
RFT with entropy regularization 0.1
0.000 0.000
0.018 0.015
RFT with entropy regularization 1
0.000 1 0.000
0.000 1 0.000
"With default hyperaarameters: learning rate
6, temperature 1, entropy regularization
Dataset: CommonGen
Train Reward
Validation Reward
RFT" SFT RFT
0.164 +: 0.004
0.209 1 0.005
0.265 +: 0.001
0.319 1 0.008
RFT with learning 2 10
0.098 0.019
0.134 + 0.015
RFT learning rate 2 10-4
0.112 0.094
0.157 + 0.072
RFT with learning rate 2 10
0.055 +: 0.020
0.106 1: 0.048
RFT with temperature 1.5
0.139 + 0.004
0.190 + 0.017
RFT with temperature 2
0.147 0.004
0.189 1: 0.000
RFT with temperature 2.5
0.140 + 0.007
0.182 + 0.006
RFT entropy regularization
0.009 0.000
0.083 0.020
RFT with entropy regularization
0.002 + 0.000
0.004 + 0.003
RFT with entropy regularization 10
0.001 0.000
0.000 + 0.000
WWith default hyperparatnteers: learning rate 2 10 temperature 1, entropy regularization 0.01
35

---

Published as a conference paper at ICLR 2024
RFT After Partial SFT Reward
RFT After Partial SFT Reward - Partial SFT Reward
RFT Alter Full SFT Reward 0.536
|RFT Reward 0.092] Fretrain Rewerd 110]]
% of SFT Ostiizztion Steps
103
0.96
1.01
1.01
1.01
1.02
1.00
6.5
5.7
5.7
5.2
na
0.95
1.00
1.00
0.99
0.99
1.01
7-2
62
64
5.9
6.0
6.6
Train
60 60
0.96
0.99
0.99
0.98
0.98
0.99
9.8
8
8.0
8.0
7.7
8.2
0.80 96
0.96
96
0.95
0.96
96
13.2
13.0
11.8
11
12.1
12.0
% SFT
0.90
0.86
0.91
0.86
0.91
0.90
18.6
17.1
18.6
17.4
18.7
18.6
0.60
0.63
0.62
0.60
0.59
0.63
11.1
12.2
11.9
11.4
10.9
12.1
20
60
60
80
100
40
60
80
100
%
SFT Samp
%
SFT Samples
RFT After Partial SFT Reward Full
RFT After Partial SFT Reward Partial SFT Reward
RFT
% of SFT Optimization Steps
[RFT
Pretrain
40 6c en 100
0.96
1.00
1.01
1.00
1.01
1.00
18.3
20.5
10.4
10.3
11.0
9.3
% Optimization
0.95
0.99
1.00
0.99
1.00
1.00
19.4
13.1
14.2
12.2
13.5
13.8
Test
0.96
0.99
0.99
0.98
0.98
0.99
28.6
20.3
19.9
19.2
18.5
20.3
0.95
0.96
0.96
0.96
0.96
0.97
40.1
39.0
34.1
33.2
34.3
34.6
0.91
0.90
0.92
71.0
66.8
69.0
66.1
69.7
52.4
56.1
50.8
20
60
80
100
60
100
Samp
Samples
Figure 16: datasets which suffers from vanishing gradients, few initial optimization steps on a small number of labeled imputs substantially boost the efficacy of RFT. This figure supplements Figure by reporting metrics based on the mean test reward, in addition to on the mean train reward. The top row is identical to Figure repeated here for case of comparison.
Pretrained Model
SFT Model: 1% Samples, 40% Steps
SFT Model: 100% Samples, 100% Steps
Reward Mean
0.8
08-
0.20
0.30
02
5.0
6.0
3.0
0.00
Sample
(by
Sample
Figure 17: A few initial SFT optimization steps small number labeled inputs effectively mitigate vanish- ing gradients the reward means and standard deviations (estimated based on ten generations per input) under the pretrained, a partial SPT, and the full (default) SFT models. The pretrained model (left) and SFT model (right) plots are identical to the respective plots in Figure 1, and are repeated here for ease of comparison. Observe that the number of samples with small reward standard deviation, i.c. with vanishing expected gradient, is considerably reduced after partial SFT. We highlight that, despite the reward mean obtained by the partial SFT model being noticeably lower than that obtained by the full SET model, after performing RFT the gap is closed almost completely. That is, as shown in Figure 4, RFT over the partial SFT model reaches roughly the same reward as RFT over the full SFT model. This highlights the effectiveness of partial initial SFT phase for addressing vanishing gradients in RFT.
the same randomly chosen NarrativeQA train samples from Figure 1, presented are
36

---

Published as a conference paper at ICLR 2024
RFT After Partial SFT Reward
RFT After Partial SFT Reward Partial SFT Reward [RFT Reward
RFT Full
% SFT Optimization Steps
Pretrain
0.024)
100
0.81
0.96
1.01
1.01
1.01
1.00
0.1
0
0.8
0.9
1.1
1.1
60
9.80
0.96
1.00
0.99
1.00
1.00
9.1
0
10
1.0
1.0
1.2
Train
62
0.80
0.97
0.98
0.98
0.98
0.98
0.1
0.8
09
1.0
1.1
1.1
40
0.80
0.94
0.94
0.97
0.96
0.95
-0.0
0.8
0.9
1.2
111
1.1
0.73
0.89
11
0.1
%
20
60
80
60
100
SFT Samp
Samples
RFT After Partial Reward
RFT After Partial SFT Reward Partial SFT Reward
0.439)
(=
(= 0.037]
% SFT Optimization Steps
100
0.76
0.96
1.00
1.00
0.99
1.00
0.1
0.0
0.0
0.0
0.0
0.1
% Optimization
8B
0.76
0.96
0.99
0.99
0.99
0.99
0.1
-0.0
0.1
-0.0
0.1
0.1
Validation
60
0.75
0.97
0.97
0.97
0.97
0.98
0.0
0.0
-0.0
40.0
-0.0
0.76
0.96
-0.0
20
0.75
0.1
60
100
Samp
Samples
Figure 18: On datasets which RFT suffers from vanishing gradients, a few initial SFT optimization steps on a small number of labeled inputs substantially boost the efficacy of RFT. This figure is identical to Figure 16, except that the experiments were carried out over the ToTTo dataset (as opposed to NarrativeQA).
RFT After Partial SFT Reward
RFT After Partial SFT Reward Partial SFT Reward
RFT Full
0.266
% of SFT Optimization Steps
|RFT
Pretrain
150
0.91
0.99
1.00
1.00
0.99
1.00
1.
2.1
2.2
2.2
2.1
2.2
0.91
0.97
1.00
0.98
0.98
0.98
1.3
2.1
2.3
2.2
2,2
2,2
Train
40 60
0.92
0.96
0.97
0.97
0.97
0.97
1.6
2
2.3
2.3
2.4
2.4
0.91
0.95
0.94
0.95
0.96
0.96
3.3
2.5
2.4
2.5
2.6
2.6
0:90
0.90
0.91
0.91
0.91
0.90
2.5
2.6
2.7
2.7
2.7
0.86
0.86
%
26
60
80
100
60
105
SFT Samp
Samples
RFT After Partial SET Reward
RFT After Partial SFT Reward Partial SET Reward
0.320)
0.179]
% SFT Opúmization Steps
60 60 100
0.94
1.00
1.01
0.99
0.99
1.00
3.6
0
0.8
0.7
0.6
0.8
Validation
% Optimization
0.93
0.98
0.99
0.99
0.98
0.97
0.5
0
0.8
0.7
0.6
0.5
0.95
0.97
0.97
0.97
1.00
0.99
0.7
0
07
0.7
1.0
0.9
96
0.97
0.97
0
0.34
60
60 Samples
100
Samp
Figure 19: On datasets which RFT suffers from vanishing gradients, a few initial SFT optimization steps on a small number of labeled inputs substantially boost the efficacy of RFT. This figure is identical to Figure 16, except that the experiments were carried out over the CommonGen dataset (as opposed to NarrativeQA).
37


Reviews:
[{"review_id": "S9dFbrCZjX", "reviewer": "Reviewer_PKCD", "summary": "In this paper, the authors identify an issue in the reinforcement finetuning phase of language models. Specifically, they first prove that the gradient norm of the value function is upper bounded by the standard deviation of the reward scores over the input sequences. They then empirically demonstrate on three datasets that small reward standard deviations are prevalent in the GRUE benchmark. Together, they suggest that the RLHF paradigm (i.e. preference reward + PPO) struggles to improve the quality on massive samples with small reward deviation. Finally, they suggest that simply using an initial supervised fine-tuning phase will greatly help mitigate the problem.\n\nOverall, the motivation of this paper is very well justified and supported. In particular, they match the theory and practice quite well in a current fast-evolving paradigm. In addition, it provides a simple and effective solution to the discovered problem. I believe the paper could be potentially impactful in RLHF for language models.\n\nMy concerns are framed as questions below.", "questions": "1. Is there a potential explanation of why the reward model generates small standard deviations on the GRUE benchmark?\n2. Could you list all the assumptions used in your proofs explicitly and discuss their realisticity?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The motivation of this paper is very well justified and supported. In particular, they match the theory and practice quite well in a current fast-evolving paradigm. In addition, it provides a simple and effective solution to the discovered problem. I believe the paper could be potentially impactful in RLHF for language models.", "weaknesses": "My concerns are framed as questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oAEVwY2L2F", "reviewer": "Reviewer_4NxS", "summary": "The paper focuses on the challenge of vanishing gradients in the context of reinforcement finetuning (RFT) of pretrained language models. The authors identify that the expected gradient for an input tends to vanish when its reward standard deviation is small, regardless of the expected reward's optimality. This phenomenon can hinder the training process. The paper suggests a solution by introducing an initial supervised finetuning (SFT) phase where a small percentage of input samples may be needed.", "questions": "Q1: Consider simple reward scaling with gradient clipping, value function learning in [1], and reward normalization in [2]. Such tricks could help the model to escape the area when it is not exactly flat but the standard deviation is small since the exact value of the reward does not matter, but the relative one matters. I wonder whether such tricks could address the problem.\n\nQ2. It could be easy to imagine that if the reward is trained on a dataset that is significantly better than the pretrained model, all outputs from the pretrained model will be bad and the reward std would be small. If we start fine-tuning with the SFT model, some of the outputs would be better after fine-tuning and the reward std would be larger. However, if the reward model is trained with human evaluation on the output of the pretrained model itself, it would be more likely that the reward std would be larger. I wonder whether the usefulness of SFT initialization is specific to some reward models trained on specific datasets.\n\n[1] DPOK: Reinforcement Learning for Fine-tuning Text-to-Image Diffusion Models\nYing Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, Kimin Lee\n\n[2] Training Diffusion Models with Reinforcement Learning\nKevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, Sergey Levine", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper pinpoints a specific, nuanced problem in the domain of reinforcement learning finetuning of language models, i.e., vanishing gradients due to small reward standard deviation.\n2. The proposed theory for small standard deviation leading to a small gradient norm is pretty intuitive, but the proof seems novel.", "weaknesses": "1. It is well-known that policy gradient could be suboptimal when the reward is pretty flat even though it's not the optimal solution. However, specific to the proposed scenario where the standard deviation is small, there could be many naive remedies not mentioned in the paper and it remains questionable whether the proposed problem could be easily solved by them. See  Q1. \n2. If we use a learned reward model, whether SFT initialization is helpful should be highly related to the dataset that the reward model is trained on, and there seems to be no discussion on this part. See Q2. \n3. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MD54xwadXy", "reviewer": "Reviewer_xsTc", "summary": "This paper studies the vanishing gradients in RLHF. Specifically, the authors argue that shrinkage in the standard deviation (std) of the reward would lead to vanishing gradients, leading to sub-optimal performance. The authors first show that norm of the gradient is upper bounded by the std of the reward plus some other factors, suggesting that as std goes to zero, training would stall due to vanishing gradients. On GRUE benchmark, the authors show that this is indeed the case; there are a substantial number of examples that have very small std under the pre-trained model. They show that RFT impacts these examples less than others; SFT, on the other hand, doesn't have this problem as ground truth labels force proper signal propagation. Using conventional heuristics such as larger learning rate, tuning temperature or entropy regularization is not effective in solving this problem. But, a few steps of SFT is shown to be effective in improving RFT performance.", "questions": "1. Can you please clarify the above example in the context of theorem-1?\n\n2. Can you discuss when and why small vanishing gradient is a problem?\n\n3. Can you discuss more on ADAM update and present more results?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper investigates a very important problem. It is well written with relevant experiments to support the claims.", "weaknesses": "Some parts of the paper need clarification.\n\n1. Theorem-1 applies more broadly beyond autoregressive models. The only part that is relevant is the Jacobian which is easily separated from the rest using the triangle inequality in the proof. Take a simple non-autoregressive example with action space A ($|A|=2$), a linear policy $\\theta^T x$ such that $p_\\theta(y=0|x)=\\theta_0 x_0$, $p_\\theta(y=1|x)=\\theta_1 x_1$, and $r(x, y=0)=r(x, y=1)=1$. The standard deviation should be zero and lead to vanishing gradients. But, the gradient is $\\nabla_\\theta(x; \\theta)=\\theta_0 x_0 [x_0, 0]/(\\theta_0x_0) + \\theta_1 x_1 [0, x_1]/(\\theta_1x_1)=[x_0, x_1]$ which doesn't necessarily have zero gradient norm. Can you clarify if the gradient estimation in this example is incorrect or how I should interpret this result?\n\n2. 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.\n\n3. 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?\n\n4. Please clarify O-notation in Appendix D.4 w.r.t. concerned variable. Otherwise, it suggests negative value inside O or instant learning when STD goes to 1 in SFT.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hdzzAzaQ7a", "reviewer": "Reviewer_voez", "summary": "This manuscript identifies a problem in policy-gradient-based Reinforcement Learning Fine-Tuning (RLFT): when the reward variance is small, the language model (LM) could be susceptible to vanishing gradients.  The authors provided both theoretical and empirical evidence to support their argument, and demonstrated that a supervised fine-tuning \"jumpstart\" could help alleviate the vanishing-gradient issue.", "questions": "Please refer to the \"Weaknesses\" section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The authors identified an interesting issue (although rather self-explanatory), and advocated for a simple and straightforward solution.\n\n- I admire the organization of this manuscript -- very clean and structured.\n- The authors included a comprehensive set of empirical results, most of which are carefully presented and scrupulously discussed.", "weaknesses": "- The writing has room for improvement.  The authors tend to make heavy use of long sentences that are difficult to parse.\n- 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.\n- 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*.\n- 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.)\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)