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 l m ( w ) p_{lm}(\mathbf{w}) plm​(w)中取样 w \mathbf{w} w的任务，其中 w = [ w 1 ⋅ ⋅ ⋅ w n ] \mathbf{w}= [w_1···w_n] w=[w1​⋅⋅⋅wn​]是离散词的序列， p l m ( w ) p_{lm}(\mathbf{w}) plm​(w)是词的序列的概率分布。可控文本生成的任务是从条件分布 p ( w ∣ c ) p(\mathbf{w}|\mathbf{c}) p(w∣c)中抽取 w \mathbf{w} w，其中 c \mathbf{c} c表示一个控制变量。对于句法控制， c \mathbf{c} c可以是一个目标句法树，而对于情感控制， c \mathbf{c} c可以是一个期望的情感标签。可控生成的目标是生成满足控制目标 c \mathbf{c} c的 w \mathbf{w} w。

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

3.2、自回归语言模型

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

3.3、连续域的扩散模型

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

为了训练扩散模型，我们定义了一个构建中间潜变量 x 1 : T \mathbf{x}_{1:T} x1:T​的前向过程。前向过程逐步向数据 x 0 \mathbf{x}_0 x0​添加高斯噪声，直到在扩散步骤 T T T，样本 x T \mathbf{x}_T xT​是近似高斯的。每个过渡 x t − 1 → x t \mathbf{x}_{t-1}→\mathbf{x}_t xt−1​→xt​的参数为 q ( x t ∣ x t − 1 ) = N ( x t ; 1 − β t x t − 1 , β t I ) q(\mathbf{x}_t|\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_t;\sqrt{1-β_t}\mathbf{x}_{t-1},β_t\mathbf{I}) q(xt​∣xt−1​)=N(xt​;1−βt​ ​xt−1​,βt​I)，其中超参数 β t β_t βt​是在扩散步骤 t t t添加的噪声量。这个前向过程 q q q的参数化不包含可训练的参数，允许我们定义一个训练目标，包括根据预先定义的前向过程 q q q生成噪声数据，并训练一个模型来逆转这个过程并重建数据。
扩散模型的训练是为了最大化数据的边际似然 E x 0 ∼ p d a t a [ log ⁡ p θ ( x 0 ) ] \mathbb{E}_{\mathbf{x}_0∼p_{data}}[\log p_θ(\mathbf{x}_0)] Ex0​∼pdata​​[logpθ​(x0​)]，目标是 log ⁡ p θ ( x 0 ) \log p_θ(\mathbf{x}_0) logpθ​(x0​)的变分下限： L v l b ( x 0 ) = E q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x T ∣ x 0 ) p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x 0 , x t ) p θ ( x t − 1 ∣ x t ) − log ⁡ p θ ( x 0 ∣ x 1 ) ] \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] Lvlb​(x0​)=Eq(x1:T​∣x0​)​[logpθ​(xT​)q(xT​∣x0​)​+t=2∑T​logpθ​(xt−1​∣xt​)q(xt−1​∣x0​,xt​)​−logpθ​(x0​∣x1​)]然而，这个目标可能是不稳定的，需要许多优化技巧来稳定。为了规避这个问题，Ho等人设计了一个简单的代理目标，对 L v l b \mathcal{L}_{vlb} Lvlb​中的每个KL-散度项进行扩展和重新加权，以获得一个均方误差损失，我们将其称为 L s i m p l e \mathcal{L}_{simple} Lsimple​。对于第 t t t个KL散度，有 E q ( x 1 : T ∣ x 0 ) [ log ⁡ q ( x t − 1 ∣ x 0 , x t ) p θ ( x t − 1 ∣ x t ) ] = E q ( x 1 : T ∣ x 0 ) [ 1 2 σ t 2 ∣ ∣ μ θ ( x t , t ) − μ ^ ( x t , x 0 ) ∣ ∣ 2 ] + C \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^2_t}||\mu_{\theta}(\mathbf{x}_t,t)-\hat{\mu}(\mathbf{x}_t,\mathbf{x}_0)||^2\right]+C Eq(x1:T​∣x0​)​[logpθ​(xt−1​∣xt​)q(xt−1​∣x0​,xt​)​]=Eq(x1:T​∣x0​)​[2σt2​1​∣∣μθ​(xt​,t)−μ^​(xt​,x0​)∣∣2]+C其中 μ ^ ( x t , x 0 ) \hat{\mu}(\mathbf{x}_t,\mathbf{x}_0) μ^​(xt​,x0​)是后验 q ( x t − 1 ∣ x 0 , x t ) q(\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t) q(xt−1​∣x0​,xt​)的平均值，它是一个封闭的高斯，而 μ θ ( x t , t ) \mu_{\theta}(\mathbf{x}_t,t) μθ​(xt​,t)是由神经网络计算的 p θ ( x t − 1 ∣ x t ) p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t) pθ​(xt−1​∣xt​)的预测平均值（假设方差为常数，已知）。将常数 C C C以及系数 1 2 σ t 2 \frac{1}{2\sigma_t^2} 2σt2​1​去掉即可得到 L s i m p l e \mathcal{L}_{simple} Lsimple​： L s i m p l e ( x 0 ) = ∑ t = 1 T E q ( x t ∣ x 0 ) ∣ ∣ μ θ ( x t , t ) − μ ^ ( x t , x 0 ) ∣ ∣ 2 \mathcal{L}_{simple}(\mathbf{x}_0)=\sum_{t=1}^T\mathbb{E}_{q(\mathbf{x}_t|\mathbf{x}_0)}||\mu_{\theta}(\mathbf{x}_t,t)-\hat{\mu}(\mathbf{x}_t,\mathbf{x}_0)||^2 Lsimple​(x0​)=t=1∑T​Eq(xt​∣x0​)​∣∣μθ​(xt​,t)−μ^​(xt​,x0​)∣∣2

虽然 L s i m p l e \mathcal{L}_{simple} Lsimple​不再是一个有效的下限，但之前的工作发现，根据经验，它使训练更加稳定，并提高了样本质量。我们将在Diffusion-LM中利用类似的简化来稳定训练和提高样本质量。

四、Diffusion-LM：连续扩散语言模型

4.1、端到端训练

为了将连续扩散模型应用于离散文本，我们定义了一个嵌入函数 E M B ( w i ) EMB(w_i) EMB(wi​)，它将每个词映射到 R d \mathbb{R}^d Rd中的一个向量。我们将长度为 n n n的序列 w \mathbf{w} w的嵌入定义为 E M B ( w ) = [ E M B ( w 1 ) , . . . , E M B ( w n ) ] ∈ R n d EMB(\mathbf{w}) = [EMB(w_1), ... , EMB(w_n)] ∈\mathbb{R}^{nd} EMB(w)=[EMB(w1​),...,EMB(wn​)]∈Rnd。

我们提出了对扩散模型训练目标的修改，联合学习扩散模型的参数和词嵌入。在预备知识实验中，我们探索了随机高斯嵌入，以及预训练的词嵌入。我们发现，与端到端训练相比，这些固定嵌入对Diffusion-LM来说是次优的。

根据图2显示，我们的方法增加了一个马尔科夫转移从离散的单词 w \mathbf{w} w到 x 0 \mathbf{x}_0 x0​的前向过程，参数符合 q ϕ ( x 0 ∣ w ) = N ( E M B ( w ) , σ 0 I ) q_{\phi}(\mathbf{x}_0|\mathbf{w})=\mathcal{N}(EMB(\mathbf{w}),\sigma_0I) qϕ​(x0​∣w)=N(EMB(w),σ0​I) ，在反向过程中，我们添加了一个可训练的rounding step， 参数满足 p θ ( w ∣ x 0 ) = ∏ i = 1 n p θ ( w i ∣ x i ) p_{\theta}(\mathbf{w}|\mathbf{x}_0)=\prod_{i=1}^np_{\theta}(w_i|x_i) pθ​(w∣x0​)=∏i=1n​pθ​(wi​∣xi​)，其中 p θ ( w i ∣ x i ) p_{\theta}(w_i|x_i) pθ​(wi​∣xi​)是一个softmax分布。第3节中介绍的训练目标现在变成了 L v l b e 2 e ( w ) = E q ϕ ( x 0 ∣ w ) [ L v l b ( x 0 ) + log ⁡ q ϕ ( x 0 ∣ w ) − log ⁡ p θ ( w ∣ x 0 ) ] \mathcal{L}_{vlb}^{e2e}(\mathbf{w})=\mathbb{E}_{q_{\phi}(\mathbf{x}_0|\mathbf{w})}\left[\mathcal{L}_{vlb}(\mathbf{x}_0)+\log q_{\phi}(\mathbf{x}_0|\mathbf{w})-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right] Lvlbe2e​(w)=Eqϕ​(x0​∣w)​[Lvlb​(x0​)+logqϕ​(x0​∣w)−logpθ​(w∣x0​)] L s i m p l e e 2 e ( w ) = E q ϕ ( x 0 : T ∣ w ) [ L s i m p l e ( x 0 ) + ∣ ∣ E M B ( w ) − μ θ ( x 1 , 1 ) ∣ ∣ 2 − log ⁡ p θ ( w ∣ x 0 ) ] \mathcal{L}_{simple}^{e2e}(\mathbf{w})=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[\mathcal{L}_{simple}(\mathbf{x}_0)+||EMB(\mathbf{w})-\mu_{\theta}(\mathbf{x}_1,1)||^2-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right] Lsimplee2e​(w)=Eqϕ​(x0:T​∣w)​[Lsimple​(x0​)+∣∣EMB(w)−μθ​(x1​,1)∣∣2−logpθ​(w∣x0​)]为了得到 L s i m p l e e 2 e \mathcal{L}_{simple}^{e2e} Lsimplee2e​，先对 L v l b e 2 e \mathcal{L}_{vlb}^{e2e} Lvlbe2e​进行化简： L v l b e 2 e = E q ϕ ( x 0 ∣ w ) [ L v l b ( x 0 ) + log ⁡ q ϕ ( x 0 ∣ w ) − log ⁡ p θ ( w ∣ x 0 ) ] = E q ϕ ( x 0 : T ∣ w ) [ log ⁡ q ( x T ∣ x 0 ) p θ ( x T ) + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x 0 , x t ) p θ ( x t − 1 ∣ x t ) − log ⁡ p θ ( x 0 ∣ x 1 ) + log ⁡ q ϕ ( x 0 ∣ w ) − log ⁡ p θ ( w ∣ x 0 ) ] = E q ϕ ( x 0 : T ∣ w ) [ log ⁡ q ( x T ∣ x 0 ) p θ ( x T ) ⏟ L T + ∑ t = 2 T log ⁡ q ( x t − 1 ∣ x 0 , x t ) p θ ( x t − 1 ∣ x t ) ⏟ L t − 1 + log ⁡ q ϕ ( x 0 ∣ w ) p θ ( x 0 ∣ x 1 ) ⏟ L 0 − log ⁡ p θ ( w ∣ x 0 ) ⏟ L r o u n d ] \begin{aligned}\mathcal{L}_{vlb}^{e2e}&=\mathbb{E}_{q_{\phi}(\mathbf{x}_0|\mathbf{w})}\left[\mathcal{L}_{vlb}(\mathbf{x}_0)+\log q_{\phi}(\mathbf{x}_0|\mathbf{w})-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right]\\&=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\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)+\log q_{\phi}(\mathbf{x}_0|\mathbf{w})-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right]\\&=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[\underbrace{\log\frac{q(\mathbf{x}_T|\mathbf{x}_0)}{p_{\theta}(\mathbf{x}_T)}}_{L_T}+\sum_{t=2}^T\underbrace{\log\frac{q(\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t)}{p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_t)}}_{L_{t-1}}+\underbrace{\log\frac{q_{\phi}(\mathbf{x}_0|\mathbf{w})}{p_{\theta}(\mathbf{x}_0|\mathbf{x}_1)}}_{L_0}-\underbrace{\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)}_{L_{round}}\right]\end{aligned} Lvlbe2e​​=Eqϕ​(x0​∣w)​[Lvlb​(x0​)+logqϕ​(x0​∣w)−logpθ​(w∣x0​)]=Eqϕ​(x0:T​∣w)​[logpθ​(xT​)q(xT​∣x0​)​+t=2∑T​logpθ​(xt−1​∣xt​)q(xt−1​∣x0​,xt​)​−logpθ​(x0​∣x1​)+logqϕ​(x0​∣w)−logpθ​(w∣x0​)]=Eqϕ​(x0:T​∣w)​⎣⎢⎢⎢⎡​LT​ logpθ​(xT​)q(xT​∣x0​)​​​+t=2∑T​Lt−1​ logpθ​(xt−1​∣xt​)q(xt−1​∣x0​,xt​)​​​+L0​ logpθ​(x0​∣x1​)qϕ​(x0​∣w)​​​−Lround​ logpθ​(w∣x0​)​​⎦⎥⎥⎥⎤​​由于 E q ϕ ( x 0 : T ∣ w ) [ L T ] → E [ ∣ ∣ E x T ∼ q [ x T ∣ x 0 ] − 0 ∣ ∣ 2 ] = E [ ∣ ∣ μ ^ ( x T ; x 0 ) ∣ ∣ 2 ]    # ( 因 为 p θ ( x T ) 服 从 标 准 正 态 分 布 ， 期 望 为 0 ) \mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[L_T\right]\to\mathbb{E}\left[||\mathbb{E}_{\mathbf{x}_T\sim q}\left[\mathbf{x}_T|\mathbf{x}_0\right]-0||^2\right]=\mathbb{E}\left[||\hat{\mu}(\mathbf{x}_T;\mathbf{x}_0)||^2\right]\space\space\#(因为p_{\theta}(\mathbf{x}_T)服从标准正态分布，期望为0) Eqϕ​(x0:T​∣w)​[LT​]→E[∣∣ExT​∼q​[xT​∣x0​]−0∣∣2]=E[∣∣μ^​(xT​;x0​)∣∣2]  #(因为pθ​(xT​)服从标准正态分布，期望为0) E q ϕ ( x 0 : T ∣ w ) [ L t − 1 ] → E [ ∣ ∣ E x t − 1 ∼ q [ x t − 1 ∣ x 0 , x t ] − E x t − 1 ∼ p θ [ x t − 1 ∣ x t ] ∣ ∣ 2 ] = E [ ∣ ∣ μ ^ ( x t ; x 0 ) − μ θ ( x t , t ) ∣ ∣ 2 ] \mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[L_{t-1}\right]\to\mathbb{E}\left[||\mathbb{E}_{\mathbf{x}_{t-1}\sim q}\left[\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t\right]-\mathbb{E}_{\mathbf{x}_{t-1}\sim p_{\theta}}\left[\mathbf{x}_{t-1}|\mathbf{x}_t\right]||^2\right]=\mathbb{E}\left[||\hat{\mu}(\mathbf{x}_t;\mathbf{x}_0)-\mu_{\theta}(\mathbf{x}_t,t)||^2\right] Eqϕ​(x0:T​∣w)​[Lt−1​]→E[∣∣Ext−1​∼q​[xt−1​∣x0​,xt​]−Ext−1​∼pθ​​[xt−1​∣xt​]∣∣2]=E[∣∣μ^​(xt​;x0​)−μθ​(xt​,t)∣∣2] E q ϕ ( x 0 : T ∣ w ) [ L 0 ] → E [ ∣ ∣ E x 0 ∼ q ϕ [ x 0 ∣ w ] − E x 0 ∼ p θ [ x 0 ∣ x 1 ] ∣ ∣ 2 ] = E [ ∣ ∣ E M B ( w ) − μ θ ( x 1 , 1 ) ∣ ∣ 2 ] \mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[L_{0}\right]\to\mathbb{E}\left[||\mathbb{E}_{\mathbf{x}_{0}\sim q_{\phi}}\left[\mathbf{x}_{0}|\mathbf{w}\right]-\mathbb{E}_{\mathbf{x}_{0}\sim p_{\theta}}\left[\mathbf{x}_{0}|\mathbf{x}_1\right]||^2\right]=\mathbb{E}\left[||EMB(\mathbf{w})-\mu_{\theta}(\mathbf{x}_1,1)||^2\right] Eqϕ​(x0:T​∣w)​[L0​]→E[∣∣Ex0​∼qϕ​​[x0​∣w]−Ex0​∼pθ​​[x0​∣x1​]∣∣2]=E[∣∣EMB(w)−μθ​(x1​,1)∣∣2]所以 L s i m p l e e 2 e = E q ϕ ( x 0 : T ∣ w ) [ ∣ ∣ μ ^ ( x T ; x 0 ) ∣ ∣ 2 + ∑ t = 2 T [ ∣ ∣ μ ^ ( x t ; x 0 ) − μ θ ( x t , t ) ∣ ∣ 2 ] ] + E q ϕ ( x 0 : T ∣ w ) [ ∣ ∣ E M B ( w ) − μ θ ( x 1 , 1 ) ∣ ∣ 2 − log ⁡ p θ ( w ∣ x 0 ) ] = E q ϕ ( x 0 : T ∣ w ) [ L s i m p l e ( x 0 ) + ∣ ∣ E M B ( w ) − μ θ ( x 1 , 1 ) ∣ ∣ 2 − log ⁡ p θ ( w ∣ x 0 ) ] \begin{aligned}\mathcal{L}_{simple}^{e2e}&=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[||\hat{\mu}(\mathbf{x}_T;\mathbf{x}_0)||^2+\sum_{t=2}^T\left[||\hat{\mu}(\mathbf{x}_t;\mathbf{x}_0)-\mu_{\theta}(\mathbf{x}_t,t)||^2\right]\right]+\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[||EMB(\mathbf{w})-\mu_{\theta}(\mathbf{x}_1,1)||^2-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right]\\&=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[\mathcal{L}_{simple}(\mathbf{x}_0)+||EMB(\mathbf{w})-\mu_{\theta}(\mathbf{x}_1,1)||^2-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right]\end{aligned} Lsimplee2e​​=Eqϕ​(x0:T​∣w)​[∣∣μ^​(xT​;x0​)∣∣2+t=2∑T​[∣∣μ^​(xt​;x0​)−μθ​(xt​,t)∣∣2]]+Eqϕ​(x0:T​∣w)​[∣∣EMB(w)−μθ​(x1​,1)∣∣2−logpθ​(w∣x0​)]=Eqϕ​(x0:T​∣w)​[Lsimple​(x0​)+∣∣EMB(w)−μθ​(x1​,1)∣∣2−logpθ​(w∣x0​)]​

我们从 L v l b e 2 e \mathcal{L}_{vlb}^{e2e} Lvlbe2e​推导出了 L s i m p l e e 2 e \mathcal{L}_{simple}^{e2e} Lsimplee2e​。由于我们是在训练嵌入函数， q ϕ q_{\phi} qϕ​现在包含可训练的参数，我们使用重新定义参数的技巧来反向传播这个采样步骤。从经验上看，我们发现所学的嵌入有意义地聚类：具有相同的词性标签（句法角色）的词往往被聚类，如图3所示。

4.2、减少Roundig Errors

学习到的嵌入定义了一个从离散文本到连续 x 0 \mathbf{x}_0 x0​的映射。我们现在描述将预测的 x 0 \mathbf{x}_0 x0​舍入到离散文本的逆过程。根据 arg max ⁡ p θ ( w ∣ x 0 ) = ∏ i = 1 n p θ ( w i ∣ x i ) \argmax p_θ(\mathbf{w}|\mathbf{x}_0)=∏_{i=1}^np_θ(w_i|x_i) argmaxpθ​(w∣x0​)=∏i=1n​pθ​(wi​∣xi​)，通过为每个位置选择最有可能的词来实现舍入。理想情况下，这种argmax舍入足以映射回离散文本，因为降噪步骤应确保 x 0 \mathbf{x}_0 x0​正好位于某个词嵌入上。然而，从经验上看，该模型未能生成致力于单个单词的 x 0 \mathbf{x}_0 x0​。

对这种现象的一个解释是，我们的目标中的 L s i m p l e ( x 0 ) \mathcal{L}_{simple}(\mathbf{x}_0) Lsimple​(x0​)项对 x 0 \mathbf{x}_0 x0​的结构建模强调不够。我们定义了 L s i m p l e ( x 0 ) = E x t ∣ ∣ μ θ ( x t , t ) − μ ^ ( x t , x 0 ) ∣ ∣ 2 \mathcal{L}_{simple}(\mathbf{x}_0)=\mathbb{E}_{\mathbf{x}_t}||\mu_θ(\mathbf{x}_t,t)-\hat{\mu}(\mathbf{x}_t,\mathbf{x}_0)||^2 Lsimple​(x0​)=Ext​​∣∣μθ​(xt​,t)−μ^​(xt​,x0​)∣∣2，其中我们的模型 μ θ ( x t , t ) \mu_θ(\mathbf{x}_t,t) μθ​(xt​,t)直接预测了每个降噪步骤 t t t的 p θ ( x t − 1 ∣ x t ） p_θ(\mathbf{x}_{t-1}|\mathbf{x}_t） pθ​(xt−1​∣xt​）的平均值。在这个目标中， x 0 \mathbf{x}_0 x0​必须致力于单字嵌入的约束只会出现在 t t t接近0的条件中，我们发现这个参数化需要仔细调整以迫使目标强调这些条件。

我们的方法对 L s i m p l e \mathcal{L}_{simple} Lsimple​进行了重新参数化，迫使Diffusion-LM在目标的每一个项中明确地建立 x 0 \mathbf{x}_0 x0​的模型。具体来说，我们推导出 L s i m p l e \mathcal{L}_{simple} Lsimple​的类似物，它通过 x 0 \mathbf{x}_0 x0​进行参数化， L x 0 − s i m p l e e 2 e ( x 0 ) = E x t ∣ ∣ f θ ( x t , t ) − x 0 ∣ ∣ 2 \mathcal{L}^{e2e}_{\mathbf{x}_0-simple}(\mathbf{x}_0)=\mathbb{E}_{\mathbf{x}_t}||f_θ(\mathbf{x}_t,t)-\mathbf{x}_0||^2 Lx0​−simplee2e​(x0​)=Ext​​∣∣fθ​(xt​,t)−x0​∣∣2，其中我们的模型 f θ ( x t , t ) f_θ(\mathbf{x}_t,t) fθ​(xt​,t)直接预测 x 0 \mathbf{x}_0 x0​。 这迫使神经网络在每个词中预测 x 0 \mathbf{x}_0 x0​，我们发现，用这个目标训练的模型很快学会 x 0 \mathbf{x}_0 x0​应该精确地以一个词嵌入为中心。根据推导，可知 q ( x t − 1 ∣ x 0 , x t ) 的 均 值 为 α ˉ t − 1 β t 1 − α ˉ t x 0 + α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t q(\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t)的均值为\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\mathbf{x}_0+\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}\mathbf{x}_t q(xt−1​∣x0​,xt​)的均值为1−αˉt​αˉt−1​ ​βt​​x0​+1−αˉt​αt​ ​(1−αˉt−1​)​xt​，可以从 q ( x t − 1 ∣ x 0 , x t ) q(\mathbf{x}_{t-1}|\mathbf{x}_0,\mathbf{x}_t) q(xt−1​∣x0​,xt​)采样 x t − 1 \mathbf{x}_{t-1} xt−1​，因此， ∣ ∣ μ ^ ( x t , x 0 ) − μ θ ( x t , t ) ∣ ∣ 2 = ∣ ∣ ( α ˉ t − 1 β t 1 − α ˉ t x 0 + α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t ) − ( α ˉ t − 1 β t 1 − α ˉ t f θ ( x t , t ) + α t ( 1 − α ˉ t − 1 ) 1 − α ˉ t x t ) ∣ ∣ 2 = ∣ ∣ α ˉ t − 1 β t 1 − α ˉ t ( x 0 − f θ ( x t , t ) ) ∣ ∣ 2 ∝ ∣ ∣ x 0 − f θ ( x t , t ) ) ∣ ∣ 2 \begin{aligned}||\hat{\mu}(\mathbf{x}_t,\mathbf{x}_0)-\mu_{\theta}(\mathbf{x}_t,t)||^2\\&=||(\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}\mathbf{x}_0+\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}\mathbf{x}_t)-(\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}f_{\theta}(\mathbf{x}_t,t)+\frac{\sqrt{\alpha_t}(1-\bar{\alpha}_{t-1})}{1-\bar{\alpha}_{t}}\mathbf{x}_t)||^2\\&=||\frac{\sqrt{\bar{\alpha}_{t-1}}\beta_t}{1-\bar{\alpha}_t}(\mathbf{x}_0-f_{\theta}(\mathbf{x}_t,t))||^2\\&\propto||\mathbf{x}_0-f_{\theta}(\mathbf{x}_t,t))||^2\end{aligned} ∣∣μ^​(xt​,x0​)−μθ​(xt​,t)∣∣2​=∣∣(1−αˉt​αˉt−1​ ​βt​​x0​+1−αˉt​αt​ ​(1−αˉt−1​)​xt​)−(1−αˉt​αˉt−1​ ​βt​​fθ​(xt​,t)+1−αˉt​αt​ ​(1−αˉt−1​)​xt​)∣∣2=∣∣1−αˉt​αˉt−1​ ​βt​​(x0​−fθ​(xt​,t))∣∣2∝∣∣x0​−fθ​(xt​,t))∣∣2​所以 L x 0 − s i m p l e e 2 e ( w ) = E q ϕ ( x 0 : T ∣ w ) [ ∣ ∣ μ ^ ( x T ; x 0 ) ∣ ∣ 2 + ∑ t = 2 T [ ∣ ∣ x 0 − f θ ( x t , t ) ∣ ∣ 2 ] ] + E q ϕ ( x 0 : 1 ∣ w ) [ ∣ ∣ E M B ( w ) − f θ ( x 1 , 1 ) ∣ ∣ 2 − log ⁡ p θ ( w ∣ x 0 ) ] \begin{aligned}\mathcal{L}_{\mathbf{x}_0-simple}^{e2e}(\mathbf{w})&=\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:T}|\mathbf{w})}\left[||\hat{\mu}(\mathbf{x}_T;\mathbf{x}_0)||^2+\sum_{t=2}^T\left[||\mathbf{x}_0-f_{\theta}(\mathbf{x}_t,t)||^2\right]\right]+\mathbb{E}_{q_{\phi}(\mathbf{x}_{0:1}|\mathbf{w})}\left[||EMB(\mathbf{w})-f_{\theta}(\mathbf{x}_1,1)||^2-\log p_{\theta}(\mathbf{w}|\mathbf{x}_0)\right]\end{aligned} Lx0​−simplee2e​(w)​=Eqϕ​(x0:T​∣w)​[∣∣μ^​(xT​;x0​)∣∣2+t=2∑T​[∣∣x0​−fθ​(xt​,t)∣∣2]]+Eqϕ​(x0:1​∣w)​[∣∣EMB(w)−fθ​(x1​,1)∣∣2−logpθ​(w∣x0​)]​

我们描述了重新参数化如何有助于模型训练，但我们也发现，同样的直觉可以在解码时使用，我们称之为clamping技巧。在 x 0 \mathbf{x}_0 x0​参数化模型的标准生成方法中，模型首先通过 f θ ( x t , t ) f_θ(\mathbf{x}_t,t) fθ​(xt​,t)计算 x 0 \mathbf{x}_0 x0​的估计值，然后以这个估计值为条件对 x t − 1 \mathbf{x}_{t-1} xt−1​进行采样，从而对 x t − 1 \mathbf{x}_{t-1} xt−1​进行降噪。 x t − 1 = α ˉ f θ ( x t , t ) + 1 − α ˉ ϵ \mathbf{x}_{t-1}=\sqrt{\bar{α}}f_θ(\mathbf{x}_t,t)+\sqrt{1-\bar{α}}\epsilon xt−1​=αˉ ​fθ​(xt​,t)+1−αˉ ​ϵ，其中 α ˉ t = ∏ s = 0 t ( 1 − β s ) \bar{α}_t=∏^t_{s=0}(1-β_s) αˉt​=∏s=0t​(1−βs​)， ϵ ∼ N ( 0 , I ) \epsilon∼\mathcal{N}(0,I) ϵ∼N(0,I)。 在clamping技巧中，模型还将预测向量 f θ ( x t , t ) f_θ(\mathbf{x}_t,t) fθ​(xt​,t)映射到其最近的词嵌入序列。现在，采样步骤变为 x t − 1 = α ˉ ⋅ C l a m p ( f θ ( x t , t ) ) + 1 − α ˉ ϵ \mathbf{x}_{t-1}=\sqrt{\bar{α}}·Clamp(f_θ(\mathbf{x}_t,t))+\sqrt{1-\bar{α}}\epsilon xt−1​=αˉ ​⋅Clamp(fθ​(xt​,t))+1−αˉ ​ϵ。clamping技巧迫使预测的向量在中间的扩散步骤中致力于一个词，使向量预测更加精确，并减少了rounding errors。

五、Diffusion-LM的解码和可控生成

5.1、可控制的文本生成

我们现在描述一个能够对Diffusion-LM进行即插即用控制的程序。我们的控制方法受到§3.1中贝叶斯公式的启发，但我们不是直接对离散文本进行控制，而是对Diffusion-LM定义的连续潜变量 x 0 : T \mathbf{x}_{0:T} x0:T​序列进行控制，并应用舍入步骤将这些潜变量转换成文本。

控制 x 0 : T \mathbf{x}_{0:T} x0:T​相当于从后验 p ( x 0 ： T ∣ c ) = ∏ t = 1 T p ( x t − 1 ∣ x t , c ) p(\mathbf{x}_{0：T}|\mathbf{c})=∏^T_{t=1}p(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{c}) p(x0：T​∣c)=∏t=1T​p(xt−1​∣xt​,c)中解码，我们将这个联合推理问题分解为每个扩散步骤的控制问题序列： p ( x t − 1 ∣ x t , c ) ∝ p ( x t − 1 ∣ x t ) ⋅ p ( c ∣ x t − 1 , x t ) p(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{c})∝p(\mathbf{x}_{t-1}|\mathbf{x}_t)·p(\mathbf{c}|\mathbf{x}_{t-1},\mathbf{x}_t) p(xt−1​∣xt​,c)∝p(xt−1​∣xt​)⋅p(c∣xt−1​,xt​) 。我们通过先前控制扩散的工作中的条件独立性假设，进一步简化 p ( c ∣ x t − 1 , x t ) = p ( c ∣ x t − 1 ) p(\mathbf{c}|\mathbf{x}_{t-1},\mathbf{x}_t)=p(\mathbf{c}|\mathbf{x}_{t-1}) p(c∣xt−1​,xt​)=p(c∣xt−1​)因此，对于第 t t t步，我们对 x t − 1 \mathbf{x}_{t-1} xt−1​运行梯度更新： ∇ x t − 1 log ⁡ p ( x t − 1 ∣ x t , c ) = ∇ x t − 1 log ⁡ p ( x t − 1 ∣ x t ) + ∇ x t − 1 p ( c ∣ x t − 1 ) \nabla_{\mathbf{x}_{t-1}}\log p(\mathbf{x}_{t-1}|\mathbf{x}_t, \mathbf{c})=\nabla_{\mathbf{x}_{t-1}}\log p(\mathbf{x}_{t-1}|\mathbf{x}_t)+\nabla_{\mathbf{x}_{t-1}}p(\mathbf{c}|\mathbf{x}_{t-1}) ∇xt−1​​logp(xt−1​∣xt​,c)=∇xt−1​​logp(xt−1​∣xt​)+∇xt−1​​p(c∣xt−1​)其中 log ⁡ p ( x t − 1 ∣ x t ) \log p(\mathbf{x}_{t-1}|\mathbf{x}_t) logp(xt−1​∣xt​)和 p ( c ∣ x t − 1 ) p(\mathbf{c}|\mathbf{x}_{t-1}) p(c∣xt−1​)都是可微的：第一项由Diffusion-LM参数化，第二项由一个神经网络分类器参数化。

与图像设置中的工作类似，我们在扩散潜变量上训练分类器，并在潜空间 x t − 1 \mathbf{x}_{t-1} xt−1​上运行梯度更新，以引导它完成控制。这些图像扩散工作在每个扩散步骤中对 ∇ x t − 1 p ( c ∣ x t − 1 ) \nabla_{\mathbf{x}_{t-1}}p(\mathbf{c}|\mathbf{x}_{t-1}) ∇xt−1​​p(c∣xt−1​)采取一个梯度步骤。为了提高文本的性能并加快解码速度，我们引入了两个关键性的修改：流畅性正则化和多个梯度步骤。

为了生成流畅的文本，我们在控制目标上运行梯度更新，并进行流畅性正则化： λ log ⁡ p ( x t − 1 ∣ x t ) + p ( c ∣ x t − 1 ) λ\log p(\mathbf{x}_{t-1}|\mathbf{x}_t) + p(\mathbf{c}|\mathbf{x}_{t-1}) λlogp(xt−1​∣xt​)+p(c∣xt−1​)，其中 λ λ λ是一个超参数，用于交换流畅性（第一项）和控制（第二项）。虽然现有的扩散可控生成方法不包括目标中的 λ log ⁡ p ( x t − 1 ∣ x t ) λ\log p(\mathbf{x}_{t-1}|\mathbf{x}_t) λlogp(xt−1​∣xt​)项，但我们发现这个项对于生成流畅的文本很有帮助。由此产生的可控生成过程可以被视为一种随机解码方法，它平衡了最大化和采样 p ( x t − 1 ∣ x t , c ) p(\mathbf{x}_{t-1}|\mathbf{x}_t,\mathbf{c}) p(xt−1​∣xt​,c)，很像流行的文本生成技术，如核采样或低温采样。为了提高控制质量，我们对每个扩散步骤采取多个梯度步骤：我们对每个扩散步骤运行3步的Adagrad更新。为了减轻增加的计算成本，我们将扩散步骤从2000步降至200步，这加快了我们的可控生成算法，而不会对样本质量造成很大伤害。

5.2、最小化贝叶斯风险解码

许多条件性文本生成任务需要单一的高质量输出序列，如机器翻译或句子填充。在这些设置中，我们应用最小贝叶斯风险（MBR）解码来聚合一组从Diffusion-LM中抽取的样本 S \mathcal{S} S，并选择在损失函数 L \mathcal{L} L（例如，负BLEU分数）下实现最小预期风险的样本： w ^ = arg min ⁡ w ∈ S ∑ w ′ ∈ S 1 ∣ S ∣ L ( w , w ′ ) \hat{\mathbf{w}}= \argmin_{\mathbf{w}∈\mathcal{S}} ∑_{\mathbf{w}'∈\mathcal{S}}\frac{1}{|\mathcal{S}| } \mathcal{L}(\mathbf{w}, \mathbf{w}') w^=w∈Sargmin​∑w′∈S​∣S∣1​L(w,w′)。我们发现，MBR解码通常会返回高质量的输出，因为低质量的样本会与其余的样本不一样，并受到损失函数的惩罚。

六、总结

我们提出了Diffusion-LM，一种基于连续扩散的新型可控语言模型，它可以实现新形式的复杂细粒度控制任务。我们证明了Diffusion-LM在6个细粒度控制任务中的成功：我们的方法几乎是先前方法控制成功率的两倍，与需要额外训练的基线微调方法相比具有竞争力。

我们发现Diffusion-LM实现的复杂控制很有说服力，而且扩散-LM是对目前离散自回归生成范式的实质性背离，这让我们感到兴奋。与任何新技术一样，我们构建的Diffusion-LMs也有缺点：（1）它的困惑度更高；（2）解码速度大大降低；（3）训练收敛更慢。我们相信，通过更多的后续工作和优化，这些问题中的许多都可以得到解决，这种方法将变成一种引人注目的大规模可控生成的方式。