模型设计

• First, we must define an embedding function that maps discrete text into a continuous space.
• Second, we require a rounding method to map vectors in embedding space back to words.

端到端训练

为了将离散文本应用到连续扩散模型上，这里设置了一个embedding函数${\mathrm{E}}_{\mathrm{MB}}(w_i)$将每个词映射到词向量${\mathbb{R}}^d$。对于长度为$n$的序列$\mathbf{w}$：

$${\mathrm{E}}_{\mathrm{MB}}(\mathbf{w})=\left[{\mathrm{E}}_{\mathrm{MB}}\left(w_1\right),\dots ,{\mathrm{E}}_{\mathrm{MB}}\left(w_n\right)\right]\in {\mathbb{R}}^{nd}$$

作者在实验中发现，使用预训练的word embedding效果不如使用随机高斯噪声初始化之后训练出来的embedding，所以加了这样一个网络，实现离散词序列$\mathbf{w}$到${\mathbf{x}}_0$的马尔科夫转换，参数化为：

$$q_{\varphi }\left({\mathbf{x}}_0\mid \mathbf{w}\right)=\mathcal{N}({\mathrm{E}}_{\mathrm{MB}}(\mathbf{w}),{\sigma }_{0}I)$$

与之对应的反向过程添加了一个可训练的舍入步骤，参数化为：

$$p_{\theta }\left(\mathbf{w}\mid {\mathbf{x}}_0\right)=\prod _{i=1}^n{p}_{\theta }\left(w_i\mid x_i\right)$$

其中$p_{\theta }\left(w_i\mid x_i\right)$是一个softmax分布。

因为增加了一个嵌入步骤和一个舍入步骤，添加的这两个网络是和扩散模型进行联合训练的，因此要对训练目标函数改进：

Lvlbe2e(w)=Eqϕ(x0∣w)[Lvlb(x0)+log⁡qϕ(x0∣w)−log⁡pθ(w∣x0)]],Lsimple e2e(w)=Eqϕ(x0:T∣w)[Lsimple (x0)+∥EMB(w)−μθ(x1,1)∥2−log⁡pθ(w∣x0)].begin{aligned} mathcal{L}_{mathrm{vlb}}^{mathrm{e2e}}(mathbf{w}) & left.=underset{q_phileft(mathbf{x}_0 mid mathbf{w}right)}{mathbb{E}}left[mathcal{L}_{mathrm{vlb}}left(mathbf{x}_0right)+log q_phileft(mathbf{x}_0 mid mathbf{w}right)-log p_thetaleft(mathbf{w} mid mathbf{x}_0right)right]right],

mathcal{L}_{text {simple }}^{mathrm{e2e}}(mathbf{w}) & =underset{q_phileft(mathbf{x}_{0: T} mid mathbf{w}right)}{mathbb{E}}left[mathcal{L}_{text {simple }}left(mathbf{x}_0right)+left|mathrm E_{mathrm{MB}} (mathbf{w})-mu_thetaleft(mathbf{x}_1, 1right)right|^2-log p_thetaleft(mathbf{w} mid mathbf{x}_0right)right] .end{aligned}
Lvlbe2e​(w)Lsimple e2e​(w)​=qϕ​(x0​∣w)E​[Lvlb​(x0​)+logqϕ​(x0​∣w)−logpθ​(w∣x0​)]],=qϕ​(x0:T​∣w)E​[Lsimple ​(x0​)+∥EMB​(w)−μθ​(x1​,1)∥2−logpθ​(w∣x0​)].​

${\mathcal{L}}_{\mathrm{simple}}^{\mathrm{e}2\mathrm{e}}(\mathbf{w})$是${\mathcal{L}}_{\mathrm{vlb}}^{\mathrm{e}2\mathrm{e}}(\mathbf{w})$的简化版，是根据DDPM设计的目标函数。

DDPM的目标函数：

$${\mathcal{L}}_{\mathrm{vlb}}\left({\mathbf{x}}_0\right)=\underset{q\left({\mathbf{x}}_{1:T}\mid {\mathbf{x}}_0\right)}{\mathbb{E}}\left[\log \frac{q\left({\mathbf{x}}_T\mid {\mathbf{x}}_0\right)}{p_{\theta }\left({\mathbf{x}}_T\right)}+\sum _{t=2}^T\log \frac{q\left({\mathrm{x}}_{t-1}\mid {\mathbf{x}}_0,{\mathbf{x}}_t\right)}{p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)}-\log p_{\theta }\left({\mathbf{x}}_0\mid {\mathbf{x}}_1\right)\right]$$
$${\mathcal{L}}_{\text{simple }}\left({\mathbf{x}}_0\right)=\sum _{t=1}^T\underset{q\left({\mathbf{x}}_t\mid {\mathbf{x}}_0\right)}{\mathbb{E}}{\Vert {\mu }_{\theta }\left({\mathbf{x}}_t,t\right)-\hat{\mu }\left({\mathbf{x}}_t,{\mathbf{x}}_0\right)\Vert }^2$$

可视化之后可以看到学到的embedding是有意义的。

减少舍入误差

embedding是将离散文本映射到连续的${\mathbf{x}}_0$上，那与之对应的反向过程就应该是将模型预测出来的${\mathbf{x}}_0$转换回离散的文本。

舍入步骤是根据$\mathrm{argmax}{p}_{\theta }\left(\mathbf{w}\mid {\mathbf{x}}_0\right)={\prod }_{i=1}^n{p}_{\theta }\left(w_i\mid x_i\right)$选择每个位置上可能性最大的词。理想状态下通过这个舍入步骤就可以将模型输出的${\mathbf{x}}_0$映射回离散文本，因为去噪步骤应该能让${\mathbf{x}}_0$恰好回到某个单词的embedding上。

然而实际情况是不行的，模型输出是不会精确到某个单词的embedding。

作者认为造成上述问题的原因，是在目标函数中${\mathcal{L}}_{\mathrm{simple}}^{\mathrm{e}2\mathrm{e}}({\mathbf{x}}_0)$对${\mathbf{x}}_0$结构的建模不够重视。

$${\mathcal{L}}_{\text{simple }}({\mathbf{x}}_0)=\sum _{t=1}^T{\mathbb{E}}_{{\mathbf{x}}_t}{\Vert {\mu }_{\theta }\left({\mathrm{x}}_t,t\right)-\hat{\mu }\left({\mathrm{x}}_t,{\mathbf{x}}_0\right)\Vert }^2$$

其中的${\mu }_{\theta }\left({\mathrm{x}}_t,t\right)$网络直接去预测时间步$t$去噪的$p_{\theta }\left({\mathbf{x}}_{t-1}\mid {\mathbf{x}}_t\right)$的均值。${\mathbf{x}}_0$对单词的约束只会出现在$t$接近于0的项中。因此需要对目标函数进行调整，去强调${\mathbf{x}}_0$。

作者对${\mathcal{L}}_{\text{simple }}$进行调整，强调模型在目标函数的每一项中都去显式地建模${\mathbf{x}}_0$。
作者推导出一个类似于${\mathcal{L}}_{\text{simple }}$的使用${\mathbf{x}}_0$参数化的公式：

$${\mathcal{L}}_{{\mathbf{x}}_0\text{-simple }}^{e2e}\left({\mathbf{x}}_0\right)=\sum _{t=1}^T{\mathbb{E}}_{{\mathbf{x}}_t}{\Vert f_{\theta }\left({\mathbf{x}}_t,t\right)-{\mathbf{x}}_0\Vert }^2$$

其中网络$f_{\theta }\left({\mathbf{x}}_t,t\right)$直接去预测${\mathbf{x}}_0$，这样会让神经网络去预测每一项的${\mathbf{x}}_0$。
使用修改之后的目标去训练模型，模型很快就会学到让${\mathbf{x}}_0$能回到以word embedding为中心。

clamping trick

作者将其称为clamping trick，在clamping trick中，模型将${\mathbf{x}}_t$降噪为${\mathbf{x}}_{t-1}$的生成过程：

1. 通过$f_{\theta }({\mathbf{x}}_t,t)$估计出一个${\mathbf{x}}_0$

2. 在这个估计的条件相爱对${\mathbf{x}}_{t-1}$进行采样

3. ${\mathbf{x}}_{t-1}=\sqrt{\bar{\alpha }}{f}_{\theta }\left({\mathbf{x}}_t,t\right)+\sqrt{1-\bar{\alpha }}ϵ$
   其中${\bar{\alpha }}_t={\prod }_{s=0}^t\left(1-{\beta }_s\right)$， $ϵ\sim \mathcal{N}(0,I)$

   这一步就是用的DDPM中的那个，因为都是高斯核，所以${\mathbf{x}}_t$可以由${\mathbf{x}}_0$一步得到。

   $\begin{array}{rl}{\mathbf{x}}_t& =\sqrt{{\alpha }_t}{x}_{t-1}+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}^{\ast }\\ & =\sqrt{{\alpha }_t}\left(\sqrt{{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_{t-1}}{ϵ}_{t-2}^{\ast }\right)+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}^{\ast }\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}{ϵ}_{t-2}^{\ast }+\sqrt{1-{\alpha }_t}{ϵ}_{t-1}^{\ast }\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}}}^2+\sqrt{1-{\alpha }_t^2}}{ϵ}_{t-2}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{{\alpha }_t-{\alpha }_t{\alpha }_{t-1}+1-{\alpha }_t}{ϵ}_{t-2}\\ & =\sqrt{{\alpha }_t{\alpha }_{t-1}}{x}_{t-2}+\sqrt{1-{\alpha }_t{\alpha }_{t-1}}{ϵ}_{t-2}\\ & =\dots \\ & =\sqrt{\prod _{i=1}^t{\alpha }_i}{x}_0+\sqrt{1-\prod _{i=1}^t{\alpha }_i}{ϵ}_0\\ & =\sqrt{{\bar{\alpha }}_t}{x}_0+\sqrt{1-{\bar{\alpha }}_t}{ϵ}_0\end{array}$

clamping trick 会将网络$f_{\theta }({\mathbf{x}}_t,t)$的预测结果映射到接近的word embedding 序列上。
现在采样步骤就变为了：

$${\mathbf{x}}_{t-1}=\sqrt{\bar{\alpha }}\cdot \mathrm{Clamp}\left(f_{\theta }\left({\mathbf{x}}_t,t\right)\right)+\sqrt{1-\bar{\alpha }}ϵ$$

clamping trick 迫使扩散模型降噪过程中每一步都去计算一个word embedding，使向量预测更为准确，以此减少舍入误差。

作者在这里提示将开始使用clamping trick的起始位置设置为超参数。具体原因看论文P5。

论文信息

论文地址：[2205.14217] Diffusion-LM Improves Controllable Text Generation (arxiv.org)

代码地址：XiangLi1999/Diffusion-LM: Diffusion-LM (github.com)

本文正在参加「金石计划」