1. 从绝对到相对为什么我们需要跨窗口的RPE在Transformer模型席卷自然语言处理领域的浪潮中Self-Attention机制无疑是其最核心的引擎。这个机制允许序列中的任意两个位置直接建立联系从而捕捉长距离依赖。然而经典的Self-Attention使用的是一个看似简单却暗藏玄机的设计绝对位置编码。简单来说就是给序列中的每个位置比如第1个词、第2个词分配一个独一无二的向量然后将这个向量加到词本身的嵌入向量上一起输入模型。这样模型在计算注意力时就能“感知”到每个词在句子中的绝对位置。这个设计在标准Transformer中运行良好但它有一个致命的缺陷它无法处理比训练时见过的序列更长的文本。想象一下你训练模型时见过的句子最长只有512个词那么模型就只学会了处理512个位置。当你试图用它来理解一篇1000个词的文档时对于第513个词及之后的词模型就“不认识”它们的位置了因为它没有学习过这些位置的编码。这极大地限制了Transformer在长文本任务如长文档摘要、书籍生成、长对话建模中的应用。为了解决这个问题研究者们提出了相对位置编码。它的核心思想不再是告诉模型“你是第几个词”而是告诉模型“你和另一个词之间相隔多远”。比如“我”和“爱”之间距离是1“我”和“深度学习”之间距离是3。这种编码方式天然具备长度外推性无论序列多长模型都学过如何处理“距离为1”、“距离为2”的关系因此理论上可以处理任意长度的序列。相对位置编码中最具代表性的工作之一就是RPE。RPE不再将位置信息加到输入上而是巧妙地修改了注意力分数的计算过程。在计算查询向量q_i和键向量k_j的注意力得分时除了常规的点积q_i·k_jRPE会额外引入一个基于相对位置偏移(i-j)的偏置项。这个偏置项是可学习的模型在训练过程中会学会当两个词距离为1时它们的注意力应该有一个怎样的基础倾向距离为2时又是另一个倾向。这样一来模型对位置的感知就从“绝对坐标”转变为了“相对关系”。然而标准的RPE实现通常假设整个序列的注意力计算是在一个“窗口”内完成的即模型一次性看到整个序列。但在处理超长序列时由于计算复杂度的限制Self-Attention的复杂度是序列长度的平方我们不得不将序列切分成多个片段或窗口然后在这些窗口内分别计算注意力。这就引出了本文要探讨的核心问题当序列被分割一个查询词在一个窗口内而它需要关注的键值词在另一个窗口内时标准的、基于窗口内相对位置的RPE就失效了。因为它无法计算跨窗口的两个词之间的相对位置关系。如何让RPE能够“看见”并正确处理这种跨窗口的相对位置就是“跨窗口的RPE”所要解决的挑战。2. 理解RPE的经典实现与跨窗口困境要理解跨窗口的挑战我们首先需要拆解一下经典RPE以Transformer-XL和后续改进工作为代表是如何工作的。它的核心公式可以简化为对注意力矩阵A的修改A_{i,j} (q_i · k_j) b_{i-j}这里b_{i-j}就是一个可学习的标量偏置其值仅依赖于查询位置i和键位置j的相对距离(i-j)。通常我们会预设一个最大相对距离k比如128对于所有|i-j| k的相对位置我们使用同一个偏置值b_{k}或b_{-k}认为超过这个距离的位置关系都“差不多远”。在标准的自回归语言建模中模型从左到右生成文本。在训练时为了模拟推理时的场景并提升效率通常会采用“片段递归”的策略也就是Transformer-XL的核心思想。模型会缓存上一个片段或窗口的隐藏状态在当前片段计算注意力时当前片段的词作为查询不仅可以关注当前片段内的词作为键还可以关注上一个片段缓存下来的词。这就在一定程度上打破了窗口的界限。但是这里的RPE计算会变得微妙。对于当前片段内的查询i和当前片段内的键j相对位置(i-j)是明确的。对于当前片段内的查询i和上一个片段缓存中的键j‘它们的绝对位置索引可能相差很远比如i是当前片段的第10个词j‘是上一个片段的第500个词。此时(i - j‘)这个值可能非常大远超预设的最大相对距离k。如果我们简单地将这个超大值作为索引去查找偏置表要么会索引越界要么会得到一个代表“非常远”的固定偏置值b_{k}。这虽然能运行但丢失了精确的相对距离信息对于模型来说“距离501”和“距离502”都被粗暴地归为了“距离k”的同一类这显然是不精确的。更复杂的情况出现在非自回归模型或双向编码器中比如BERT。在这些模型中一个窗口内的词需要同时关注窗口内所有其他词。如果我们简单地将长序列切成不重叠的窗口例如每512个词一个窗口那么窗口1中的词完全无法与窗口2中的词建立注意力连接RPE就更无从谈起了。这种硬切割会破坏序列的连贯性对于理解跨句、跨段落的语义关系是灾难性的。因此跨窗口RPE的核心目标就是设计一种机制使得在计算窗口化注意力时模型能够准确地获知并利用任意两个词无论它们是否在同一个计算窗口内之间的相对位置信息。这不仅仅是技术实现上的挑战更关乎模型能否真正具备处理长程、细粒度依赖关系的能力。3. 滑动窗口、扩张注意力与相对位置索引的重新校准为了解决跨窗口的RPE问题社区和业界提出了几种主流的思路它们从不同的角度对计算过程进行了改造。3.1 滑动窗口注意力这是最直观的一种方法。它不完全将序列切成独立的块而是让一个固定大小的窗口在序列上滑动。对于序列中的每个位置i其注意力范围是[i-w, iw]这样一个固定大小的邻域其中w是窗口半径。这样位于窗口边缘的词比如位置i自然可以“看到”相邻窗口中的词位置i-w到i-1如果i靠近当前块末尾的话。在这种设置下实现跨窗口RPE的关键在于统一所有词的位置坐标系。我们不能再用每个窗口内部的局部位置索引0到2w来计算相对位置因为不同窗口的局部索引0对应着序列中完全不同的绝对位置。解决方案是使用全局绝对位置索引。我们为序列中的每个词分配一个唯一的、连续的绝对位置编号0, 1, 2, …。在计算注意力时对于查询i和键j我们使用它们的绝对位置索引来计算相对距离d i - j。然后将这个距离d映射到RPE的偏置表中。由于滑动窗口保证了i和j的距离不会超过窗口大小2w1因此d的范围是有限的[-w, w]完全在预设的最大相对距离k的覆盖范围内通常k w。注意这里有一个重要的实现细节。在训练非常长的序列时我们可能无法一次性将整个序列的绝对位置比如10000都编码进模型。常用的技巧是使用循环位置编码或相对位置桶。例如可以将绝对位置对某个模数如512取余或者将相对距离d通过一个函数如对数桶映射到有限的几个桶中每个桶对应一个可学习的偏置。这样模型学到的不是“距离137的偏置”而是“距离在128-256这个桶范围内的偏置”在保证外推性的同时降低了参数量。3.2 扩张注意力与局部敏感哈希滑动窗口虽然解决了邻近窗口的问题但对于需要超长程依赖的任务窗口大小w可能仍然不够。扩张注意力是一种受空洞卷积启发的方法。它不是在每个位置都计算注意力而是每隔一定的步长扩张率选取一个位置进行计算。这样在不增加计算量的情况下每个位置的实际感受野变大了。结合RPE时我们需要处理的不再是连续位置之间的距离而是稀疏采样位置之间的距离。此时相对位置d i - j可能是一个很大的值并且是扩张率的倍数。RPE偏置表需要能够适应这种稀疏的、大间隔的距离模式。一种实践是将距离除以扩张率后再进行映射或分桶让模型学会对“扩张后的距离”进行建模。更激进的方法是使用局部敏感哈希等近似注意力机制将相似的词哈希到同一个桶中无论它们的位置多远。在这种情况下RPE的设计需要与哈希策略协同我们可能不再需要精确的相对距离而是需要一个能表示“是否在同一个哈希桶内”或“哈希桶之间的相对关系”的偏置。这为RPE的设计打开了新的思路即从基于数值距离的编码转向基于语义或结构分组的编码。3.3 长序列建模框架中的RPE集成以Longformer和BigBird为例一些专门为长序列设计的Transformer变体如Longformer和BigBird本身就采用了混合的注意力模式滑动窗口注意力全局注意力。它们天然需要处理跨窗口的RPE。以Longformer为例它的注意力模式是对于大多数词采用滑动窗口局部注意力对于少量预先选定的“全局词”如[CLS]标记或某些关键实体则赋予其关注整个序列的能力。在实现这种混合注意力时RPE需要被灵活地应用。对于局部滑动窗口部分使用上述的全局绝对位置索引方法计算RPE。对于全局注意力部分当一个全局词作为查询需要关注序列中所有键时它需要计算与每一个键的相对位置。由于序列可能极长这里必须使用分桶策略。将所有可能的巨大相对距离通过一个函数如对数函数映射到几十个或几百个有限的桶中。例如距离1-2映射到桶0距离3-4映射到桶1距离5-8映射到桶2以此类推。模型为每个桶学习一个偏置。这样全局词在关注远处一个词时使用的RPE偏置是基于“距离桶”的而非精确距离这是一种在计算效率和模型能力之间的有效折衷。在实际编码中这通常体现为一个庞大的相对位置偏置矩阵的查找过程。我们需要预先计算好序列中所有位置对之间的相对距离桶索引形成一个索引矩阵。在注意力计算时根据查询i和键j从这个索引矩阵中取出对应的桶编号再去查找一个小的、可学习的偏置嵌入表获得标量偏置b_{bucket}。# 伪代码示意基于分桶的RPE偏置获取 def get_relative_position_bucket(relative_position, max_distance512, num_buckets32): 将相对位置映射到桶中。 relative_position: 标量或矩阵表示相对距离 (i-j) max_distance: 超过此距离的视为同一类远距离 num_buckets: 桶的总数 ret 0 n -relative_position if max_distance is not None: # 将距离限制在[-max_distance, max_distance]内 n torch.clamp(n, -max_distance, max_distance) # 是否为远距离负方向 is_negative n 0 n torch.abs(n) else: is_negative n 0 n torch.abs(n) # 对数分桶近距离区分细致远距离区分粗糙 max_exact num_buckets // 2 if n max_exact: ret n else: # 对数空间分桶 val torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) ret max_exact val.to(torch.long) ret torch.min(ret, torch.tensor(num_buckets - 1)) if is_negative: ret -ret return ret # 假设我们有相对位置矩阵 rel_pos [batch, seq_len, seq_len] bucket_indices get_relative_position_bucket(rel_pos, max_distance128, num_buckets64) # bucket_indices 的形状也是 [batch, seq_len, seq_len]值在 [-63, 63] 之间包含0 relative_bias relative_bias_embedding(bucket_indices) # 查表relative_bias_embedding 是一个 nn.Embedding(2*num_buckets, 1) # relative_bias 就是最终要加到注意力分数上的偏置矩阵4. 实践中的挑战、解决方案与效果评估将跨窗口RPE理论付诸实践尤其是在现有深度学习框架和模型架构中会遇到一系列工程和算法上的挑战。4.1 计算与内存开销的平衡引入跨窗口的、基于全局位置的RPE最直接的影响是计算开销。在标准的窗口内RPE中偏置矩阵的大小是[窗口大小, 窗口大小]。而在跨窗口设置下如果我们为一个长度为L的序列计算所有位置对的RPE理论上需要一个[L, L]的矩阵这在内 存上是不可接受的例如L8192时单是存储这个FP16矩阵就需要约512MB。因此稀疏计算和高效查找是关键。我们不会真的去计算和存储这个L x L的稠密矩阵。而是利用以下特性注意力模式本身的稀疏性例如在滑动窗口中每个查询只关注固定范围内的键所以只需要计算一个带状矩阵。相对位置的可重复性相对位置(i-j)只依赖于差值而不是具体的i和j。我们可以预先计算一个长度为(2*L-1)的偏置向量b其中b[k]对应相对距离为(k - L 1)的偏置。在计算注意力时通过索引技巧生成偏置矩阵。这种方法在Transformer-XL中就有体现它高效地生成了相对位置偏置。对于更复杂的模式如Longformer的混合注意力需要根据具体的注意力掩码mask来动态地计算或查找RPE偏置。这通常需要定制化的CUDA内核来实现高效操作因为标准的矩阵操作库难以处理这种不规则的模式。4.2 训练稳定性与初始化RPE的可学习偏置参数需要谨慎初始化。通常这些偏置会被初始化为零或很小的随机数。这是因为在训练初期我们希望注意力机制主要由内容相似性q·k主导位置偏置作为一个微调项慢慢加入。如果初始化过大可能会淹没内容信息导致模型难以收敛。另一个陷阱是分桶边界的不连续性。在对数分桶中距离4和距离5可能被分到不同的桶从而对应完全不同的可学习偏置b4和b5。这可能导致模型对距离的微小变化过于敏感。为了缓解这个问题有些实现会采用“软分桶”或给偏置嵌入表加上平滑正则鼓励相邻桶的偏置值变化平缓。4.3 长文本任务上的效果验证跨窗口RPE的有效性最终需要在长文本任务上进行检验。常见的评测基准包括长文本语言建模如PG-19书籍语料、arXiv数据集评测模型在长上下文下的困惑度。长文档摘要如GovReport、SummScreen评测生成摘要的质量。长文档问答如HotpotQA需要多文档推理、NarrativeQA基于故事全文。代码生成与理解代码文件往往很长且依赖关系复杂。在这些任务上配备了有效跨窗口RPE的模型如Longformer、BigBird相比仅使用绝对位置编码或标准窗口RPE的基线模型通常能展现出显著优势。例如困惑度更低生成的摘要更连贯、覆盖更多关键点问答的准确率更高。这证明了让模型准确感知长程相对位置关系对于理解文档级语义结构至关重要。4.4 一个简化的代码示例为滑动窗口注意力添加跨窗口RPE假设我们使用PyTorch并已有一个基础的滑动窗口注意力函数。以下是如何集成一个基于全局绝对位置和分桶的RPE的简化流程import torch import torch.nn as nn import math class SlidingWindowAttentionWithCrossWindowRPE(nn.Module): def __init__(self, embed_dim, num_heads, window_size, max_distance1024, rpe_buckets64): super().__init__() self.embed_dim embed_dim self.num_heads num_heads self.window_size window_size self.max_distance max_distance self.rpe_buckets rpe_buckets # 标准的Q, K, V投影层 self.q_proj nn.Linear(embed_dim, embed_dim) self.k_proj nn.Linear(embed_dim, embed_dim) self.v_proj nn.Linear(embed_dim, embed_dim) self.out_proj nn.Linear(embed_dim, embed_dim) # RPE偏置嵌入表我们为每个桶学习一个标量偏置每个注意力头可以不同这里简化为共享 # 桶的数量是 2 * rpe_buckets (正负距离) self.relative_bias_table nn.Parameter(torch.zeros(2 * rpe_buckets, num_heads)) # 初始化 nn.init.trunc_normal_(self.relative_bias_table, std0.02) def _get_relative_position_bucket(self, relative_position): 将相对位置映射到桶索引同前面的函数此处略去详细实现 # 返回的索引范围在 [0, 2*rpe_buckets-1] pass def forward(self, x, global_positions): x: 输入序列 [batch, seq_len, embed_dim] global_positions: 全局绝对位置索引 [batch, seq_len] 或 [seq_len] batch, seq_len, _ x.shape # 1. 计算Q, K, V q self.q_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) # [B, H, L, D_h] k self.k_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) v self.v_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) # 2. 计算滑动窗口内的内容注意力分数 (QK^T) # 这里简化了滑动窗口的掩码生成实际中可能需要更复杂的实现如使用banded matrix attn_scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1)) # [B, H, L, L] # 3. 计算跨窗口的RPE偏置 # 生成相对位置矩阵 [L, L] if global_positions.dim() 1: global_positions global_positions.unsqueeze(0).expand(batch, seq_len) # rel_pos[i, j] position_i - position_j rel_pos global_positions.unsqueeze(2) - global_positions.unsqueeze(1) # [B, L, L] # 将相对位置映射到桶索引 bucket_idx self._get_relative_position_bucket(rel_pos) # [B, L, L] # 查找RPE偏置表得到每个位置对的偏置 [B, L, L, H] rpe_bias self.relative_bias_table(bucket_idx) # 假设bucket_idx已适配嵌入层输入 rpe_bias rpe_bias.permute(0, 3, 1, 2) # 调整维度为 [B, H, L, L] # 4. 将RPE偏置加到注意力分数上 attn_scores attn_scores rpe_bias # 5. 应用滑动窗口掩码将窗口外的注意力分数设为负无穷 # 生成一个带状掩码仅保留中心带宽度为 (2*window_size1) 的区域 mask self._create_sliding_window_mask(seq_len, self.window_size).to(attn_scores.device) attn_scores attn_scores.masked_fill(mask 0, float(-inf)) # 6. 计算注意力权重和输出 attn_weights torch.softmax(attn_scores, dim-1) context torch.matmul(attn_weights, v) # [B, H, L, D_h] context context.transpose(1, 2).contiguous().view(batch, seq_len, self.embed_dim) output self.out_proj(context) return output def _create_sliding_window_mask(self, seq_len, window_size): 创建一个带状掩码矩阵 mask torch.ones(seq_len, seq_len, dtypetorch.bool) i torch.arange(seq_len).view(-1, 1) j torch.arange(seq_len).view(1, -1) mask torch.abs(i - j) window_size return mask.unsqueeze(0).unsqueeze(0) # [1, 1, L, L] 便于广播这个示例展示了核心思想使用全局位置计算相对距离通过分桶映射到可学习的偏置并将其与滑动窗口注意力结合。在实际的复杂模型如Longformer中注意力掩码和RPE偏置的生成逻辑会更加复杂需要处理局部、全局等多种注意力模式。5. 超越距离RPE的未来演进方向跨窗口RPE解决了距离计算的问题但当前主流的RPE仍然建立在“相对距离”这个单一维度上。然而位置关系远不止线性距离这么简单。未来的RPE可能会朝着更丰富、更结构化的方向发展。5.1 二维及高维位置编码对于图像、视频、图结构数据位置关系是多维的。在图像Transformer中相对位置通常用二维向量 (Δx, Δy) 表示。此时的RPE偏置表可能是一个二维查找表或者将二维向量编码成一个标量。这可以看作是跨窗口RPE在二维空间上的自然延伸其中“窗口”可能是图像的一个局部区块。5.2 基于内容的相对位置偏置当前的RPE是静态的、与内容无关的只要相对距离相同偏置就相同。但事实上两个词之间的位置重要性可能取决于它们本身是什么词。例如“因为”和“所以”之间的位置关系比两个普通名词之间的位置关系更重要。未来的RPE可能会动态化让偏置b_{i-j}不仅仅依赖于距离还依赖于查询q_i和键k_j的内容或者它们的交互结果。这相当于让模型自己学习在何种语义情境下距离因素应该如何被加权。5.3 与其它长程增强技术的结合跨窗口RPE是增强长程建模能力的一种手段。它可以与其它技术结合使用形成更强大的解决方案与记忆机制结合如Transformer-XL将过去片段的隐藏状态作为可延伸的上下文。RPE需要处理当前查询与记忆库中键的相对位置。与层次化注意力结合先对句子或段落进行粗粒度编码再在粗粒度表示上进行注意力。RPE需要在不同粒度层次上定义相对位置如词间距离、句间距离。与稀疏激活专家模型结合如Mixture of Experts (MoE)。RPE的设计可能需要考虑不同专家所处理的子序列之间的相对位置关系。在我个人的实验和项目应用中尤其是在处理法律长文档、学术论文和长篇对话时一个稳定且高效的跨窗口RPE模块是模型能否“读懂”全文结构的关键。初期最容易踩的坑就是错误地混用了局部和全局位置索引导致模型在窗口边界处行为异常。我的经验是在实现任何复杂的注意力模式时一定要先可视化出前向传播过程中生成的注意力掩码和相对位置偏置矩阵确保它们与你的设计意图完全一致。例如可以检查对于一个靠近片段末尾的查询它是否能正确地对前一片段开头的键赋予一个负的、较大绝对值的相对位置偏置。这种细致的调试往往比盲目调整超参数更能带来性能的提升。 Transformer跨窗口相对位置编码：突破长序列建模瓶颈的关键技术
1. 从绝对到相对为什么我们需要跨窗口的RPE在Transformer模型席卷自然语言处理领域的浪潮中Self-Attention机制无疑是其最核心的引擎。这个机制允许序列中的任意两个位置直接建立联系从而捕捉长距离依赖。然而经典的Self-Attention使用的是一个看似简单却暗藏玄机的设计绝对位置编码。简单来说就是给序列中的每个位置比如第1个词、第2个词分配一个独一无二的向量然后将这个向量加到词本身的嵌入向量上一起输入模型。这样模型在计算注意力时就能“感知”到每个词在句子中的绝对位置。这个设计在标准Transformer中运行良好但它有一个致命的缺陷它无法处理比训练时见过的序列更长的文本。想象一下你训练模型时见过的句子最长只有512个词那么模型就只学会了处理512个位置。当你试图用它来理解一篇1000个词的文档时对于第513个词及之后的词模型就“不认识”它们的位置了因为它没有学习过这些位置的编码。这极大地限制了Transformer在长文本任务如长文档摘要、书籍生成、长对话建模中的应用。为了解决这个问题研究者们提出了相对位置编码。它的核心思想不再是告诉模型“你是第几个词”而是告诉模型“你和另一个词之间相隔多远”。比如“我”和“爱”之间距离是1“我”和“深度学习”之间距离是3。这种编码方式天然具备长度外推性无论序列多长模型都学过如何处理“距离为1”、“距离为2”的关系因此理论上可以处理任意长度的序列。相对位置编码中最具代表性的工作之一就是RPE。RPE不再将位置信息加到输入上而是巧妙地修改了注意力分数的计算过程。在计算查询向量q_i和键向量k_j的注意力得分时除了常规的点积q_i·k_jRPE会额外引入一个基于相对位置偏移(i-j)的偏置项。这个偏置项是可学习的模型在训练过程中会学会当两个词距离为1时它们的注意力应该有一个怎样的基础倾向距离为2时又是另一个倾向。这样一来模型对位置的感知就从“绝对坐标”转变为了“相对关系”。然而标准的RPE实现通常假设整个序列的注意力计算是在一个“窗口”内完成的即模型一次性看到整个序列。但在处理超长序列时由于计算复杂度的限制Self-Attention的复杂度是序列长度的平方我们不得不将序列切分成多个片段或窗口然后在这些窗口内分别计算注意力。这就引出了本文要探讨的核心问题当序列被分割一个查询词在一个窗口内而它需要关注的键值词在另一个窗口内时标准的、基于窗口内相对位置的RPE就失效了。因为它无法计算跨窗口的两个词之间的相对位置关系。如何让RPE能够“看见”并正确处理这种跨窗口的相对位置就是“跨窗口的RPE”所要解决的挑战。2. 理解RPE的经典实现与跨窗口困境要理解跨窗口的挑战我们首先需要拆解一下经典RPE以Transformer-XL和后续改进工作为代表是如何工作的。它的核心公式可以简化为对注意力矩阵A的修改A_{i,j} (q_i · k_j) b_{i-j}这里b_{i-j}就是一个可学习的标量偏置其值仅依赖于查询位置i和键位置j的相对距离(i-j)。通常我们会预设一个最大相对距离k比如128对于所有|i-j| k的相对位置我们使用同一个偏置值b_{k}或b_{-k}认为超过这个距离的位置关系都“差不多远”。在标准的自回归语言建模中模型从左到右生成文本。在训练时为了模拟推理时的场景并提升效率通常会采用“片段递归”的策略也就是Transformer-XL的核心思想。模型会缓存上一个片段或窗口的隐藏状态在当前片段计算注意力时当前片段的词作为查询不仅可以关注当前片段内的词作为键还可以关注上一个片段缓存下来的词。这就在一定程度上打破了窗口的界限。但是这里的RPE计算会变得微妙。对于当前片段内的查询i和当前片段内的键j相对位置(i-j)是明确的。对于当前片段内的查询i和上一个片段缓存中的键j‘它们的绝对位置索引可能相差很远比如i是当前片段的第10个词j‘是上一个片段的第500个词。此时(i - j‘)这个值可能非常大远超预设的最大相对距离k。如果我们简单地将这个超大值作为索引去查找偏置表要么会索引越界要么会得到一个代表“非常远”的固定偏置值b_{k}。这虽然能运行但丢失了精确的相对距离信息对于模型来说“距离501”和“距离502”都被粗暴地归为了“距离k”的同一类这显然是不精确的。更复杂的情况出现在非自回归模型或双向编码器中比如BERT。在这些模型中一个窗口内的词需要同时关注窗口内所有其他词。如果我们简单地将长序列切成不重叠的窗口例如每512个词一个窗口那么窗口1中的词完全无法与窗口2中的词建立注意力连接RPE就更无从谈起了。这种硬切割会破坏序列的连贯性对于理解跨句、跨段落的语义关系是灾难性的。因此跨窗口RPE的核心目标就是设计一种机制使得在计算窗口化注意力时模型能够准确地获知并利用任意两个词无论它们是否在同一个计算窗口内之间的相对位置信息。这不仅仅是技术实现上的挑战更关乎模型能否真正具备处理长程、细粒度依赖关系的能力。3. 滑动窗口、扩张注意力与相对位置索引的重新校准为了解决跨窗口的RPE问题社区和业界提出了几种主流的思路它们从不同的角度对计算过程进行了改造。3.1 滑动窗口注意力这是最直观的一种方法。它不完全将序列切成独立的块而是让一个固定大小的窗口在序列上滑动。对于序列中的每个位置i其注意力范围是[i-w, iw]这样一个固定大小的邻域其中w是窗口半径。这样位于窗口边缘的词比如位置i自然可以“看到”相邻窗口中的词位置i-w到i-1如果i靠近当前块末尾的话。在这种设置下实现跨窗口RPE的关键在于统一所有词的位置坐标系。我们不能再用每个窗口内部的局部位置索引0到2w来计算相对位置因为不同窗口的局部索引0对应着序列中完全不同的绝对位置。解决方案是使用全局绝对位置索引。我们为序列中的每个词分配一个唯一的、连续的绝对位置编号0, 1, 2, …。在计算注意力时对于查询i和键j我们使用它们的绝对位置索引来计算相对距离d i - j。然后将这个距离d映射到RPE的偏置表中。由于滑动窗口保证了i和j的距离不会超过窗口大小2w1因此d的范围是有限的[-w, w]完全在预设的最大相对距离k的覆盖范围内通常k w。注意这里有一个重要的实现细节。在训练非常长的序列时我们可能无法一次性将整个序列的绝对位置比如10000都编码进模型。常用的技巧是使用循环位置编码或相对位置桶。例如可以将绝对位置对某个模数如512取余或者将相对距离d通过一个函数如对数桶映射到有限的几个桶中每个桶对应一个可学习的偏置。这样模型学到的不是“距离137的偏置”而是“距离在128-256这个桶范围内的偏置”在保证外推性的同时降低了参数量。3.2 扩张注意力与局部敏感哈希滑动窗口虽然解决了邻近窗口的问题但对于需要超长程依赖的任务窗口大小w可能仍然不够。扩张注意力是一种受空洞卷积启发的方法。它不是在每个位置都计算注意力而是每隔一定的步长扩张率选取一个位置进行计算。这样在不增加计算量的情况下每个位置的实际感受野变大了。结合RPE时我们需要处理的不再是连续位置之间的距离而是稀疏采样位置之间的距离。此时相对位置d i - j可能是一个很大的值并且是扩张率的倍数。RPE偏置表需要能够适应这种稀疏的、大间隔的距离模式。一种实践是将距离除以扩张率后再进行映射或分桶让模型学会对“扩张后的距离”进行建模。更激进的方法是使用局部敏感哈希等近似注意力机制将相似的词哈希到同一个桶中无论它们的位置多远。在这种情况下RPE的设计需要与哈希策略协同我们可能不再需要精确的相对距离而是需要一个能表示“是否在同一个哈希桶内”或“哈希桶之间的相对关系”的偏置。这为RPE的设计打开了新的思路即从基于数值距离的编码转向基于语义或结构分组的编码。3.3 长序列建模框架中的RPE集成以Longformer和BigBird为例一些专门为长序列设计的Transformer变体如Longformer和BigBird本身就采用了混合的注意力模式滑动窗口注意力全局注意力。它们天然需要处理跨窗口的RPE。以Longformer为例它的注意力模式是对于大多数词采用滑动窗口局部注意力对于少量预先选定的“全局词”如[CLS]标记或某些关键实体则赋予其关注整个序列的能力。在实现这种混合注意力时RPE需要被灵活地应用。对于局部滑动窗口部分使用上述的全局绝对位置索引方法计算RPE。对于全局注意力部分当一个全局词作为查询需要关注序列中所有键时它需要计算与每一个键的相对位置。由于序列可能极长这里必须使用分桶策略。将所有可能的巨大相对距离通过一个函数如对数函数映射到几十个或几百个有限的桶中。例如距离1-2映射到桶0距离3-4映射到桶1距离5-8映射到桶2以此类推。模型为每个桶学习一个偏置。这样全局词在关注远处一个词时使用的RPE偏置是基于“距离桶”的而非精确距离这是一种在计算效率和模型能力之间的有效折衷。在实际编码中这通常体现为一个庞大的相对位置偏置矩阵的查找过程。我们需要预先计算好序列中所有位置对之间的相对距离桶索引形成一个索引矩阵。在注意力计算时根据查询i和键j从这个索引矩阵中取出对应的桶编号再去查找一个小的、可学习的偏置嵌入表获得标量偏置b_{bucket}。# 伪代码示意基于分桶的RPE偏置获取 def get_relative_position_bucket(relative_position, max_distance512, num_buckets32): 将相对位置映射到桶中。 relative_position: 标量或矩阵表示相对距离 (i-j) max_distance: 超过此距离的视为同一类远距离 num_buckets: 桶的总数 ret 0 n -relative_position if max_distance is not None: # 将距离限制在[-max_distance, max_distance]内 n torch.clamp(n, -max_distance, max_distance) # 是否为远距离负方向 is_negative n 0 n torch.abs(n) else: is_negative n 0 n torch.abs(n) # 对数分桶近距离区分细致远距离区分粗糙 max_exact num_buckets // 2 if n max_exact: ret n else: # 对数空间分桶 val torch.log(n.float() / max_exact) / math.log(max_distance / max_exact) * (num_buckets - max_exact) ret max_exact val.to(torch.long) ret torch.min(ret, torch.tensor(num_buckets - 1)) if is_negative: ret -ret return ret # 假设我们有相对位置矩阵 rel_pos [batch, seq_len, seq_len] bucket_indices get_relative_position_bucket(rel_pos, max_distance128, num_buckets64) # bucket_indices 的形状也是 [batch, seq_len, seq_len]值在 [-63, 63] 之间包含0 relative_bias relative_bias_embedding(bucket_indices) # 查表relative_bias_embedding 是一个 nn.Embedding(2*num_buckets, 1) # relative_bias 就是最终要加到注意力分数上的偏置矩阵4. 实践中的挑战、解决方案与效果评估将跨窗口RPE理论付诸实践尤其是在现有深度学习框架和模型架构中会遇到一系列工程和算法上的挑战。4.1 计算与内存开销的平衡引入跨窗口的、基于全局位置的RPE最直接的影响是计算开销。在标准的窗口内RPE中偏置矩阵的大小是[窗口大小, 窗口大小]。而在跨窗口设置下如果我们为一个长度为L的序列计算所有位置对的RPE理论上需要一个[L, L]的矩阵这在内 存上是不可接受的例如L8192时单是存储这个FP16矩阵就需要约512MB。因此稀疏计算和高效查找是关键。我们不会真的去计算和存储这个L x L的稠密矩阵。而是利用以下特性注意力模式本身的稀疏性例如在滑动窗口中每个查询只关注固定范围内的键所以只需要计算一个带状矩阵。相对位置的可重复性相对位置(i-j)只依赖于差值而不是具体的i和j。我们可以预先计算一个长度为(2*L-1)的偏置向量b其中b[k]对应相对距离为(k - L 1)的偏置。在计算注意力时通过索引技巧生成偏置矩阵。这种方法在Transformer-XL中就有体现它高效地生成了相对位置偏置。对于更复杂的模式如Longformer的混合注意力需要根据具体的注意力掩码mask来动态地计算或查找RPE偏置。这通常需要定制化的CUDA内核来实现高效操作因为标准的矩阵操作库难以处理这种不规则的模式。4.2 训练稳定性与初始化RPE的可学习偏置参数需要谨慎初始化。通常这些偏置会被初始化为零或很小的随机数。这是因为在训练初期我们希望注意力机制主要由内容相似性q·k主导位置偏置作为一个微调项慢慢加入。如果初始化过大可能会淹没内容信息导致模型难以收敛。另一个陷阱是分桶边界的不连续性。在对数分桶中距离4和距离5可能被分到不同的桶从而对应完全不同的可学习偏置b4和b5。这可能导致模型对距离的微小变化过于敏感。为了缓解这个问题有些实现会采用“软分桶”或给偏置嵌入表加上平滑正则鼓励相邻桶的偏置值变化平缓。4.3 长文本任务上的效果验证跨窗口RPE的有效性最终需要在长文本任务上进行检验。常见的评测基准包括长文本语言建模如PG-19书籍语料、arXiv数据集评测模型在长上下文下的困惑度。长文档摘要如GovReport、SummScreen评测生成摘要的质量。长文档问答如HotpotQA需要多文档推理、NarrativeQA基于故事全文。代码生成与理解代码文件往往很长且依赖关系复杂。在这些任务上配备了有效跨窗口RPE的模型如Longformer、BigBird相比仅使用绝对位置编码或标准窗口RPE的基线模型通常能展现出显著优势。例如困惑度更低生成的摘要更连贯、覆盖更多关键点问答的准确率更高。这证明了让模型准确感知长程相对位置关系对于理解文档级语义结构至关重要。4.4 一个简化的代码示例为滑动窗口注意力添加跨窗口RPE假设我们使用PyTorch并已有一个基础的滑动窗口注意力函数。以下是如何集成一个基于全局绝对位置和分桶的RPE的简化流程import torch import torch.nn as nn import math class SlidingWindowAttentionWithCrossWindowRPE(nn.Module): def __init__(self, embed_dim, num_heads, window_size, max_distance1024, rpe_buckets64): super().__init__() self.embed_dim embed_dim self.num_heads num_heads self.window_size window_size self.max_distance max_distance self.rpe_buckets rpe_buckets # 标准的Q, K, V投影层 self.q_proj nn.Linear(embed_dim, embed_dim) self.k_proj nn.Linear(embed_dim, embed_dim) self.v_proj nn.Linear(embed_dim, embed_dim) self.out_proj nn.Linear(embed_dim, embed_dim) # RPE偏置嵌入表我们为每个桶学习一个标量偏置每个注意力头可以不同这里简化为共享 # 桶的数量是 2 * rpe_buckets (正负距离) self.relative_bias_table nn.Parameter(torch.zeros(2 * rpe_buckets, num_heads)) # 初始化 nn.init.trunc_normal_(self.relative_bias_table, std0.02) def _get_relative_position_bucket(self, relative_position): 将相对位置映射到桶索引同前面的函数此处略去详细实现 # 返回的索引范围在 [0, 2*rpe_buckets-1] pass def forward(self, x, global_positions): x: 输入序列 [batch, seq_len, embed_dim] global_positions: 全局绝对位置索引 [batch, seq_len] 或 [seq_len] batch, seq_len, _ x.shape # 1. 计算Q, K, V q self.q_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) # [B, H, L, D_h] k self.k_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) v self.v_proj(x).view(batch, seq_len, self.num_heads, -1).transpose(1, 2) # 2. 计算滑动窗口内的内容注意力分数 (QK^T) # 这里简化了滑动窗口的掩码生成实际中可能需要更复杂的实现如使用banded matrix attn_scores torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(q.size(-1)) # [B, H, L, L] # 3. 计算跨窗口的RPE偏置 # 生成相对位置矩阵 [L, L] if global_positions.dim() 1: global_positions global_positions.unsqueeze(0).expand(batch, seq_len) # rel_pos[i, j] position_i - position_j rel_pos global_positions.unsqueeze(2) - global_positions.unsqueeze(1) # [B, L, L] # 将相对位置映射到桶索引 bucket_idx self._get_relative_position_bucket(rel_pos) # [B, L, L] # 查找RPE偏置表得到每个位置对的偏置 [B, L, L, H] rpe_bias self.relative_bias_table(bucket_idx) # 假设bucket_idx已适配嵌入层输入 rpe_bias rpe_bias.permute(0, 3, 1, 2) # 调整维度为 [B, H, L, L] # 4. 将RPE偏置加到注意力分数上 attn_scores attn_scores rpe_bias # 5. 应用滑动窗口掩码将窗口外的注意力分数设为负无穷 # 生成一个带状掩码仅保留中心带宽度为 (2*window_size1) 的区域 mask self._create_sliding_window_mask(seq_len, self.window_size).to(attn_scores.device) attn_scores attn_scores.masked_fill(mask 0, float(-inf)) # 6. 计算注意力权重和输出 attn_weights torch.softmax(attn_scores, dim-1) context torch.matmul(attn_weights, v) # [B, H, L, D_h] context context.transpose(1, 2).contiguous().view(batch, seq_len, self.embed_dim) output self.out_proj(context) return output def _create_sliding_window_mask(self, seq_len, window_size): 创建一个带状掩码矩阵 mask torch.ones(seq_len, seq_len, dtypetorch.bool) i torch.arange(seq_len).view(-1, 1) j torch.arange(seq_len).view(1, -1) mask torch.abs(i - j) window_size return mask.unsqueeze(0).unsqueeze(0) # [1, 1, L, L] 便于广播这个示例展示了核心思想使用全局位置计算相对距离通过分桶映射到可学习的偏置并将其与滑动窗口注意力结合。在实际的复杂模型如Longformer中注意力掩码和RPE偏置的生成逻辑会更加复杂需要处理局部、全局等多种注意力模式。5. 超越距离RPE的未来演进方向跨窗口RPE解决了距离计算的问题但当前主流的RPE仍然建立在“相对距离”这个单一维度上。然而位置关系远不止线性距离这么简单。未来的RPE可能会朝着更丰富、更结构化的方向发展。5.1 二维及高维位置编码对于图像、视频、图结构数据位置关系是多维的。在图像Transformer中相对位置通常用二维向量 (Δx, Δy) 表示。此时的RPE偏置表可能是一个二维查找表或者将二维向量编码成一个标量。这可以看作是跨窗口RPE在二维空间上的自然延伸其中“窗口”可能是图像的一个局部区块。5.2 基于内容的相对位置偏置当前的RPE是静态的、与内容无关的只要相对距离相同偏置就相同。但事实上两个词之间的位置重要性可能取决于它们本身是什么词。例如“因为”和“所以”之间的位置关系比两个普通名词之间的位置关系更重要。未来的RPE可能会动态化让偏置b_{i-j}不仅仅依赖于距离还依赖于查询q_i和键k_j的内容或者它们的交互结果。这相当于让模型自己学习在何种语义情境下距离因素应该如何被加权。5.3 与其它长程增强技术的结合跨窗口RPE是增强长程建模能力的一种手段。它可以与其它技术结合使用形成更强大的解决方案与记忆机制结合如Transformer-XL将过去片段的隐藏状态作为可延伸的上下文。RPE需要处理当前查询与记忆库中键的相对位置。与层次化注意力结合先对句子或段落进行粗粒度编码再在粗粒度表示上进行注意力。RPE需要在不同粒度层次上定义相对位置如词间距离、句间距离。与稀疏激活专家模型结合如Mixture of Experts (MoE)。RPE的设计可能需要考虑不同专家所处理的子序列之间的相对位置关系。在我个人的实验和项目应用中尤其是在处理法律长文档、学术论文和长篇对话时一个稳定且高效的跨窗口RPE模块是模型能否“读懂”全文结构的关键。初期最容易踩的坑就是错误地混用了局部和全局位置索引导致模型在窗口边界处行为异常。我的经验是在实现任何复杂的注意力模式时一定要先可视化出前向传播过程中生成的注意力掩码和相对位置偏置矩阵确保它们与你的设计意图完全一致。例如可以检查对于一个靠近片段末尾的查询它是否能正确地对前一片段开头的键赋予一个负的、较大绝对值的相对位置偏置。这种细致的调试往往比盲目调整超参数更能带来性能的提升。