论文阅读[精读]-Diffusion-LM Improves Controllable Text Generation(2)

之前我们聊了

• DDPM的原理 DDPM笔记
• 也聊了 score based models 原理 博客阅读笔记

今天我们满足了所有前置知识，可以来看这篇的本体了。如何把diffusion models 用到可控文本生成中。虽然是老本行，但这篇还是精读一下。

摘要

• 可控文本生成是一个重要的任务
• 已有的方法对粗粒度的控制尚可(情感)，但对细粒度的控制效果不好(句法树)
• 本文设计了一个非自回归的、连续的diffusion model来解决这个问题
• 去噪过程中的多个中间量可以指导可控的文本生成过程
• 在6个可控生成任务中，Diffusion-LM模型大幅战胜了SOTA

Introduction

自回归的语言模型的生成能力很强，但可控性很弱。如果想做可控的生成：

• 直观方法是做有监督的训练，但这样需要每个任务炼一个模型，成本高。
• plug-and-play的方法，冻住模型。但此类方法效果差，难以实现细粒度的控制

传统的连续diffusion model不适用于文本，因为word是离散的。本文通过添加一个embedding步骤和一个rounding步骤解决这个问题。同时，本文通过一种可控的gradient update指导模型可以平衡fluency和controllable两种属性。

本文在6个可控任务中大幅战胜plug-and-play SOTA，同时战胜了fine-tune的方法。甚至，在不用classifier的情况下和自回归的模型有的一拼。

Related work

Diffusion Models for Text： 已有方法都是在离散的空间做这件事，提到的论文我也看了，有时间给大家写一篇笔记。本文用了连续的方法来做。

Autoregressive and Non-autoregressive LMs： 已有的模型大多是是从左到右训练的，但是对于infilling等任务来说，这限制了模型的发挥。这个问题，有一些解决方法。本文的diffusion model是classifier指导的，所以可以看到句子的全部。

Plug-and-Play Controllable Generation：冻结模型，训一个类似分类器的东西。已有方法大多基于自回归模型：

• FUDGE：通过对生成质量的估计来reweight LM prediction
• GeDi, DExperts： 通过fine—tune一个小的LM来reweight LM prediction
  • GeDi的论文没准后面也写个笔记

PPLM方法和本文最接近，通过一个分类器做梯度上升。但由于是自回归的，只能从左到右，不能实现细粒度的控制。

Problem Statement and Background

Generative Models and Controllable Generation for Text

如何在控制条件c的情况下生成w，最大化$P(w|c)$.由于

$$\begin{aligned} P(w|c) & = \frac{P(w)\times P(c|w)}{P(c)} \\ & \propto P(w)\times P(c|w) \end{aligned}$$

如果对w的生成求导：

$$\bigtriangledown_w \log P(w|c) = \bigtriangledown_w \log P(w)+ \bigtriangledown_w \log P(c|w)$$

可以发现，右边是正常模型的输出的梯度，后面是某个分类器的输出的梯度。这个其实就是经典的在图像领域做可控生成的diffusion model的方法

Autoregressive Language Models

$$P_{lm}(w) = P_{lm}(w_1) \sum_{i=2}^N P_{lm}(w_i| w_{<i})$$

这是一个经典的从左到右的解码过程。一般大家都是用一个transformer结构来做$P_{lm}(w)$

Diffusion Models for Continuous Domains

这一部分是DDPM的子集，可以参考

• DDPM笔记
• what are diffusion models

结论是：

$$\mathcal{L}_{simple} = \sum_{t=1}^T \mathop{\mathbb{E}} \limits_{q(x_t| x_0)} || \mu_\theta(x_t,t) - \hat{\mu}(x_t,x_0)||^2 \\ \hat{\mu}(x_t,x_0) = \text{mean} [q(x_{t−1}|x_0, x_t)]$$

大概就是估计模型对每一个$x_{t-1}$都做估计，把估计和实际值的偏差叠加在一起。所谓的”mean”是指后面来自一个高斯分布，这个是期望

Diffusion-LM: Continuous Diffusion Language Modeling

其实这一部分开始才是本文的方法：

• 首先用一个EMB层把离散的word映射到连续空间
• 定义一个端到端的training objectives
• 定义一个rounding method再把生成的连续向量映射回word
• 同时，为了rounding的效果，提出了专门的训练策略

End-to-end Training

$$EMB(w) = [EMB(w_1), . . . , EMB(w_n)] \in R^{nd}.$$

是一个词向量映射，把n个word映射成n个d维向量，这个是可学习的。

接下来把

$$q_\phi (x_0|w) = \mathcal{N} (EMB(w), \sigma_0I)$$

视做diffusion process的第一步(其实$\phi$就是embedding层模型……)

此外，添加rounding过程

$$p_\theta(w | x_0) = \sum^n_{i=1} p_\theta(w_i | x_i)$$

$p_\theta(w_i | x_i)$就是一个经典的softmax变换。

最后：

$$\begin{aligned} \mathcal{L}_{vlb}(x_0) & = \mathbb{E}{(x_0\sim p_{data})} [\log p_\theta (x_0 )] \\ & = \mathop{\mathbb{E}}\limits_{x_{1:T}|x_0} \left[ \log \frac{P_\theta(x_T | x_0)}{P_\theta (x_T)} + \sum_{t=2}^T \log \frac{q(x_{t−1}|x_0, x_t)}{P_\theta(x_{t-1} | x_t)} -\log P_\theta(x_{0} | x_1) \right] \end{aligned}$$ $$\begin{aligned} \mathcal{L}_{vlb}^{e2e}(w) & = \mathop{\mathbb{E}}\limits_{x_0|w}\left[\mathcal{L}_{vlb}(x_0) +\log q_\phi(x_0 | w) - \log p_\theta(w | x_0) \right] \\ & = \mathop{\mathbb{E}}\limits_{x_{0:T}|w} \left[ \mathcal{L}_{simple}(x_0) + ||EMB(w)−\mu_\theta (x_1 ,1) ||^2 -\log p_\theta(w | x_0) \right] \end{aligned}$$

这是一个端到端的训练loss。只需要给出很多w，用DDPM的算法，就可以慢慢让模型学习了。作者也发现，这样模型和embedding一起学出来，最后word embedding也可以表现出一定的聚合性

Reducing Rounding Errors

rounding的过程其实是对每一个位置选取最接近的embedding。

作者提到，如果像一般的DDPM那样用模型预测$p\theta(x{t−1} | x_t) $的均值，会导致输入长度很短的时候效果很差(图像领域没出问题，因为最短也是32x32)。

解决办法是用模型直接预测$x_0$.

$$\mathcal{L}_{x0-simple}^{e2e}(x_0) = \sum_{t=1}^T \mathbb{E}_{x_t} ||f_\theta(x_t,t) - x_0||^2\\ f_\theta(x_t,t) \to x_0$$

这里有点怪：换模型的目标好像会影响sample方法？还是说同样是预测$x{t-1}$，但是把$x_0$视为$x{t-1},x_t$的线性组合

这样模型可以很快发现$x_0$的每一个位置都应该对应某个word的embedding。

同时，修改原来的sample方法：

$$x_{t−1} =\sqrt{\overline{\alpha}}f_\theta(x_t,t)+ \sqrt{1- \overline{\alpha}} \epsilon \\ \overline{\alpha} = \sum_{s=0}^t (1 -\beta_t) ,\quad \epsilon \sim \mathcal{N}(0,I)$$

(能成立是因为连续的高斯采样等价于从最开始只做一次高斯采样)

上面的方法随着t的递减，逐渐从随机噪声恢复成一个无噪音的数据

如果希望模型可以显式地对齐某个embedding，可以在每一步都加一个clamp。取和$f_\theta(x_t,t)$最接近的word embedding

$$x_{t−1} =\sqrt{\overline{\alpha}} \text{Clamp}(f_\theta(x_t,t))+ \sqrt{1- \overline{\alpha}} \epsilon$$

同时作者也提到，对于t很大时做clamp是一个负优化，因为这时$f_\theta(x_t,t)$还没有对应到某些embedding，因此选取clamp的开始时间是一个超参。

Decoding and Controllable Generation with Diffusion-LM

Controllable Text Generation

正如前面提到的：

$$\bigtriangledown_{x_{t-1}} \log P(x_{t-1}|x_t,c) = \bigtriangledown_{x_{t-1}} \log P(x_{t-1} | x_t)+ \bigtriangledown_{x_{t-1}} \log P(c|x_{t-1})$$

其中式子右边的第一项是Diffusion-LM的输出。右边是一种分类器的输出。如果对于目标任务，我们在目标数据集训练一个分类器，在sample时每一步都根据这个输出做更新，那么最终就会生成符合条件c的结果。

同时作者为了细粒度的控制”可控性“，定义一个超参$\lambda$来平衡流畅度和控制程度

$$p(x_{t−1} | x_t, c) = \lambda \log P(x_{t-1} | x_t)+\log P(c|x_{t-1})$$

作为最终用来指导梯度的东西。作者提到每次得到$x_t$之后：

• sample $x_{t-1}$

• 再做三轮Adagrad更新

• 把sample步骤从2000步变成200步来加速

这个方法我理解类似于之前微分方程那篇文章提到的PC更新方式

Minimum Bayes Risk Decoding

这个方法可以得到一个高质量的结果。作者用下面的最小化risk函数：

$$\hat{w} = \mathop{\text{argmin}}_{w \in S} \sum_{w' \in S} \frac{1}{|S|} \mathcal{L}_{bleu}(w,w')$$

这种方法可以参考 MBR解码博客

Experimental Setup

作者在很多个任务上使用了Diffusion-LM方法：

• classifier-guided
• classifier-free：不需要分类器，比如说句子长度。只要最开始的初始化向量变成需要长度即可(控制率显然是100%)

Diffusion model里面的内部解码层作者使用了一个非常传统的transformer结构，80M参数，可以说是很少:

• 步骤T = 2000
• 词向量维度 d = 16-128
• 句长n=64

输入就像上图input所示：通过分类器来在解码中每一次去噪都提供梯度

Result

主要结果就如上图所示。

• 基本上把SOTA得分翻倍了，但这个SOTA是指plug-and-play的方法
• 在Syntax Tree， Syntax Spans任务上，甚至比fine-tune的模型效果还好

同时作者也做了综合的控制实验，比如说同时限定两个条件。可以看出，效果很好，方法就是单纯的把控制的梯度转换成两个classifier梯度的综合。

在infilling任务中，Diffusion-LM在没有做额外训练的情况下就表现很好

我的思考

这篇工作很新颖，很有启发，我觉得可以做的点还有以下两个方面：

• 这个方法好像不能实现平行的转换，只能”根据某些控制生成特定的句子“。像之前DALL.E都是可以把特定的句子翻译成对应的图像。
• 感觉短文本领域，只有更加平行，才能更加泛用。不过在长文本领域，可能这种没有平行的生成会更好，但作者并没有测试。