Diffusion-LM模型 - 论文阅读

Title: Diffusion-LM Improves Controllable Text Generation

Date: 2022/5/27

Paper: http://arxiv.org/abs/2205.14217

Code: https://github.com/XiangLi1999/Diffusion-LM

这篇论文针对的问题是在不重新训练模型的的情况下，可以多语言模型（Language Model）的行为控制，即可控的文本生成(Controllable Text Generation)。

论文中提出Diffusion-LM模型，即上面的架构，：

• embedding和对应的rounding步骤，放在了$x_0$之前
• 并且设计了embedding损失函数，用来学习embedding信息
• 也提出了一种技术改进rounding的准确性，这一个步骤是把连续的变量转换成离散的word序列。
• 提出一种基于梯度的方法，增强文本生成的可控性。其迭代地对$Diffusion-LM$的连续潜在变量进行梯度更新，以平衡流畅性和控制满意度。

本论文还通过设计6个可控范围（从细粒度到复杂的结构，如语法树）实验来验证模型。

论文特点

1. 相较于之前研究在离散空间上的扩散模型论文研究，该论文是在连续潜在空间上应用扩散模型，这使得基于梯度的有效方法有利于可控生成。
2. 针对大部分语言模型都是基于从左到右的自回归模型，这对文本可控并没有利，这是由于自回归LMs不能直接以正确的上下文为条件。$Diffusion-LM$可以以观察句子的复杂、全局属性的任意分类器作为条件。
3. Diffusion-LM有即插即用的可控生成能力，这是通过对潜在变量的连续序列应用分类器引导(classifier-guided) 的梯度更新(gradient update) 来完成的。

可控文本生成

文本生成：从一个训练好的语言模型$p_{lm}(\vec w)$中获取抽取样本$\vec w$，其中$\vec w = [w_1, ..., w_n]$是离散的单词序列，$p_{lm}(\vec w)$是单词序列上的概率分布。 可控文本生成：从一个条件概率分布$p(\vec w \ | \ \vec c)$上抽取样本$\vec w$，其中$\vec c$表示控制变量。

论文中采用即插即用(plug-and-play) 的可控生成方式，即先从一个无标签的文本数据 上训练出一个语言模型$p_{lm}(\vec w)$；然后针对不同生成控制表达式使用不同的分类器，即针对每一个任务，给出一个在较少的标记文本数据上训练得到的对应分类器$p(\vec c \ | \ \vec w)$，最后通过这两种模型在后验$p(\vec w \ | \ \vec c)$上进行抽样，其中通过贝叶斯公式，可以得到：

$p_{lm}(\vec w)$表示$\vec w$是否流畅，$p(\vec c \ | \ \vec w)$表示$\vec w$是否完成控制。

Diffusion-LM模型

标准的扩散模型：

• 正向过程：通过一系列随$t$的超参数$\beta$和噪音来确定一个个$x_{t}$。即$x_{t-1} \to x_t$通过$q(x_t \ | \ x_{t-1}) = \mathcal N(x_t \ ; \sqrt{1-\beta_t}x_{t-1}\ , \beta_t I)$。这个过程中没有需要训练的参数。同时，这得到向前扩散中添加的噪音数据信息以及添加噪音后的数据分布，并且通过反转这个过程来预测这个噪音或者数据分布来定义损失函数。
• 方向过程：即$x_{t} \to x_{t-1}$，通过$p_{\theta}(x_{t-1} \ | \ x_t) = \mathcal N (x_{t-1}\ ; \mu_{\theta}(x_t, t), \sum_{\theta}(x_t, t))$，其中$p_{\theta}(x_t) \approx \mathcal N(0, I)$。

  $\mu_{\theta}和\sum_{\theta}$都是可以通过U-Net或者Transformer获取。

• 损失函数：最大化边际似然$\mathbb E_{x_0 \sim p_{data}}[log\ p_{\theta}(x_0)]$:
  $$\begin{align} \mathcal L_{vlb}(x_0) = \mathbb E_{q(x_{1:T}\ | \ x_0)} [log\ \frac{q(x_T \ | \ x_0)}{p_{\theta}(x_T)} + \sum_{t=2}^T\ log\ \frac{q(x_{t-1} \ | \ x_0, x_1)}{p_{\theta}(x_{t-1} \ | \ x_t)} - log\ p_{\theta}(x_0\ | \ x_1)] \end{align}$$

  这个损失函数是不稳定的，需要很多技巧来保证稳定性

• 简化的损失函数：重新权重化$\mathcal L_{vlb}$中的KL散度项得到：
  $$\begin{align} \mathcal L_{simple}(x_0) = \sum_{t=1}^T \mathbb E_{q(x_t\ | \ x_0)}\| \mu_{\theta}(x_t,\ t) - \hat\mu(x_t,\ x_0) \|^2 \end{align}$$

  其中$\hat\mu(x_t,\ x_0)$是$q(x_{t-1} \ | \ x_0, x_t)$的后验均值，$\mu_{\theta}(x_t,\ t)$是$p_{\theta}(x_{t-1} | x_t)$模型的预测均值。 这个损失函数不再是有效的下界，但是却更稳定，并且也能提高样本质量。

Diffusion-LM模型

Diffusion-LM模型相对于标准的扩散模型所做的修改：

• 定义embedding function处理离散文本到连续空间的映射。完成这个映射是通过end-to-end训练中学习特征表示，即把特征表示的学习加入到训练目标中(training object)
• 需要一个roundding method把特征表示映射回离散的文本。

End-to-end Training

定义一个$Emb(w_i)$函数，把每个单词映射到一个向量$v_i \in R^d$； $n$个词序列的向量$\vec w$可以表示$Emb(\vec w) = [Emb(w_1), ...,Emb(w_n)] \in R^{n \times d}$

这个embedding function被添加在标准扩散模型的前向过程$x_0$前，即单词序列通过Emb得到$x_0$，之后再通过$x_0$计算其他的$x_t$完成标准的扩散过程。其中$\vec w \to x_0$可以通过：

方向过程公式中的$p_{\theta}(\vec w \ | \ \vec x_0)$是softmax分布。

通过上面的定义，则可以确定新的目标函数:

Reducing Rounding Errors

rounding过程就是把连续的$x_0$转成离散的文本ids。理想情况下，就是在$\vec x_0$的情况下获取最大的$\vec w$，即$\ argmax \ p_{\theta} (\vec w \ | \ x_0) = \prod_{i=1}^n\ p_{\theta}\ (w_i \ | \ x_i)$，可是这并不能保证得到好的words。一种解释是：公式(7)中的$\mathcal L_{simple}(x_0)$项不够重视$x_0$的结构建模。

论文中提出解决该问题的方法是re-parametrizes $\mathcal L_{simple}$，强制使模型Diffusion-LM明确地在公式(7)中的每一项对$x_0$建模：

其中$f_{\theta}(x_t, t)$模型直接预测$x_0$。这迫使神经网络预测目标函数每项中的$x_0$，发现用这个目标训练的模型很快就知道$x_0$应该精确地以单词嵌入为中心。

相同的思路，在解码时(decoding time)也可以使用，论文中称为clamping trick。通过$f_{\theta}(x_t, t)$预测得到$x_0$，然后获取$x_{t-1}$可以通过公式(9)替换标准公式，公式(10)是使用clamping trick后的公式，clamping trick强制预测的向量用于中间扩散步骤的单词，从而使向量预测更加精确并减少舍入误差：

通过这个技巧，模型另外也能映射预测的向量$f_{\theta}(x_t, t)$到他自己附近的单词序列特征表示。

Diffusion-LM的解码及可控生成

这里是说明Diffusion-LM怎么控制文本生成的。基本原理就是使用公式(1)代替直接在离散的文本数据上。该论文就是通过控制在Diffusion-LM上定义的连续潜在变量$x_{0:T}$序列，然后应用rounding转换潜在变量到文本数据。

控制$x_{0:T}$生成等同于从后验公式$p(x_{0:T} \ | \ c) = \prod_{t=1}^T \ p(x_{t-1} \ | \ x_t, c)$解码；该论文中通过将这个联合推理问题分解为每个扩散步骤的一系列控制问题：

这样我们就可以在第$t$步，通过在$x_{t-1}$上建立梯度更新：

其中$log \ p(x_{t-1}\ | \ x_t)$和$log \ p(c \ | \ x_{t-1})$都是可微的；前者是通过Diffusion-LM对其参数化，第二项是通过一个神经网络分类器参数化。

训练过程：在扩散潜在变量(diffusion latent variables)上训练分类器，并且在$x_{t-1}$的潜在空间上执行梯度更新，从而达到令人满意的控制。（这个过程是类似可控图像的方式）

不过论文中针对文本数据以及加速解码过程，又提出了两种方法：

• fluency regularization：引入超参数$\lambda$，权衡fluenncy和control：
  $$\begin{align} \text{公式(13)可得}\qquad log\ p(x_{t-1} \ | \ x_t, c) \propto & \ log\ p(x_{t-1} \ | \ x_t) + log\ p(c \ | \ x_{t-1}) \nonumber \\ \nonumber \\ \text{引入超参数$\lambda$}\qquad log\ p(x_{t-1} \ | \ x_t, c) \propto & \ \lambda \ log\ p(x_{t-1} \ | \ x_t) + log\ p(c \ | \ x_{t-1}) \end{align}$$

  虽然现有的扩散可控生成方法在目标中不包括$\lambda\ log\ p(x_{t−1}\ | \ x_t)$项，但我们发现该项有助于生成流畅的文本。

• multiple gradient steps：这个是为了提高控制质量（the control quality）。通过在每个扩散步骤中执行多次梯度步骤：论文中使用的是每个扩散步骤执行Adgrad更新的3个步骤。

其中由于计算量太大，论文中使用downsample方法，将扩散步长从2000降到200，这加快了可控生成算法，并且也不会对样本质量造成太大影响。

针对单样本高质量生成要求

面对机器翻译（machine translation） 或者句子填充（sentence infilling） 等任务，这类任务要求输出一个单独且高质量的序列。所以论文中针对这类问题，应用了**Minimum Bayes Risk (MBR)**进行解码，通过聚合从Diffusion-LM模型抽取出来的一系列样本集合$S$，然后选择在损失函数$\mathcal L$下可以满足最小期待风险(the minimum expected risk) 的样本，从而达到较好的结果：

参考文献

1. Li, Xiang Lisa, John Thickstun, Ishaan Gulrajani, Percy Liang和Tatsunori B. Hashimoto. 《Diffusion-LM Improves Controllable Text Generation》. arXiv, 2022年5月27日. http://arxiv.org/abs/2205.14217.