1. 概述

Transformer^[1] 框架的提出，为后续大预言模型 LLM 的产生提供了基础，回顾 Transformer 的网络结构，在 Transformer 网络中，是由多层编码器和解码器堆叠而成，Transformer 的网络结构如下图所示：

在 Transformer 网络结构中，编码器和解码器都是由多头注意力和前馈网络模块组成，在多头注意力机制中，最重要的是注意力的计算，其一般化的计算公式为：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^\top}{\sqrt{d_k}})V$$

其中，$Q,K,V$ 分别是查询、键、值向量。前馈网络通常由两个线性层和非线性激活组成。

随着大语言模型的发展，逐渐从这套 Transformer 的 Encoder + Decoder 框架中，发展出只使用 Decoder 模块的生成式模型^[2]，结构如下图所示：

其中，xN 表示的是 N 次重复这个部分。Decoder 模型在训练和推理的过程中也不是说可以处理无限的长度，当然，对于模型来说是希望上下文的长度是越长越好，越长对于接受到的信息就会越多，但是越长的上下文带来的最直接的问题就是计算复杂度的上升，因此，一般对于超长文本，会采用分段处理的方法，也就是将超长的文本切分成多段，每一段单独处理。

这种方法固然使得 Decoder 能处理这种超长文本，但是也引入了一个问题：把文本切成多段，每段独立过模型，段与段之间没有了信息的流动，为了解决这个问题，Transformer-XL^[3] 中采用的不是简单地分段，而是提出片段级循环的概念，即缓存前一段的隐状态，处理新段时复用，目的是让新段能看到前一段的信息。因为缓存了前一段的隐状态，在 Decoder 中使用的正弦、余弦绝对位置编码将失去作用，因此，在 Transformer-XL 中又针对性的提出了相对位置编码的概念，下面将针对这两处的优化详细介绍。

注意：Transformer-XL 虽然名称中带有 Transformer，但是其要解决的是 Transformer 中 Decoder 部分的超长文本以及编码问题。

2. Transformer-XL 的算法原理

2.1. 为什么需要位置编码

在上面的分析中，我们发现，Self-Attention 的计算是整个 Transformer 的核心，而在 Self-Attention 的计算中，最重要的是 $Q,K,V$，这三者都是由输入的 Token 的向量经过线性变换得到的，我们假设 Token 的长度为$m$，每个 Token 的维度为 $n$，输入可以表示为 $X\in \mathbb{R}^{m\times n}$，这时候有三个权重矩阵，分别为 $W_Q\in \mathbb{R}^{n\times d}$，$W_K\in \mathbb{R}^{n\times d}$ 和 $W_V\in \mathbb{R}^{n\times d}$，得到如下的线性变换的结果：

$$Q=XW_Q,\quad K=XW_K,\quad V=XW_V$$

假设此时，我们交换了输入 $X$ 的第 $i$ 行和第 $j$ 行，得到 $X'$，实际上相当于左乘一个 $P\in \mathbb{R}^{m\times m}$ 矩阵，得到 $X'=PX$，其中 $P$ 具体为：

$$P=\begin{pmatrix} 1 & 0 & \cdots & 0 & \cdots & 0 & \cdots & 0 \\ 0 & 1 & \cdots & 0 & \cdots & 0 & \cdots & 0 \\ \vdots & \vdots & \cdots & \vdots &\cdots &\vdots &\cdots &\vdots \\ 0 & 0 & \cdots & 0 & \cdots & 1& \cdots &0 \\ \vdots & \vdots & \cdots & \vdots &\cdots &\vdots &\cdots &\vdots \\ 0 & 0 & \cdots & 1 & \cdots & 0& \cdots &0 \\ \vdots & \vdots & \cdots & \vdots &\cdots &\vdots &\cdots &\vdots \\ 0 & 0 & \cdots & 0 & \cdots & 0& \cdots &1 \end{pmatrix}$$

交换后，重新计算 $Q,K,V$，得到：

$$Q'=X'W_Q=PXW_Q,\quad K'=XW_K=PXW_K,\quad V'=XW_V=PXW_V$$

再计算注意力得到：

$$\begin{aligned} \mathrm{Attention}(Q',K',V') &= \mathrm{softmax}(\frac{Q'K'^\top}{\sqrt{d_k}})V'\\ & \\ &= \mathrm{softmax}(\frac{\left( PQ\right)\left(PK\right)^\top}{\sqrt{d_k}})\left(PV\right) \\ & \\ &=P\mathrm{softmax}(\frac{ QK^\top}{\sqrt{d_k}})P^T\left(PV\right) \\ &\\ &=P\mathrm{softmax}(\frac{ QK^\top}{\sqrt{d_k}})V \\ &\\ &=P\mathrm{Attention}(Q,K,V) \end{aligned}$$

这说明什么，如果只是交换了第 $i$ 行和第 $j$ 行，实际上对于最终的输出也就是交换了第 $i$ 行和第 $j$ 行。我们的目的是希望第 $i$ 个位置和 $j$ 个位置不同，这个时候就需要加入位置信息，我们注意到在 Input/Output Embedding 和 $Q,K,V$ 之间，有一个 Positional Encoding，这就是位置编码。格式为：

$$x_i = \text{TokenEmbed}(w_i) + \text{PositionEmbed}(i)$$

这个时候 $X'\neq PX$，也就没有了后面的 $\mathrm{Attention}(Q',K',V')=P\mathrm{Attention}(Q,K,V)$。

2.2. 正弦、余弦绝对位置编码

在原始的 Transformer 中的 Decoder 阶段采用的是正弦、余弦绝对位置编码，这是如何实现的呢？假设 Token 的词向量是 $x_i$，位置 $i$ 对应的位置向量是 $PE_i$，那么输入 Decoder 的向量为：

$$h_i=x_i+PE_i$$

位置编码的各个维度按照不同频率计算：

$$PE_{(pos,2i)} =\sin\left(\frac{pos}{10000^{2i/d}}\right)$$

$$PE_{(pos,2i+1)} =\cos\left(\frac{pos}{10000^{2i/d}}\right)$$

其中，$pos$ 是 Token 的位置，$d$ 是 Embedding 维度。这里的位置编码是固定公式算出来的，不需要训练额外的位置参数。

但是在 Decoder 中，为了控制计算量，一般会将超长文本截断成固定长度的片段，并单独处理每一个片段，如下图所示：

生成式模型的工作原理是根据上文不断生成 Token 的过程，在训练阶段和推理阶段，针对超长上下文的情况，会略有不同。其中，在训练阶段，Decoder 模型会把长文本切分成固定长度的片段，然后独立地处理每一个片段，如上图中的 Segment 1、Segment 2 等等，它们彼此之间没有任何连接。这种方式存在两个很大的问题，第一，文本被强行切断，没有了语义的信息；第二，依赖的长度有限，都是固定的。

在推理阶段，因为要不断根据上文生成新的 Token，通常可以采用的是滑动窗口的策略，这样虽然保证了每次的预测有了足够的上下文，但是每一次的预测都需要重新计算之前已经计算过的 Token，存在大量的重复计算。

这一系列的问题也就导致急需要一个优化的方案来处理超长文本的问题。

2.3. Transformer-XL

Transformer-XL 中提出了片段级循环机制（Segment-Level Recurrence）和相对位置编码（Relative Position Encoding）两个策略协同解决超长文本的问题。

2.3.1. 片段级循环机制（Segment-Level Recurrence）

片段级循环，简单来说就是在处理下一个片段时，如何重复使用上一个片段的信息，这样才能使得在处理下一个片段时保留上文的信息。

上面也介绍到在 Decoder 模块中包含了$N$个 Masked Multi-Self-Attention + FFN 的模块，每一个模块的输出作为下一个模块的输入，假设输入片段为 $\tau$，这个过程如下所示：

$$\mathbf{h}^{0}_{\tau} \;\xrightarrow{\text{第1层}}\; \mathbf{h}^{1}_{\tau} \;\xrightarrow{\text{第2层}}\; \mathbf{h}^{2}_{\tau} \;\cdots\; \xrightarrow{\text{第}N\text{层}}\; \mathbf{h}^{N}_{\tau}$$

其中，$\mathbf{h}^{0}_{\tau}\in \mathbb{R}^{L\times d}$ 表示的是第 $0$ 层，也就是输入层，这一层由 Token Embedding + Position Embedding 组成，$\mathbf{h}^{n}_{\tau}\in \mathbb{R}^{L\times d}$ 是第 $n$ 层的输出，每一行对应片段中一个位置的隐藏向量，其中，第 $n$ 层的输入是第 $n-1$ 层的输出 $\mathbf{h}^{n-1}_{\tau}\in \mathbb{R}^{L\times d}$。

现在回到 Tranformer-XL 中，以处理第 $\tau+1$ 个片段为例，首先，假设前一个片段 $\tau$ 在第 $n-1$ 层的隐藏状态的输出为 $\mathbf{h}^{n-1}_{\tau}$，当前片段 $\tau+1$ 在同一层的状态为 $\mathbf{h}^{n-1}_{\tau+1}$。模型会先将前者的缓存状态与后者拼接，形成扩展后的隐藏序列：

$$\widetilde{\mathbf{h}}^{n-1}_{\tau+1} = \left[\text{SG}(\mathbf{h}^{n-1}_{\tau}) \circ \mathbf{h}^{n-1}_{\tau+1}\right]$$

其中，$\text{SG}(\cdot)$ 表示停止梯度（stop-gradient）。停止梯度的含义是在训练时被视为常量，不参与反向传播，这保证了计算效率。符号 $\circ$ 表示沿序列长度维度的拼接。这样就使得 Tranformer-XL 在处理片段 $\tau$ 时能够同时使用上片段 $\tau+1$ 的有效信息。因为 $\mathbf{h}^{n-1}_{\tau}$ 和 $\mathbf{h}^{n-1}_{\tau+1}$ 两者的维度相同，都为 $L\times d$，因此沿序列维度拼接后的 $\widetilde{\mathbf{h}}^{n-1}_{\tau+1}$ 的维度为 $\widetilde{\mathbf{h}}^{n-1}_{\tau+1}\in \mathbb{R}^{2L\times d}$。同时，拼接后的 $\widetilde{\mathbf{h}}^{n-1}_{\tau+1}$ 作为第 $n$ 层注意力的 $K$ 和 $V$ 的来源，而 $Q$ 只来自于当前的片段 $\tau$：

$$Q = \mathbf{h}^{n-1}_{\tau+1} W_q$$

$$K = \widetilde{\mathbf{h}}^{n-1}_{\tau+1} W_k$$

$$V = \widetilde{\mathbf{h}}^{n-1}_{\tau+1} W_v$$

这样，通过循环，信息可以跨越多个片段向后传递，扩展了有效上下文的长度。

2.3.2. 相对位置编码（Relative Positional Encoding）

由于在上面的片段 $\tau+1$ 在第 $n-1$ 层的输出时，会将片段 $\tau$ 在第 $n-1$ 层的隐藏状态的输出 $\mathbf{h}^{n-1}_{\tau}$ 拼接到 $\mathbf{h}^{n-1}_{\tau+1}$ 中，如果还是使用正弦、余弦绝对位置编码，那么就会产生位置语义的混乱，举例来说，假设，片段 $\tau$ 的片段位置为 [0, 1, 2, 3]，片段 $\tau+1$ 的片段位置也为 [0, 1, 2, 3]，当两者拼接时，就会出现 [0, 1, 2, 3, 0, 1, 2, 3] 这样的位置序列，此时，片段 $\tau$ 中的位置 3 和片段 $\tau+1$ 中的位置 0，在绝对位置编码的情况下就不好区分了，也不好判断他们的真实距离。那是不是有一种相对距离的概念，这就是相对位置编码，简单来说就是对于一个 $Q$，它不需要知道每个 $k$ 的绝对位置，只需要知道它与每个 $k$ 的相对距离 $i-j$ 即可，具体如何计算呢？我们还是从绝对位置编码开始。

首先，在 Decoder 模块中，每个 Token 是由自身的编码和位置编码共同组成的：

$$h_i = E_{x_i} + U_i$$

其中，$E_{x_i}$ 是 Token $x_i$ 的语义编码，$U_i$ 是位置 $i$ 的绝对位置编码。$Q$ 和 $K$ 经由线性变换得到：

$$q_i = (E_{x_i}+U_i)W_q$$

$$k_j = (E_{x_j}+U_j)W_k$$

假设我们这里只计算注意力得分，也就是 $QK^\top$，那么 $A_{i,j}^{\text{abs}}$ 为：

$$A_{i,j}^{\text{abs}} = q_i k_j^\top = (E_{x_i}+U_i)W_q \, W_k^\top (E_{x_j}+U_j)^\top$$

这实际上是一个值，对其转置实际上并没有变化，可以得到：

$$A_{i,j}^{\text{abs}} = (E_{x_i}+U_i)^\top W_q^\top W_k (E_{x_j}+U_j)$$

对其展开，得到：

$$A_{i,j}^{\text{abs}} = \underbrace{E_{x_i}^\top W_q^\top W_k E_{x_j}}_{(a)} + \underbrace{E_{x_i}^\top W_q^\top W_k U_j}_{(b)} + \underbrace{U_i^\top W_q^\top W_k E_{x_j}}_{(c)} + \underbrace{U_i^\top W_q^\top W_k U_j}_{(d)}$$

其中，$(a)$ 部分是 $Q$ 的语义和 $K$ 的语义相互计算，$(b)$ 部分是 $Q$ 的语义和 $K$ 的绝对位置相互计算，$(c)$ 部分是 $Q$ 的绝对位置和 $K$ 的语义相互计算，$(d)$ 部分是 $Q$ 的绝对位置和 $K$ 的绝对位置相互计算。

此时，如果每一段都单独计算，那么 $Q$ 和 $K$ 中的位置信息都相互不影响，但是在 Transformer-XL 中 $K$ 和 $V$ 会拼接上一个片段的输出，也就是 $U_j$ 中包含了上一个片段的绝对位置信息，因此对上上述公式中的 $(b)$、$(c)$ 和 $(d)$ 会出现片段 $\tau$ 和片段 $\tau+1$ 的位置信息的冲突。

为了解决这个问题，在 Transformer-XL 中将注意力分数改成只依赖于相对距离 $R_{i-j}$ 的形式，因此，对上式中的 $(b)$ 修改为：

$$E_{x_i}^\top W_q^\top W_{k,R} R_{i-j}$$

其中，$R_{i-j}$ 是 $Q$ 中位置 $i$ 和 $K$ 中位置 $j$ 之间的相对位置编码，具体为：

$$R_{d,2k} = \sin\left(\frac{d}{10000^{2k/d_{\text{model}}}}\right)$$

$$R_{d,2k+1} = \cos\left(\frac{d}{10000^{2k/d_{\text{model}}}}\right)$$

其中，$d=i-j$ 表示 $Q$ 和 $K$ 的相对距离，$d_{\text{model}}$ 表示的是编码维度，$k$ 为维度的索引：$k=0,1,\dots,d_{\text{model}}/2-1$。

对于上式中的 $(c)$，由于这一项依赖 $Q$ 的绝对位置 $U_i$，而 $Q$ 来自当前片段，$U_i$ 的坐标系和 $K$ 的坐标系不一致。在 Transformer-XL 中，不再让 $Q$ 的位置参与计算，而是用一个可学习的向量 $u$ 代替 $U_i^\top W_q^\top$，最后，上式中的 $(c)$ 修改为：

$$u^\top W_{k,E} E_{x_j}$$

其中，$u$ 对所有 $Q$ 位置共享。

同样，在上式中的 $(d)$ 中也存在 $Q$ 的绝对位置，同样换成可学习的 $v$，最后，上式中的 $(d)$ 修改为：

$$v^\top W_{k,R} R_{i-j}$$

其中，$v$ 是另一个可学习向量。

最终，融合了 $(a)$、$(b)$、$(c)$ 和 $(d)$ 后，上式变成：

$$A_{i,j}^{\text{rel}} = \underbrace{E_{x_i}^\top W_q^\top W_{k,E} E_{x_j}}_{(a)} + \underbrace{E_{x_i}^\top W_q^\top W_{k,R} R_{i-j}}_{(b)} + \underbrace{u^\top W_{k,E} E_{x_j}}_{(c)} + \underbrace{v^\top W_{k,R} R_{i-j}}_{(d)}$$

最值得注意的是对于 $Q$ 的绝对位置被可学习全局偏置 $u$ 和 $v$ 替代了，$Q$ 中不再携带绝对位置，避免片段 $\tau+1$ 中的 $Q$ 和片段 $\tau$ 中的 $K$ 的坐标系冲突。

3. 结论

在 Decoder 框架下，处理超长文本通常会采用分段的方式，如果独立地处理每一段势必会丢失掉很多的信息，此时就引入了将上文信息融入到当前片段的处理手段，也就是片段级循环机制，有了这个机制后，原本的绝对位置编码就会失效了，因为拼接后会出现两段相同的位置编码，此时，就可以使用相对位置编码，去除所有会相互冲突的地方，改用相对位置编码和全局可学习向量代替。

参考文献

[1] Vaswani A , Shazeer N , Parmar N ,et al.Attention Is All You Need[J].arXiv, 2017.DOI:10.48550/arXiv.1706.03762.

[2] Radford A, Narasimhan K, Salimans T, et al. Improving language understanding by generative pre-training[J]. 2018.

[3] Dai Z, Yang Z, Yang Y, et al. Transformer-xl: Attentive language models beyond a fixed-length context[C]//Proceedings of the 57th annual meeting of the association for computational linguistics. 2019: 2978-2988.