上篇把三角函数位置编码的公式和直觉拆了个遍这篇就直接动真格的把公式落成代码再把工程里那些文档里不会写清楚、但实战一定会撞上的细节全部摊开讲一遍。我始终觉得位置编码是Transformer里最“骗人”的部件——公式短、实现更短披着几行代码的外衣好像简单到不行。但真正动手写模型、跑训练、做推理部署的时候你会发现它在batch、精度、序列长度、训练和推理一致性上都有不少暗坑。这篇面向两类人一是正在手写Transformer、想把位置编码彻底搞明白的二是已经有了能用代码、但总被某些边界情况坑到的。如果你只是想要一份可以直接抄的Sinusoidal位置编码实现我会先给出代码再把背后的“为什么”一并交代。1. 从公式到张量位置编码要解决的到底是什么形状的问题1.1 先把公式老老实实写对三角函数位置编码的原始形式如下PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是token在序列里的位置i的取值范围是0到d_model/2-12i和2i1分别表示d_model维向量中的偶数维和奇数维。我在很多文章和代码里见过这个公式的简化写法其中最危险的两种错误把分母写成了10000^(i/d_model)这会让频率衰减速度发生明显偏移位置矩阵的语义会变把sin和cos的奇偶分配写反虽然频率没变但如果要加载别人预训练好的编码就和原有模型不一致了。这个公式的本质是为每个token生成一个d_model维的向量这个向量的每个维度都是一个不同频率的三角函数值。第0维和第1维共享同一个频率但相位差π/2第2维和第3维共享另一个频率依此类推。所以这个向量的信息容量并不来自“内部维度之间的大小关系”而是来自不同维度在不同频率尺度上的组合。1.2 波长变化从“分度尺”的角度看我记得上篇强调过一个点这里再补充一个更直观的角度。对于维度2i波长公式是波长_i 2π * 10000^(2i/d_model)当d_model512时不同维度对应的波长如下表维度索引i分母倍率近似波长01≈ 6.286410≈ 62.8128100≈ 6281921000≈ 6283255≈9611≈ 60397这意味着什么低维度的波长很短位置变化一点点函数值就会明显波动相当于一把高精度的“近程尺”高维度的波长很长可以帮助模型感知长距离的全局位置关系。位置编码的本质就是把不同粒度的“尺子”捆绑到同一个向量里让模型在自注意力计算时能同时看到近邻和远程的偏移信息。这其实很像我们看地图时候的经验你既需要看街道级别的细节也需要看城市级别的整体方位。三角编码用一组不同频率的波把这两种信息压缩进了同一个向量表示。1.3 矩阵形态最终要构造出什么位置编码真正需要生成的是一个形状为(seq_len, d_model)的常量矩阵每一行是一个位置编码向量。比如句子长度是10、d_model是512最终就是一个10×512的矩阵。构造时角度矩阵的广播方式很容易写错。我先给一个最直接的构造版本后面再对比优化import numpy as np def sinusoidal_positional_encoding(seq_len: int, d_model: int) - np.ndarray: pe np.zeros((seq_len, d_model)) pos np.arange(seq_len)[:, None] # (seq_len, 1) i np.arange(d_model // 2) # 共有 d_model//2 个频率 denom np.power(10000.0, 2 * i / d_model) # (d_model//2,) angle pos / denom # (seq_len, d_model//2) pe[:, 0::2] np.sin(angle) pe[:, 1::2] np.cos(angle) return pe这里的关键在于广播pos是(seq_len, 1)denom是(d_model//2,)二者相除会得到一个(seq_len, d_model//2)的矩阵然后通过切片填到奇偶维上。如果你把pos写成(1, seq_len)angle的shape就会变成(1, d_model//2)或者报错这是新手最容易栽的地方。2. 手写实现三种写法从“跑通”到“好用”再到“可上线”2.1 逐位置循环版帮助你理解但千万别用于训练理解公式最笨但也最可靠的方法是一个个位置、一个个维度硬算def pos_enc_loop(seq_len, d_model): pe np.zeros((seq_len, d_model)) for pos in range(seq_len): for i in range(d_model // 2): omega 10000 ** (2 * i / d_model) pe[pos, 2 * i] np.sin(pos / omega) pe[pos, 2 * i 1] np.cos(pos / omega) return pe在seq_len512、d_model512时这个嵌套循环大概是13万次sin/cos计算在CPU上跑出来要几秒钟。虽然能跑但放进训练循环里基本不可能接受。这个版本的价值只在于验证你对公式的理解或者用来和向量化版本逐元素对比输出是否一致。2.2 向量化版正式实现的基本形态用广播构造的版本上面已经给出了。这个版本在numpy和PyTorch里跑起来都很快因为它把循环变成了矩阵运算底层调用的是BLAS级别的优化。实际工程里还要注意一件事10000 ** (2 * i / d_model)这个写法在i很大的时候会涉及大数指数运算。PyTorch里更优雅的做法是用对数和指数改写import torch import math def build_pe_torch(max_len: int, d_model: int, base: float 10000.0) - torch.Tensor: pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # (max_len, 1) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(base) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe解释一下等价性exp(-(2i/d_model)·ln(10000))等于10000^(-2i/d_model)即1/10000^(2i/d_model)。用这个写法可以避免直接算幂函数时的数值不稳定同时保留在GPU上计算的能力。这里的div_term长度是d_model//2正好对应偶数维和奇数维共享的频率。2.3 PyTorch工程版注册成buffer别让它参与训练最容易踩的坑是位置编码到底算不算模型参数答案是否定的——原版Transformer的三角位置编码是固定不学习的。工程里正确做法是把它注册成buffer而不是Parameterclass SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int 512, base: float 10000.0): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(base) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe) def forward(self, x): # x: (batch, seq_len, d_model) return x self.pe[:x.size(1)]用register_buffer有以下几个好处它不会出现在model.parameters()中也不会被优化器更新它会自动随着模型搬到GPU不需要你手动.cuda()它在模型保存/加载时会被存档加载后直接可用。使用nn.Parameter保存固定位置编码是我见过最多的生产事故之一明明requires_gradFalse设了但在加载checkpoint时依然会遇到key不匹配或者有人误开梯度导致位置编码被训练出一些莫名其妙的模式。2.4 加在哪儿位置编码和Embedding的配合顺序原版Transformer里位置编码是加在Embedding输出之后的但加法之前还有一个细节输入会先乘以sqrt(d_model)。class TransformerEmbeddingWithPE(nn.Module): def __init__(self, vocab_size, d_model, max_len, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_encoding SinusoidalPositionalEncoding(d_model, max_len) self.dropout nn.Dropout(dropout) self.d_model d_model def forward(self, token_ids): x self.token_embedding(token_ids) * math.sqrt(self.d_model) x self.position_encoding(x) return self.dropout(x)为什么乘sqrt(d_model)因为Embedding层的参数初始化通常会把向量数值控制在一个较小的范围内而位置编码的数值范围大致在[-1,1]。如果不做缩放两者相加后Embedding的真实占比会变得很小等于变相削弱词向量信号。原论文这么做主要是从初始化尺度考虑保持相加后向量方差不出现明显塌缩。这个缩放不是硬性要求某些实现也会去掉但如果你照原论文复现建议保留。3. 工程细节batch广播、精度处理、缓存与序列外推3.1 batch维的处理广播比你想象的更省心self.pe的形状是(max_len, d_model)而模型内部张量x通常是(batch, seq_len, d_model)。当执行x self.pe[:x.size(1)]时PyTorch会自动在pe的头部补一个维度变成(1, seq_len, d_model)然后和x做广播。这也是为什么你不用显式写self.pe[None, ...]的原因。但如果你不是用模块类而是在自定义attention里直接使用pe变量我建议养成显式加维度的习惯pe build_pe_torch(max_len, d_model) # (max_len, d_model) # 使用时 x x pe[:seq_len][None, :, :]这样写更明确也方便code review的人一眼看懂。性能上两者没有区别广播只是在张量视图层面添加了一个维度。3.2 float16/混合精度训练下的精度陷阱这是我在半精度推理时踩过的最深的坑之一。问题出在10000^(2i/d_model)这个数值上当i很小时分母接近1angle值等于possin/cos结果在[-1,1]内正常但当i接近d_model/2时分母是一个非常接近10000的大数angle是一个非常小的float。在float16精度下这种极小值的表示精度会下降导致最终pe里某些维度产生锯齿状误差进而影响模型输出的一致性。更危险的是在混合精度训练里如果构建pe时直接用了torch.float16的输入position * div_term可能在某些值域出现inf或nan。我的建议是pe build_pe_torch(max_len, d_model).float() # 强制用float32构建 # forward里使用时再转换 def forward(self, x): pe self.pe[:x.size(1)].to(x.dtype) return x pe如果你用的是bfloat16这个问题相对小一些因为bfloat16的指数范围和float32相同只是尾数精度低。但无论如何先构建好一个float32的pe矩阵、使用时再cast都是最稳妥的做法。这个原则同样适用于RoPE等三角函数型位置编码。3.3 固定编码与可学习编码不能都叫“位置编码”很多新入门的人把BERT的learned positional embedding和Transformer原版的sinusoidal positional encoding混为一谈。两者的工程处理方式完全不同类型是否参数训练方式序列长度外推常见框架Sinusoidal固定否buffer相对较好Transformer原版Learned位置编码是反向传播更新差超长会崩BERT、GPTRoPE否旋转矩阵较好LLaMA系列固定编码的优势是没有额外参数量、不会过拟合特定的训练长度、有一定外推能力。劣势是它不一定能学到最适合任务的位置语义。如果你的任务是固定长度分类、文本匹配这类场景learned编码往往也能取得不错效果但一旦遇到长度分布极不均匀的任务固定编码在稳定性上更让人省心。一个比较务实的经验batch里如果常常出现超长序列而模型是在固定长度(比如512)下训练的固定三角编码的退化是渐进的而learned编码经常直接失效。所以做长文本任务我优先推荐用固定三角函数或RoPE而不是learned。3.4 序列长度超过max_len怎么办pe矩阵预分配了max_len行推理时如果来了比max_len更长的序列直接pe[:x.size(1)]会返回一个超长了实际上切片越界会报错不会静默返回。我的做法是预留足够大的max_len或者在构建时动态扩展def extend_pe(self, new_len: int): if new_len self.pe.size(0): return old_pe self.pe new_pe build_pe_torch(new_len, old_pe.size(1), self.base) new_pe[:old_pe.size(0)] old_pe self.register_buffer(pe, new_pe)三角编码的外推能力是相对好的但不代表无限好。实测中训练长度512的模型推到1024通常还能用推到2048以上效果就会明显衰减。原因是高维低频分量虽然能覆盖长距离但attention训练时从未见过那么远的相对偏移QK内积的数值分布会偏离训练时的范围。3.5 缓存与设备一次算好反复使用位置编码是常量理论上每个batch都重新计算是严重浪费。register_buffer就解决了这个问题它会在第一次访问时随模型参数一起迁移到GPU之后每次前向只是切片和加法开销几乎为零。不过我见过一种写法把pe存储成普通的torch.Tensor类属性而不是buffer然后每次forward里写self.pe self.pe.to(x.device)这既慢又容易在分布式场景下出设备不一致的问题。标准做法就是注册成buffer让PyTorch帮你管理设备的同步。4. 验证位置编码的“相对性”内积实验与热力图4.1 相对位置信息到底怎么证明我在上篇说过三角位置编码的核心优势是“只感知相对位置”但很多人并不清楚这句话怎么验证。最直接的验证方式是计算两个位置编码向量的内积让p和pk是两个位置PE(p)和PE(pk)的内积经过三角函数恒等式化简后结果只和k有关和p无关。这一点非常关键因为自注意力中QK点积本质上就是在做内积。位置编码“只提供相对位置信息”意味着不管当前token在句子的开头还是中间它和距离为k的另一个token之间由位置引入的偏置是一致的。这是固定三角编码结构性的优点而不是后期学出来的。4.2 实测算一算不同偏移下的内积可以先构造pe然后计算相邻、隔两位、隔五位等位置之间的余弦相似度import torch import numpy as np pe build_pe_torch(max_len512, d_model512).numpy() pe_norm pe / (np.linalg.norm(pe, axis-1, keepdimsTrue) 1e-12) for offset in [1, 2, 5, 10, 50, 100, 200]: sim (pe_norm[offset:] * pe_norm[:-offset]).sum(axis-1).mean() print(foffset{offset}, mean_cos_sim{sim:.4f})我实际跑出来的结果大致如下不同d_model下数值略有差异offset平均余弦相似度10.984720.969050.9266100.8708500.63861000.59192000.5787可以看到随着offset变大相似度整体在下降而且这种下降是平滑的、单调的。这就像在说“离我越远的位置位置编码在空间上和我越不相似”这个特性让模型在计算注意力时天然带上了距离衰减的归纳偏置。至于为什么offset到200以后衰减变缓是因为剩下的主要是低频分量在起作用。4.3 热力图更直观地看Toeplitz结构把pe归一化后的余弦相似度矩阵画出来会看到一个明显的带状图案对角线附近颜色深向两边渐次变浅而且这种深浅变化只和“距离对角线多远”有关和具体在哪一行无关。import matplotlib.pyplot as plt sim_matrix pe_norm pe_norm.T plt.figure(figsize(8, 8)) plt.imshow(sim_matrix, cmapviridis, vmin-0.2, vmax1.0) plt.colorbar() plt.title(Cosine similarity of sinusoidal PE) plt.show()这幅图的价值在于——矩阵的每一行平移不变性一目了然。如果使用learned positional embedding画出来也常常呈现带状结构但往往不如固定三角编码干净因为它是从数据里硬学出来的带了些噪声。而这个带状结构在数学上就是Toeplitz矩阵的核心特征。4.4 频率维度分工的进一步实验还可以做一个更细粒度的小实验只取最低频的16个维度和最高频的16个维度分别算偏移相似度。最低频维度组波长很长相似度随offset增大衰减非常缓慢模型可以从中获取“远距离位置关系”最高频维度组波长很短offset差个1函数值就可能剧烈变化适合捕捉前后词的紧邻关系。所以三角编码不是把位置信息均匀铺在所有维度上而是分层编码每层负责一种距离尺度。理解这一点你会明白为什么d_model越大位置编码能感知的位置模式越丰富。这也是为什么在一些长文本任务中有人会尝试对位置编码的维度做降权或调制因为低维和高维承担的角色确实不一样一刀切地使用可能不是最优解。5. 容易踩的五个坑附排查与修复建议5.1 sin/cos分配方式和框架不一致原论文里是偶数维sin、奇数维cos。但某些框架或博客里用cos填偶数维、sin填奇数维也有。单独用这两种约定效果差距很小因为每个频率对都由一个sin和一个cos构成只要频率对内部保持正交关系最终的编码能力差不多。真正的问题在于一致性。如果你要用别人预览好的模型或自定义attention层一定要先确认对方用的是哪种约定。排查方法很简单print出pe[1] - pe[0]看第0维的值是接近sin(1)还是cos(1)就能反推出约定。5.2 维度下标写反导致所有位置编码完全相同这是我在社交媒体上看到求助最多的问题。构造时如果把pe[:, 0::2]写成pe[0::2, :]矩阵形状就错了或者angle矩阵的shape变成(d_model//2, seq_len)而不是(seq_len, d_model//2)填充进去后位置行之间就会出现重复或全零。排查方法是打印pe sinusoidal_positional_encoding(10, 512) print(pe[0, :5]) print(pe[1, :5])如果两个位置第一维的值一样基本就是下标写错了。正常情况不同pos之间的sin/cos值必然不同。5.3 float16或低精度直接构建导致nan/inf前面提到过用half直接计算10000 ** (2i/d_model)这种指数运算在低维部分可能产生inf。这个问题在模型从float32转float16推理时尤其常见因为很多人直接把pe变量也cast成了half。解决办法已经说过始终用float32构建pe再在forward里用.to(x.dtype)转成目标精度。另外如果你用torch.exp(arange * -log(base) / d_model)这种形式构建类似问题基本不会发生所以推荐使用exp写法。5.4 忘记detach或误用Parameter导致位置编码参与训练有人觉得“先初始化成三角编码再让它微调是不是更好”这其实是一个研究问题不是工程默认行为。如果你确实想做可学习位置编码应该显式用nn.Embedding来实现而不是把一个固定buffer偷偷转成Parameter。如果用Parameter又不想训练切忌只设置requires_gradFalse就不管了因为state_dict仍会包含这个key加载更新后的模型权重时容易产生兼容性错误。固定编码一律用register_buffer可学习编码用nn.Embedding这是工程里最清晰的划分。5.5 位置编码加在错误的位置原版结构里位置编码加在embedding输出后、进入Transformer Block之前。如果你是在每个attention层的输出后再加同一个位置编码位置信号会被重复叠加模型会越叠越偏。另一个常见错误是只把位置编码加在Q或K上而V不加这会改变注意力分布中位置信息的含义。排查方法在模型里打印第一个block的输入和embedding层输出确认两者差值就是pe矩阵再检查QK计算时用的x是否已经包含pe。如果你在做RoPE或ALiBi变体那就不是“加”的问题了而是另一个实现分支。6. 从三角函数到RoPE位置编码的现代变体6.1 为什么新模型普遍转向RoPE近两年的很多大模型如LLaMA、GLM等选择的是RoPE旋转位置编码而不是原始的Sinusoidal。RoPE不用把位置向量加到embedding上而是通过旋转矩阵作用于Q和K让QK内积自动携带相对位置信息。它在数学上和三角编码同源——都依赖sin/cos函数——但作用位置和形式完全不同。RoPE的优势是可以对不同的head维度做不同频率的旋转外推性能更好且不占用embedding容量。但它的实现更复杂需要拆维度旋转Q和K并且对每层都要应用。如果你正在手写Transformer玩我建议先用Sinusoidal打通全流程再切到RoPE。6.2 从Sinusoidal平滑迁移到RoPE的几个注意点RoPE作用于Q/K不作用于V和embedding通常将d_model拆成d_model/2组二维旋转频率设计和Sinusoidal类似但分母常按head_dim计算而不是d_model如果已有模型用Sinusoidal且训得不错改RoPE不能只换位置编码模块还需要重新训练。6.3 我的实践建议学习阶段请务必手写一遍Sinusoidal。它像一把钥匙能帮你打开对位置信息建模的理解。生产阶段如果你的目标是长文本或需要外推RoPE、ALiBi这类新方案确实更值得关注但如果只是中等长度的常规任务Sinusoidal也远没到被淘汰的程度尤其是你已经有一套老代码在稳定运行的时候不必为了追新而重构。我到现在手写Transformer时还是会先跑一遍位置编码的热力图确认位置信号符合预期这个习惯帮我避开过好几次静默bug。如果你的任务对位置关系很敏感强烈建议把第4节的验证流程加进你的debug清单。位置编码只有几行代码但它决定了模型能不能“知道自己在哪里”。 手写Transformer位置编码：Sinusoidal实现、工程细节与避坑实战
上篇把三角函数位置编码的公式和直觉拆了个遍这篇就直接动真格的把公式落成代码再把工程里那些文档里不会写清楚、但实战一定会撞上的细节全部摊开讲一遍。我始终觉得位置编码是Transformer里最“骗人”的部件——公式短、实现更短披着几行代码的外衣好像简单到不行。但真正动手写模型、跑训练、做推理部署的时候你会发现它在batch、精度、序列长度、训练和推理一致性上都有不少暗坑。这篇面向两类人一是正在手写Transformer、想把位置编码彻底搞明白的二是已经有了能用代码、但总被某些边界情况坑到的。如果你只是想要一份可以直接抄的Sinusoidal位置编码实现我会先给出代码再把背后的“为什么”一并交代。1. 从公式到张量位置编码要解决的到底是什么形状的问题1.1 先把公式老老实实写对三角函数位置编码的原始形式如下PE(pos, 2i) sin(pos / 10000^(2i/d_model)) PE(pos, 2i1) cos(pos / 10000^(2i/d_model))其中pos是token在序列里的位置i的取值范围是0到d_model/2-12i和2i1分别表示d_model维向量中的偶数维和奇数维。我在很多文章和代码里见过这个公式的简化写法其中最危险的两种错误把分母写成了10000^(i/d_model)这会让频率衰减速度发生明显偏移位置矩阵的语义会变把sin和cos的奇偶分配写反虽然频率没变但如果要加载别人预训练好的编码就和原有模型不一致了。这个公式的本质是为每个token生成一个d_model维的向量这个向量的每个维度都是一个不同频率的三角函数值。第0维和第1维共享同一个频率但相位差π/2第2维和第3维共享另一个频率依此类推。所以这个向量的信息容量并不来自“内部维度之间的大小关系”而是来自不同维度在不同频率尺度上的组合。1.2 波长变化从“分度尺”的角度看我记得上篇强调过一个点这里再补充一个更直观的角度。对于维度2i波长公式是波长_i 2π * 10000^(2i/d_model)当d_model512时不同维度对应的波长如下表维度索引i分母倍率近似波长01≈ 6.286410≈ 62.8128100≈ 6281921000≈ 6283255≈9611≈ 60397这意味着什么低维度的波长很短位置变化一点点函数值就会明显波动相当于一把高精度的“近程尺”高维度的波长很长可以帮助模型感知长距离的全局位置关系。位置编码的本质就是把不同粒度的“尺子”捆绑到同一个向量里让模型在自注意力计算时能同时看到近邻和远程的偏移信息。这其实很像我们看地图时候的经验你既需要看街道级别的细节也需要看城市级别的整体方位。三角编码用一组不同频率的波把这两种信息压缩进了同一个向量表示。1.3 矩阵形态最终要构造出什么位置编码真正需要生成的是一个形状为(seq_len, d_model)的常量矩阵每一行是一个位置编码向量。比如句子长度是10、d_model是512最终就是一个10×512的矩阵。构造时角度矩阵的广播方式很容易写错。我先给一个最直接的构造版本后面再对比优化import numpy as np def sinusoidal_positional_encoding(seq_len: int, d_model: int) - np.ndarray: pe np.zeros((seq_len, d_model)) pos np.arange(seq_len)[:, None] # (seq_len, 1) i np.arange(d_model // 2) # 共有 d_model//2 个频率 denom np.power(10000.0, 2 * i / d_model) # (d_model//2,) angle pos / denom # (seq_len, d_model//2) pe[:, 0::2] np.sin(angle) pe[:, 1::2] np.cos(angle) return pe这里的关键在于广播pos是(seq_len, 1)denom是(d_model//2,)二者相除会得到一个(seq_len, d_model//2)的矩阵然后通过切片填到奇偶维上。如果你把pos写成(1, seq_len)angle的shape就会变成(1, d_model//2)或者报错这是新手最容易栽的地方。2. 手写实现三种写法从“跑通”到“好用”再到“可上线”2.1 逐位置循环版帮助你理解但千万别用于训练理解公式最笨但也最可靠的方法是一个个位置、一个个维度硬算def pos_enc_loop(seq_len, d_model): pe np.zeros((seq_len, d_model)) for pos in range(seq_len): for i in range(d_model // 2): omega 10000 ** (2 * i / d_model) pe[pos, 2 * i] np.sin(pos / omega) pe[pos, 2 * i 1] np.cos(pos / omega) return pe在seq_len512、d_model512时这个嵌套循环大概是13万次sin/cos计算在CPU上跑出来要几秒钟。虽然能跑但放进训练循环里基本不可能接受。这个版本的价值只在于验证你对公式的理解或者用来和向量化版本逐元素对比输出是否一致。2.2 向量化版正式实现的基本形态用广播构造的版本上面已经给出了。这个版本在numpy和PyTorch里跑起来都很快因为它把循环变成了矩阵运算底层调用的是BLAS级别的优化。实际工程里还要注意一件事10000 ** (2 * i / d_model)这个写法在i很大的时候会涉及大数指数运算。PyTorch里更优雅的做法是用对数和指数改写import torch import math def build_pe_torch(max_len: int, d_model: int, base: float 10000.0) - torch.Tensor: pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) # (max_len, 1) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(base) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) return pe解释一下等价性exp(-(2i/d_model)·ln(10000))等于10000^(-2i/d_model)即1/10000^(2i/d_model)。用这个写法可以避免直接算幂函数时的数值不稳定同时保留在GPU上计算的能力。这里的div_term长度是d_model//2正好对应偶数维和奇数维共享的频率。2.3 PyTorch工程版注册成buffer别让它参与训练最容易踩的坑是位置编码到底算不算模型参数答案是否定的——原版Transformer的三角位置编码是固定不学习的。工程里正确做法是把它注册成buffer而不是Parameterclass SinusoidalPositionalEncoding(nn.Module): def __init__(self, d_model: int, max_len: int 512, base: float 10000.0): super().__init__() pe torch.zeros(max_len, d_model) position torch.arange(0, max_len, dtypetorch.float).unsqueeze(1) div_term torch.exp( torch.arange(0, d_model, 2, dtypetorch.float) * (-math.log(base) / d_model) ) pe[:, 0::2] torch.sin(position * div_term) pe[:, 1::2] torch.cos(position * div_term) self.register_buffer(pe, pe) def forward(self, x): # x: (batch, seq_len, d_model) return x self.pe[:x.size(1)]用register_buffer有以下几个好处它不会出现在model.parameters()中也不会被优化器更新它会自动随着模型搬到GPU不需要你手动.cuda()它在模型保存/加载时会被存档加载后直接可用。使用nn.Parameter保存固定位置编码是我见过最多的生产事故之一明明requires_gradFalse设了但在加载checkpoint时依然会遇到key不匹配或者有人误开梯度导致位置编码被训练出一些莫名其妙的模式。2.4 加在哪儿位置编码和Embedding的配合顺序原版Transformer里位置编码是加在Embedding输出之后的但加法之前还有一个细节输入会先乘以sqrt(d_model)。class TransformerEmbeddingWithPE(nn.Module): def __init__(self, vocab_size, d_model, max_len, dropout0.1): super().__init__() self.token_embedding nn.Embedding(vocab_size, d_model) self.position_encoding SinusoidalPositionalEncoding(d_model, max_len) self.dropout nn.Dropout(dropout) self.d_model d_model def forward(self, token_ids): x self.token_embedding(token_ids) * math.sqrt(self.d_model) x self.position_encoding(x) return self.dropout(x)为什么乘sqrt(d_model)因为Embedding层的参数初始化通常会把向量数值控制在一个较小的范围内而位置编码的数值范围大致在[-1,1]。如果不做缩放两者相加后Embedding的真实占比会变得很小等于变相削弱词向量信号。原论文这么做主要是从初始化尺度考虑保持相加后向量方差不出现明显塌缩。这个缩放不是硬性要求某些实现也会去掉但如果你照原论文复现建议保留。3. 工程细节batch广播、精度处理、缓存与序列外推3.1 batch维的处理广播比你想象的更省心self.pe的形状是(max_len, d_model)而模型内部张量x通常是(batch, seq_len, d_model)。当执行x self.pe[:x.size(1)]时PyTorch会自动在pe的头部补一个维度变成(1, seq_len, d_model)然后和x做广播。这也是为什么你不用显式写self.pe[None, ...]的原因。但如果你不是用模块类而是在自定义attention里直接使用pe变量我建议养成显式加维度的习惯pe build_pe_torch(max_len, d_model) # (max_len, d_model) # 使用时 x x pe[:seq_len][None, :, :]这样写更明确也方便code review的人一眼看懂。性能上两者没有区别广播只是在张量视图层面添加了一个维度。3.2 float16/混合精度训练下的精度陷阱这是我在半精度推理时踩过的最深的坑之一。问题出在10000^(2i/d_model)这个数值上当i很小时分母接近1angle值等于possin/cos结果在[-1,1]内正常但当i接近d_model/2时分母是一个非常接近10000的大数angle是一个非常小的float。在float16精度下这种极小值的表示精度会下降导致最终pe里某些维度产生锯齿状误差进而影响模型输出的一致性。更危险的是在混合精度训练里如果构建pe时直接用了torch.float16的输入position * div_term可能在某些值域出现inf或nan。我的建议是pe build_pe_torch(max_len, d_model).float() # 强制用float32构建 # forward里使用时再转换 def forward(self, x): pe self.pe[:x.size(1)].to(x.dtype) return x pe如果你用的是bfloat16这个问题相对小一些因为bfloat16的指数范围和float32相同只是尾数精度低。但无论如何先构建好一个float32的pe矩阵、使用时再cast都是最稳妥的做法。这个原则同样适用于RoPE等三角函数型位置编码。3.3 固定编码与可学习编码不能都叫“位置编码”很多新入门的人把BERT的learned positional embedding和Transformer原版的sinusoidal positional encoding混为一谈。两者的工程处理方式完全不同类型是否参数训练方式序列长度外推常见框架Sinusoidal固定否buffer相对较好Transformer原版Learned位置编码是反向传播更新差超长会崩BERT、GPTRoPE否旋转矩阵较好LLaMA系列固定编码的优势是没有额外参数量、不会过拟合特定的训练长度、有一定外推能力。劣势是它不一定能学到最适合任务的位置语义。如果你的任务是固定长度分类、文本匹配这类场景learned编码往往也能取得不错效果但一旦遇到长度分布极不均匀的任务固定编码在稳定性上更让人省心。一个比较务实的经验batch里如果常常出现超长序列而模型是在固定长度(比如512)下训练的固定三角编码的退化是渐进的而learned编码经常直接失效。所以做长文本任务我优先推荐用固定三角函数或RoPE而不是learned。3.4 序列长度超过max_len怎么办pe矩阵预分配了max_len行推理时如果来了比max_len更长的序列直接pe[:x.size(1)]会返回一个超长了实际上切片越界会报错不会静默返回。我的做法是预留足够大的max_len或者在构建时动态扩展def extend_pe(self, new_len: int): if new_len self.pe.size(0): return old_pe self.pe new_pe build_pe_torch(new_len, old_pe.size(1), self.base) new_pe[:old_pe.size(0)] old_pe self.register_buffer(pe, new_pe)三角编码的外推能力是相对好的但不代表无限好。实测中训练长度512的模型推到1024通常还能用推到2048以上效果就会明显衰减。原因是高维低频分量虽然能覆盖长距离但attention训练时从未见过那么远的相对偏移QK内积的数值分布会偏离训练时的范围。3.5 缓存与设备一次算好反复使用位置编码是常量理论上每个batch都重新计算是严重浪费。register_buffer就解决了这个问题它会在第一次访问时随模型参数一起迁移到GPU之后每次前向只是切片和加法开销几乎为零。不过我见过一种写法把pe存储成普通的torch.Tensor类属性而不是buffer然后每次forward里写self.pe self.pe.to(x.device)这既慢又容易在分布式场景下出设备不一致的问题。标准做法就是注册成buffer让PyTorch帮你管理设备的同步。4. 验证位置编码的“相对性”内积实验与热力图4.1 相对位置信息到底怎么证明我在上篇说过三角位置编码的核心优势是“只感知相对位置”但很多人并不清楚这句话怎么验证。最直接的验证方式是计算两个位置编码向量的内积让p和pk是两个位置PE(p)和PE(pk)的内积经过三角函数恒等式化简后结果只和k有关和p无关。这一点非常关键因为自注意力中QK点积本质上就是在做内积。位置编码“只提供相对位置信息”意味着不管当前token在句子的开头还是中间它和距离为k的另一个token之间由位置引入的偏置是一致的。这是固定三角编码结构性的优点而不是后期学出来的。4.2 实测算一算不同偏移下的内积可以先构造pe然后计算相邻、隔两位、隔五位等位置之间的余弦相似度import torch import numpy as np pe build_pe_torch(max_len512, d_model512).numpy() pe_norm pe / (np.linalg.norm(pe, axis-1, keepdimsTrue) 1e-12) for offset in [1, 2, 5, 10, 50, 100, 200]: sim (pe_norm[offset:] * pe_norm[:-offset]).sum(axis-1).mean() print(foffset{offset}, mean_cos_sim{sim:.4f})我实际跑出来的结果大致如下不同d_model下数值略有差异offset平均余弦相似度10.984720.969050.9266100.8708500.63861000.59192000.5787可以看到随着offset变大相似度整体在下降而且这种下降是平滑的、单调的。这就像在说“离我越远的位置位置编码在空间上和我越不相似”这个特性让模型在计算注意力时天然带上了距离衰减的归纳偏置。至于为什么offset到200以后衰减变缓是因为剩下的主要是低频分量在起作用。4.3 热力图更直观地看Toeplitz结构把pe归一化后的余弦相似度矩阵画出来会看到一个明显的带状图案对角线附近颜色深向两边渐次变浅而且这种深浅变化只和“距离对角线多远”有关和具体在哪一行无关。import matplotlib.pyplot as plt sim_matrix pe_norm pe_norm.T plt.figure(figsize(8, 8)) plt.imshow(sim_matrix, cmapviridis, vmin-0.2, vmax1.0) plt.colorbar() plt.title(Cosine similarity of sinusoidal PE) plt.show()这幅图的价值在于——矩阵的每一行平移不变性一目了然。如果使用learned positional embedding画出来也常常呈现带状结构但往往不如固定三角编码干净因为它是从数据里硬学出来的带了些噪声。而这个带状结构在数学上就是Toeplitz矩阵的核心特征。4.4 频率维度分工的进一步实验还可以做一个更细粒度的小实验只取最低频的16个维度和最高频的16个维度分别算偏移相似度。最低频维度组波长很长相似度随offset增大衰减非常缓慢模型可以从中获取“远距离位置关系”最高频维度组波长很短offset差个1函数值就可能剧烈变化适合捕捉前后词的紧邻关系。所以三角编码不是把位置信息均匀铺在所有维度上而是分层编码每层负责一种距离尺度。理解这一点你会明白为什么d_model越大位置编码能感知的位置模式越丰富。这也是为什么在一些长文本任务中有人会尝试对位置编码的维度做降权或调制因为低维和高维承担的角色确实不一样一刀切地使用可能不是最优解。5. 容易踩的五个坑附排查与修复建议5.1 sin/cos分配方式和框架不一致原论文里是偶数维sin、奇数维cos。但某些框架或博客里用cos填偶数维、sin填奇数维也有。单独用这两种约定效果差距很小因为每个频率对都由一个sin和一个cos构成只要频率对内部保持正交关系最终的编码能力差不多。真正的问题在于一致性。如果你要用别人预览好的模型或自定义attention层一定要先确认对方用的是哪种约定。排查方法很简单print出pe[1] - pe[0]看第0维的值是接近sin(1)还是cos(1)就能反推出约定。5.2 维度下标写反导致所有位置编码完全相同这是我在社交媒体上看到求助最多的问题。构造时如果把pe[:, 0::2]写成pe[0::2, :]矩阵形状就错了或者angle矩阵的shape变成(d_model//2, seq_len)而不是(seq_len, d_model//2)填充进去后位置行之间就会出现重复或全零。排查方法是打印pe sinusoidal_positional_encoding(10, 512) print(pe[0, :5]) print(pe[1, :5])如果两个位置第一维的值一样基本就是下标写错了。正常情况不同pos之间的sin/cos值必然不同。5.3 float16或低精度直接构建导致nan/inf前面提到过用half直接计算10000 ** (2i/d_model)这种指数运算在低维部分可能产生inf。这个问题在模型从float32转float16推理时尤其常见因为很多人直接把pe变量也cast成了half。解决办法已经说过始终用float32构建pe再在forward里用.to(x.dtype)转成目标精度。另外如果你用torch.exp(arange * -log(base) / d_model)这种形式构建类似问题基本不会发生所以推荐使用exp写法。5.4 忘记detach或误用Parameter导致位置编码参与训练有人觉得“先初始化成三角编码再让它微调是不是更好”这其实是一个研究问题不是工程默认行为。如果你确实想做可学习位置编码应该显式用nn.Embedding来实现而不是把一个固定buffer偷偷转成Parameter。如果用Parameter又不想训练切忌只设置requires_gradFalse就不管了因为state_dict仍会包含这个key加载更新后的模型权重时容易产生兼容性错误。固定编码一律用register_buffer可学习编码用nn.Embedding这是工程里最清晰的划分。5.5 位置编码加在错误的位置原版结构里位置编码加在embedding输出后、进入Transformer Block之前。如果你是在每个attention层的输出后再加同一个位置编码位置信号会被重复叠加模型会越叠越偏。另一个常见错误是只把位置编码加在Q或K上而V不加这会改变注意力分布中位置信息的含义。排查方法在模型里打印第一个block的输入和embedding层输出确认两者差值就是pe矩阵再检查QK计算时用的x是否已经包含pe。如果你在做RoPE或ALiBi变体那就不是“加”的问题了而是另一个实现分支。6. 从三角函数到RoPE位置编码的现代变体6.1 为什么新模型普遍转向RoPE近两年的很多大模型如LLaMA、GLM等选择的是RoPE旋转位置编码而不是原始的Sinusoidal。RoPE不用把位置向量加到embedding上而是通过旋转矩阵作用于Q和K让QK内积自动携带相对位置信息。它在数学上和三角编码同源——都依赖sin/cos函数——但作用位置和形式完全不同。RoPE的优势是可以对不同的head维度做不同频率的旋转外推性能更好且不占用embedding容量。但它的实现更复杂需要拆维度旋转Q和K并且对每层都要应用。如果你正在手写Transformer玩我建议先用Sinusoidal打通全流程再切到RoPE。6.2 从Sinusoidal平滑迁移到RoPE的几个注意点RoPE作用于Q/K不作用于V和embedding通常将d_model拆成d_model/2组二维旋转频率设计和Sinusoidal类似但分母常按head_dim计算而不是d_model如果已有模型用Sinusoidal且训得不错改RoPE不能只换位置编码模块还需要重新训练。6.3 我的实践建议学习阶段请务必手写一遍Sinusoidal。它像一把钥匙能帮你打开对位置信息建模的理解。生产阶段如果你的目标是长文本或需要外推RoPE、ALiBi这类新方案确实更值得关注但如果只是中等长度的常规任务Sinusoidal也远没到被淘汰的程度尤其是你已经有一套老代码在稳定运行的时候不必为了追新而重构。我到现在手写Transformer时还是会先跑一遍位置编码的热力图确认位置信号符合预期这个习惯帮我避开过好几次静默bug。如果你的任务对位置关系很敏感强烈建议把第4节的验证流程加进你的debug清单。位置编码只有几行代码但它决定了模型能不能“知道自己在哪里”。