在[上一篇]中，我们提到 RPE 提出后出现了分裂式的发展趋势，而按时间来讲，对 RPE 的初次改进出自 19 年的论文： Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context ，即超长上下文的 Transformer 。
这这篇论文其实并不是专门对 RPE 的改进，而是为了解决 Transformer 本身的结构问题：
上下文窗口分割让长文本序列无法实现真正的全局注意力，如何实现跨窗口的注意力？
而其对 RPE 的改进可以说是一个“配套方案”，因此，本篇主要展开 Transformer-XL 的主干逻辑，它可以说是现代 AI 记忆系统的最初形态。
在我们之前介绍 Transformer时，我们提到了位置编码，提到了因果掩码，但并没有对数据本身展开太多。
这里我们先补充一些 Transformer 对文本数据处理细节， 展开 Transformer 的上下文窗口逻辑。
我们知道，注意力本身就是是矩阵运算，假设一个 batch 中有三句话：
| 文本序列 | Tokenize 后 | Token 数量 |
|---|---|---|
| "I love AI" | ["I", "love", "AI"] | 3 |
| "Transformer is powerful" | ["Transformer", "is", "powerful"] | 3 |
| "Hi" | ["Hi"] | 1 |
显然，这种不规则长度，根本无法直接组成矩阵。
因此，Transformer 会使用经典方法：Padding（填充），来把所有序列补齐到同一个长度。
例如补到当前 batch 的最长长度：
| 原始序列 | Padding 后 |
|---|---|
| [I, love, AI] | [I, love, AI] |
| [Transformer, is, powerful] | [Transformer, is, powerful] |
| [Hi] | [Hi, PAD, PAD] |
这里的 <PAD> 就是填充符： 它没有任何语义，仅仅用于对齐矩阵，来满足模型输入要求。
这里很容易产生一个问题：
既然填充符也进入了序列，那模型会不会真的去关注这些“伪 token”？
答案是会：如果不做额外处理，Attention 会把填充符成正常 token 一样参与计算。
这显然是无意义的。因此，Transformer 引入 Attention Mask，也就是注意力掩码，它和我们之前提到的解码器阶段的因果掩码不同。它的作用就是：
在填充的同时生成一组掩码，告诉模型输入的哪些位置是填充符，不用计算。
就像这样：
\[[Hi, PAD, PAD]\rightarrow[1,0,0] \]
显然，1 表示有效 token，0 表示 PAD.
随后，在 Attention 内部计算完 \(QK^T\) 得到注意力分数后，模型会根据掩码把 PAD 位置对应的分数变成 \(-\infty\) ，这样在进行 Softmax 之时：
\[e^{-\infty}=0 \]
这样，PAD 的注意力权重就会归零，不污染真正的注意力信息。
但是，这只是解决了逻辑问题，这里还有一个问题：计算量问题。
在有了 PAD 后，从理论上来说： Transformer 可以处理无限长的序列，只要把该 batch 中的其他序列都填充至最长长度就好了。
但显然，这只是理论。
原因在于注意力计算中每一个 token 都要计算和所有 token 的注意力分数,其复杂度为：$$O(n^2)$$因此，计算量会随着序列长度增加而暴涨，得到结论：Transformer 不可能无限扩展上下文。
所以，大多数 Transformer 都会设定一个固定长度上下文窗口，例如：512、1024、2048 等等，这就是现在的AI模型记忆的最初形态。
当序列长度超过这个窗口长度时，Transformer 就会进行切段（Segment）。
例如，一个最大上下文窗口长度为 \(512\) 的 Transformer，如果输入了一篇长度为 1500 tokens 的长文章，模型通常会将其切分为以下三段：
之后，这三个 Segment 会只在自己的序列范围内计算注意力，从而解决长序列带来的计算量问题。
但这也同时带来了新的问题：
Segment A 中的 token，无法看到 Segment B 和 Segment C 中的内容。
因为它们不在一个上下文窗口里，这就是“看到结尾，忘了开头”的原因。
这便是 Transformer-XL 的核心改进点，同时，作者提出了一种改进后的 RPE，下面就来详细展开。
现在，我们已经理解了原始 Transformer 的核心问题：
上下文窗口是固定的，不同 Segment 之间完全隔离。
而 Transformer-XL 的核心思想其实非常直观：
既然当前窗口看不到历史内容，那就把历史窗口缓存下来。
下面来分点展开其详细逻辑：
在展开这部分前，我们要先强调一点设计：
同一个序列的 Segment A、B、C 不会出现在同一个 batch 内，而是按顺序出现在前后不同 batch。
现在我们来展开 Memory 的内容，你就会明白这么设计的原因：
如图所示，我们知道序列数据进入 Transformer block 后会进行注意力计算、融合等处理，得到编码信息，进行下一步堆叠或其他处理：
\[Y = \text{TransformerLayer}(X) \]
而现在，我们新增了一个缓存窗口，我们会这这一批次的编码信息存入缓存窗口，而下一批次的输出后，就会把下一个批次的输出再存入其中。
值得一提的是，缓存窗口不一定要大于等于上下文窗口。
如果设置缓存窗口为 256，而上下文窗口是 512 ，那么只会切分编码信息的最后 256 部分进入缓存窗口。
这一步，我们是把现在的存下来，其作用自然就是在下一步使用，来实现“记忆功能”。
现在，Memory 已经建立好了。接下来真正关键的问题来了：
历史编码信息到底是如何参与下一轮计算的？
在这里你会发现这种逻辑非常类似于 RNN ，因此，我们在这里也称 Memory 为上一轮的隐藏状态 \(h^{n-1}\) ，而其传递逻辑是这样的：
这段是递归传递的核心，我们来详细展开：先回到 Attention 的核心公式：
\[\text{Attention}(Q,K,V)=softmax\left(\frac{QK^T}{\sqrt d}\right)V \]
我们之前说过：
因此，在当前批次中，我们首先要使用自身信息 \(h\) 来生成需求，即 \(Q\):
\[Q=hW_Q \]
而为了实现记忆，我们就要从当前和缓存中确认供给和真正信息，即\(K\)，\(V\)，因此，XL 的处理是现将当前批次信息和缓存进行拼接：
\[\widetilde h=[h^{n-1};h] \]
注意，这是序列长度维度的拼接，所以不会破坏的 Q/K/V 投影计算。
展开一下，对于：
\[h \in \mathbb{R}^{L \times d} \]
其中：
如果当前 segment：
\[h\in \mathbb{R}^{L_c\times d} \]
而历史 memory：
\[h^{n-1}\in \mathbb{R}^{L_m\times d} \]
拼接后的结构是这样的：
\[\widetilde h \in \mathbb{R}^{(L_m+L_c)\times d} \]
相当于：
memory token1
memory token2
...
current token1
current token2
...
就像把两段句子接起来。
所以后面的 KV 投影计算是完全合法的：
\[K = \widetilde h W_K \quad K \in \mathbb{R}^{(L_m+L_c)\times d_k} \]
\[V = \widetilde h W_V \quad V \in \mathbb{R}^{(L_m+L_c)\times d_v} \]
这样，在后续的归一化融合部分里，当前 token 就可以同时注意当前 segment 和历史 memory，实现跨 segment 的注意力。
到这里，你会发现：Transformer-XL 已经开始出现一种“循环结构”了。
因为融合了上轮 memory 的当前 segment 的输出，会成为下一轮的 memory，继续和下一个 segment 融合，在各个block 内相继进行：
\[h^{n-1} \rightarrow h^{n} \rightarrow h^{n+1} \]
这种机制就被称为：Segment-Level Recurrence（段级递归）。
这里的“递归”并不是 RNN 那种 token-by-token 的时间递归，而是在 Segment 层级上的递归。
这种递归保留了 Transformer 的并行计算优势，因为 segment 内部仍然是并行 Attention，只有 segment 之间是递归的。
前面我们已经把正向传播的框架建立完成，但一个新的问题就出现了：
在这种设计里反向传播会让梯度会无限跨 segment 传播。
而这会导致反向传播链无限增长，训练极度缓慢，出现类似 RNN 的长距离依赖问题。
因此，Transformer-XL 的应对措施是：Stop Gradient，停止梯度。
含义是：
Memory 可以参与前向传播， 但不会参与反向传播。
即历史 memory 只作为“只读缓存”， 当前 segment 才参与训练更新。
当然，这也意味着： 历史 memory 不会被后续 segment 反向修正。
这便催生了后来的研究方向。
前面我们用 Memory 和段级递归解决了跨段信息传递的问题，但这同样催生了一个新问题：
如何进行跨窗口的位置编码？
举个简单例子，假设模型上下文窗口长度为 4，一段长文本被切成了两个 segment：
| Segment 1 | Segment 2 | |
|---|---|---|
| Position 1 | A | E |
| Position 2 | B | F |
| Position 3 | C | G |
| Position 4 | D | H |
引入 Memory 后，Segment 2 中的 token 可以看到 Segment 1 的编码信息。这时模型会面对这样一个问题：
位置 4 到底是指 D（第一个 segment 的末尾），还是 H（第二个 segment 的末尾）？
显然，模型无法区分。
因为这时的位置编码 \(p\) 只依赖在 segment 内部的绝对位置，不依赖在原始文本中的实际位置。
当两个 segment 的位置编码范围完全相同时，模型就出现了"位置混淆"，这就是绝对编码在段级递归场景下的根本缺陷。
因此，Transformer-XL 在 RPE 的基础上进行了改进：跨窗口的 PRE。
我们在下一篇再详细展开。
此内容由惯性聚合(RSS阅读器)自动聚合整理，仅供阅读参考。 原文来自 — 版权归原作者所有。