Fine-tuning policies with self-generated data improves learning speed
Fine-Tuning on Self-Generated and Reward-Weighted Data: Learning Dynamics, Convergence Rates, and Benefits of Off-Policyness
Machine Learning
Summary
This work looks at how computer programs can improve themselves by practicing, using their own attempts and feedback to learn better strategies. The authors studied a method called RE(S) that changes how often these programs update their approach during learning. They found that updating less frequently can still lead to the best result and sometimes helps avoid getting stuck with poor strategies. This finding challenges the usual belief that frequent updates are always best and helps explain successes in training large language models.
What this means in practice
- •For machine learning engineers: Optimize training schedules for reinforcement learning models by tuning update frequency to improve convergence without sacrificing final performance.
- •For language model trainers: Guide fine-tuning strategies for post-training large language models using off-policy reward-guided self-training to accelerate learning.
Tested on simulated data.
Authors
Zhiwei Wang, Yanxi Chen, Yaliang Li, Bolin Ding
Abstract
We study the learning dynamics of fine-tuning a policy model on self-generated and reward-weighted data, with particular focus on a generalized version of REINFORCE -- referred to as RE(S) -- that updates the rollout distribution once every $S \ge 1$ gradient steps. Prior work in bandits and reinforcement learning has developed rich theory for policy gradient methods, and on-policy sampling (i.e., a small $S$, ideally $1$) is often viewed as crucial to their success; yet in prominent application like post-training large language models, reward-guided self-training has proved to be effective even when the rollout distribution is updated infrequently, but theoretical understanding remains limited for the convergence properties of these off-policy methods. To bridge these gaps, we develop a unified theory for RE(S) that covers the full spectrum of $S \ge 1$: it can be interpreted as a stage-wise optimization process, where each stage takes $S$ gradient steps for minimizing the Kullback-Leibler distance to a fixed reward-weighted rollout distribution. For multi-arm bandits with softmax policies, our in-depth analysis and numerical experiments reveal three key findings: (1) for any fixed $S$, RE(S) enjoys global convergence to the optimal policy as the number of rollout distribution updates $B = \lfloor T / S \rfloor \rightarrow \infty$, where $T$ denotes the number of gradient steps; (2) we prove tight two-sided bounds showing that the suboptimality gap of RE(S) achieves an asymptotic $Θ(1 / T)$ convergence rate, while $S$ only affects the length of a burn-in phase; (3) when initialized at a weak policy with a small optimal-action probability, RE(1) gets trapped around suboptimal policies for a long period, whereas RE(S) with a suitable $S$ avoids the detour and achieves significantly faster convergence to the global optimum, highlighting the benefits of off-policyness in this case.