Supervised Fine Tuning on Curated Data is Reinforcement Learning (and can be improved) — part 2/4
Cleaned-text code points [15691, 31563). Bibliography omitted here; source.md includes it.
We can derive a lower bound on this RL objective via a well known inequality [ Kober and Peters, 2008 , Roux, 2016a ] . Assuming strictly positive rewards, we can utilize the inequality x ≥ 1 + log ( x ) x\geq 1+\log(x) applied to the importance weights, to show that (for more details see Appendix A ): J ( θ ) ≥ ∫ ℬ π ref ( τ ) log p ( τ , θ ) R ( τ ) 𝑑 τ + cst = 𝔼 π ref ( τ ) [ R ( τ ) log p ( τ , θ ) ] + cst , \displaystyle J(\theta)\geq\int_{\mathcal{B}}\pi_{\mathrm{ref}}(\tau)\log p(\tau;\theta)R(\tau)\mathrm{d}\tau+\mathrm{cst}=\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}\Big[R(\tau)\log p(\tau;\theta)\Big]+\mathrm{cst}, (4) where cst \mathrm{cst} denotes constant terms independent of θ \theta . This is a well known objective in the RL literature [ Dayan and Hinton, 1997 , Peters and Schaal, 2007 , Kober and Peters, 2008 , Roux, 2016a ] often referred to as reward weighted regression. We can relate this objective to supervised learning in two cases (i) we assume access to a dataset of filtered “good data” (e.g. a natural setting in cases where we have success indicators assigned by humans); (ii) we additionally assume access to quality scores for all “good data” (i.e. we assume each τ \tau can be given a quality score on an ordinal scale). For case (i) we can find a relation by assuming binary, sparse, rewards in the RL objective. Then the return simplifies to an indicator function, R ( τ ) = 𝕀 ( S ( τ ) > 0 ) R(\tau)=\mathbb{I}(S(\tau)>0) (where S S is some score function indicating success), collapsing Equation 4 to the maximum likelihood supervised fine tuning objective: J ( θ ) ≥ 𝔼 π ref ( τ ) [ 𝕀 ( S ( τ ) > 0 ) log p ( τ , θ ) ] = c ref 𝔼 τ ∈ 𝒟 + [ log p ( τ , θ ) ] ⏟ 𝒥 SFT ( θ ) , \displaystyle J(\theta)\geq\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}\Big[\mathbb{I}(S(\tau)>0)\log p(\tau;\theta)\Big]=c_{\mathrm{ref}}\underbrace{\mathbb{E}_{\tau\in\mathcal{D}^{+}}\Big[\log p(\tau;\theta)\Big]}_{\mathcal{J}_{\text{SFT}}(\theta)}, (5) where 𝒟 + = { τ i | τ i ∼ π ref ( τ ) , 𝕀 ( S ( τ i ) > 0 ) } i = 1 N \mathcal{D}^{+}=\{\tau_{i}|\tau_{i}\sim\pi_{\mathrm{ref}}(\tau),\mathbb{I}(S(\tau_{i})>0)\}_{i=1}^{N} is a dataset of trajectories sampled from π ref \pi_{\textrm{ref}} and filtered by the binary reward indicator and c ref = 𝔼 π ref ( τ ) [ 𝕀 ( S ( τ ) > 0 ) ] c_{\mathrm{ref}}=\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}[\mathbb{I}(S(\tau)>0)] is constant independent of θ \theta . And we start optimization of J ( θ ) J(\theta) by initializing to the reference distribution, setting π θ 0 = π ref ( τ ) \pi_{\theta_{0}}=\pi_{\mathrm{ref}}(\tau) . If π ref \pi_{\mathrm{ref}} is not available in closed form we approximate it by first learning it via standard behavior cloning on the full dataset 𝒟 \mathcal{D} . Thus we can interpret SFT on filtered data as maximizing a lower bound on the RL objective .
Analogously, for case (ii) we can assume rewards R ( τ ) = ∑ i = 0 C 𝕀 ( S ( τ ) > c i ) R(\tau)=\sum_{i=0}^{C}\mathbb{I}(S(\tau)>c_{i}) , with C C (ordinal) quality classes. Then RL simplifies to maximum likelihood of trajectories proportional to quality: J ( θ ) ≥ 𝔼 π ref ( τ ) [ ∑ i = 0 C 𝕀 ( S ( τ ) > c i ) log p ( τ , θ ) ] = c ref Q 𝔼 τ ∈ 𝒟 Q + [ log p ( τ , θ ) ] ⏟ 𝒥 SFT(Q) ( θ ) , \displaystyle J(\theta)\geq\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}\Big[\sum_{i=0}^{C}\mathbb{I}(S(\tau)>c_{i})\log p(\tau;\theta)\Big]=c_{\mathrm{ref}}^{Q}\underbrace{\mathbb{E}_{\tau\in\mathcal{D}_{Q}^{+}}\Big[\log p(\tau;\theta)\Big]}_{\mathcal{J}_{\text{SFT(Q)}}(\theta)}, (6) where 𝒟 Q + = 𝒟 c 1 + ∪ ⋯ ∪ 𝒟 c C + \mathcal{D}_{Q}^{+}=\mathcal{D}_{c_{1}}^{+}\cup\dots\cup\mathcal{D}_{c_{C}}^{+} with 𝒟 c + = { τ i | τ i ∼ π ref ( τ ) , 𝕀 ( S ( τ i ) > c i ) } i = 1 N \mathcal{D}_{c}^{+}=\{\tau_{i}|\tau_{i}\sim\pi_{\mathrm{ref}}(\tau),\mathbb{I}(S(\tau_{i})>c_{i})\}_{i=1}^{N} . Note c ref Q = 𝔼 π ref ( τ ) [ ∑ i = 0 C 𝕀 ( S ( τ ) > c i ) ] c_{\mathrm{ref}}^{Q}=\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}[\sum_{i=0}^{C}\mathbb{I}(S(\tau)>c_{i})] . And we have the relation ∀ i : c i > c i − 1 \forall i:c_{i}>c_{i-1} . Hence, we can interpret SFT over quality curated data as maximizing a lower bound on the RL objective.
It is important to realize, however, that both lower bounds, become looser throughout training, as the distance between p ( τ , θ ) p(\tau;\theta) and π ref ( τ ) \pi_{\mathrm{ref}}(\tau) increases. This connection between RL via reward weighted regression and maximum likelihood is not new [ Roux, 2016a , Roux, 2016b , Ferbach et al., 2024 ] but we believe it is underexplored in the context of SFT and RL for large models.
### 3.3 Improving SFT from the perspective of RL
The bound SFT optimizes becomes looser as we train, thus weakening the connection to RL. We show that by re-framing the RL problem we can introduce a tighter bound that only requires a minimal change to the maximum likelihood objective used in SFT and can achieve better results.
To derive this variant we first introduce an auxiliary distribution q ( τ ) q(\tau) that we are free to choose. And we will choose it later to trade off the tension between staying close to the reference policy for stable optimization and tightening the bound. Using this auxiliary we can re-write the RL objective as: J ( θ ) = ∫ ℬ π ref ( τ ) q ( τ ) π ref ( τ ) p ( τ , θ ) q ( τ ) R ( τ ) 𝑑 τ . \displaystyle J(\theta)=\int_{\mathcal{B}}\pi_{\mathrm{ref}}(\tau)\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}\frac{p(\tau;\theta)}{q(\tau)}R(\tau)\mathrm{d}\tau. (7) By applying the inequality x ≥ 1 + log ( x ) x\geq 1+\log(x) to x = p ( τ , θ ) / q ( τ ) x=\nicefrac{{p(\tau;\theta)}}{{q(\tau)}} , as previously considered in [ Roux, 2016b , Kober and Peters, 2008 ] we can obtain an analogous bound to ( 4 ): J ( θ ) \displaystyle J(\theta) ≥ ∫ ℬ π ref ( τ ) q ( τ ) π ref ( τ ) log p ( τ , θ ) R ( τ ) 𝑑 τ + cst = 𝔼 π ref ( τ ) [ q ( τ ) π ref ( τ ) R ( τ ) log p ( τ , θ ) ] + cst . \displaystyle\geq\int_{\mathcal{B}}\pi_{\mathrm{ref}}(\tau)\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}\log p(\tau;\theta)R(\tau)\mathrm{d}\tau+\mathrm{cst}=\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}\Big[\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}R(\tau)\log p(\tau;\theta)\Big]+\mathrm{cst}. For simplicity we drop the constant (wlog). If we assume sparse rewards as above we can further turn this into a training objective over filtered data that results in an importance weighted version of SFT: J ( θ ) \displaystyle J(\theta) ≥ 𝔼 π ref ( τ ) [ q ( τ ) π ref ( τ ) 𝕀 ( S ( τ ) ) log p ( τ , θ ) ] = c ref 𝔼 τ ∈ 𝒟 + [ q ( τ ) π ref ( τ ) log p ( τ , θ ) ] ⏟ 𝒥 iw-SFT , \displaystyle\geq\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}\Big[\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}\mathbb{I}(S(\tau))\log p(\tau;\theta)\Big]=c_{\mathrm{ref}}\underbrace{\mathbb{E}_{\tau\in\mathcal{D}^{+}}\Big[\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}\log p(\tau;\theta)\Big]}_{\mathcal{J}_{\text{iw-SFT}}}, (8) where 𝒟 + = { τ i | τ i ∼ π ref ( τ ) , 𝕀 ( S ( τ i ) ) > 0 } i = 1 N \mathcal{D}^{+}=\{\tau_{i}|\tau_{i}\sim\pi_{\mathrm{ref}}(\tau),\mathbb{I}(S(\tau_{i}))>0\}_{i=1}^{N} and c ref = 𝔼 π ref ( τ ) [ 𝕀 ( S ( τ ) > 0 ) ] c_{\mathrm{ref}}=\mathbb{E}_{\pi_{\mathrm{ref}}(\tau)}[\mathbb{I}(S(\tau)>0)] . Here, the auxiliary distribution q ( τ ) q(\tau) is now used to importance weight the examples in the dataset. And an analogous bound for quality scored data 𝒟 Q + = 𝒟 1 + ∪ ⋯ ∪ 𝒟 C + \mathcal{D}_{Q}^{+}=\mathcal{D}_{1}^{+}\cup\dots\cup\mathcal{D}_{C}^{+} as in Eq. ( 6 ) can be found: J ( θ ) \displaystyle J(\theta) ≥ c ref Q 𝒥 iw-SFT(Q) = 𝔼 τ ∈ 𝒟 Q + [ q ( τ ) π ref ( τ ) log p ( τ , θ ) ] . \displaystyle\geq c_{\mathrm{ref}}^{Q}\mathcal{J}_{\text{iw-SFT(Q)}}=\mathbb{E}_{\tau\in\mathcal{D}_{Q}^{+}}\Big[\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}\log p(\tau;\theta)\Big]. (9) Critically, both objectives have the property that as q ( τ ) → p ( τ , θ ) q(\tau)\rightarrow p(\tau;\theta) we have 𝒥 iw-SFT ( θ ) → J ( θ ) \mathcal{J}_{\text{iw-SFT}}(\theta)\rightarrow J(\theta) . Namely, as q q gets closer to our policy distribution, the bound tightens. However, at the same time, tightening the bound will increase variance in the importance weights and it is thus not generally safe to just set q ( τ ) = p ( τ , θ ) q(\tau)=p(\tau,\theta) . That being said, and in contrast to standard importance weighting for policy gradients [ Jiang and Li, 2016 , Metelli et al., 2018 ] , we are free to choose q q to smooth out the bound. We can thus side-step common problems with large variance in the importance weights/blow-up by enforcing importance weights to not be too extreme. Optimizing 𝒥 iwSFT \mathcal{J}_{\text{iwSFT}} or 𝒥 iwSFT(Q) \mathcal{J}_{\text{iwSFT(Q)}} then can be achieved by sampling from the filtered data distribution and computing importance weights on the fly (requiring one extra reference model to be kept in memory) as shown in Algorithm 1 in the appendix.
### 3.4 The importance of choosing q ( τ ) q(\tau)
As mentioned, we are free to choose q ( τ ) q(\tau) in our method. We would generally like to achieve q ( τ ) → p ( τ , θ ) q(\tau)\rightarrow p(\tau;\theta) over the course of training so that the bound becomes tight. However, as q ( τ ) q(\tau) moves away from π ref \pi_{\textrm{ref}} the variance of the bound (and thus the variance of its gradient) increases. As a result we want to bound the KL divergence K L ( q ( τ ) ∥ π ref ( τ ) ) < β KL(q(\tau)\|\pi_{\textrm{ref}}(\tau))<\beta but need to trade this off with tightening the bound; a well known problem when dealing with importance sampling based objectives.
In contrast to standard off-policy policy gradient algorithms, choosing q ( τ ) q(\tau) such that it strikes this trade-off (e.g. via clipping or other means as discussed below) will not bias our algorithm–but may nonetheless lead to different optima and thus better (or worse) policies. In practice, we choose q ( τ ) = p ( s 0 ) ∏ t = 0 T − 1 p ( s t + 1 | s t , a t ) π ( a t | s t ; θ q ) \displaystyle q(\tau)=p(s_{0})\prod_{t=0}^{T-1}p(s_{t+1}|s_{t},a_{t})\pi(a_{t}|s_{t};\theta_{q}) (10) where θ q → θ \theta_{q}\rightarrow\theta towards the end of training such that q ( τ ) → p π ( τ , θ ) q(\tau)\rightarrow p_{\pi}(\tau;\theta) .
For controlling the variance, we consider two different variants of bounding the difference between q ( τ ) q(\tau) and π ref \pi_{\textrm{ref}} by controlling the importance weights. In either variant we set θ q \theta_{q} to be a time-lagged version of θ \theta , the parameters we are optimizing (see experiments section for details) and then constrain the importance weights to control the evolution of p ( τ , θ ) p(\tau,\theta) during training (and hence indirectly q ( θ ) q(\theta) ). We note that for each variant we always consider importance weights over full trajectories .
#### Clipping per step importance weights
A simple way to control the importance weights is to simply bound the per step importance weights to be within a certain range. That is we choose: q ( τ ) = p ( s 0 ) ∏ t = 0 T − 1 ( p ( s t + 1 | s t , a t ) clip ( π ( a t | s t ; θ ) , α min , α max ) CLOSE q(\tau)=p(s_{0})\prod_{t=0}^{T-1}(p(s_{t+1}|s_{t},a_{t})\text{clip}(\pi(a_{t}|s_{t};\theta),\alpha_{\text{min}},\alpha_{\text{max}}) , where we typically set α min \alpha_{\text{min}} and α max \alpha_{\text{max}} fixed, i.e. in the language model experiments we will choose ( α min , α max ) = ( 1 − 0.8 , 1 + 0.8 ) (\alpha_{\text{min}},\alpha_{\max})=(1-0.8,1+0.8) . We will then do an overall clipping of the importance weight: clip ( q ( τ ) / π ref ( τ ) , β min , β max ) \text{clip}(q(\tau)/\pi_{\text{ref}}(\tau),\beta_{\min},\beta_{\max}) .
#### Smoothing importance weights
As a second, more general, option we consider ensuring low variance by smoothing the importance weights on the trajectory level. In particular we consider transforming the weights as: q ( τ ) π ref ( τ ) = exp ( ∑ i g ( ρ i ) ) , ρ i = log ( π q ( a i | s i ) ) − log ( π ref ( a i | s i ) ) \displaystyle\frac{q(\tau)}{\pi_{\mathrm{ref}}(\tau)}=\exp\left(\sum_{i}g(\rho_{i})\right),\quad\rho_{i}=\log(\pi_{q}(a_{i}|s_{i}))-\log(\pi_{\mathrm{ref}}(a_{i}|s_{i})) (11) Here, i i denotes the i i -th token in the sequence τ \tau , and g ( x ) g(x) is a monotonically increasing function with respect to x x . We use a simple temperature function for continuous control experiments 1 1 1 But the above clipping can also be expressed using this formulation as g ( x ) = clip ( x , min , max ) g(x)=\mathrm{clip}(x,\min,\max) . : g ( x ) = k x g(x)=kx , where k k is a hyperparameter, for all our experiments in continuous control domains. By setting k = 1 / / ∑ i M i k=\nicefrac{{1}}{{/}}{\sum_{i}M_{i}} (where M M is the token mask for a transformer LLM) we can implement using the averaged log importance weight. We instead use k = α ∑ i M i k=\alpha\sum_{i}M_{i} which gives us an extra degree of freedom to smooth out the importance weights. This way we can smoothly move from no importance weighting ( α → 0 \alpha\rightarrow 0 ) to more and more “peaked” weights. We ablate the choice of α \alpha in the supplementary material.
### 3.5 A Toy Example to build Intuition for SFT on curated/filtered data
We consider a toy example to get an intuition for when the bound optimized via SFT is problematic–and thus when the importance weighted version can be beneficial. We construct this toy example to answer the following questions: Q1: Does the loose lower bound lead to suboptimal policies? Q2: What is the effect of the importance weighting? Can it help learn better policies? Q3: What are the failure modes of the importance weighted variant, and how can they be avoided?
Consider a simple bandit optimization problem in which an agent has a choice of two actions 𝒜 = { pull-left , pull-right } \mathcal{A}=\{\text{pull-left},\text{pull-right}\} . The episode terminates after picking a single action. R ( τ ) R(\tau) is either 0 or 1; the probability distribution of R ( τ ) R(\tau) is defined by: ℙ ( R ( τ ) = 1 ) = { 0.5 if τ = pull-left 1 if τ = pull-right \displaystyle\mathbb{P}(R(\tau)=1)=\left\{\begin{array}[]{l l}0.5&\text{if }\tau=\text{pull-left}\\ 1&\text{if }\tau=\text{pull-right}\end{array}\right.
Thus the optimal strategy is to always pull the right arm. We assume that we can observe the behavior of the system only by watching (collecting data from) a reference policy π ref ( a ) = 𝒰 { pull-left , pull-right } \pi_{\mathrm{ref}}(a)=\mathcal{U}\{\text{pull-left},\text{pull-right}\} that picks between the two arms uniformly at random.
We collect a dataset 𝒟 \mathcal{D} using this policy which will contain 50 % 50\% pull-left and 50 % 50\% pull-right actions. Filtering this dataset by the binary success reward to produce 𝒟 + = { τ i | τ i ∼ π ref ( τ ) , R ( τ i ) = 1 } i = 1 N \mathcal{D}^{+}=\{\tau_{i}|\tau_{i}\sim\pi_{\mathrm{ref}}(\tau),R(\tau_{i})=1\}_{i=1}^{N} then results in a dataset with 2 times more pull-right than pull-left actions. Subsequently training via SFT on this data results in a policy π SFT \pi_{\text{SFT}} that chooses pull-left one third of the time and pull-right two thirds of the time. The resulting average reward is 5 / 6 5/6 . See Figure 1 for a numerical experiment. This policy improves on the reference policy, but is also not optimal. The key problem here is that the SFT procedure cannot incorporate knowledge from failures. For Q1 : while SFT on curated data can be seen as optimizing a lower bound on the RL objective and will improve on the reference policy it can result in sub-optimal behavior.