深度学习进阶（二十一）跨窗口的 RPE

合集 - 深度学习进阶(31)

1.深度学习进阶（一）从注意力到自注意力03-312.深度学习进阶（二）多头自注意力机制（Multi-Head Attention）04-023.深度学习进阶（三）Transformer Block04-044.深度学习进阶（四）Transformer 整体结构04-065.深度学习进阶（五）Vision Transformer04-086.深度学习进阶（六）归纳偏置与蒸馏04-107.深度学习进阶（七）Data-efficient Image Transformer04-138.深度学习进阶（八）Swin Transformer04-159.深度学习进阶（九）池化技术的初步改进：RoI Pooling04-1910.深度学习进阶（十） RoI Align04-2111.深度学习进阶（十一）Position-Sensitive RoI Pooling04-2412.深度学习进阶（十二）可变形池化 deformable RS RoI Pooling04-2713.深度学习进阶（十三）可变形卷积 DCN04-2914.深度学习进阶（十四）ConvNeXt04-3015.深度学习进阶（十五）通道注意力 SE05-0316.深度学习进阶（十六） 混合注意力 CBAM05-0417.深度学习进阶（十七）高效通道注意力 ECA05-0518.深度学习进阶（十八）坐标注意力 CA05-0619.深度学习进阶（十九）相对位置编码 RPE05-0920.深度学习进阶（二十）Transformer-XL05-11

21.深度学习进阶（二十一）跨窗口的 RPE05-13

22.深度学习进阶（二十二）T5：NLP任务的首次大一统05-1523.深度学习进阶（二十三）偏置型 RPE05-1824.深度学习进阶（二十四）Swin 的二维 RPE05-2025.深度学习进阶（二十五）RoPE：现代 NLP 的位置编码范式05-2926.深度学习进阶（二十六）现代 LLM 的核心架构设计其一：RMSNorm06-0127.深度学习进阶（二十七）现代 LLM 的核心架构设计其二：SwiGLU06-0428.深度学习进阶（二十八）现代 LLM 的核心架构设计其三：Decoder-Only 下的 KV Cache06-0829.深度学习进阶（二十九）现代 LLM 的核心架构设计其四：GQA06-0930.深度学习进阶（三十）从 Transformer 到 LLaMA：现代 LLM 架构总览06-1631.深度学习进阶（三十一）FlashAttention：IO 感知的精确注意力06-19

收起

上一篇我们介绍了 Transformer-XL 的段级递归和记忆缓存机制，也留下了一个问题：

Memory 让不同 segment 可以互相看到，但配套的位置编码逻辑却并不完善。

这一篇就来展开 Transformer-XL 配套的改进方案：跨窗口的相对位置编码。

1. 为什么要提出跨窗口的 RPE？#

1.1 正余弦绝对编码的局限#

我们还是用上一篇的例子来展开这个问题：假设模型上下文窗口长度为 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 的末尾）。

那可不可以这样：

利用正余弦的无限外推能力，直接让位置编码跨 segment 递增不就行了？Segment 1 用位置 1~4，Segment 2 用位置 5~8？

但问题在于：

正余弦位置编码虽然可以理论上无限生成，却并不意味着模型真的具备无限长度泛化能力。

要理解这一点，我们需要明白的是：PE 只是帮助学习的一环。
注意力真正的学习是词嵌入、位置编码、QKV 矩阵等综合实现的，这些实际上都是在训练长度范围内形成的。

假设模型训练的上下文长度为 512，那么模型长期接触的位置范围其实只有 0 ~ 512，它会在这个空间内形成比较稳定的分布。

但如果采用跨 segment 递增，出现了一个 PE(512000)，模型没有在这种位置分布上训练过，导致无法泛化。
从专业名词上说，这叫：extrapolation degradation（外推退化）。
这时就像让一个每天练习正常考试题的学生没有任何防备的去做竞赛题。巨大的差异会直接影响原本已经稳定的“解题思路”。
从而导致KQV 矩阵投影异常、注意力得分漂移等问题。
总结来说就是：

强行使用看似合理实际却不适配的配件，表面上解决了问题，但实际却带来了更多 bug。

因此，单纯使用全局递增的绝对位置编码，实际上并不能真正跨段的长上下文建模。

1.2 可学习绝对编码的局限#

相比正余弦绝对编码，可学习绝对编码的局限就更明显了：

完全没有外推能力。

展开来说，Transformer 中的可学习绝对位置编码本质上其实是一张位置查找表，当我们设置上下文长度为 512 时， 它就为每个位置初始化一组向量作为位置编码并注入，在反向传播中不断更新。

它不是位置越大，自动推导，而只是：“训练时记住了这个位置应该长什么样”。
所以，当采用跨 segment 递增，出现了一个 PE(512000），但可学习绝对编码只到PE（512）。而要增加就又回到了最初的窗口大小和计算量问题。

最终，可学习绝对位置编码因为本质上无法支持无限长度扩展，不适用于 Transformer-XL 。

1.3 原始 RPE 的局限#

在原始 RPE中，我们介绍了 Shaw 等人提出的相对位置编码：引入可学习的 ${\mathbf{r}}^K$ 和 ${\mathbf{r}}^V$，分别加到 Key 和 Value 上。

从逻辑上说，相对位置天然就是跨窗口的，但过去的水救不了现在的火，原始 PRE 存在两个待优化问题：

1. 可学习的相对位置表有长度限制。 如果训练时最大相对距离是 10，而需要分段的长序列相对距离可能达到几千，这时候查表查不到，只能截断，信息丢失严重。
2. "K/V 加法注入"的方式过于粗糙。 它只是简单地把位置信息线性叠加到语义表示上，无法形成更丰富的交互。

因此，Transformer-XL 并没有直接用原始 RPE，而是重新设计了一套相对位置编码方案。

2.正余弦相对位置编码#

现在我们知道了 Transformer-XL 必须使用 RPE ，在正式展开其注入逻辑前，我们需要先了解 Transformer-XL 本身对编码逻辑的设计：正余弦相对位置编码。

一个问题是：

“继续沿用原始 RPE 的相对位置表不可以吗？”

就像这样：

相对距离	编码
-1	$r_{-1}$
-2	$r_{-2}$
-3	$r_{-3}$

原因就是我们上面提到的：

可学习的相对位置表有长度限制，外推能力不足。

于是 Transformer-XL 做了一个非常关键的决定：

重新使用正余弦函数生成相对编码。

具体来说，对于相对距离：

$$\mathrm{\Delta }=i-j$$

Transformer-XL 会直接使用正余弦函数生成对应位置编码：

$$R_{\mathrm{\Delta }}^{(2k)}=\sin \left(\frac{\mathrm{\Delta }}{{10000}^{2k/d}}\right)$$

$$R_{\mathrm{\Delta }}^{(2k+1)}=\cos \left(\frac{\mathrm{\Delta }}{{10000}^{2k/d}}\right)$$

公式本身的计算逻辑和原始 Transformer的正余弦绝对编码没有任何区别，只是把绝对位置 $pos$ 变成了相对距离 $\mathrm{\Delta }=i-j$ 。

假设位置编码维度：$d=4$，$\mathrm{\Delta }=i-j=-2$：

维度	公式	计算结果
$R_{-2}^{(0)}$	$\sin \left(\frac{-2}{{10000}^0}\right)=\sin (-2)$	$\approx -0.909$
$R_{-2}^{(1)}$	$\cos \left(\frac{-2}{{10000}^0}\right)=\cos (-2)$	$\approx -0.416$
$R_{-2}^{(2)}$	$\sin \left(\frac{-2}{{10000}^{2/4}}\right)=\sin (-0.02)$	$\approx -0.020$
$R_{-2}^{(3)}$	$\cos \left(\frac{-2}{{10000}^{2/4}}\right)=\cos (-0.02)$	$\approx 0.999$

最终得到：

$$R_{-2}\approx [-0.909,-0.416,-0.020,0.999]$$

这就是 “相对距离为 -2” 对应的相对位置编码，它完善了长序列中的位置编码，相比绝对编码具有更好的泛化性。

不过这里需要注意的是，这种编码仍然存在上面提到的正余弦绝对编码的局限。
因为模型中的QKV 投影矩阵、Attention 模式等本质上仍然是在有限上下文范围内训练得到的，不能通过单独的编码设计就完全解决这些问题。
因此也催生了后续的更多改进。

3. 跨窗口的 RPE#

现在，我们通过正余弦相对位置编码解决了原始 RPE 可学习的相对位置表有长度限制的问题，一下个问题就是："K/V 加法注入"的方式过于粗糙。

而 Transformer-XL 的做法并不是继续修补原始 RPE，而是重新推导 Relative Attention 的计算形式。

还是先回忆一下普通 Self-Attention 公式：

$$Attention(Q,K,V)=softmax\left(\frac{Q{K}^T}{\sqrt{d}}\right)V$$

其中计算注意力分数的部分是：

$$A_{i,j}=Q_i^T{K}_j$$

但问题在于原始的它只包含“内容”，没有真正包含“位置关系”，所以原始 RPE 才会把 RPE 注入到 Attention 中。
而现在，这部分再次被重新改写为：

$$A_{i,j}^{rel}=Q_i^T{K}_j+Q_i^T{R}_{i-j}+u^T{K}_j+v^T{R}_{i-j}$$

这也是 Transformer-XL 最核心的 Relative Attention 公式。
看起来很复杂，但它其实只是把“内容”和“位置”进行了分开建模。
我们分点来看：

3.1 内容项和位置项#

先来看前两项：

$$Q_i^T{K}_j+Q_i^T{R}_{i-j}$$

其中：

$$Q_i^T{K}_j$$

其实就是普通 Self-Attention 的逻辑：Query 和 Key 的语义相似度。

但还是老问题：

Attention 不仅需要“内容信息”，还需要“位置信息”。

于是 Transformer-XL 额外加入了：

$$Q_i^T{R}_{i-j}$$

这项的语义是： “当前 Query 更倾向关注什么距离的位置”。
通过反向传播，模型会逐渐把这些 “距离规律” 也学习进 Attention 中。

3.2 内容偏置和位置偏置#

在上一项完成后还没结束，因为作者发现：

仅仅让 Query 感知位置，其实还不够。

于是 Transformer-XL 又进一步加入了：

$$u^T{K}_j+v^T{R}_{i-j}$$

这两项全局偏置。
先来看：

$$u^T{K}_j$$

这里的 $u$ 是一个可学习向量。它不再依赖于个具体位置的 Query，而是在像偏置一样，是计算中的全局偏好。
可以把它理解成：

这个自注意力头天然更喜欢关注什么类型的信息。

同理，对于：

$$v^T{R}_{i-j}$$

这里的 $v$ 同样是一个可学习向量，它学习的是：

这个自注意力头天然更偏好什么距离。

这便是跨窗口的 RPE 的完整逻辑。
首先，它通过正余弦相对编码改善了原始 RPE 的建表大小问题，让编码拥有从理论上无限外推的能力，又相对绝对编码提高了泛化能力。
其次，它不再把位置信息简单的注入 K/V ，而是让其真正参与注意力计算，让位置不再只是辅助信息，而是真正开始影响注意力组成。

此外，Transformer-XL 删除了原始 RPE 中对 Value 的位置注入。原因在于作者认为位置真正需要影响的，其实是“关注谁”，而不是“传递什么内容”。
但其实后续也有别的模型又把 V 加回去了，所以这只是选择和讲故事问题，不再多说了。

最后，还有一点需要提到的是，原论文还提到了一种工程优化技巧，叫做 Relative Shift，是为了在实际运算中高效计算 $Q_i^T{R}_{i-j}$，但因为其只是一种计算技巧，在后续也出现了许多代替方案，就不再展开了。