大模型位置编码全解析：从三角函数到RoPE再到YaRN

本文系统介绍大模型中的位置编码技术，包括三角位置编码、相对位置编码、旋转位置编码(RoPE)及其优化版本YaRN。详细阐述各种方法的数学原理、优缺点和实现方式，特别关注位置编码如何影响模型长序列处理能力。RoPE通过旋转变换显式建模相对位置关系，而YaRN通过多种插值策略增强RoPE外推能力，有效解决了大模型上下文窗口扩展问题。

零、基础理论

笔者在梳理位置编码的过程中看到很多数学推导，很多基础理论都忘了，这里把用到的基础理论前置在这里温习一遍：

1.三角函数和差公式

αβαβαβ

另外，对于三角函数ωϕ，周期T=波长λπ，频率

二维旋转变换

对于向量逆时针旋转角度φ得到，那么有如下公式成立：

反过来对于二维向量v左乘上面的三角矩阵，可以认为对向量v执行了逆时针旋转角度φ的操作。

正交矩阵

定义：对于一个的实方阵 A（元素均为实数），如果它的转置矩阵等于其逆矩阵，即，也就是则称 A 为正交矩阵。正交矩阵有几个性质：

两个正交矩阵的乘积还是正交矩阵
任何向量左乘正交矩阵后都不改变模长（这就是RoPE提到的数值稳定性）

上面的旋转矩阵就是一个正交矩阵。此外，还可以证明旋转矩阵的转置矩阵就是反方向旋转矩阵（RoPE就是用到了这个性质）。将平面向量逆时针旋转角的旋转矩阵为：

它的转置矩阵（行和列互换）是：

顺时针旋转角的旋转矩阵 —— 本质就是逆时针旋转，记为。根据三角函数的奇偶性：代入旋转矩阵定义，得到：

对比和的表达式，两者完全相等。

复数的几何意义

复数，这里a和b都是实数，a称为实部，b称为虚部，i为虚数单位，满足。复数z对应复平面上一个向量(a,b)，实部a对应横坐标，虚部b对应纵坐标。

复数也可以写成极坐标的形式： , r 是向量的模长，是向量与实轴的夹角（幅角）。

共轭复数：对任意复数，其共轭复数定义为：，典型性质是复数和共轭复数的乘积是实数

欧拉公式

也就是是一个模长恒为 1 的复数（因为 )，且与实轴的夹角是φ。

这里有一个重要的推论：一个复数乘以相当于将这个复数逆时针旋转角度φ。证明如下：

ϕϕϕϕϕϕϕ

也就是乘积后的复数对应的复平面的向量可以表示成如下形式，结合上面旋转变换的矩阵可以发现正好是原来的向量旋转了角度φ。

一、三角位置编码

论文《Attention Is All You Need》

公式如下，pos表示位置，i表示向量维度索引，作用位置上是位置编码和input层的emb相加：

优势

纵向看，固定位置下不同向量维度的位置编码都对应一条不同尺度的正弦/余弦曲线，且波长成等比数列 (2π到10000 * 2π)，这可以让不同维度分别捕捉局部位置（短波长，相邻 token 差异大）和全局位置（长波长，序列整体位置趋势）信号；
横向看，可以捕捉相对位置信号，以偶数位为例结合三角函数和差公式有如下推导。也就是对于任意固定相对位置偏移k，是的线性函数，模型易学习相对位置注意力。
基于上述两点，可以实现推理长度外推，也就是训短推长。直观理解对于训练max_len=10000，那么推理的时候对于第10001个位置的位置编码，由于三角位置编码的周期重复性，该位置的部分维度的位置编码是之前出现过(短波长维度)或者跟第10000个位置的位置编码差异不大(长波长维度)，此外模型还知道该位置和第10000个位置的相对位置是1，可以用第10000个位置的位置编码通过训练好的线性变换矩阵来学到该位置的位置编码。

待改进点

位置向量和emb加法结合破坏了原始emb的模长，而模长代表了语义强度，模型需要花费额外的力气来解耦它们
虽然论文中提到PE(pos+k)理论上可以表示为PE(pos)的线性组合，但是并不直接建模相对位置，需要模型隐性学习，也就是对相对位置的表征能力较弱。
外推性没有那么强，主要是因为相对位置学习的比较弱，所以也限制了训短推长的效果。

代码实现

复制自tensor2tensor的tensor2tensor/layers/common_attention.py

def get_timing_signal_1d(length, channels, min_timescale=1.0, max_timescale=1.0e4, start_index=0): position = tf.to_float(tf.range(length) + start_index) num_timescales = channels // 2 log_timescale_increment = ( math.log(float(max_timescale) / float(min_timescale)) / tf.maximum(tf.to_float(num_timescales) - 1, 1)) inv_timescales = min_timescale * tf.exp( tf.to_float(tf.range(num_timescales)) * -log_timescale_increment) scaled_time = tf.expand_dims(position, 1) * tf.expand_dims(inv_timescales, 0)# Please note that this slightly differs from the published paper.# See a discussion here: https://github.com/tensorflow/tensor2tensor/pull/177 signal = tf.concat([tf.sin(scaled_time), tf.cos(scaled_time)], axis=1) signal = tf.pad(signal, [[0, 0], [0, tf.mod(channels, 2)]]) signal = tf.reshape(signal, [1, length, channels])return signal

二、相对位置编码

论文《Self-Attention with Relative Position Representations》，这篇论文是相对位置编码的开山之作。

1.公式

改造的是下面两个公式，也就是分别在计算value加权求和、计算attn score即的时候加上了相对位置的向量，注意K和V的的table是不同的，不同head/不同layer之间的position emb的table是共享的。另外这里的位置编码并不是作用在input层的，而是每层transformer在计算attn的时候都会加上。

2 位置变量的表示

这部分介绍上面的具体怎么算的。考虑超过一定距离后，精确的相对位置信息便不再具有实用价值，因此对最大距离加以截断到k<n，这能让模型对训练过程中未见过的序列长度具备泛化能力。对于位置i，左右分别有最大距离-k和k，因此，位置emb size是2k+1。具体公式如下，其中w将索引映射成向量。

3 优缺点

优点就是直接建模了相对位置，缺点就是引入了额外的位置emb，增大的资源利用，另外因为有最大距离阶段机制，对于长距离的依赖不够友好。

4 代码实现

通过调整计算顺序可以节省一些计算复杂度，上面的公式(3)和(4) 都可以使用如下公式分开进行矩阵运算。参考tensor2tensor的实现，相对位置编码的tf代码实现如下：

def _generate_relative_positions_matrix(length, max_relative_position, cache=False):"""生成相对位置矩阵，比如长度为4，cache=False的结果如下： [[ 0, 1, 2, 3], [-1, 0, 1, 2], [-2, -1, 0, 1], [-3, -2, -1, 0]]， 这里cache用于区分训练还是推理阶段，训练阶段需要计算所有token的相对位置； 推理阶段只需要计算「1 个新 token」与「length 个历史 token」的相对位置，而非 length×length 的全量矩阵。 长度为4，cache=True的结果如下： [[-3, -2, -1, 0]] # shape=(1, 4)"""if not cache: range_vec = tf.range(length) range_mat = tf.reshape(tf.tile(range_vec, [length]), [length, length]) distance_mat = range_mat - tf.transpose(range_mat)else: distance_mat = tf.expand_dims(tf.range(-length+1, 1, 1), 0) distance_mat_clipped = tf.clip_by_value(distance_mat, -max_relative_position, max_relative_position)# Shift values to be >= 0. Each integer still uniquely identifies a relative# position difference. final_mat = distance_mat_clipped + max_relative_positionreturn final_matdef _generate_relative_positions_embeddings(length, depth, max_relative_position, name, cache=False):"""根据相对位置矩阵查表生成相对位置emb，size： [1 if cache else length, length, depth].""" with tf.variable_scope(name): relative_positions_matrix = _generate_relative_positions_matrix( length, max_relative_position, cache=cache) vocab_size = max_relative_position * 2 + 1 # Generates embedding for each relative position of dimension depth. embeddings_table = tf.get_variable("embeddings", [vocab_size, depth]) embeddings = tf.gather(embeddings_table, relative_positions_matrix) return embeddingsdef _relative_attention_inner(x, y, z, transpose):""" 这里就是计算的上面公式(5)的分子部分。 Relative position-aware dot-product attention inner calculation. This batches matrix multiply calculations to avoid unnecessary broadcasting. Args: x: Tensor with shape [batch_size, heads, length or 1, length or depth]. y: Tensor with shape [batch_size, heads, length or 1, depth]. z: Tensor with shape [length or 1, length, depth]. transpose: Whether to transpose inner matrices of y and z. Should be true if last dimension of x is depth, not length. Returns: A Tensor with shape [batch_size, heads, length, length or depth]. """ batch_size = tf.shape(x)[0] heads = x.get_shape().as_list()[1] length = tf.shape(x)[2] dz = tf.shape(x)[-1]# xy_matmul is [batch_size, heads, length or 1, length or depth] xy_matmul = tf.matmul(x, y, transpose_b=transpose)# x_t is [length or 1, batch_size, heads, length or depth] x_t = tf.transpose(x, [2, 0, 1, 3])# x_t_r is [length or 1, batch_size * heads, length or depth] x_t_r = tf.reshape(x_t, [length, heads * batch_size, -1])# x_tz_matmul is [length or 1, batch_size * heads, length or depth] x_tz_matmul = tf.matmul(x_t_r, z, transpose_b=transpose)# x_tz_matmul_r is [length or 1, batch_size, heads, length or depth] x_tz_matmul_r = tf.reshape(x_tz_matmul, [length, batch_size, heads, -1])# x_tz_matmul_r_t is [batch_size, heads, length or 1, length or depth] x_tz_matmul_r_t = tf.transpose(x_tz_matmul_r, [1, 2, 0, 3])# 这里原版tensor2tensor代码看似没有乘以(dz**-0.5)，其实是前置对q进行缩放过了，效果是等价的。return (xy_matmul + x_tz_matmul_r_t)*(dz**-0.5) def dot_product_attention_relative(q, k, v, bias=None, max_relative_position, dropout_rate=0.0, image_shapes=None, save_weights_to=None, name=None, make_image_summary=True, cache=False):"""总的调用入口。 Calculate relative position-aware dot-product self-attention. The attention calculation is augmented with learned representations for the relative position between each element in q and each element in k and v. Args: q: a Tensor with shape [batch, heads, length, depth]. k: a Tensor with shape [batch, heads, length, depth]. v: a Tensor with shape [batch, heads, length, depth]. bias: bias Tensor. max_relative_position: an integer specifying the maximum distance between inputs that unique position embeddings should be learned for. dropout_rate: a floating point number. image_shapes: optional tuple of integer scalars. save_weights_to: an optional dictionary to capture attention weights for visualization; the weights tensor will be appended there under a string key created from the variable scope (including name). name: an optional string. make_image_summary: Whether to make an attention image summary. cache: whether use cache mode Returns: A Tensor. Raises: ValueError: if max_relative_position is not > 0. """if not max_relative_position: raise ValueError("Max relative position (%s) should be > 0 when using " "relative self attention." % (max_relative_position)) with tf.variable_scope( name, default_name="dot_product_attention_relative", values=[q, k, v]) as scope: # This calculation only works for self attention. # q, k and v must therefore have the same shape. if not cache: q.get_shape().assert_is_compatible_with(k.get_shape()) q.get_shape().assert_is_compatible_with(v.get_shape()) # Use separate embeddings suitable for keys and values. depth = k.get_shape().as_list()[3] length = common_layers.shape_list(k)[2] relations_keys = _generate_relative_positions_embeddings( length, depth, max_relative_position, "relative_positions_keys", cache=cache) relations_values = _generate_relative_positions_embeddings( length, depth, max_relative_position, "relative_positions_values", cache=cache) # Compute self attention considering the relative position embeddings. logits = _relative_attention_inner(q, k, relations_keys, True) if bias is not None: logits += bias weights = tf.nn.softmax(logits, name="attention_weights") if save_weights_to is not None: save_weights_to[scope.name] = weights save_weights_to[scope.name + "/logits"] = logits weights = tf.nn.dropout(weights, 1.0 - dropout_rate) if not tf.get_variable_scope().reuse and make_image_summary: attention_image_summary(weights, image_shapes) return _relative_attention_inner(weights, v, relations_values, False)

三、旋转位置编码(RoPE)

论文《ROFORMER: ENHANCED TRANSFORMER WITH ROTARY POSITION EMBEDDING》

代码：https://github.com/lucidrains/performer-pytorch

方案

1.1方向

如下公式，显式建模两个token 、的相对位置关系，就是找到将原始emb和位置编码融合的函数、，使得融合后的两个向量的内积(也就是attn的logit)的值可以表示成这样的形式：即两个token原始的emb（这里是指和）以及相对位置m-n的函数。注意这样的g函数是跟m和n绝对位置无关的，只跟m-n有关。

1.2方案

作者找到的函数就是向量旋转，旋转的角度跟位置m成正比，也就是逆时针旋转θ角度，公式和示意图如下（以维度=2为例）：

1.3证明

旋转变换后向量：，其中R为旋转矩阵。

目标：证明，即点积仅与原始向量、相对位置有关。

推导步骤:

step1：写出旋转后向量的点积表达式，根据矩阵转置的乘积性质：，展开得

step2：利用正交矩阵性质化简 RoPE 的旋转矩阵是正交矩阵，满足，且正交矩阵的乘积仍为正交矩阵。同时，旋转矩阵满足乘法加法性：，其逆矩阵满足（反向旋转）。因此：这里的正是相对位置k = n-m对应的旋转矩阵，其表达式为

step3：代入点积公式，展开最终形式将代入点积表达式：

1.4一般形式

上面是以维度=2的形式来说明的，扩展到高纬空间的形式可以两两分组，旋转矩阵如下(不要纠结这个高纬向量整体旋转了多少角度，对比固定某个维度下不同位置下的位置向量才有意义)：

这里作者也借鉴经典三角位置编码的方案，为不同的维度设置了不同的旋转角度θ来实现不同尺度的位置信息感知能力，公式如下：

因为上面的大旋转矩阵太稀疏，直接用会导致算力浪费，因此作者提出了等价的计算方式如下，也就是在每层transformer计算attn之前先对q和k的向量做如下计算来融合相对位置信息：

优点

2.1长距离衰减

如下图所示，随着相对位置变大，两个向量的内积呈现衰减趋势，这一特性与我们的直观认知高度吻合：相对距离较远的 token 对，彼此间的关联程度理应更弱。

2.2灵活适配任意序列长度

因为是显式建模的相对位置关系，因为外推能力更强。

代码

import tensorflow as tfimport numpy as npdef apply_rotary_pos_emb(x, cos, sin): """ 应用 RoPE 旋转：将 x 按照 cos/sin 进行旋转 Args: x: [batch, seq_len, num_heads, head_dim] 或 [batch, seq_len, dim] cos: [1, seq_len, 1, head_dim] sin: [1, seq_len, 1, head_dim] """ # 核心旋转公式: # rope(x) = [x0, x1, x2, x3, ...] * cos + [-x1, x0, -x3, x2, ...] * sin # 即把向量两两分组 (x_2i, x_2i+1)，然后做旋转 head_dim = tf.shape(x)[-1] # 1. 将 x 切分为两半：x1 (偶数位), x2 (奇数位) # 假设 head_dim 是偶数。 # x1: [..., 0, 2, 4...] # x2: [..., 1, 3, 5...] # 为了高效，通常不直接 gather，而是 reshape 后 split # [..., head_dim/2, 2] x_reshaped = tf.reshape(x, tf.concat([tf.shape(x)[:-1], [head_dim // 2, 2]], axis=0)) # x1: [..., head_dim/2], x2: [..., head_dim/2] x1, x2 = tf.unstack(x_reshaped, axis=-1) # 2. 构造旋转后的 x # rotate_x = [-x2, x1] # 拼接回去: [..., head_dim/2, 2] -> reshape -> [..., head_dim] rotate_x = tf.stack([-x2, x1], axis=-1) rotate_x = tf.reshape(rotate_x, tf.shape(x)) # 3. 应用公式: x * cos + rotate_x * sin # 注意：这里的 cos 和 sin 应该已经包含了频率信息 return (x * cos) + (rotate_x * sin)def get_rope_cos_sin(max_seq_len, dim, base=10000): """ 预计算 RoPE 的 cos 和 sin 表 Args: max_seq_len: 最大序列长度 dim: head_dim (注意不是 hidden_size，是每个 head 的维度) base: 频率基数，LLaMA 1/2 是 10000，CodeLlama 是 1000000 Returns: cos, sin: shape [1, max_seq_len, 1, dim] """ # 1. 计算频率 (inv_freq) # theta_i = base^(-2i/d) # 只计算前半部分维度 (dim/2)，因为 RoPE 是两两配对的 inv_freq = 1.0 / (base ** (np.arange(0, dim, 2).astype(np.float32) / dim)) # 2. 生成位置索引 t # t: [0, 1, ..., max_seq_len-1] t = np.arange(max_seq_len, dtype=np.float32) # 3. 计算位置 * 频率 (freqs) # outer product: [seq_len, dim/2] freqs = np.einsum("i,j->ij", t, inv_freq) # 4. 拼接回完整维度 [seq_len, dim] # LLaMA 的实现通常是 [freqs, freqs] 拼接，对应 x 的 [x_0, x_1, x_2...] 排列方式 # 为了配合 apply_rotary_pos_emb 中的 reshape split 逻辑 (两两配对): # 我们需要让 cos 的排列是 [theta_0, theta_0, theta_1, theta_1 ...] # 这样 reshape 后刚好对应 x 的 [x_0, x_1] 共享同一个 theta_0 # 这里有点 trick： # 有两种常见的 RoPE 实现风格： # 风格 A (Google/T5): x 拆分为 [x_0, x_1, ...] 和 [x_d/2, x_d/2+1, ...] # 风格 B (LLaMA/HuggingFace): x 视为 [(x_0, x_1), (x_2, x_3)...] # 上面的 apply_rotary_pos_emb 是按 **风格 B** (两两相邻) 写的。 # 所以 freqs 需要 repeat 一次: [theta_0, theta_1] -> [theta_0, theta_0, theta_1, theta_1] freqs_repeated = np.repeat(freqs, 2, axis=-1) # [seq_len, dim] cos = np.cos(freqs_repeated) sin = np.sin(freqs_repeated) # 5. 增加 Batch 和 Head 维度以便广播 # [1, seq_len, 1, dim] return ( tf.constant(cos[None, :, None, :], dtype=tf.float32), tf.constant(sin[None, :, None, :], dtype=tf.float32) )# --- 使用示例 ---def demo(): batch_size = 2 seq_len = 10 num_heads = 4 head_dim = 64 # 模拟 Q 向量 q = tf.random_normal([batch_size, seq_len, num_heads, head_dim]) # 1. 预计算 cos, sin (通常在类初始化时做一次) cos, sin = get_rope_cos_sin(seq_len, head_dim) # 2. 应用 RoPE # 注意：实际推理时需要根据当前 seq_len 切片 cos/sin q_rope = apply_rotary_pos_emb(q, cos, sin) with tf.Session() as sess: out = sess.run(q_rope) print("Output shape:", out.shape) # 验证模长不变性 (RoPE 是旋转，不改变模长) q_norm = np.linalg.norm(sess.run(q), axis=-1) out_norm = np.linalg.norm(out, axis=-1) print("模长误差:", np.max(np.abs(q_norm - out_norm))) # 应该非常小，接近 1e-6demo()

四、YaRN

论文《YaRN: Efficient Context Window Extension of Large Language Models》

代码： https://github.com/jquesnelle/yarn

这个是RoPE的外推增强版本。

方案

1.0 符号定义

定义缩放因子s如下，其中L是模型训练的上下文长度，L‘是推理的上下文长度,且L’>L：

重新定义RoPE的原始语义emb和位置向量融合的函数f’如下：

语义向量size是D，d表示维度的索引，也就是上面RoPE的i，在RoPE中(这里b=10000)，值得注意的是这里d的最大值是|D|/2-1:

θθ

波长公式如下（本质是角度θ的定义）：

下面会经常出现高频和低频，这里对应不同维度上实行不同尺度(也即是不同维度使用了不同θ)的位置编码。高频是指短波长的三角函数，有助于捕捉临近token的位置关系，因为距离变化一点点位置向量就变化很大；低频是长波长的三角函数，有助于捕捉长距离依赖，位置向量感知不到小距离的变化，距离变化很大时位置向量才会有明细差异。

1.1 线性插值方案PI（Position Interpolation）

直接将超出距离的token缩放回来并且用小数据量finetune下，对应上面的g(m)=m/s，

1.2 “NTK-aware” 插值

上述PI插值的方法对于所有维度无差别缩放(因为是作用在位置m上的)会导致高频信息受损(相当于高频旋转的角度更小了，感知能力会减弱)，进而损伤短上下文的效果。

为解决 RoPE 嵌入插值过程中高频信息丢失的问题，NTK 感知插值方法摒弃了对 RoPE 所有维度进行等比例因子 s 缩放的策略，转而通过 “高频少缩放、低频多缩放”的方式，将插值带来的 “缩放压力” 分散到不同维度上。NTK感知插值定义如下：

注意这里s的指数的分母是|D|-2，这样做的目的是想对于最低频的维度（）希望能正好缩放倍。这里相当于在原来的θ的基础上乘以了项，也就是高频基本不受影响，越低频受影响越大。比如对于d=0，这里就不缩放，对于，这里乘以了实现了缩放。该方法在不微调的情况下超过PI的方法，但是微调后效果比PI方法差，原因是因为这里其实存在微小数量的旋转角度越界。比如|D|=512，根据波长公式，d=|D|/2-2=254，那么该维度波长=58469.56574841605，假设训练长度是50000，推理长度是100000，s=2，那么训练阶段该维度最大旋转角度是，推理的时候最大旋转角度是，也就是推理的时候出现了训练阶段没见过的角度，即越界。

(NTK是Neural Tangent Kernel的缩写，对我们理解这种插值原理关系不大，可以只当成是一个符号就行。)