Diffusion-LM Improves Controllable Text Generation 扩散语言模型改善可控文本生成

论文链接：https://arxiv.org/abs/2205.14217
代码链接：https://github.com/XiangLi1999/Diffusion-LM

一、简介

无需重训模型控制语言模型的表现是自然语言生成（NLG）的一个重要开放问题，近期工作已经在简单句子属性控制生成上取得了一定的进展，但是在复杂的、细粒度的控制上的进展还远远不够。作者指的更复杂，更细粒度的控制：

为了解决这一挑战，我们提出了一个基于连续扩散模型的非自回归语言模型，我们称之为Diffusion-LM。Diffusion-LM是在连续空间上使用扩散模型，将随机噪声逐渐还原成词向量。连续的层级式的中间隐空间可以通过简单的梯度回传方法实现复杂的可控生成任务。实验表明，Diffusion-LM在六个细粒度控制任务上取得领先的效果。

传统的自回归语言模型可以生成高质量文本，但是生成的结果一般都是随机不可控的，现实情况下往往需要生成可控的文本。一个最直接的方法就是通过有标签数据微调语言模型，但是这种方法针对不同的可控任务需要不同的标注数据和重训模型。所以针对这一缺陷，应运而生了即插即用式的可控文本生成任务。

扩散模型已经在图像生成上取得了优秀的性能，但是很少使用在NLP领域，主要是因为文本模态的离散性。为了解决这一问题，文章提出在隐空间上进行扩散模型的加噪和去噪，在隐空间和词向量之间采用round的方法，将隐空间估计到词向量空间。

实验证明使用扩散模型进行即插即用的文本可控生成可以取得更优的性能，并且支持更多复杂的可控生成任务。

至于选择扩散模型的原因，根据文章总结为以下几点：

1. 中间变量的连续且分层的特性使得简单的梯度算法就可以实现复杂且可控的生成任务。

2. 由于文本通常是离散的，因此运用离散的扩散模型具有极大的困难，因此作者对标准的扩散模型进行了改造: 在传统的步骤上添加了embedding step和rounding step，设计了学习embedding的训练目标，并且提出了提升rounding表现的方法。

3. 之前的控制生成方法，通常基于自回归语言模型，只能从左至右生成。这导致PPLM无法修复在之前步中的错误。

二、相关工作

文本领域扩散模型。扩散模型被证明在连续空间上有较好的生成效果，但是文本模态天然的与图像、音频不同，是连续的语言符号表征的语义信息。为了处理离散数据，之前的工作有研究离散状态空间的扩散模型，方法是在离散数据上采用一个破坏过程（每个token有一定概率被删除或者被替换）。本文聚焦在文本领域上的连续扩散模型，文章提出的Diffusion-LM在连续的隐空间上进行扩散过程，其可以有效地使用梯度回传方法进行可控生成。
自回归和非自回归语言模型。很多的预训练语言模型都是从左至右自回归形式的，如GPT-3、PaLM等，固定的生成顺序将会限制可控生成的灵活性，尤其那种左右上下文全局都相关的可控任务。因为自回归模型无法利用右侧上下文信息作为先验条件，之前很多工作在这些控制任务上采用了特殊的训练和解码技术，比如Qin等人提出了将离散的语言模型输出到连续的变量中，然后采用右侧上下文回传梯度信息。Diffusion-LM的优势就是可以在复杂的、全局的句子属性上以任意的分类器为条件进行可控生成。此外，还有一些在机器翻译和语音识别任务上的非自回归语言模型，但是这些方法都是对机器翻译和语言识别任务定制的，都无法成功对语言建模。
即插即用的可控文本生成。即插即用的可控文本生成旨在将语言模型冻结，使用势函数引导输出的方向。通过给定概率势函数来衡量生成文本是否符合控制要求，并且生成的文本优化的方向不仅要符合控制要求，还要保证语言的通顺性。自回归模型上已有很多种即插即用的可控生成方法。比如FUDGE通过每个token对整个句子的部分控制进行重新赋权重，根据候选挑选最好的控制结果，GeDi对每个token采用一个更小的语言模型微调来达到可控生成的效果。与本文方法最接近的方法是PPLM，它是一种在隐空间上升高梯度以满足控制方向的一种方法，但是PPLM控制是根据自回归语言模型方向从左到右的控制token的生成，所以在某些复杂的任务上表现并不理想，比如句法结构控制，Diffusion-LM在隐空间上进行可控生成缓解了PPLM的这一问题。

三、问题陈述和背景

3.1、文本的生成模型和可控生成

文本生成是指从训练好的语言模型$p_{lm}(\mathbf{w})$中取样$\mathbf{w}$的任务，其中$\mathbf{w}=[w_1\cdot \cdot \cdot w_n]$是离散词的序列，$p_{lm}(\mathbf{w})$是词的序列的概率分布。可控文本生成的任务是从条件分布$p(\mathbf{w}|\mathbf{c})$中抽取$\mathbf{w}$，其中$\mathbf{c}$表示一个控制变量。对于句法控制，$\mathbf{c}$可以是一个目标句法树，而对于情感控制，$\mathbf{c}$可以是一个期望的情感标签。可控生成的目标是生成满足控制目标$\mathbf{c}$的$\mathbf{w}$。

考虑即插即用的可控生成设置：我们得到一个从大量未标注文本数据中训练出来的语言模型$p_{lm}(\mathbf{w})$，对于每个控制任务，我们得到一个从较小数量的标注文本数据中训练出来的分类器$p(\mathbf{c}|\mathbf{w})$（例如，对于句法控制，分类器是一个概率分析器）。目标是利用这两个模型，通过贝叶斯规则$p(\mathbf{w}|\mathbf{c})\propto p_{lm}(\mathbf{w})\cdot p(\mathbf{c}|\mathbf{w})$从后验$p(\mathbf{w}|\mathbf{c})$中近似取样。这里，$p_{lm}(\mathbf{w})$鼓励w流畅，而$p(\mathbf{c}|\mathbf{w})$鼓励$\mathbf{w}$履行控制。

3.2、自回归语言模型

语言模型的典型方法将$p_{lm}$纳入自回归左至右的mannar，$p_{lm}(\mathbf{w})=p_{lm}(w_1){\prod }_{i=2}^n{p}_{lm}(x_i|x_{<i})$ 。在这种情况下，文本生成被简化为以迄今为止生成的部分序列为条件反复预测下一个token的任务。下一个token的预测$p_{lm}(x_i|x_{<i})$通常由Transformer架构来参数化。

3.3、连续域的扩散模型

扩散模型是一个潜变量模型，它将数据${\mathbf{x}}_0\in {\mathbb{R}}^d$建模为马尔可夫链${\mathbf{x}}_T...{\mathbf{x}}_0$，每个变量的维度是$d$，${\mathbf{x}}_T$是一个高斯分布。扩散模型对潜在变量${\mathbf{x}}_{T:1}$的序列进行增量降噪，以接近目标数据分布的样本（图2）。初始状态$p_{\theta }({\mathbf{x}}_T)\approx \mathcal{N}(0,I)$，每个降噪过渡${\mathbf{x}}_t\to {\mathbf{x}}_{t-1}$由模型$p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)=\mathcal{N}({\mathbf{x}}_{t-1};{\mu }_{\theta }({\mathbf{x}}_t,t),{\Sigma }_{\theta }({\mathbf{x}}_t,t))$来参数化。例如，${\mu }_{\theta }$和${\Sigma }_{\theta }$可以由U-Net或Tranformer来计算。

为了训练扩散模型，我们定义了一个构建中间潜变量${\mathbf{x}}_{1:T}$的前向过程。前向过程逐步向数据${\mathbf{x}}_0$添加高斯噪声，直到在扩散步骤$T$，样本${\mathbf{x}}_T$是近似高斯的。每个过渡${\mathbf{x}}_{t-1}\to {\mathbf{x}}_t$的参数为$q({\mathbf{x}}_t|{\mathbf{x}}_{t-1})=\mathcal{N}({\mathbf{x}}_t;\sqrt{1-{\beta }_t}{\mathbf{x}}_{t-1},{\beta }_t\mathbf{I})$，其中超参数${\beta }_t$是在扩散步骤$t$添加的噪声量。这个前向过程$q$的参数化不包含可训练的参数，允许我们定义一个训练目标，包括根据预先定义的前向过程$q$生成噪声数据，并训练一个模型来逆转这个过程并重建数据。
扩散模型的训练是为了最大化数据的边际似然${\mathbb{E}}_{{\mathbf{x}}_0\sim p_{data}}[\log p_{\theta }({\mathbf{x}}_0)]$，目标是$\log p_{\theta }({\mathbf{x}}_0)$的变分下限：$${\mathcal{L}}_{vlb}({\mathbf{x}}_0)={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_T|{\mathbf{x}}_0)}{p_{\theta }({\mathbf{x}}_T)}+\sum _{t=2}^T\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0,{\mathbf{x}}_t)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}-\log p_{\theta }({\mathbf{x}}_0|{\mathbf{x}}_1)\right]$$然而，这个目标可能是不稳定的，需要许多优化技巧来稳定。为了规避这个问题，Ho等人设计了一个简单的代理目标，对${\mathcal{L}}_{vlb}$中的每个KL-散度项进行扩展和重新加权，以获得一个均方误差损失，我们将其称为${\mathcal{L}}_{simple}$。对于第$t$个KL散度，有$${\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\log \frac{q({\mathbf{x}}_{t-1}|{\mathbf{x}}_0,{\mathbf{x}}_t)}{p_{\theta }({\mathbf{x}}_{t-1}|{\mathbf{x}}_t)}\right]={\mathbb{E}}_{q({\mathbf{x}}_{1:T}|{\mathbf{x}}_0)}\left[\frac{1}{2{\sigma }_t^2}||{\mu }_{\theta }({\mathbf{x}}_t,t)-\hat{\mu }({\mathbf{x}}_t,{\mathbf{x}}_0)|{|}^2\right]+C$$其中$\hat{\mu }({\mathbf{x}}_t,{\mathbf{x}}_0)$是后验 q ( x t − 1 ∣