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

发布时间 2023-03-27 17:37:06作者: 馒头and花卷

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\) 会产生相当不错的结果.