Training Diffusion with RL

**RLHF** 将决策当成二分类问题，人认为action1比action2更好的概率以及他们对应reward之间的关系为: $$ p\left(a_{1} \succ a_{2} \mid s\right)=\frac{\exp \left(r\left(s, a_{1}\right)\right)}{\exp \left(r\left(s, a_{1}\right)\right)+\exp \left(r\left(s, a_{2}\right)\right)} $$ 通过大量的HF，我们可以用一个网络从人类偏好的数据集当中拟合出reward model，优化的loss如下： $$ \mathcal{L}_{R}\left(r_{\phi}, \mathcal{D}\right)=-\mathbb{E}_{\left(x, y_{w}, y_{l}\right) \sim \mathcal{D}}\left[\log \sigma\left(r_{\phi}\left(x, y_{w}\right)-r_{\phi}\left(x, y_{l}\right)\right)\right] $$ 其中$x$表示状态$s$,$y$表示动作$a$,$\sigma$是sigmoid函数，$y_{w}$是人类认为更好的数据，$y_{l}$是人类认为更差的数据。当获取到reward model之后，对baseline进行finetune的公式为: $$ \max _{\pi_{\theta}} \mathbb{E}_{x \sim \mathcal{D}, y \sim \pi_{\theta}(y \mid x)}\left[r_{\phi}(x, y)\right]-\beta \mathbb{D}_{\mathrm{KL}}\left[\pi_{\theta}(y \mid x) \| \pi_{\mathrm{ref}}(y \mid x)\right] $$ 其中ref是原始模型，加KL散度是为了避免OOD。 **DPO** > https://arxiv.org/abs/2305.18290 对于上面的公式，给定reward function $r(x,y)$,这个优化目标的最优解为: $$ \pi^*(y|x) = \frac{1}{Z(x)}\pi_\text{ref}(y|x)\exp(\frac{1}{\beta}r(x,y) ) $$ 因此$r(x,y)$可以被表示为 $$ r(x,y)=\beta \log\frac{\pi^*(y|x)}{\pi_\text{ref}(y|x)}+\beta \log Z(x) $$ 则偏好的公式可以写成： $$ p^{*}\left(y_{1} \succ y_{2} \mid x\right)=\frac{1}{1+\exp \left(\beta \log \frac{\pi^*\left(y_{2} \mid x\right)}{\pi_{\mathrm{ref}}\left(y_{2} \mid x\right)}-\beta \log \frac{\pi^*\left(y_{1} \mid x\right)}{\pi_{\mathrm{ref}}\left(y_{1} \mid x\right)}\right)} $$ 我们要优化的目标从学习reward model变成了优化$\pi_\theta(y|x)$，此时进行finetune的优化目标为 $$ L_{DPO}(\pi_\theta;\pi_\text{ref}) = -E_{(x,y_w,y_l)}[\log \sigma(\beta\log\frac{\pi_\theta(y_w|x) }{\pi_\text{ref}(y_w|x)}-\beta\log\frac{\pi_\theta(y_l|x) }{\pi_\text{ref}(y_l|x)} )] $$ 此时可以发现不再需要reward model即可开始训练，最终的效果与RLHF的效果没有很大的差别。 **Training diffusion with rl** Inference的时候，$x_{t-1}$与$x_{t}$的关系为: $$ p_\theta(x_{t-1}|x_t,c) = N(x_{t-1}|\mu_\theta(x_t,t,c),\sigma^2_t\bf{I}) $$ 如果看成MDP过程，则有如下对应关系: $$ s_t\doteq (c,t,x_t)\\ \pi(a_t|s_t)\doteq p_\theta(x_{t-1}|x_t,c)\\ a_t \doteq x_{t-1}\\ R(s_t,a_t) \doteq r(x_0,c) \;\;\text{if t = 0 else 0} $$ 其中$r(x_0,c)$可以是用LLAVA对$x_0$与$c$之间匹配程度的打分作为reward,或者其他可压缩性等指标作为reward 如果使用偏好作为reward，则训练的loss可以使用上述的DPOloss，此时需要对输出的图片进行人类偏好选择即可。  **Theory** 对于MDP问题，此时不能只最大化 $r(x,y)$，而是需要考虑期望回报$Q(x,y)$。对于价值学习的方法来说，如果不考虑epsilon greedy探索，RL的优化目标可以表示为 $$\max_\pi\mathbb{E}_{x\sim D,y\sim \pi(y|x)}[Q_*(x,y)]$$ 其中$Q_*(x,y)$表示当前reward下的最优动作价值函数。最终学习到的策略为: $$\pi(y|x)={1 \;if\; y =argmax_y[Q_*(x,y)]\; else \; 0}$$ 为了考虑更新之后模型与原始模型之间的OOD问题，我们定义新的RL优化目标为: \begin{align*} &\max_\pi\mathbb{E}_{x\sim d^\pi,y\sim \pi(y|x)}[Q_*(x,y)] - \beta \mathbb{D}_{KL}[\pi(y|x)\|\pi_{ref}(y|x)] \\ &= \max_\pi\mathbb{E}_{x\sim d^\pi,y\sim \pi(y|x)}[Q_*(x,y)-\beta\log\frac{\pi(y|x)}{\pi_{ref}(y|x)}] \\ & = \min_\pi\mathbb{E}_{x\sim d^\pi,y\sim \pi(y|x)}[\log\frac{\pi(y|x)}{\pi_{ref}(y|x)}-\frac{1}{\beta}Q_*(x,y)] \\ & = \min_\pi\mathbb{E}_{x\sim d^\pi,y\sim \pi(y|x)}[\log\frac{\pi(y|x)}{\pi_{ref}(y|x)\exp(\frac{1}{\beta}Q_*(x,y))}]\\ & = \min_\pi\mathbb{E}_{x\sim d^\pi}[\mathbb{D}_{KL}[\pi(y|x)\|\tilde \pi(y|x)]] \end{align*} 其中 $\tilde \pi(y|x) = \pi_{ref}(y|x)\exp(\frac{1}{\beta}Q_*(x,y))$. 优化上述目标的最终结果为 $$ \pi(y|x) =\tilde \pi(y|x) = \pi_{ref}(y|x)\exp(\frac{1}{\beta}Q_*(x,y)) $$ 所以可得$Q_*(x,y) = \beta\log\frac{ \pi(y|x)}{\pi_{ref}(y|x)}$. 对于任意一个给定的reward model $r(x,y)$，在去噪过程中环境的动力学我们已知为： $$ p(s_{t+1}|s_t,a_t) = 1 \;\; if \;\;s_{t+1}=a_t \;\;else\;\; 0 $$ 可以通过值迭代算法或者策略迭代算法计算出$Q_*(x,y)$,具体的关系式为$Q^{*}(s_t, a_t)=r(s_t, a_t)+\gamma \sum_{s_{t+1} \subset S} P\left(s_{t+1} \mid s_t, a_t\right) \max _{a^{\prime} \subset A} Q^{*}\left(s_{t+1}, a^{\prime}\right)$，因此可以把给定reward function $r$前提下的$Q_*(x,y)$记为$Q_r(x,y)$。对于HF问题，我们把人类的policy当成是最优的policy，因此人类的policy为$\pi_*(y|x)$, 这个策略对应的reward model记作最优，即r=r*. 此时r*对应的$Q$为 $$ Q_{r_*}(x,y) = \beta\log\frac{\pi_*(y|x)}{\pi_{ref}(y|x)} $$ 对于一段片段$\sigma$,基于人类偏好的公式可以写为 $$ p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)=\frac{\exp \sum_{t=k}^{M}\left(r^*\left(s_t^1, a_{t}^1\right)\right)}{\exp \sum_{t=k}^{M}\left(r^*\left(s_t^1, a_{t}^1\right)\right)+\exp \sum_{t=k}^{M}\left(r^*\left(s_t^0, a_{t}^0\right)\right)} $$ 其中$\sigma=\{s_k,a_k,s_{k+1},a_{k+1},...,s_M,a_M\}$表示MDP过程的轨迹，在实验中即为从高斯噪声开始到生成图像结束的总片段。 ***Assumption:*** > 1. 由于$\mathbb{E}[\sum_{t=k}^{}\left(r^*\left(S_t, A_{t}\right)\right)]=Q_{r^*}(s_k,a_k)$, 我们近似认为$\sum_{t=k}^{}\left(r^*\left(s_t, a_{t}\right)\right) \approx Q_{r^*}(s_k,a_k)$ > 2. 对于两个片段来说，如果生成出来的最终图像A比B更好，则A片段比B片段好 ***Loss*** $$ \tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)=\frac{\exp \left(Q_{r^*}\left(s_k^1, a_{k}^1\right)\right)}{\exp \left(Q_{r^*}\left(s_k^1, a_{k}^1\right)\right)+\exp \left(Q_{r^*}\left(s_k^0, a_{k}^0\right)\right)} $$ 利用*Assumption 1* 让 $\tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)$近似代替$p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)$,使用$\pi_\theta(y|x)$代替$\pi^*(y|x)$作为待优化的目标，结合DPO的公式以及给定片段人类偏好的公式，最终优化的loss为： $$ L_{DPO}(\pi_\theta;\pi_\text{ref}) = -E_{(x\sim d^{\pi_{\theta}},\sigma_w\sim \pi_{\theta},\sigma_l\sim \pi_{\theta})}[\log \sigma(\beta\log\frac{\pi_\theta(a_k^w|s_k^w) }{\pi_\text{ref}(a_k^w|s_k^w)}-\beta\log\frac{\pi_\theta(a_k^l|s_k^l) }{\pi_\text{ref}(a_k^l|s_k^l)} )] $$ Like TRPO, 当$\pi_\theta和\pi_{ref}$之间的差距不大时,我们认为状态访问分布变化很小，即$d^{\pi_\theta}\approx d^{\pi_{ref}}$. 我们同样使用重要性采样的方法基于$\pi_{ref}$来更新$\pi_{\theta}$，此时Loss可以表示为: $$ L_{DPO}(\pi_\theta;\pi_\text{ref}) = -E_{(x\sim d^{\pi_{ref}},\sigma_w\sim \pi_{ref},\sigma_l\sim \pi_{ref})}[\frac{\prod_{t=k}^M\pi_\theta(a_t^w|s_t^w)\pi_\theta(a_t^l|s_t^l)}{\prod_{t=k}^M\pi_{ref}(a_t^w|s_t^w)\pi_{ref}(a_t^l|s_t^l)}\log \sigma(\beta\log\frac{\pi_\theta(a_k^w|s_k^w) }{\pi_\text{ref}(a_k^w|s_k^w)}-\beta\log\frac{\pi_\theta(a_k^l|s_k^l) }{\pi_\text{ref}(a_k^l|s_k^l)} )] $$ 为了保证$\pi_\theta和\pi_{ref}$之间的差距不要太大，我们可以采取PPO的clip方法将$\frac{\prod_{t=k}^M\pi_\theta(a_t^w|s_t^w)\pi_\theta(a_t^l|s_t^l)}{\prod_{t=k}^M\pi_{ref}(a_t^w|s_t^w)\pi_{ref}(a_t^l|s_t^l)}$ 约束到$[1-\epsilon,1+\epsilon]$之间，然而这种方式需要预先把轨迹当中所有的概率都计算出来，比较费时。如果进行进一步近似，即认为这个值近似为1, 此时我们得到简化的loss: $$ L_{DPO_{simple}}(\pi_\theta;\pi_\text{ref}) = -E_{(x\sim d^{\pi_{ref}},\sigma_w\sim \pi_{ref},\sigma_l\sim \pi_{ref})}[\log \sigma(\beta\log\frac{\pi_\theta(a_k^w|s_k^w) }{\pi_\text{ref}(a_k^w|s_k^w)}-\beta\log\frac{\pi_\theta(a_k^l|s_k^l) }{\pi_\text{ref}(a_k^l|s_k^l)} )] $$ 在实验当中我们选择了后者作为loss。我们每一轮采样少量的数据的方式更新$\pi_\theta$使得其与每一轮的$\pi_{ref}$尽量相似，并且选择$\epsilon=0.1$保证上述近似为1的假设是合理的。 ***Increase sample utilization*** 对于两个片段来说，我们只能通过最后一个图片进行打分，认为生成图像好的片段中所有的action都是好的，而差的图片对应的片段中所有的action都是差的。然而使用上述的loss只能更新$\pi_\theta(a_k,s_k)$,在实验当中由于轨迹是从开始去噪到结束去噪，因此只更新了k=0情况下的概率，并没有更新片段当中其他的概率。利用*Assumption 2*，我们可以取出任意包含最后图像的连续子片段，两个子片段之间仍然满足偏好关系。因此对于总序列长度为$M$的好片段$\sigma^w$和差片段$\sigma^l$，我们分别定义$M$条子序列: \begin{align} \sigma_i^w &= \{s_i^w,a_i^w,s_{i+1}^w,a_{i+1}^w,...s_{M-1}^w,a_{M-1}^w\} \;\;\;\; 0 \le i \le M-1 \\ \sigma_i^l &= \{s_i^l,a_i^l,s_{i+1}^l,a_{i+1}^l,...s_{M-1}^l,a_{M-1}^l\} \;\;\;\; 0 \le i \le M-1 \end{align} 此时对于$0\le i\le M$,片段$\sigma_i^w$和$\sigma_i^l$对应的DPO loss为:  $$ L_{DPO_{simple}}(\pi_\theta;\pi_\text{ref}) = -E_{(x\sim d^{\pi_{ref}},\sigma_w\sim \pi_{ref},\sigma_l\sim \pi_{ref})}[\log \sigma(\beta\log\frac{\pi_\theta(a_i^w|s_i^w) }{\pi_\text{ref}(a_i^w|s_i^w)}-\beta\log\frac{\pi_\theta(a_i^l|s_i^l) }{\pi_\text{ref}(a_i^l|s_i^l)} )] $$ 此时一段序列可以更新$M$次，大大提高了样本的利用效率. **Experiment** > 实验目前是使用的脚本代替人进行偏好选择，因此有标准的reward项使用上述的方法训练diffusion，认为亮的图片更好，最终训练的结果为: <div style="text-align:center"> <img src="https://hackmd.io/_uploads/S1RMLClh3.png" width="200" height="200"> </div> 其中reward是计算img.mean(),在实际训练当中并不存在reward model。上述的训练lr=3e-4,训练后完全看不出狗的形状，模型被改变的很彻底；当lr=1e-4时的结果为： <div style="text-align:center"> <img src=https://hackmd.io/_uploads/BkM4_Cgh2.png width="200" height="200"> </div> 此时模型仍然保留着原本给定prompt生成的能力，只不过会趋向于生成亮色的狗如果认为红色的图片更好，最终训练的结果为: <div style="text-align:center"> <img src=https://hackmd.io/_uploads/rkyfuAe3h.png width="200" height="200"> </div> 如果认为压缩后的图像大小越大越好，最终的训练结果如下： <div style="text-align:center"> <img src=https://hackmd.io/_uploads/SkHtIzQ3n.png width="200" height="200"> </div> <div style="text-align:center"> <img src=https://hackmd.io/_uploads/rJP0pIQ3n.png width="200" height="2enj00"> </div> 可以看到上述的方法是有效的，可以实现根据人类的反馈进行打分从而使得模型朝着我们想要的方向进行finetune。 **Advantage** 无需学习reward model，避免了使用网络拟合reward出现的一些问题。并且这种方式生成的图像完全是由人类偏好进行选择的，不会出现生成奇怪的图像 **Drawback** 需要模型生成一些图片然后再给人类进行偏好选择，非常费时。 **使用期望替代具体值的合理性分析** **Lemma**: Using assumption 1, for $i\in\{0,1\}$,suppose the expected return satisfy normal distribution, i.e, $\sum_{t}^{}r^*\left(s_t^i, a_{t}^i\right)\sim N(Q_{r^*}(s_0^i,a_0^i), \sigma^2)$. Given $Q_{r^*}^i(s,a)\in[Q_\min,Q_\max]$, then $P(|p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)-\tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)|< \frac{(\xi^2+1)(\exp(\sigma^2)-1)}{16\xi \delta})>1-\delta$ where $\xi=\frac{\exp Q_\max}{\exp Q_\min}.$ **Proof**:For simplicity, we define $Q_i=Q_*(s_0^i,a^i_0),X_i=\sum_{t}^{}r^*\left(s_t^i, a_{t}^i\right) \quad i\in\{0,1\}$. $E[p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)]=\frac{E[\exp(X_1)]}{E[\exp(X_1)+\exp(X_0)}=\frac{\exp(Q_1+1/2\sigma)}{\exp(Q_1+1/2\sigma)+\exp(Q_0+1/2\sigma)}=\frac{\exp(Q_1)}{\exp(Q_1)+\exp(Q_0)}=E[\tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)],\\\\ E[(p\left(\sigma_{1} \succ \sigma_{0} \mid s\right))^2]=\frac{E[\exp(2X_1)]}{E[\exp(2X_1)]+E[\exp(2X_0)]+E[2\exp(X_0)\exp(X_1)]}=\frac{\exp(2Q_1+2\sigma^2)}{\exp(2Q_1+2\sigma^2)+\exp(2Q_0+2\sigma^2)+\exp(Q_0+Q_1+\sigma^2)}\\\\= \frac{\exp(2Q_1+\sigma^2)}{\exp(2Q_1+\sigma^2)+\exp(2Q_0+\sigma^2)+2\exp(Q_0+Q_1)}\\\\ Var[p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)]=E[(p\left(\sigma_{1} \succ \sigma_{0} \mid s\right))^2]-(E[p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)])^2=\\\\\frac{2\exp(3Q_1+Q_0)(\exp(\sigma^2)-1)}{[\exp(2Q_1+\sigma^2)+\exp(2Q_0+\sigma^2)+2\exp(Q_0+Q_1)][\exp(Q_1)+\exp(Q_0)]^2} \;\;\le\frac{2\exp(3Q_1+Q_0)(\exp(\sigma^2)-1)}{[\exp(Q_1)+\exp(Q_0)]^4}$ Similarly we have $Var[p\left(\sigma_{0} \succ \sigma_{1} \mid s\right)]\le\frac{2\exp(Q_1+3Q_0)(\exp(\sigma^2)-1)}{[\exp(Q_1)+\exp(Q_0)]^4}$ Note that $Var[p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)]=Var[1-p\left(\sigma_{0} \succ \sigma_{1} \mid s\right)]=Var[p\left(\sigma_{0} \succ \sigma_{1} \mid s\right)]$，considering these two inequalities, we have: $Var[p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)]\le \frac{[\exp(Q_1+3Q_0)+\exp(Q_0+3Q_1)](\exp(\sigma^2)-1)}{[\exp(Q_1)+\exp(Q_0)]^4}\le\frac{[\exp(Q_1+3Q_0)+\exp(Q_0+3Q_1)](\exp(\sigma^2)-1)}{16[\exp(2Q_1)\exp(2Q_0)]}\\\\=\frac{[\exp(Q_0-Q_1)+\exp(Q_1-Q_0)](\exp(\sigma^2)-1)}{16}\le\frac{(\xi+\frac{1}{\xi})(\exp(\sigma^2)-1)}{16}.$ By using the Chebyshev inequality, we can obtaine that $P(|p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)-\tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)< t)>1-\frac{(\xi^2+1)(\exp(\sigma^2)-1)}{16\xi t}.$ We choose $t=\frac{(\xi^2+1)(\exp(\sigma^2)-1)}{16\xi \delta}$ so that $P(|p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)-\tilde p\left(\sigma_{1} \succ \sigma_{0} \mid s\right)|< \frac{(\xi^2+1)(\exp(\sigma^2)-1)}{16\xi \delta})>1-\delta$