155. Vanishing Gradients in Reinforcement Finetuning of Language Models
Invited abstract in session WD-3: Optimization in neural architectures II, stream Optimization in neural architectures: convergence and solution characterization.
Wednesday, 11:25 - 12:40Room: M:J
Authors (first author is the speaker)
| 1. | Noam Razin
|
| Computer Science, Tel Aviv University |
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. In this talk, I will present a recent work identifying a fundamental optimization obstacle in RFT. Namely, 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 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 of inputs whose expected gradient vanishes, as measured by the reward standard deviation, is crucial for the successful execution of RFT.
Keywords
- Artificial intelligence based optimization methods and appl
- Data driven optimization
- Large- and Huge-scale optimization
Status: accepted
Back to the list of papers