CoCon: A Self-Supervised Approach for Controlled Text Generation
Motivation
在使用语言模型做文本生成的任务时，给定提示文字 $x_{:t-1}=\{x_1,\cdots,x_{t-1}\}$，后续的文本 $\{x_{t}, \cdots, x_{l}\}$ 是通过自回归的方式生成的：
\[p\left(x_{t}, \ldots, x_{l} \mid x_{1}, \ldots, x_{t-1}\right)=\prod_{i=t}^{l} p\left(x_{i} \mid x_{1}, \ldots, x_{i-1}\right)\]之前的利用语言模型做可控文本生成任务的研究工作，如PPLM和CTRL等，都认为 $p(x)$ 可以以目标属性或控制代码为条件，来控制文本的情感或主题：
\[p\left(x_{t}, \ldots, x_{l} \mid x_{1}, \ldots, x_{t-1}\right)=\prod_{i=1}^{l} p\left(x_{i} \mid \mathbf{a},\left\{x_{1}, \ldots, x_{i-1}\right\}\right)\]其中，$\mathbf{a}$ 是目标属性。这些控制生成的方法都是基于一个全局的属性（情感或主题），而不是更加局部的内容（词或词组）。因此，这激发了一种以输入目标内容 $\mathbf{c}$ 为条件的方法，以便对文本生成进行更细粒度的控制：
\[p\left(x_{t}, \ldots, x_{l} \mid x_{1}, \ldots, x_{t-1}\right)=\prod_{i=1}^{l} p\left(x_{i} \mid \mathbf{c},\left\{x_{1}, \ldots, x_{i-1}\right\}\right)\]Model Architecture
语言模型的生成过程可以分为两个部分：编码器 $enc$ 和解码器 $dec$ 。编码器用作特征提取器，它接受输入序列的embedding，并输出中间表示 $\mathbf{h}_{: t-1}=\operatorname{enc}\left(x_{: t-1}\right)$。然后，解码器根据这一个中间表示产生下一个词：
\[\mathbf{o}\_{t}=\mathrm{LM}\left(x_{: t-1}\right)=\operatorname{dec}\left(\operatorname{enc}\left(x_{: t-1}\right)\right)=\operatorname{dec}\left(\mathbf{h}_{: t-1}\right)\]作者提出一个CoCon块，将 $\mathbf{h}$ 转换为以目标内容 $\mathbf{c}$ 为条件的表示：
\[\mathbf{h}_{: t-1}^{\prime}=\operatorname{CoCon}\left(\mathbf{h}_{: l_{c}}^{(\mathbf{c})}, \mathbf{h}_{: t-1}\right)\]其中，$\mathbf{h}_{: l_{c}}^{(\mathbf{c})}=\operatorname{enc}(\mathbf{c})$ 是目标内容的表示，$l_c$ 是目标文本的长度。作者使用一个Transformer块来作为CoCon块。与传统的语言模型中的注意力层相似， $\mathbf{Q}, \mathbf{K}, \mathbf{V} \in \mathbb{R}^{(t-1) \times d}$ 矩阵通过中间表示的线性变换得到。为了加入内容表示 $\mathbf{h}_{: l_{r}}^{(\mathbf{c})}$，关于内容的 $\mathbf{K}^{(\mathbf{c})}, \mathbf{V}^{(\mathbf{c})} \in \mathbb{R}^{l_{c} \times d}$ 矩阵也需要计算，然后与矩阵拼接起来：
\[\mathbf{K}^{\prime}=\left[\mathbf{K}^{(\mathbf{c})} ; \mathbf{K}\right], \quad \mathbf{V}^{\prime}=\left[\mathbf{V}^{(\mathbf{c})} ; \mathbf{V}\right], \quad \mathbf{A}=\operatorname{Softmax}\left(\mathbf{Q} \mathbf{K}^{\prime \top}\right) \mathbf{V}^{\prime}=\operatorname{Softmax}(\mathbf{W}) \mathbf{V}^{\prime}\]其中，$\mathbf{A}=\{\mathbf{a}_{1}, \ldots, \mathbf{a}_{t-1}\}$ 和 $\mathbf{W} \in \mathbb{R}^{(t-1) \times\left(l_{c}+t-1\right)}$ 代表了注意力权重。CoCon块的输出由如下得到：
\[\mathbf{h}_{i}^{\prime}=\operatorname{FF}\left(\mathbf{a}_{i}\right), \quad \tilde{\mathbf{o}}_{t}=\operatorname{dec}\left(\left[\mathbf{h}_{: t-2} ; \mathbf{h}_{t-1}^{\prime}\right]\right), \quad p_{\theta, \psi}\left(\tilde{x}_{t} \mid \mathbf{c}, x_{: t-1}\right)=\operatorname{Softmax}\left(\tilde{\mathbf{o}}_{t}\right)\]Multiple Content Inputs：CoCon允许在同一个生成过程中输入多个目标内容。假设有 $N$ 个目标内容 $\left(\mathbf{c^1},\cdots,\mathbf{c^N}\right)$，输出文本可以以注意力的键值矩阵为条件得到：
\[\mathbf{K}^{\prime}=\left[\mathbf{K}^{\left(\mathbf{c}^{1}\right)} \ldots \mathbf{K}^{\left(\mathbf{c}^{N}\right)} ; \mathbf{K}\right], \quad \mathbf{V}^{\prime}=\left[\mathbf{V}^{\left(\mathbf{c}^{1}\right)} \ldots \mathbf{V}^{\left(\mathbf{c}^{N}\right)} ; \mathbf{V}\right], \quad \mathbf{A}=\operatorname{Softmax}\left(\mathbf{Q} \mathbf{K}^{\prime \top}\right) \mathbf{V}^{\prime}\]Strength of Content Conditioning： 可以为内容注意力权重 $\mathbf{W}_{:,: l_{c}} \in \mathbb{R}^{(t-1) \times l_{c}}$ 添加偏置项 $\tau_{content}$。正的偏置项可以增强内容控制，负的偏置会减轻控制作用。
Self-Supervised Learning
作者采用自我监督的学习方法训练 CoCon。给定长度为 $l$ 的自然语言文本 $\mathbf{x}=\{x_{1}, \ldots, x_{t-1}, x_{t}, \ldots, x_{l}\}$，可以将其切分为两个连续的片段 $\mathbf{x}^{a}=\{x_{1}, \ldots, x_{t-1}\}$ 和 $\mathbf{x}^{b}=\{x_{t}, \ldots, x_{l}\}$，于是有 $\mathbf{x}=\left[\mathbf{x}^{a} ; \mathbf{x}^{b}\right]$。
Self Reconstruction Loss：将 $\mathbf{x}^{b}$ 作为目标内容的输入，由 $\mathbf{x}^{a}$ 重构出完整的句子 $\mathbf{x}$。具体来说，先计算输入文本 $\mathbf{x}$ 和 $\mathbf{c}$ 的中间表示： \(\mathbf{h}_{: l}=\operatorname{enc}(\mathbf{x})=\operatorname{enc}\left(x_{: l}\right), \quad \mathbf{h}_{: l_{c}}^{(\mathbf{c})}=\operatorname{enc}(\mathbf{c})=\operatorname{enc}\left(x_{t: l}\right)\)
其中，$l_c=l-t+1$ 是 $\mathbf{c}$ 的长度。基于内容 $\mathbf{c}$ 的中间表示可以通过CoCon块计算得到：
\[\mathbf{h}_{i}^{\prime}=\operatorname{CoCon}\left(\mathbf{h}_{: l_{\mathrm{r}}}^{(\mathbf{c})}, \mathbf{h}_{: i}\right), \quad \forall i \geq t-1\] \[\tilde{\mathbf{o}}_{i+1}=\operatorname{dec}\left(\left[\mathbf{h}_{: t-2} ; \mathbf{h}_{t-1: i}^{\prime}\right]\right), \quad p_{\theta, \psi}\left(\tilde{x}_{i+1} \mid \mathbf{c}, x_{: i}\right)=\operatorname{Softmax}\left(\tilde{\mathbf{o}}_{i+1}\right), \quad \forall i \geq t-1\]通过语言模型的训练目标，我们得出了 self reconstruction loss，训练CoCon通过将 $\mathbf{x}^b$ 本身作为内容输入，来预测产生 $\mathbf{x}^b$ 自身的词：
\[\mathcal{L}_{\mathrm{self}}=-\sum_{i=t}^{l} \log p_{\theta, \psi}\left(x_{i} \mid\left(\mathbf{c}=\mathbf{x}^{b}\right),\left\{x_{1}, \ldots, x_{i-1}\right\}\right)\]作者还在CoCon的注意力层中应用了 $\mathbf{c}$-mask，使得 $\mathbf{h}_{i}^{\prime}$ 中不存在由 $\mathbf{h}_{i+1}$ 计算得到的值。
Null Content Loss：为了鼓励 CoCon 的输出遵循提示文本 $\mathbf{x}^{a}$，而无需依赖 $\mathbf{x}^{b}$，作者提出了一个新的损失函数，它将 self reconstruction loss 中的内容输入用空字符代替：
\[\mathcal{L}_{\mathrm{null}}=-\sum_{i=t}^{l} \log p_{\theta, \psi}\left(x_{i} \mid(\mathbf{c}=\varnothing),\left\{x_{1}, \ldots, x_{i-1}\right\}\right)\]Cycle Reconstruction Loss：在推理阶段，CoCon 的目标内容 $\mathbf{c}$ 不太可能与提示文字 $\mathbf{p}$ 共现。因此，为了鼓励 $\mathbf{c}$ 和 $\mathbf{p}$ 来自不同来源的情况下的泛化性，作者引入了一个 cycle reconstruction 训练。CoCon 的自回归生成可以写为：
\[\mathbf{y}=f_{\theta, \psi}(\mathbf{c}, \mathbf{p})\]然后 $[\mathbf{p};\mathbf{y}]$ 可以作为一个流畅的句子，$\mathbf{y}$ 的内容由 $\mathbf{c}$ 决定。
第一步，以从 $\mathbf{x}$ 中得到的目标内容 $(\mathbf{c})$ 和从 $\mathbf{x^{\prime}}$ 中得到的提示文本 $(\mathbf{p})$ 为条件，计算CoCon的输出：
\[\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}=f_{\theta, \psi}\left(\left(\mathbf{c}=\mathbf{x}^{b}\right),\left(\mathbf{p}=\mathbf{x}^{\prime a}\right)\right)\]其中，$\mathbf{x}=\left[\mathbf{x}^{a} ; \mathbf{x}^{b}\right] $ 和 $\mathbf{x}^{\prime}=\left[\mathbf{x}^{\prime a} ; \mathbf{x}^{\prime b}\right]$。由于CoCon使用了预训练的语言模型，$\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}$ 应该是一个流畅的文本序列，它跟随着提示文本 $\mathbf{x}^{\prime a}$，并试图合并 $\mathbf{x}^{\prime b}$ 的内容。
第二步，以 $\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}$ 为目标内容的输入，以 $\mathbf{x}^{a}$ 为提示文本：
\[\mathbf{y}_{\text {cycle }}=f_{\theta, \psi}\left(\left(\mathbf{c}=\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}\right),\left(\mathbf{p}=\mathbf{x}^{a}\right)\right)\]由于 $\mathbf{x}=\left[\mathbf{x}^{a} ; \mathbf{x}^{b}\right]$，$\mathbf{x}^{b}$ 是提示文本 $\mathbf{x}^{a}$ 的符合逻辑的下文，$\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}$ 在第一步中以 $\mathbf{x}^{b}$ 为条件。它将 $\mathbf{x}^{b}$ 假定为 $\mathbf{y}_{\text {cycle}}$ 的训练标签，于是 cycle reconstruction loss 为：
\[\mathcal{L}_{\text {cycle }}=-\sum_{i=t}^{l} \log p_{\theta, \psi}\left(\mathbf{y}_{\text {cycle }}=\mathbf{x}^{b} \mid\left(\mathbf{c}=\mathbf{y}_{\mathbf{x}, \mathbf{x}^{\prime}}\right),\left(\mathbf{p}=\mathbf{x}^{a}\right)\right)\]Adversarial Loss：对抗训练的目标是帮助生成真实的文本。作者引入了对抗损失，通过最小化损失来鼓励输出文本的表示 $(enc(\mathbf{y}))$ 与训练样本 $(dec(\mathbf{x}))$ 相匹配：
\[\mathcal{L}_{\mathrm{adv}}=\mathbb{E}_{\mathbf{x}}\left[\log f_{\mathrm{disc}}(\operatorname{enc}(\mathbf{x}))\right]+\mathbb{E}_{\mathbf{y}}\left[\log \left(1-f_{\mathrm{disc}}(\operatorname{enc}(\mathbf{y}))\right],\right.\]其中，$f_{disc}$ 是一个判别器网络。它用于分类表示是否是CoCon生成的文本。判别器的训练目标是最大化 $\mathcal{L}_{\mathrm{adv}}$：
\[\phi^{*}=\underset{\phi}{\arg \max } \mathcal{L}_{\mathrm{adv}}\]Full Training：完整的学习目标是通过随机梯度下降来最小化四个损失项：
\[\theta^{*}=\arg \min \left(\lambda_{\text {self }} \mathcal{L}_{\text {self }}+\lambda_{\text {null }} \mathcal{L}_{\text {null }}+\lambda_{\text {cycle }} \mathcal{L}_{\text {cycle }}+\lambda_{\text {adv }} \mathcal{L}_{\text {adv }}\right),\]