深度学习进阶（二十四）Swin 的二维 RPE

为什么需要二维 RPE#在 T5 中相对位置是一个标量−。因为文本是一维序列两个 token 之间的关系只需要一个数字就能描述。但图像数据不同。一张图像被划分为 × 的 patch 网格后两个 patch 之间的相对位置是二维的。一个 patch 到另一个 patch 的偏移不仅有水平方向的距离还有垂直方向的距离。具体来说对于图像中的位置 (1,1) 和 (2,2)相对位置是分成下面两部分Δ1−2,Δ1−2这时如果再用一维标量来描述这个二维偏移必然丢失方向信息。好在Swin 本身的 Window Attention 设计其实已经为 RPE 减了负注意力只在 7×7 的窗口内进行。这种设计让我们可以不再过多考虑 NLP 中的编码外推问题但相应的在这个局部范围内精确的二维相对关系对建模视觉结构至关重要。因此Swin 设计了一套二维的相对位置编码方案。2. 二维 RPE 如何构造#我们直接来看 Swin 在窗口注意力中使用的公式Attention(,,)Softmax()公式本身在形式上和 T5 是完全相同的关键在于偏置矩阵 的构造上。我们分点来展开2.1 直接将 RPE 推广到二维#我们先来看看最直接的方法对于一个 × 的窗口直接设计 ∈2×2其中 表示窗口内第 个 patch 和第 个 patch 之间的偏置值。我们用一个简单的例子来演示为什么是 2×2 假设窗口大小2 那么窗口就是[1234]现在每个 token 都要和另外所有 token 建立关系。那么 计算的注意力得分矩阵形状就是这样的[1→11→21→31→42→12→22→32→43→13→23→33→44→14→24→34→4]偏置矩阵必须和注意力矩阵一一对应。所以 ∈2×2。这种方法当然是可以跑通的但我们要考虑二维带来的参数问题如果直接学习一个 2×2 的参数矩阵那每个注意力头就得维护 4 个参数。一个 Swin 有多个头和多个层累计下来参数巨大。因此 Swin 自然有对应的改进。2.2 空间关系的平移不变性#在 NLP 中我们只针对每种相对位置设计偏置但是在上面方案里你会发现直接推广会带来很多无意义的参数核心是因为在二维数据中相对逻辑更加凸显窗口内大量位置对其实拥有相同的相对偏移。比如patch (0,0) 和 (1,0) 之间的偏移是 (Δ1,Δ0)而 patch (2,0) 和 (3,0) 之间的偏移同样是 (Δ1,Δ0)。它们本质上描述的是同一种空间关系理应共享同一个偏置值。于是 Swin 的做法是推广相对逻辑不直接学习 而是学习一个小得多的偏置表再通过二维索引从中查值。3. 紧凑偏置表与查表逻辑#3.1 二维相对位置的计算#首先对于一个 × 的窗口给每个位置一个坐标 (,)显然,∈[0,−1]对于任意两个 patch 二维相对偏移是Δ−,Δ−那么Δ 的取值范围就是 [−(−1),−1]一共 2−1 种可能。Δ 同理这部分的计算逻辑和 T5 是完全相同的。现在我们知道了所有可能的 (Δ,Δ) 组合一共有 (2−1)2 种也就是说我们只需要一个 (2−1)×(2−1) 的偏置表就能覆盖窗口内所有可能的位置关系。这就是 Swin 的紧凑偏置表 ^^∈(2−1)×(2−1)建表本身的逻辑到此结束但现在还有一个小问题 和 ^ 大小不一对于每组注意力计算我要如何查表注入相应偏置3.2 查表过程#其实这步可以理解为如何将 ^ 内的值映射到总公式里的 中首先前面我们已经知道了∈2×2因此真正参与 Attention 计算的偏置矩阵 也必须是 2×2。但我们刚刚学习的紧凑偏置表只有^∈(2−1)×(2−1)不难理解为了让二者适配Swin 的设计是这样的对于 Attention Matrix 中的每一个元素都先计算两个 patch 的相对位移再去 ^ 中查对应 bias。展开来说 中的每一个元素本质上都对应“一对 patch 的关系”而每一对 patch 都有自己的 (Δ,Δ)因此我们可以计算相对位移实现查表取值^[Δ,Δ]这就实现了相同相对位移的 patch 对共享同一个偏置。不过这在实现中还有一个问题数组索引没有负数负偏移并不能和其索引直接对应。而 Δ,Δ∈[−(−1),−1]因此 Swin 会先做一次平移去寻找正确索引Δ′Δ(−1)Δ′Δ(−1)现在Δ′,Δ′∈[0,2−2]于是查表过程就变成^[Δ′,Δ′]字母还是有些抽象我们再举一个实例设 3 那么 patch 网格可以就是[(0,0)(0,1)(0,2)(1,0)(1,1)(1,2)(2,0)(2,1)(2,2)]此时 2−15 因此 ^∈5×5如果当前 patch 为 (0,0)它去关注 (2,2) 那么Δ0−2−2Δ0−2−2现在我们需要查^[−2,−2]显然数组索引不能为负数。所以进行平移−12于是Δ′Δ2Δ′Δ2原本的 [−2,−2] 就被平移成 [0,0]^[−2,−2]⇒^[0,0]这里可能容易疑惑的一点是^∈5×5 中存储的并不是“偏移坐标本身”而是“对应相对位移的偏移参数”。展开来说数学意义上的 (,)(−2,−2) 会被映射到数组索引(0,0)因此^[0,0] 实际存储的就是相对位移为 (−2,−2) 时对应的偏置。这样所有原本可能为负数的二维位移都被映射到了合法数组索引,可以稳定完成查表。最终所有 patch 两两之间都会完成一次查表从而动态构造出完整的偏置矩阵∈2×2随后即可完成二维相对位置信息的注入。值得一提的是在具体实现中二维紧凑表会被展平成一维以类似“编号”的逻辑取值根本逻辑没变明白即可。3.3 参数对比#来看看两种方式的参数对比方式7Swin 默认14暴力直接法49×492401196×19638416Swin 紧凑法13×1316927×27729压缩比约 14 倍约 53 倍很明显随着窗口增大紧凑表的优势会更加明显。