Transformer-XL

Transformer 的上下文窗口问题#在我们之前介绍 Transformer时我们提到了位置编码提到了因果掩码但并没有对数据本身展开太多。这里我们先补充一些 Transformer 对文本数据处理细节 展开 Transformer 的上下文窗口逻辑。1.1 长短不一的文本如何统一输入#我们知道注意力本身就是是矩阵运算假设一个 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]这里的 就是填充符 它没有任何语义仅仅用于对齐矩阵来满足模型输入要求。1.2 Attention 如何处理填充符#这里很容易产生一个问题既然填充符也进入了序列那模型会不会真的去关注这些“伪 token”答案是会如果不做额外处理Attention 会把填充符成正常 token 一样参与计算。这显然是无意义的。因此Transformer 引入 Attention Mask也就是注意力掩码它和我们之前提到的解码器阶段的因果掩码不同。它的作用就是在填充的同时生成一组掩码告诉模型输入的哪些位置是填充符不用计算。就像这样显然1 表示有效 token0 表示 PAD.随后在 Attention 内部计算完得到注意力分数后模型会根据掩码把 PAD 位置对应的分数变成这样在进行 Softmax 之时这样PAD 的注意力权重就会归零不污染真正的注意力信息。但是这只是解决了逻辑问题这里还有一个问题计算量问题。1.3 Transformer 如何处理长文本序列#在有了 PAD 后从理论上来说 Transformer 可以处理无限长的序列只要把该 batch 中的其他序列都填充至最长长度就好了。但显然这只是理论。原因在于注意力计算中每一个 token 都要计算和所有 token 的注意力分数,其复杂度为因此计算量会随着序列长度增加而暴涨得到结论Transformer 不可能无限扩展上下文。所以大多数 Transformer 都会设定一个固定长度上下文窗口例如512、1024、2048 等等这就是现在的AI模型记忆的最初形态。当序列长度超过这个窗口长度时Transformer 就会进行切段Segment。例如一个最大上下文窗口长度为的 Transformer如果输入了一篇长度为 1500 tokens 的长文章模型通常会将其切分为以下三段Segment A: [1 ~ 512] tokensSegment B: [513 ~ 1024] tokensSegment C: [1025 ~ 1500] tokens之后这三个 Segment 会只在自己的序列范围内计算注意力从而解决长序列带来的计算量问题。但这也同时带来了新的问题Segment A 中的 token无法看到 Segment B 和 Segment C 中的内容。因为它们不在一个上下文窗口里这就是“看到结尾忘了开头”的原因。这便是 Transformer-XL 的核心改进点同时作者提出了一种改进后的 RPE下面就来详细展开。2.Transformer-XL#现在我们已经理解了原始 Transformer 的核心问题上下文窗口是固定的不同 Segment 之间完全隔离。而 Transformer-XL 的核心思想其实非常直观既然当前窗口看不到历史内容那就把历史窗口缓存下来。下面来分点展开其详细逻辑2.1 记忆缓存 Memory#在展开这部分前我们要先强调一点设计同一个序列的 Segment A、B、C 不会出现在同一个 batch 内而是按顺序出现在前后不同 batch。现在我们来展开 Memory 的内容你就会明白这么设计的原因image.png如图所示我们知道序列数据进入 Transformer block 后会进行注意力计算、融合等处理得到编码信息进行下一步堆叠或其他处理而现在我们新增了一个缓存窗口我们会这这一批次的编码信息存入缓存窗口而下一批次的输出后就会把下一个批次的输出再存入其中。值得一提的是缓存窗口不一定要大于等于上下文窗口。如果设置缓存窗口为 256而上下文窗口是 512 那么只会切分编码信息的最后 256 部分进入缓存窗口。这一步我们是把现在的存下来其作用自然就是在下一步使用来实现“记忆功能”。2.2 Memory 如何参与下一轮 Attention#现在Memory 已经建立好了。接下来真正关键的问题来了历史编码信息到底是如何参与下一轮计算的在这里你会发现这种逻辑非常类似于 RNN 因此我们在这里也称 Memory 为上一轮的隐藏状态而其传递逻辑是这样的image.png这段是递归传递的核心我们来详细展开先回到 Attention 的核心公式我们之前说过Query表示“当前需要什么”Key表示“当前可以提供什么”Value表示“真正携带的信息”因此在当前批次中我们首先要使用自身信息来生成需求即:而为了实现记忆我们就要从当前和缓存中确认供给和真正信息即因此XL 的处理是现将当前批次信息和缓存进行拼接注意这是序列长度维度的拼接所以不会破坏的 Q/K/V 投影计算。展开一下对于其中token 数序列长度embedding 维度如果当前 segment而历史 memory拼接后的结构是这样的相当于memory token1memory token2…current token1current token2…就像把两段句子接起来。所以后面的 KV 投影计算是完全合法的这样在后续的归一化融合部分里当前 token 就可以同时注意当前 segment 和历史 memory实现跨 segment 的注意力。2.3 Segment-Level Recurrence 段级递归#到这里你会发现Transformer-XL 已经开始出现一种“循环结构”了。因为融合了上轮 memory 的当前 segment 的输出会成为下一轮的 memory继续和下一个 segment 融合在各个block 内相继进行这种机制就被称为Segment-Level Recurrence段级递归。image.png这里的“递归”并不是 RNN 那种 token-by-token 的时间递归而是在 Segment 层级上的递归。这种递归保留了 Transformer 的并行计算优势因为 segment 内部仍然是并行 Attention只有 segment 之间是递归的。2.4 Stop Gradient 停止梯度#前面我们已经把正向传播的框架建立完成但一个新的问题就出现了在这种设计里反向传播会让梯度会无限跨 segment 传播。而这会导致反向传播链无限增长训练极度缓慢出现类似 RNN 的长距离依赖问题。因此Transformer-XL 的应对措施是Stop Gradient停止梯度。含义是Memory 可以参与前向传播 但不会参与反向传播。即历史 memory 只作为“只读缓存” 当前 segment 才参与训练更新。image.png当然这也意味着 历史 memory 不会被后续 segment 反向修正。这便催生了后来的研究方向。3.Transformer-XL 里的位置编码问题#前面我们用 Memory 和段级递归解决了跨段信息传递的问题但这同样催生了一个新问题如何进行跨窗口的位置编码举个简单例子假设模型上下文窗口长度为 4一段长文本被切成了两个 segmentSegment 1 Segment 2Position 1 A EPosition 2 B FPosition 3 C GPosition 4 D H引入 Memory 后Segment 2 中的 token 可以看到 Segment 1 的编码信息。这时模型会面对这样一个问题位置 4 到底是指 D第一个 segment 的末尾还是 H第二个 segment 的末尾显然模型无法区分。因为这时的位置编码只依赖在 segment 内部的绝对位置不依赖在原始文本中的实际位置。当两个 segment 的位置编码范围完全相同时模型就出现了位置混淆这就是绝对编码在段级递归场景下的根本缺陷。