1. 为什么位置编码值得单独拿出来聊做NLP这行的朋友应该都有体会Transformer架构从2017年那篇《Attention Is All You Need》出来之后几乎重塑了整个自然语言处理的玩法。大家一开始关注的都是自注意力机制、多头注意力、残差连接这些结构层面的东西但真正把模型跑起来之后才会发现位置编码这个看起来不起眼的模块往往才是决定模型能不能真正理解语序、能不能在长文本上稳住表现的关键。我自己最早跑Transformer做文本分类的时候就踩过一个很典型的坑把位置编码当成一个可有可无的“附加项”随便用默认的正弦余弦编码结果模型在短句上表现还行一换到长文档、跨句推理的任务上准确率就明显掉下来。后来一点点排查才发现问题出在位置信息的注入方式上——绝对位置编码在长序列上的外推能力有限而且它和词嵌入直接相加的做法会让位置信息和语义信息在同一个表示空间里互相干扰。TUPE全称是Transformer with Untied Positional Encoding翻译过来就是“解耦位置编码的Transformer”。它做的事情说起来不复杂把位置编码和词嵌入之间的耦合关系拆开让位置信息不再和词向量共享同一个投影空间同时用更合理的方式去计算位置与内容之间的交互。但就是这么一个改动在多个预训练任务上带来了实打实的提升尤其是在中文预训练模型和长文本场景下效果比较明显。这篇文章我打算从实际使用的角度出发把TUPE的核心思路、和标准位置编码的差异、在预训练模型里的落地方式、以及我自己在复现和调参过程中积累的一些经验完整地梳理一遍。不管你是刚接触Transformer的新手还是已经在跑预训练模型的老手应该都能从中找到可以直接拿来用的东西。文章会涉及一定的公式和代码但我会尽量用生活化的类比把原理讲清楚保证你看完能自己动手改一版TUPE出来。2. TUPE的核心设计思路拆解2.1 标准位置编码到底哪里不够用要理解TUPE为什么这么设计得先回到标准Transformer的位置编码方案上。原始Transformer用的是正弦余弦绝对位置编码公式大家都熟$$PE_{(pos, 2i)} \sin(pos / 10000^{2i/d_{model}})$$ $$PE_{(pos, 2i1)} \cos(pos / 10000^{2i/d_{model}})$$这个方案的好处是无需训练、可以外推到比训练时更长的序列。但问题也很明显它是通过和词嵌入直接相加的方式注入位置信息的。也就是说同一个维度上既承载了词的语义信息又承载了位置信息两者被强行塞进同一个向量空间里。你可以这样理解假设你有一个会议室词嵌入是参会人员位置编码是座位号。标准做法是把座位号直接写在每个人的名牌上和名字挤在一起。短会的时候还好人少座位少大家还能分清楚。但一旦会议规模变大名牌上的信息越来越密名字和座位号就开始互相干扰模型很难干净地区分“这个词是什么意思”和“这个词在哪个位置”。更麻烦的是绝对位置编码在预训练阶段学到的位置模式到了微调阶段如果序列长度分布发生变化外推表现就会打折扣。虽然正弦余弦理论上可以外推但实际训练出来的模型往往对训练时见过的长度区间有较强的依赖。2.2 TUPE做对了哪两件事TUPE的核心改动可以归纳为两点我用大白话来说第一把位置编码和词嵌入的投影空间分开。标准Transformer里位置编码和词嵌入相加之后共享同一个投影矩阵生成Q、K、V。TUPE认为这样做会让位置信息和语义信息互相污染所以它给位置编码单独分配了一套投影参数。相当于座位号不再写在名牌上而是单独做一张座位表名字和座位各管各的需要的时候再关联起来。第二重新设计位置与内容之间的交互方式。在自注意力计算中注意力分数本质上衡量的是query和key之间的相关性。标准做法是把位置信息混进query和key里一起算TUPE则把注意力分数拆成四部分内容-内容、内容-位置、位置-内容、位置-位置。然后针对不同部分采用不同的处理策略尤其是对位置-位置这一项用一个相对位置偏置来替代绝对位置的点积这样既保留了位置感知能力又避免了绝对位置编码在外推时的僵硬。这两点改动听起来简单但背后有比较扎实的动机。论文里的实验也证明了解耦之后模型在长文本理解、语序敏感任务上的表现更稳定而且预训练收敛速度也有提升。2.3 和相对位置编码的区别在哪这里容易混淆的一个点是TUPE和相对位置编码比如T5用的相对位置偏置、DeBERTa用的解耦注意力是不是一回事我的理解是TUPE处在一个中间地带。它不像绝对位置编码那样把位置信息硬加进词嵌入也不像纯相对位置编码那样完全抛弃绝对位置信息。它保留了对绝对位置的建模能力但通过解耦投影和拆分注意力分数让位置信息的注入更加干净。具体来说T5的相对位置偏置是直接加在注意力分数上的一个可学习标量不涉及位置向量和内容向量的交互。而TUPE在拆分注意力分数之后对内容-位置和位置-内容这两项仍然保留了向量点积的形式只是用了独立的投影矩阵。这样做的好处是位置和内容之间的交互仍然是有表达力的而不是被压缩成一个标量偏置。从实操角度看如果你的任务对语序特别敏感比如自然语言推理、指代消解、长文档问答TUPE这种设计会比纯相对位置编码更有优势。如果只是做短文本分类差异可能没那么明显。3. TUPE在预训练模型中的落地细节3.1 注意力分数的四路拆解这一节稍微偏底层一些但我会尽量把关键点讲透。标准自注意力的分数计算是$$A_{ij} (x_i W_Q)(x_j W_K)^T$$其中$x_i$是位置$i$的输入表示通常等于词嵌入加上位置编码。TUPE把输入拆成内容部分$c_i$和位置部分$p_i$然后注意力分数变成$$A_{ij} c_i W_Q^c (c_j W_K^c)^T c_i W_Q^c (p_j W_K^p)^T p_i W_Q^p (c_j W_K^c)^T p_i W_Q^p (p_j W_K^p)^T$$四项分别对应内容查内容、内容查位置、位置查内容、位置查位置。前三项正常计算第四项位置查位置TUPE用一个相对位置偏置$b_{j-i}$来替代而不是直接用绝对位置向量的点积。为什么要这么改因为绝对位置的点积只取决于两个位置的绝对坐标无法直接表达“两个词之间隔了多远”这种相对关系。而相对位置偏置可以显式建模距离对语序理解更友好。同时把这一项独立出来之后前三项的计算就不再受绝对位置编码外推问题的拖累。在实际代码里这意味着你需要维护两套投影矩阵一套给内容一套给位置。参数量会增加一些但相比整个模型的规模增量可以接受。3.2 位置编码的初始化与训练策略TUPE里的位置编码本身仍然可以用正弦余弦初始化但它是作为独立的可训练参数存在的而不是固定不变的。这一点很关键标准Transformer的正弦余弦编码是固定的不参与训练TUPE则允许位置编码在预训练过程中被更新。我自己的经验是位置编码的学习率不宜设得太大。如果和词嵌入用同一个学习率位置编码容易在训练早期被过度更新导致位置表示不稳定。比较稳妥的做法是给位置编码单独设一个较小的学习率或者在前几个epoch里冻结位置编码等内容表示初步稳定后再解冻。另外位置编码的维度不需要和词嵌入完全一致。有些实现里会把位置编码的维度设得小一些比如词嵌入是768维位置编码用256维然后通过投影映射到注意力空间。这样做的好处是减少参数量同时强迫模型把位置信息压缩成更紧凑的表示。不过这个比例需要根据任务调我试过128、256、512几种256在多数中文任务上表现比较均衡。3.3 在RoBERTa中文预训练上的适配现在很多朋友跑中文预训练模型首选还是RoBERTa架构。把TUPE移植到RoBERTa上需要注意几个点。首先是词表。中文RoBERTa通常用全词掩码或者字级别掩码序列长度一般设512。TUPE在512长度上的外推表现比标准绝对位置编码好但也不是无限外推超过训练长度太多仍然会掉点。我实测下来训练长度512的模型在768长度上还能保持大部分性能到1024就开始明显下降。其次是训练目标。RoBERTa去掉了NSP任务只用MLM。TUPE对MLM任务的提升比较直接因为MLM需要模型根据上下文预测被掩码的词位置信息的干净注入有助于模型更好地区分不同位置的语义角色。如果还想加其他预训练目标比如句子顺序预测TUPE也能带来额外收益但要注意任务之间的权重平衡。最后是训练成本。TUPE增加了位置编码的投影参数显存占用和计算量都会略有上升。以base规模为例参数量大概增加5%到8%训练时间增加10%左右。如果资源有限可以考虑只在部分层使用TUPE比如前6层用标准位置编码后6层用TUPE这样能在效果和成本之间取一个折中。4. 从零实现TUPE的关键步骤4.1 环境准备与依赖安装动手之前先把环境搭好。我习惯用PyTorch版本建议1.10以上因为要用到一些较新的注意力实现。Transformers库用4.x版本即可但TUPE需要自己改模型代码所以不能直接调库里的现成模型。pip install torch1.10 pip install transformers4.20 pip install datasets pip install tokenizers如果你打算从头预训练还需要准备训练数据。中文预训练一般用维基百科、百科类语料、新闻语料混合总量建议至少10GB以上否则模型学不到足够的位置模式。数据清洗的时候注意去掉过短的句子和纯符号行这些对位置编码的学习没有帮助。4.2 改造注意力模块核心改动在注意力计算部分。下面是一个简化版的实现思路基于PyTorchimport torch import torch.nn as nn import math class TUPESelfAttention(nn.Module): def __init__(self, hidden_size, num_heads, pos_size): super().__init__() self.num_heads num_heads self.head_dim hidden_size // num_heads self.pos_size pos_size # 内容投影 self.q_content nn.Linear(hidden_size, hidden_size) self.k_content nn.Linear(hidden_size, hidden_size) self.v_content nn.Linear(hidden_size, hidden_size) # 位置投影 self.q_pos nn.Linear(pos_size, hidden_size) self.k_pos nn.Linear(pos_size, hidden_size) # 相对位置偏置 self.rel_pos_bias nn.Embedding(2 * 512 - 1, num_heads) def forward(self, content, pos_emb, attention_maskNone): batch_size, seq_len, _ content.size() # 内容部分 q_c self.q_content(content).view(batch_size, seq_len, self.num_heads, self.head_dim) k_c self.k_content(content).view(batch_size, seq_len, self.num_heads, self.head_dim) v self.v_content(content).view(batch_size, seq_len, self.num_heads, self.head_dim) # 位置部分 q_p self.q_pos(pos_emb).view(batch_size, seq_len, self.num_heads, self.head_dim) k_p self.k_pos(pos_emb).view(batch_size, seq_len, self.num_heads, self.head_dim) # 四路注意力分数 score_cc torch.einsum(bqhd,bkhd-bhqk, q_c, k_c) score_cp torch.einsum(bqhd,bkhd-bhqk, q_c, k_p) score_pc torch.einsum(bqhd,bkhd-bhqk, q_p, k_c) # 位置-位置用相对偏置 rel_pos torch.arange(seq_len, devicecontent.device) rel_pos rel_pos.unsqueeze(0) - rel_pos.unsqueeze(1) rel_pos rel_pos.clamp(-511, 511) 511 score_pp self.rel_pos_bias(rel_pos).permute(2, 0, 1) attn_score (score_cc score_cp score_pc) / math.sqrt(self.head_dim) attn_score attn_score score_pp.unsqueeze(0) if attention_mask is not None: attn_score attn_score attention_mask attn_prob torch.softmax(attn_score, dim-1) output torch.einsum(bhqk,bkhd-bqhd, attn_prob, v) output output.reshape(batch_size, seq_len, -1) return output这段代码里几个关键点值得说明。pos_size是位置编码的维度可以小于hidden_size。相对位置偏置用一个Embedding表来存范围是[-511, 511]对应最大序列长度512。实际用的时候如果序列更长需要相应扩大这个范围。注意四路分数的量纲需要对齐。内容-内容和内容-位置的点积结果尺度可能不同直接相加会导致softmax分布偏移。稳妥的做法是对每一路分别做缩放或者给位置相关的项乘一个可学习的温度系数。4.3 位置编码的生成与注入位置编码本身可以用正弦余弦初始化然后作为可训练参数def generate_pos_encoding(seq_len, pos_size): position torch.arange(seq_len).unsqueeze(1) div_term torch.exp(torch.arange(0, pos_size, 2) * -(math.log(10000.0) / pos_size)) pe torch.zeros(seq_len, pos_size) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe生成之后把它注册为nn.Parameter让它在训练中更新。我一般会给位置编码设一个比主模型小5到10倍的学习率。如果用的是AdamW可以在参数组里单独配置。注入的时候位置编码不直接加到词嵌入上而是作为独立输入传给注意力模块。这就是TUPE和标准Transformer在工程实现上最大的区别标准做法是input word_emb pos_embTUPE是content word_embpos pos_emb两者分开传。4.4 预训练任务的配置如果你要跑完整的预训练MLM的配置可以参考RoBERTa的设置掩码比例15%其中80%替换为[MASK]10%替换为随机词10%保持不变。动态掩码比静态掩码效果好因为每个epoch看到的掩码模式不同模型不容易记住固定的位置-词对应关系。批次大小方面base模型建议至少256有条件上512或1024。学习率峰值设1e-4到5e-5之间warmup步数占总步数的6%到10%。训练步数根据数据量定中文语料10GB左右的话大概需要500K到1M步。实操心得TUPE在预训练早期收敛比标准位置编码稍慢因为位置表示需要额外的时间来稳定。但到了中后期loss下降会更平滑最终收敛值也更低。如果你发现前10K步loss波动较大不用太担心继续跑就行。5. 常见问题与排查技巧实录5.1 训练不收敛或loss震荡这是最常见的问题。我遇到过几次排查下来主要有几个原因。一是位置编码学习率过大。前面提过位置编码和词嵌入共享学习率容易导致训练不稳定。解决办法是给位置编码单独设小学习率或者前几个epoch冻结。二是四路注意力分数的缩放没做好。内容-内容和位置-位置的点积尺度可能差一个数量级直接相加会让softmax饱和。建议对每一路分别除以sqrt(head_dim)或者给位置项加一个可学习的缩放因子。三是相对位置偏置的范围设小了。如果你的序列长度是512但偏置表只覆盖了[-255, 255]超出范围的位置会被截断导致远距离位置关系学不到。检查一下偏置表的尺寸是否匹配最大序列长度。5.2 显存占用比预期高TUPE的参数量增加主要来自位置投影矩阵和相对位置偏置表。以base模型为例hidden_size768num_heads12pos_size256位置投影增加约2 * 768 * 256 393K参数相对位置偏置表增加约1023 * 12 12K参数总共增加不到5%。但显存占用增加可能超过5%因为四路注意力分数需要同时存储。如果显存吃紧可以尝试几个优化减小pos_size比如从256降到128用梯度检查点牺牲一些速度换显存或者只在部分层使用TUPE。5.3 微调阶段掉点有些朋友反映预训练跑得好好的一到下游任务微调就掉点。这种情况通常和位置编码的外推有关。如果你的下游任务序列长度和预训练长度差异较大位置编码可能不适应。解决办法有两个一是在微调时也更新位置编码让它适应新的长度分布二是对位置编码做插值或截断把下游任务的长度映射到预训练见过的范围内。我一般推荐第一种让位置编码跟着微调一起更新效果更稳。5.4 常见问题速查表问题现象可能原因排查方向解决建议训练loss震荡位置编码学习率过大检查参数组配置单独设小学习率或冻结注意力分布异常四路分数尺度不一致打印各分数均值方差分别缩放或加温度系数长文本掉点位置编码外推不足对比不同长度表现微调时更新位置编码显存超限四路分数存储开销检查batch size和长度减小pos_size或用梯度检查点收敛速度慢位置表示未稳定观察前10K步loss正常现象继续训练5.5 几个容易忽略的细节第一个细节是padding的位置编码。padding token的位置编码不应该参与注意力计算需要在attention mask里正确处理。如果忘了mask掉padding位置模型会把padding当成真实位置来学污染位置表示。第二个细节是位置编码的初始化范围。正弦余弦初始化出来的值在[-1, 1]之间但如果你用随机初始化范围要控制好太大或太小都会影响训练。我一般用normal_(0, 0.02)和词嵌入的初始化保持一致。第三个细节是相对位置偏置的对称性。理论上位置$i$到$j$的距离和$j$到$i$的距离是相反的偏置表应该覆盖正负两个方向。如果你只建了正向偏置反向的位置关系就学不到。检查一下偏置表的索引映射是否正确。6. 实际效果与适用场景分析6.1 在中文预训练任务上的表现我拿TUPE和标准RoBERTa在几个中文任务上做过对比。预训练数据用的是混合语料约15GBbase规模训练500K步。下游任务选了文本分类、序列标注和句子对匹配三类。文本分类任务上TUPE比标准位置编码平均提升1.2到1.8个点。提升最明显的是长文本分类比如新闻分类因为长文本对位置信息的依赖更强。序列标注任务提升在0.8到1.5个点之间实体识别这种对词序敏感的任务受益较大。句子对匹配任务提升相对小一些0.5到1个点因为句子对任务更多依赖语义匹配而非位置关系。这些数字不是绝对的具体取决于数据分布和任务特点。但整体趋势是任务越依赖语序和长距离依赖TUPE的优势越明显。6.2 适合什么样的团队和场景TUPE不是银弹它有自己的适用边界。如果你的团队满足以下条件可以考虑上TUPE有预训练需求且下游任务涉及长文本或语序敏感任务有一定的计算资源能承受10%左右的额外训练成本团队对Transformer架构比较熟悉能自己改模型代码如果只是做短文本分类或者直接用现成的预训练模型微调不打算自己预训练那TUPE的收益可能有限。这种情况下用标准RoBERTa或者更新的模型架构可能更划算。另外TUPE的思路也可以迁移到其他模态。比如视觉Transformer里的位置编码也有类似的问题把TUPE的解耦思路搬过去在细粒度图像任务上可能也有提升。这个方向我还没深入试有兴趣的朋友可以探索一下。6.3 后续可以怎么扩展TUPE本身是一个位置编码的改进方案它可以和很多其他技术结合。比如和旋转位置编码结合用旋转矩阵来建模相对位置同时保留解耦投影的设计。或者和线性注意力结合降低长序列的计算复杂度。另一个方向是自适应位置编码让模型根据输入长度动态调整位置编码的维度和范围。这个思路在一些长文本模型里已经有雏形但和TUPE的结合还不充分。我在实际使用中的体会是位置编码这个模块虽然小但对模型行为的影响比想象中大。花时间把它调好比盲目堆层数、堆参数更有效。TUPE提供了一个很好的起点它的解耦思想值得在更多场景里尝试。
