Difformer: Empowering Diffusion Models on the Embedding Space for Text Generation

目录

• 概
• 符号说明
• 主要内容

Gao Z., Guo J., Tan X., Zhu Y., Zhang F., Bian J. and Xu L. Difformer: Empowering diffusion models on the embedding space for text generation. arXiv preprint arXiv:2212.09412, 2022.

概

本文发现并改进了原先 Diffusion 模型的缺点 (针对文本生成任务).

符号说明

• $\bm{x} = [x_1, x_2, \cdots, x_m]$, source sentence;
• $\bm{y} = [y_1, y_2, \cdots, y_n]$, target sentence;
• 前向过程:
  $$q(\bm{z}_t|\bm{z}_{t-1}) = \mathcal{N}(\bm{z}_t; \sqrt{1 - \beta_t} \bm{z}_{t-1}, \beta_t \bm{I}), \\ \bm{z}_0 \sim \mathcal{N}(\bm{e(y)}; \beta_0 \bm{I}).$$
  其中 $\bm{e}(\cdot)$ 将离散的 token $y$ 映射到 embedding 空间.
• 反向过程:
  $$p_{\theta}(\bm{z}_{t-1}|\bm{z}_t) = \mathcal{N}(\bm{z}_{t-1}; \bm{\mu}_{\theta}(\bm{z}_t, t), \bm{\Sigma}_{\theta}(\bm{z}_t, t)).$$
  一般, 方差和前向保持一致, 然后我们实际上拟合的是
  $$\hat{\bm{z}}_0 = f_{\theta}(\bm{z}_t, t).$$

主要内容

• Diffusion 在文本生成之上的应用大抵为优化如下的损失:

  $$\underbrace{\mathbb{E}_{\bm{z}_0, t} [\|\hat{\bm{z}}_0 - \bm{z}_0\|^2]}_{\mathcal{L}_{vlb}} + \underbrace{(- \mathbb{E}_{\bm{z}_0, \bm{y}} \log p_{\theta}(\bm{y}|\bm{z}_0))}_{\mathcal{L}_{round}}.$$

  前者是普通的 diffusion 损失, 后者是用于离散到连续的对齐.

• 问题1: 可学习的 embeddings: 和 image, video 不同的是, embedding 是可学习的, 这可能导致训练的不稳定.

• 解决办法1: 如上图所示, 我们用不同时刻 t 的预测 $\hat{\bm{z}}_0$ 来靠近 $\bm{y}$, 即

  $$\mathcal{L}_{anchor} = -\log p_{\theta}(\bm{y}|\hat{\bm{z}}_0)$$

  来替代 $\mathcal{L}_{round}$.

• 问题2: embedding 的 norm 不一致: 由于词频不一致, 作者发现训练过程中高频词的 norm 会普遍大一点. 此时加入相同量级的噪声对于信号的干扰程度就不一致了:
  

• 解决办法2: 故作者在 embedding 之后加了一个 embedding normalization 层 (不过看起来对实际效果的提升不是特别大):

  $$\text{LN}(\bm{e}(y_i)) = \frac{\bm{e}(y_i) - \mathbb{E}[\bm{e}(y_i)]}{\sqrt{\mathbb{V}[\bm{e}(y_i)] + \epsilon}} \odot \bm{\gamma} + \bm{\eta}.$$

• 问题3: noise schedule 导致扩散过程的大部分阶段都是'无意义'的: 作者发现, 对于一般的前向过程和 noise schedule 而言, 整个过程的大部分时间的干扰效果都是很差的, 此时很难充分训练 (特别是对于 sqrt 这种 schedule):
  

• 解决办法3: 作者对整体的扰动的方差进行一个扩张:

  $$q(\bm{z}_t|\bm{z}_{t-1}) = \mathcal{N}(\bm{z}_t; \sqrt{1 - \beta_t} \bm{z}_{t-1}, \beta_t F^2 \bm{I}).$$

  一般的扩散 $F=1$, 作者发现 $F=4$ 会产生相当不错的结果.