Position Encoding on Jiang Yi(姜祎)'s Homepage

Notes on NoPE

Wed, 24 Dec 2025 15:19:42 +0800

Introduction

【参考文献 1】中系统性对比了 AliBi, RoPE , T5 提出的 T5 bias 以及 Transformer 提出的绝对位置编码 (APE).

作者发现，常用的方法在 length generalization 上表现并不是最好的，而 NoPE 不需要额外的计算开销反而效果最好。

【参考文献 2】进一步探究了 NoPE 长度外推的泛化性。作者有三点发现：

NoPE 相比于 RoPE, 其长度外推泛化能力更强
对于 NoPE 来说，模型会在还没有到达预训练上下文长度之前，表现就出现下降的情况
通过调整 softmax 的温度超参数，我们可以提高 NoPE 的长度外推泛化性能力。

Method

【参考文献 1】对比了不同 position encoding 的相似度，结果如下图所示

实验结果表明，NoPE 与 T5 提出的 T5 bias 最相似。

作者在理论上推导出了 NoPE 的两个性质：

Theorem 1 (Absolute Encoding) Let $x$ be an input sequence of length $T + 1$ to the model. Then, the first layer of $f_θ$ can recover absolute positions $[1, . . . , T + 1]$ in the hidden state $H^{(1)}$. That is, there exist $W_Q, W_K , W_V , W_O, W_1$, and $W_2$ such that the self-attention and feedforward operations in the first layer compute absolute positions and write it to the next hidden state.

Theorem 2 (Relative Encoding) Suppose that the hidden state $H^{(1)}$ contains absolute positional information, as stated in Theorem 1, and assume that it is not overwritten by any subsequent layers. Then, the self-attention in all subsequent layers can implement a relative positional encoding: there exists a parameterization of fθ such that, for $\ell ≥ 2$, the attention dot product between query $q_n$ and key $k_m$ at positions n and m can be expressed as:

$$ \langle q_n, k_m\rangle = f_{cnt}(q, k) + f_{rel}(n − m) $$

where $f_{cnt}$ is a function of their content, and $f_{rel}$ is a function of their relative distance.

【参考文献 2】探究了 softmax 中 normalization factor 对模型表现的影响，作者定义 attention 为

$$ \mathrm{Attn}(q,k,v) = \mathrm{softmax}\left(\lambda q^Tk\right)v $$

实验结果如下图所示

结果说明，通过调整 $\lambda$ 我们可以有效提高 NoPE 的上下文扩展泛化能力

Conclusion

NoPE 说明在 transformer 中我们可以不需要加入位置编码模块，这两篇论文均验证了 NoPE 的有效性。

References

Notes on ALiBi

Wed, 24 Dec 2025 15:10:55 +0800

meta 等提出了 ALiBi, 一个通过 linear biases 来实现位置编码的方法来提高 LLM 在推理阶段的外推能力。

Introduction

当下，有若干种位置编码的方式：

Sinusoidal position embeddings: Transformer 提出的正弦位置编码
RoPE: RoPE 提出的旋转位置编码
T5 bias: T5 提出的相对位置编码

作者通过实验对比了不同的位置编码方法，发现这些方法在推理阶段的外推能力都比较差。

为了解决这个问题，作者提出了 ALiBi (attention with linear biases), 一个几乎不增加计算和内存开销的位置编码方法，来提高 LLM 在推理阶段的外推能力。

Method

作者将外推能力定义为

a model’s ability to continue performing well as the number of input tokens during validation increases beyond the number of tokens on which the the model was trained.

计 $L$ 为训练阶段的上下文长度， $L_{valid}$ 为推理阶段的上下文长度。

作者首先对比了不同的位置编码方法的外推能力，结果如下图所示

结果显示，不同位置编码在推理阶段扩展模型的上下文能力均有限。

Context Length	$L$	$L_{valid}$
Sinusoidal	512	50
	1024	50
RoPE	512	200
	1024	100
T5 bias	512	600
	1024	800
ALiBi	512	-
	1024	-

为了解决这个问题，作者提出了 AliBi, 其表达式为

$$ \mathrm{softmax}(q_iK^T+m\cdot [-(i-1),\dots,-2,-1,0]) $$

其中 $m$ 是一个和 heads 相关的超参数。如果我们有 8 个 heads, 则对应的 scaling 值分别为 $[1/2^1,1/2^2,\dots,1/2^8]$, 如果我们有 16 个 heads, 则我们对 8 个 heads 的结果进行插值，得到 $[1/2^{0.5},1/2^1,\dots,1/2^8]$. ALiBi 的示意图如下所示

ALiBi 通过 bias 惩罚了较远的 query-key pairs, 并且不同的 heads 的惩罚项也不同，从而每个 head 对距离的信息敏感度也不尽相同。

Experiments

ALiBi 在 WikiText-103 上的实验结果如下图所示

Conclusion

作者分析了已有的 position embedding 方法，发现已有的方法在推理阶段均不能有效扩展模型的上下文长度。因此，作者提出了 AliBi, 一个通过 linear bias 来增加位置信息的方法，作者通过实验验证了 ALiBi 的有效性。

References

Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation

Notes on T5

Wed, 24 Dec 2025 15:07:08 +0800

google 在 2020 年发表了 T5 (Text-to-Text Transfer Transformer), 一个使用统一框架来将所有 NLP 任务转换为 text-to-text 格式的迁移学习框架。

Introduction

作者首先回顾了迁移学习和 pre-training, 迁移学习是提高模型在下游任务上表现的一类方法，但是目前还没有一个能够对比各种方法的框架。pre-training 通过在大量数据上进行预训练然后再进行微调，可以有效提高模型在下游任务上的表现。

为了解决这两个问题，作者首先将所有的文本处理任务统一为 “text-to-text” 的形式，这样我们就可以对比不同架构，训练方式以及数据对模型表现的影响

作者提到，本文并不是提供一个新的方法，而是详细对比不同方法，为后续研究提供基础。

Method

Model

在架构上，作者使用了 Transformer 的 encoder-decoder 架构，但是作者做了几点修改

作者提出了 T5 bias, 一个用于替换原始 transformer 绝对位置编码的相对位置编码形式
作者使用了 RMSNorm 替换了 Transformer 中的 LayerNorm.

Data

作者基于 Common Crawl 构建训练数据集，作者对数据进行了清洗，最终数据集大小为 750GB. 作者将这个数据集记为 C4 (Clean Crawled Corpus).

Downstream Tasks

下游任务包括：

text classification: GLUE, SuperGLUE
abstractive summarization: CNN/Daily Mail
question answering: SQuAD
translation: WMT English to German, French and Romanian

Input and Output Format

所有任务的输入输出都被转换为 text-to-text 格式。

Experiments

作者使用的 baseline 模型是一个基于 encoder-decoder 架构的 transformer 模型，其大小以及 configuration 与 BERT base 差不多，最终模型参数量为 220M。

field	num layers	hidden size	MLP hidden size	num heads	head dim	dropout	seq len
value	12	768	3072	12	64	0.1	512

训练时，作者使用了 AdaFactor 优化器，batch size 为 512，训练使用了 34B token. 学习率作者使用了 inverse square root learning schedule: $1/\sqrt{\max(n,k)}$, $n$ 和 $k$ 分别代表当前 step 和 warming up steps.

作者基于 sentencepiece (见 LLM tokenizer) 构建了 Tokenizer, 覆盖 English, German, French 和 Romanian 四种语言。

模型训练的目标函数为 BERT 使用的 “masked language modeling”, 格式如下所示

1
2
3
4
5
6
7
8


# original text
Thank you for inviting me to your party last week.

# inputs
Thank you <X> me to your party <Y> week.

# targets
<X> for inviting <Y> last <Z>

首先，作者对比了不同的架构。作者对比了如下三种 transformer 的变体：

实验结果如下图所示

结果显示，encoder-decoder 架构，denoising 训练目标的效果最好。并且，当 layers 减少一半之后，模型的表现大幅度下降。共享参数的 encoder-decoder 架构表现比 prefix LM 效果更好

接下来作者针对 denoising 的配置进行了测试，实验结果发现 BERT-style 的训练目标效果最好，并且 corruption 比例对模型的表现影响有限，作者使用了 BERT 的配置，即 $15\%$ 的 token 被 masked 掉。对于 span length, 作者通过实验发现不同的 span length 对结果影响不大。因此，作者将 span length 设置为 $3$.

在数据上，作者发现：

对数据进行过滤可以提高模型的表现
使用 in-domain 的数据可以提高模型在该 domain 上的表现，但是问题在于 In-domain 的数据往往比较少
数据量过少时，模型会出现 memorization，也就是过拟合的情况

Overall

作者总结前面的发现，构建了 5 个 size 的模型：

Model	layers	hidden size	FFN hidden size	num heads	head dimension
Small	6	512	2048	8	64
Base	12	768	3072	12	64
Large	24	1024	4096	16	64
3B	24	1024	16384	32	128
11B	24	1024	65536	128	128

Conclusion

作者在本文中提出了 T5, 一个统一所有文本处理任务的迁移学习框架，作者系统性探究了架构，数据以及训练对模型最终表现的影响。最终作者基于 encoder-decoder transformer 架构以及 denoising training objective 训练得到了 T5 系列大语言模型。

References

Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

Base of RoPE Bounds Context Length

Mon, 22 Dec 2025 11:34:42 +0800

Introduction

RoPE 已经成为了大多数 LLM 使用的 position encoding 范式，但是，RoPE 与 LLM long context 之间的关系还没有被探索清楚。在本文中，作者就探究了 base frequency 与 LLM context capability 之间的关系，并给出了一个达到指定上下文长度所需要的 base frequency 的 lower bound.

Method

首先作者回顾了 attention 与 RoPE 的定义，关键就是 RoPE 这部分，如下所示

$$ A_{ij} = (R_{i,\theta}q_i)^T(R_{j,\theta}k_j) = q_i^TR_{i-j,\theta}k_j $$

这里 $\theta$ 就是 base frequency, 作者总结不同模型的 base frequency 配置如下表所示

Model	Llama-7B	Llama2-7B	Llama3-8B	Mistral-7B	Baichuan2-7B
Base frequency	10,000	10,000	500,000	1,000,000	10,000
Context length	2048	4096	8192	32,768	4,096

接下来，作者回顾了 YARN. 其核心思想在于，预训练阶段所有可能的 $\cos(t-s)\theta_i$ 都见过，才能保证模型的 OOD 表现

作者认为 base frequency 的设置应该满足两个条件：

The closer token gets more attention: 当前的 token 应该给邻近的 token 更高的注意力
The similar token gets more attention: 当前的 token 应该给相似的 token 更高的注意力

在 RoPE 中，作者已经给出了 $A_{ij}$ 与相对距离 $|i-j|$ 之间的关系。因此第一个性质已经满足了。

接下来，作者分析了相似 token 的性质，作者定义 token 的相似性如下：

$$ \mathbb{E}_{q,k^*}[q^TR_{m,\theta}k^*] - \mathbb{E}_{q,k}[q^TR_{m,\theta}k] $$

这里 $k^*=q+\epsilon$ 代表了相似的 token, 而 $k$ 是一个随机 token. 作者给出的结论如下

Theorem 假设 $q,k\in\mathbb{R}^d$ 独立同分布，它们的标准差为 $\sigma\in\mathbb{R}$, 则对于 $k^*=q+\epsilon$, $\epsilon$ 是一个随机变量满足 $\mathbb{E}[\epsilon]=0$, 则我们有

$$ \mathbb{E}_{q,k^*}[q^TR_{m,\theta}k^*] - \mathbb{E}_{q,k}[q^TR_{m,\theta}k] = 2\sigma^2\sum_{i=0}^{d/2-1}\cos(m\theta_i) $$

作者定义 $B_{m,\theta}=\sum_{i=0}^{d/2-1}\cos(m\theta_i)$, 作者认为给定 $\theta$, 模型的上下文长度 $L_\theta$ 满足

$$ L_\theta = \sup\{L\mid B_{m,\theta}\geq 0, \forall m\in[L]\} $$

也就是说，base frequency 决定了 LLM 的上下文长度。作者给出了不同的上下文长度对应的 base frequency 如下表所示

Context Len.	1k	2k	4k	8k	16k	32k	64k	128k	256k	512k	1M
Lower Bound	4.3e3	1.6e4	2.7e4	8.4e4	3.1e5	6.4e5	2.1e6	7.8e6	3.6e7	6.4e7	5.1e8

总的来说，远距离衰减性保证了模型会更关注邻近的 token, 而相似 token 保证了模型能够区分出真正有意义的 token.

Experiments

作者首先分析了 base frequency 在 fine-tuning 阶段对模型上下文能力的影响，实验结果如下图所示

从实验结果可以看到，当 base frequency 低于阈值时，模型的表现急剧下降。

作者进一步探讨了 base frequency 对于模型 pre-training 阶段的影响，结果也是一样的，即非常小的 base frequency 会限制模型的 context 能力，结果下图所示（三行分别代表了 base frequency 为 1e2, 1e4 和 1e6 的情况）

可以看到，尽管 perplexity 都差不多，但是使用更大的 base frequency 其长上下文能力明显更好。

作者进一步分析了为什么较小的 base frequency 会影响模型的长上下文能力。作者认为较小的 base frequency 会导致 $B_{m,\theta}$ 接近于 0，从而模型难以区分随机 token 和相似 token, 这样模型只能依赖于邻近 token 进行学习，这样就限制了模型的长上下文能力

作者还进一步对比了提高 base frequency 与 Interpolation 两种做法，实验结果如下表所示

实验结果说明，Interpolation 在上下文超过 30K 之后，其 $B_{m,\theta}\leq0$ 的次数显著增加，表明了其和上下文能力之间的关系。

Conclusion

作者在本文中，探究了 RoPE 中 base frequency 与 LLM 上下文能力之间的关系，发现了提高模型的上下文能力需要关注 RoPE 的 base frequency 超参数，并给出了对应的 lower bound. 作者通过实验验证了这个观点。

References

arxiv

Notes on ViT

Thu, 04 Dec 2025 11:00:44 +0800

Introduction

Transformer 架构在 NLP 领域已经成为了事实上的标准。而在 CV 领域，目前还是 CNN 架构占据了主导。

在本文中，作者就探究是否可以应用 Transformer 架构来完成图像识别的任务。

作者发现，当在小规模数据集上进行训练时，Transformer 架构的表现是不如 CNN 架构的，作者认为原因是 Transformer 架构缺乏如 translation equivariance 和 locality 等 inductive biases. 但是当数据集规模上去之后，作者发现这种 inductive bias 可以通过大规模训练抵消掉。其中，最好的模型在 ImageNet 上达到了 $88.55\%$ 的准确率。

Method

ViT 的架构如下图所示

为了能够处理图片，对于输入的图片 $x\in\mathbb{R}^{H\times W\times C}$, 其中 $(H,W)$ 是图片大小，$C$ 是 channel 的个数。作者将其展开为一个 2D patch 序列, $x_p\in\mathbb{R}^{N\times (P^2\times C)}$, 其中 $(P, P)$ 是 image patch 的大小，$N=HW/P^2$ 是 image patch 的的个数，也是 Transformer 输入 token 的个数，Transformer 的 hidden size 为 $D$, 作者使用了一个 linear layer 来将 image patch 转换为 Transformer 的输入。

与 BERT 一致，作者使用了一个 [class] token 来作为输入 patch 序列的 embedding, 即 $z_0^0=x_{class}$, 其对应的输出 $z_L^0$ 作为该图片的表示。作者还使用了 1D 的 position encoding. 最终，ViT 表达式如下所示

$$ \begin{aligned} z_0 &= [x_{class};x_p^1\mathbf{E};x_p^2\mathbf{E};\cdots;x_p^N\mathbf{E};]+\mathbf{E}_{pos}, &\mathbf{E}\in\mathbb{R}^{(P^2\cdot C)\times D},\mathbf{E}_{pos}\in\mathbb{R}^{(N+1)\times D}\\ z_{\ell}'&=\mathrm{MultiHeadAttention}(\mathrm{LayerNorm}(z_{\ell-1}))+z_{\ell-1},&\ell=1,\dots,L\\ z_{\ell} &= \mathrm{MLP}(\mathrm{LayerNorm}(z_{\ell}'))+z_{\ell}',&\ell=1,\dots,L\\ y&=\mathrm{LayerNorm}(z_L^0) \end{aligned} $$

作者分析认为，相比于 CNN 架构，Transformer 架构缺乏 inductive bias, 这是因为每个 patch 对其周围的 2D 临近信息用的非常少，ViT 架构必须从零开始学习这些位置以及空间关系。

作者还介绍了混合架构，即 patch embedding 由一个 linear layer 替换为一个 CNN 架构

在 finetune 时，作者移除了 pre-trained 的 prediction head, 然后加入了一个 $D\times K$ 的 feed forward layer, 其中 $K$ 是分类的类别数。

Experiments

作者使用的数据集如下所示

Dataset	classes	images
ImageNet	1K	1.3M
ImageNet-21K	21K	14M
JFT	18K	303M

作者在 ImageNet, CIFAR-10/100, Oxford-IIIT Pets, Oxford Flowers-102 这些数据集上进行评测。

模型配置如下表所示

Model	Layers	$D$	$D_{FFN}$	# heads	# params
ViT-Base	12	768	3072	12	86M
ViT-Large	24	1024	4096	16	307M
ViT-Huge	32	1280	5120	16	632M

为了简便，作者在后续实验过程中，对模型名字进行了改写，比如 ViT-L/16 就代表了 ViT-Large model, 其对应的 patch size 为 16.

ViT 的实验结果如下图所示

可以看到相比于 CNN 架构，ViT 所需要的训练时间更少，且效果更好。

为了验证 ViT 模型具有更高的训练效率，作者在不同的数据集上进行的实验，结果如下图所示

可以看到，当数据集比较小时，大部分 ViT 模型表现都比 ResNet 差，但是当数据集规模增加之后，ViT 的表现逐渐超过了 ResNet. 并且，随着数据集规模的增加，大模型的表现也比小模型好。

作者进一步在四个不同大小的随机数据集上进行了验证，结果如下图所示

实验结果显示，inductive bias 在小规模数据集上比较有效，但是当数据集规模上去之后，直接从数据集中学习相关的 pattern 更加高效和有效。

由于不同的模型参数可能不太一致，为了更加公平地对比，作者使用算力来进行对比，结果如下图所示

实验结果显示，ViT 的训练效率比 ResNet 更高，其次，hybrid model 在算力比较小的时候，其表现比 ViT 更好，但是当算力提升之后，两者表现差不多。最后，ViT 随着算力的提升，还有进一步提升的空间。

作者还对 ViT 的 attention 进行了进一步的分析，结果有以下几点：

row-column structure 很明显，同一行同一列的 patches 有比较相似的 embedding
ViT 学习到的 position embedding 对于比较近的 patches 其相似度也更高
对于 attention 来说，某些 heads 在 early layer 就已经开始关注全局信息，并且随着 layer 的提升，attention distance 也在提升，说明模型关注全局信息的能力逐渐增强

作者尝试了对 Transformer 进行 scaling up, 结果如下图所示

实验结果显示，使用更多的 layer 可以有效提高模型的表现，但是当 layer 数大于 16 之后，其提升有限。另一方面，提高宽度对模型表现影响不大。作者发现，降低 patch size 也可以提高模型的表现。

作者还对比了不同类型的 position encoding, 结果发现，使用/不使用 position encoding 对模型表现影响非常大，但是不同的 position encoding 总体来说表现差别不是很大。

Conclusion

作者在本文中提出了 ViT, 一个基于 Transformer 架构的图像识别模型，作者通过在大规模数据集上进行训练验证了 ViT 架构的有效性。

作者发现目前还存在如下挑战：

如何使用自监督的方式进行训练，比如类似 BERT 的 masked patch prediction
将 ViT 应用于其他的视觉任务，比如检测和分割
进一步 scaling ViT

References

openreview

Notes on CoMP

Thu, 04 Dec 2025 10:58:30 +0800

Introduction

作者首先提到已有的 vision foundation model (VFM) 存在两个问题：

不能处理动态分辨率图片输入，尽管我们可以使用 Bilinear interpolation 和 multi-resolution training 等方法，但是模型对于动态分辨率图片输入处理能力仍然不足
VFM 和 LLM 之间存在 representation gap

针对这两个问题，作者提出了 CoMP, 一个 continual pre-training pipeline, CoMP 包含两个模块：

C-RoPE, 一个针对 VFM 的 continual RoPE, 用于帮助 VFM 处理动态分辨率图片输入
Alignment Loss, 用于对齐 VFM 和 LLM 的 representation

Method

CoMP 整体的架构图如下所示

可以看到，CoMP 本质上就是一个多模态大模型，知识我们训练的目标为 VFM

C-RoPE

C-RopE 的核心思想是结合绝对位置编码以及相对位置编码来使得 pre-trained ViT 可以接受任意精度图片输入，对于输入图片 $X_V\in\mathbb{R}^{H\times W}$, 首先经过 patchify 得到 $N=HW/P^2$ 个 patch, 这里 $P$ 是 patch size, 每个 patch 大小为 $x_p\in\mathbb{R}^{N\times (P^2\cdot C)}$, $C$ 是 channels. 然后 VFM 每一个 layer 的计算过程为

$$ \begin{aligned} z_0 &= [x_p^1E;\dots;x_p^NE] + \mathrm{Int}(E_{pos})\\ q_i,k_i,v_i &= \mathrm{Proj}_q(z_i), \mathrm{Proj}_k(z_i), \mathrm{Proj}_v(z_i)\\ y_i &= z_i + \mathrm{Proj}_o(\mathrm{Softmax}\left((Rq_i)^T(Rk_i) / D_v\right)v_i)\\ z_{i+1} &= y_i + \mathrm{FFN}(y_i) \end{aligned} $$

这里 $E\in\mathbb{E}^{P^2\cdot C\times D_v}$, $E_{pos}\in\mathbb{R}^{N\times D_V}$ 分别是 patch embedding 和 learnable position embedding, $\mathrm{Int}(\cdot)$ 是 bilinear interpolation.

相比于 ViT, C-RoPE 做出了两点改动：

使用了 Interpolation 来支持动态分辨率图片输入
使用了 RoPE 来高效利用位置信息

Text-supervised Generative Pre-training

作者还是用了 LLM 的 cross-entropy loss 来进行对齐，text decoding loss 定义为

$$ \mathcal{L}_{dec} = -\frac1T\sum_{i=V+1}^{V+T}\log P(X_i\mid X_{其中 $H_v=\mathcal{F}_{\psi}(\mathcal{V}(X_V))$ 是视觉特征，$T, V$ 分别代表了文本 token 和视觉 token 个数

Vision-language Representation Alignment

text-decoding loss 可以对其视觉特征和文本特征，但是对于使用自监督训练方式的 VLM, 比如 DINOv2, 其预训练目标与 text-decoding loss 之间存在较大 gap, 为了解决这个问题，作者提出了 alignment loss.

具体做法就是先计算出文本和视觉特征：

$$ F_v = \mathrm{Pool}(H_v), F_t = \mathrm{Pool}(P_{\theta}(X_t)) $$

然后作者将 $F_v, F_t$ 映射到语言空间，

$$ C_v = W^TF_v, C_t=W^TF_t $$

这里 $W\in\mathbb{R}^{D_t\times K}$, $D_t, K$ 分别为 LLM 的 hidden size 以及 vocabulary size.

接下来作者使用了 iterative Sinkhorn-Knopp 算法来归一化 $C_t$

$$ p_t=\mathrm{Diag}(u_W)\exp(\frac{C_t}{\epsilon})\mathrm{Diag}(v) $$

这里 $u_W\in\mathbb{R}^K$ 是 words 的 prior marginal distribution,$v\in\mathbb{R}^B$ 是 renormalization vector

最终 alignment loss 定义为

$$ \mathcal{L}_{align} = -p_t\log p_v $$

这里 $p_v=\mathrm{softmax}(C_v)$, 作者对 LLM 使用了 stop gradient 操作，仅训练 VFM

Training

CoMP 训练分为三个阶段：

Stage 1: warming up, 仅训练 adapter
Stage 2: 所有模型参数参与训练，使用了 RoPE 2D 和高精度图片输入
Stage 3: instruction tuning, 使用 RoPE 2D 和动态分辨率图片输入

训练的损失为

$$ \mathcal{L} = \begin{cases} \mathcal{L}_{dec} + \alpha \mathcal{L}_{align}, &\text{Stage 1 and Stage 2}\\ \mathcal{L}_{dec}, &\text{Stage 3} \end{cases} $$

Experiments

作者基于 CLIP, SigLIP, SigLip-2, DINOv2, AIMv2 等 VFM 进行了实验，主要结果如下图所示

作者还对 training recipe 进行的消融实验，结果如下，可以看到，RoPE-2D 对模型表现提升最大，数据量和动态分辨率图片输入次之

作者还对 C-RoPE 以及 alignment loss 进行的消融实验，结果如下图所示

Conclusion

作者在本文中提出了 CoMP，一个 continual pre-training VFM 的方法，通过 C-RoPE 以及 alignment loss, 作者提高了 VFM 在 MLLM 中的表现

References

arxiv

Notes on YaRN

Thu, 03 Jul 2025 14:40:49 +0800

Introduction

YaRN (Yet Another RoPE extentionN method) 时23年9月EleutherAI等提出来的一个扩展LLM上下文长度的方法，后来被Qwen系列模型所应用。

Preliminary

作者首先回顾了一下RoPE, 具体内容请参见上一篇blog。并使用了 $f_{W}(x_m, m, \theta_d)$ 来表示RoPE：

$$ f_{W}(x_m, m, \theta_{d}) = \Theta_{\theta, m}^d W x_m $$

其中 $\Theta_{\theta, m}^d\in\mathbb{R}^{d\times d}$ 是多维旋转矩阵, $\theta_d=[\theta_{0,d},\dots,\theta_{(d-2)/2,d}]\in\mathbb{R}^{d/2}$, $\theta_{i,d}=\theta^{2i/d}$, $\theta>0$ 是一个超参数，RoPE中设置为 $\theta=10000$, $W\in\{W_q,W_k\}\subset\mathbb{R}^{d\times d}$ 是对应的query/key projection layer的权重矩阵， $x\in\mathbb{R}^{d}$ 是输入的hidden states.

接下来，作者定义了两个新的变量：

scaling factor 假设预训练的上下文长度为 $L$, 扩展的上下文长度为 $L'>L$, 则我们定义scaling factor 为

$$ s = \frac{L'}{L} $$

易知 $s>1$.

wavelength 我们将 $\lambda_d$ 定义为 $i$-th 维的RoPE embedding对应的 wavelength：

$$ \lambda_{i,d} = \frac{2\pi}{\theta_{i,d}} = 2\pi\theta^{2i/d}, \ i=0,\dots,(d-2)/2 $$

wavelength描述了对于第$i$ 个维度，RoPE旋转一周 ($2\pi$) 所需要的上下文长度。

基于 $f_{W}(x_m, m, \theta_d)$, 作者统一了已有的扩展上下文长度的方法，作者将不同的扩展方法使用一个通用函数 $f_W(x_m,g(m), h(\theta_d))$ 来表示，这里 $g(m)$ 和 $h(\theta_d)$ 分别代表了不同的长度外推方法所使用的变换。

Position Interpolation

Position Interpolation (PI) 的核心思想在于，我们可以通过Interpolation，将超过预训练长度的文本给压缩到当前最大长度，以此来避免RoPE外推产生的问题，其对应的公式为

$$ f'_W(x_m,m,\theta_d) = f_W(x_m, \frac{mL}{L'}, \theta_d) $$

其中 $L'>L$ 为我们扩展之后的上下文长度， $L$ 为我们预训练的上下文长度。使用通用函数表示的话，我们有

$$ g(m) = \frac{m}{s},\quad h(\theta_d) = \theta_d $$

NTK-aware Interpolation

PI的问题是，并没有考虑不同维度的wavelength。基于NTK理论，DNN当输入维度比较低，且embedding缺少高频内容时，模型就会很难学习到高频信息。对应到RoPE里面，输入的token position id是低位信息（1维），而输出的RoPE是一个 $d$ 维的复杂向量。因此，当输入token非常相似却距离非常近时，RoPE就会丢失高频的细节信息

因此，为了解决这个问题，作者对不同的维度使用了不同的缩放策略：维度比较小时，其缩放的更多，维度比较大是，其缩放的更少。

基于这个策略，作者提出了NTK-aware interpolation，其定义如下：

$$ g(m) = m, \quad h(\theta_{i,d}) = \theta'^{-2i/d}, i=0,\dots,(d-2)/2 $$

其中

$$ \theta' = \theta\cdot s^{\frac{d}{d-2}} $$

实际中，这种方法会产生out-of-bound的值，因此最终结果会比PI要差一点，为了解决这个问题，一般会使用比 $s$ 更大的scaling factor.

上式的推导基于一个简单的假设：我们希望最后一个维度的wavelength在scaling之后，是线性变化的，即 $\theta'_{(d-2)/2,d}=s\theta_{(d-2)/2,d}$, 求解之后，我们就得到了上面的定义

PI 和NTK-aware interpolation的问题在于，我们对不同的维度的处理都是一样的。这类不在乎wavelength的方法被称为blind interpolation methods, 接下来要介绍的就是基于wavelength的方法，即target interpolation methods.

NTK-by-parts Interpolation

与NTK-aware interpolation， NTK-by-parts interpolation基于wavelength来考虑不同维度上所做的变换。

对于低维度，其 $\theta_{i,d}$ 非常大，因此旋转一周所需要的上下文长度也非常大。实际上就导致某些维度的embedding并不是均匀分布的，（比如说只有 $0\sim\pi$ 这个区间的embedding），这个时候，模型就只能访问到绝对位置信息，而访问不到相对位置信息。另外，当我们对所有的维度都进行scale的时候，所有的token都会与彼此更加靠近，这损害了模型对于局部信息的获取能力。基于这些认知，作者基于wavelength，对不同的维度分别进行处理：

如果wavelength远小于上下文长度 $L$，则我们不做任何处理
如果wavelength等于或者大于上下文长度 $L$，则我们使用NTK-aware interpolation 进行处理
对于中间的其他维度，我们进行了一个trade off

作者定义了一个ratio $r$ 来描述原始上下文长度 $L$ 和 wavelength 之间的关系

$$ r(i,d) = \frac{L}{\lambda_{i,d}} = \frac{L}{2\pi\theta'^{2i/d}} $$

基于这个ratio，我们可以定义上面的三种处理方式对应的权重

$$ \gamma(r) = \begin{cases} 0, &\text{if } r < \alpha\\ 1, &\text{if } r > \beta\\ \frac{r-\alpha}{\beta-\alpha}, &\text{otherwise} \end{cases} $$

其中， $\beta>\alpha>0$ 是超参数， $r<\alpha$, $r<\beta$ 分别代表了上面的第1种，第2种情况。

最后，NTK-by-parts interpolation的定义如下

$$ g(m) = m, \quad h(\theta_i) = \left(1-\gamma(r(i,d)\right)\frac{\theta_i}{s} + \gamma(r(i,d))\theta_i $$

作者通过实验发现，在LLaMA上，$\alpha=1$ 和 $\beta=32$ 是一个比较好的选择

Dynamic NTK Interpolation

在实际中，一个经常遇到的场景就是sequence length会从1逐步上升到最大上下文长度，比如说inference的时候。对于这种情况，我们有两种解决方法：

在整个inference周期中，RoPE的scaling factor都设置为 $s=K'/L$, 其中$L'$ 是扩展后的上下文长度
在每次foward的过程汇总，都更新sclaing factor $s=\max(1, \ell'/L)$, 这里 $\ell'$ 是当前sequence的长度

作者发现，方案1在sequence长度小于 $L$的时候性能会下降，并且当上下文长度超过 $L'$ 时，性能下降的更快。但是，方案2可以让模型的性能下降曲线更平缓。因此，作者将这种inference-time方法称为 Dynamic Scaling method, 当其与NTK-aware方法结合时，就得到了 Dynamic NTK interpolation

作者通过实验发现，Dynamic NTK interpolation在$L'=L$ 时，效果非常好

YaRN

在YaRN中，作者针对Dynamic NTK interpolation做了进一步改进，也就是在计算attention softmax时，加入了一个温度参数 $t>0$, 这样attention的计算就变成了

$$ \mathrm{Attn}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{t\sqrt{d}}\right)V $$

作者发现，通过这种scaling的方式，YaRN可以在不改变代码的前提下，更改attention的机制。并且，其不增加训练和推理的cost

作者将YaRN定义为结合了NTK-by-parts interpolation和上述scaling技巧的方法

对于LLaMA，作者推荐使用如下参数：

$$ \sqrt{\frac{1}{t}} = 0.1\ln(s) + 1. $$

实验结果如下

作者发现：

对于合适的 $t$, 扩展上下文之后，perplexity会变的更好
最好的$t$ 对于不同的位置和样本提升都是一样的

Evaluation

实验结果如下

Extension Method	Trained Tokens	Context Window			Evaluation Context Window Size
			2048	4096	6144	8192	10240
PI (s = 2)	1B	8k	3.92	3.51	3.51	3.34	8.07
NTK ($\theta$ = 20k)	1B	8k	4.20	3.75	3.74	3.59	6.24
YaRN (s = 2)	400M	8k	3.91	3.50	3.51	3.35	6.04

可以看到，YaRN使用的数据更少，并且当模型扩展到10240的时候，其表现下降的最慢，这说明了YaRN在扩展上下文长度时的有效性

原始RoPE，dynamic-PI和dynamic-YaRN的对比

可以看到，RoPE的上下文扩展能力很差，Dynamic-YaRN的表现最好。

Code

YaRN的实现在HuggingFace/src/transformers/modeling_rope_utils.py 里的 _compute_yarn_parameters 函数里，其返回 inv_freq 以及 attention_factor 两个量，前者代表了 $\theta_d$, 后者代表 $t\sqrt{d}$.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46



def get_mscale(scale):
 if scale <= 1:
 return 1.0
 return 0.1 * math.log(scale) + 1.0

def find_correction_dim(num_rotations, dim, base, max_position_embeddings):
 """Inverse dimension formula to find the dimension based on the number of rotations"""
 return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (2 * math.log(base))

def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings):
 """Find dimension range bounds based on rotations"""
 low = math.floor(find_correction_dim(low_rot, dim, base, max_position_embeddings))
 high = math.ceil(find_correction_dim(high_rot, dim, base, max_position_embeddings))
 return max(low, 0), min(high, dim - 1)

def linear_ramp_factor(min, max, dim):
 if min == max:
 max += 0.001 # Prevent singularity

 linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min)
 ramp_func = torch.clamp(linear_func, 0, 1)
 return ramp_func


class YaRNRotaryEmbedding(nn.Module):
 def __init__(self, config):
 beta_fast = config.get("beta_fast") or 32
 beta_slow = config.get("beta_slow") or 1
 dim = config["head_dim"]
 factor = config.max_position_embeddings / original_max_position_embeddings
 pos_freqs = base ** (torch.arange(0, dim, 2) / dim)
 inv_freq_extrapolation = 1.0 / pos_freqs
 inv_freq_interpolation = 1.0 / (factor * pos_freqs)

 low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings)

 # Get n-dimensional rotational scaling corrected for extrapolation
 inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2)
 inv_freq = (inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + inv_freq_extrapolation * inv_freq_extrapolation_factor)

 attention_factor = get_mscale(factor)

 def forward(self, x, position_ids):
 ...
 return cos, sin

实际的transformers代码中，Qwen使用的还是默认的RoPE，在inference时如果我们需要扩展上下文，可以通过修改config的形式：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18


from transformers import pipeline

model_name_or_path = "Qwen/Qwen3-8B"

generator = pipeline(
 "text-generation",
 model_name_or_path,
 torch_dtype="auto",
 device_map="auto",
 model_kwargs={
 "max_position_embeddings": 131072,
 "rope_scaling": {
 "rope_type": "yarn",
 "factor": 4.0,
 "original_max_position_embeddings": 32768,
 },
 }
)

Conclusion

在本文中，作者首先构建了一个统一的表征不同上下文长度扩展的形式，接下来作者分析了不同上下文长度扩展的不足，并提出了YaRN这种上下文长度扩展方式，结果发现，YaRN不仅在短上下文长度下面表现很好，当上下文长度扩展之后，其表现依然非常优秀。

References

Notes on Position encoding

Mon, 19 May 2025 10:46:39 +0800

本文前半部分参考参考文献1，推荐大家看博客原文。

Position encoding总结

在上一篇blog 中, 我们介绍了 Attention 的两个性质，也就是在不加 position encoding 的情况下，Attention 对于 query 是 permutation equivariant 的，对于 key 和 value 是 permutation invariant 的。

但是“我爱你”和“你爱我”这两句话所表示的含义应该是不一样的，我们将这两句话作为 key 和 value 的时候，我们发现模型的输出是一致的，这显然是不能接受的。因此，我们就需要加入 position encoding，让模型学习到语序信息，从而明白不同的语序有不同的含义。

下面是测试代码（来自参考文献1）

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27


import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModel

model_id = "meta-llama/Llama-3.2-1B"
tok = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id)

text = "The dog chased another dog"
tokens = tok(text, return_tensors="pt")["input_ids"]
embeddings = model.embed_tokens(tokens)
hdim = embeddings.shape[-1]

W_q = nn.Linear(hdim, hdim, bias=False)
W_k = nn.Linear(hdim, hdim, bias=False)
W_v = nn.Linear(hdim, hdim, bias=False)
mha = nn.MultiheadAttention(embed_dim=hdim, num_heads=4, batch_first=True)

with torch.no_grad():
 for param in mha.parameters():
 nn.init.normal_(param, std=0.1) # Initialize weights to be non-negligible

output, _ = mha(W_q(embeddings), W_k(embeddings), W_v(embeddings))

dog1_out = output[0, 2]
dog2_out = output[0, 5]
print(f"Dog output identical?: {torch.allclose(dog1_out, dog2_out, atol=1e-6)}") #True

Position encoding 可以分为绝对位置编码 (absolute position encoding, APE)，相对位置编码 (relative position encoding, RPE) 以及可学习的位置编码。可学习位置编码主要是 BERT 类的模型在使用，其训练成本比较高，本文不做讨论。绝对位置编码是原始 transformer 里提出的编码模式，现在的大多数基于 transformer 模型使用的都是相对位置编码。

本文中，我们先介绍位置编码应该具有的性质，然后我们分别介绍绝对位置编码和相对位置编码，我们将着重关注苏剑林老师提出来的 RoPE。

位置编码

在介绍位置编码之前，我们首先应该关注位置编码的性质，位置编码的目标是为输入的 token embedding 增加位置信息，那么理想的位置编码应该是怎么样的呢？

我们这里直接引用参考文献1 中给定的性质：

性质 1: token sequence 中每个位置的位置编码都是唯一的。这个很好理解，如果不唯一的话，那么根据前面推导的性质，这两个位置的 attention 输出就完全一致了
性质 2: 线性相关性。也就是说，如果我们知道了位置 $p$ 处的位置编码，那么理想情况下，我们应该能比较简单地得到 $p+k$ 处的位置编码，理想情况下，我们应该有 $PE(p+k)=W_kPE(p)$.
性质 3: 泛化到长上下文中去。我们希望位置编码不仅在 8K 的上下文起作用，还希望位置编码能够泛化到 32K 的上下文
性质 4: 生成模式是固定的。固定的模式有助于模型更好地学习位置相关的信息
性质 5: 可以扩展到多维。我们希望位置编码可以从文本扩展到图片再到视频，也就是从 $1D$ 到 $nD$.

绝对位置编码

绝对位置编码依照其名称，其思想就是为每个位置的 token 分配一个固定的位置信息，也就是对于输入的 hidden states $\bm{x}=[\bm{x}_1,\dots,\ \bm{x}_m]\in\mathbb{R}^{m\times d}$, 我们有

$$ \bm{x}_i' = \bm{x}_i + p_i, i=1,\dots, m $$

这里，$p_i\in\mathbb{R}^d$. 我们的 attention 就变成了

$$ \mathrm{Attn}(X) = \mathrm{softmax}\left(\frac{(Q+P)(K+P)^T}{\sqrt{d}}\right)V\in\mathbb{R}^{m\times d} $$

这里

$$ P = [p_1,\dots,p_m]\in\mathbb{R}^{m\times d}， Q= W_QX\in\mathbb{R}^{m\times d}, K=W_KX, V=W_VX\in\mathbb{R}^{n\times d} $$

整数位置编码

一个最简单的想法就是我们使用正整数来标记 token 所在的位置，也就是

$$ PE(i) = [i, \dots, i]=i\mathbf{1}_{d\times 1}\in\mathbb{R}^d,\ i=1,\dots,m $$

可以看到，这个简单的设计满足性质 1，性质 2，性质 3，性质 4.

但是，注意到 attention 的输入 $X$ 通常是经过 Layer Normalization 处理过后的，因此其按列符合正态分布，并且均值和方差一般较小。当我们加上整数位置编码之后，其 token 本身的信息就会被污染，也就是信噪比非常低。一个解决方法就是我们对 $PE(i)$ 进行 normalization，即

$$ PE(i)' = \frac{1}{m}PE(i) = \frac{i}{m}\mathbf{1}_{d\times 1} $$

现在所有的位置编码的值都比较小，但是我们发现新的位置编码不满足性质 2 了，这是因为现在位置编码还和 sequence 长度有关，我们从位置 $p$ 到位置 $p+k$ 不仅取决于 $k$ 还取决于 sequence 长度 $m$

二进制位置编码

既然整数位置编码的主要问题是对输入影响太大，我们能否找一个不影响输入的整数位置编码方式呢？参考文献1 提出了二进制位置编码，因为每个 token 是 $d$ 维的，因此我们可以使用 $d$ 位二进制来表示 $i$. 比如说，当 $d=3$, $m=4$ 时，我们的位置编码分别为

$$ PE(0) =p_{(000)_2} = [0, 0, 0],\ PE(1) =p_{(001)_2}= [0, 0, 1],\ PE(2) =p_{(010)_2} = [0, 1, 0],\ PE(3) =p_{(011)_2} = [0, 1, 1] $$

现在，我们二进制位置编码满足性质 1，性质 2. 对于性质 3，由于 $d$ 位二进制的表示范围为 $[0, 2^d-1]$，因此其泛化性受到 $d$ 的影响。

参考文献1 画出了不同位置的值的变化情况。我们这里也模仿绘制出类似的曲线图

我们发现，二进制位置编码高位，也就是 $PE(i)_{d}$ 的变化很慢，而低位，也就是 $PE(i)_{0}$ 变化很快，

二进制位置编码解决了整数位置编码的信噪比过低和线性相关性。但是其问题是其对不同位置的 token embedding 产生的影响是不一样的。比如位置 1 和位置 2 的相同的 token embedding 之间的区别是：

$$ (\bm{x}_2 + PE(2)) - (\bm{x}_1 + PE(1)) = (\bm{x}_2-\bm{x}_1)+ [0, 1, -1] $$

一般来说, $\bm{x}_2-\bm{x}_1$ 比较小，因此使用二进制位置编码的问题是输入位置的微小变化（增加一个 token 或减少一个 token）都会对最终结果产生巨大影响。因此，我们需要想办法解决这个问题。

Sinusoidal

前面提到二进制位置编码的问题是相邻 token 之间变化太大，不够光滑。因此我们想要增加一个光滑性质，也就是说我们希望：

位置编码值在 $[-1, 1]$ 之间，防止对 token embedding 产生影响
相邻 token 的位置编码尽可能相近，即 $|PE(k+p)-PE(p)| \leq \delta |k|$, 其中 $\delta>0$ 是一个比较小的数。
与二进制一样，高位的变化比较慢，低位的变化比较快

一个想法就是利用三角函数 $\sin$ 或者 $\cos$，三角函数满足前两个性质，对于第三个性质，我们可以通过控制频率来满足。这样我们得到的位置编码就具有如下形式：

$$ PE(p, i) = \sin\left(\frac{p}{\theta^{i/d}}\right) $$

其中 $\theta$ 是我们的超参数。

我们现在来推导一下上面位置编码的线性相关性：

$$ PE(p+k) = \sin\left(\frac{p+k}{\theta^{i/d}}\right)=PE(p)\cos\left(\frac{k}{\theta^{i/d}}\right) + \cos\left(\frac{p}{\theta^{i/d}}\right)\sin\left(\frac{k}{\theta^{i/d}}\right) $$

我们发现，$\sin$ 位置编码不满足线性相关性。但是出现的 $\cos$ 给了我们启发，也就是我们可以同时使用 $\sin$ 和 $\cos$ 来完成位置编码，这也是原始 transformer 里提出来的 Sinusoidal 位置编码，其形式为：

$$ \begin{aligned} PE(p, 2i) &= \sin\left(\frac{p}{\theta^{2i/d}}\right)\\ PE(p, 2i+1) &= \cos\left(\frac{p}{\theta^{2i/d}}\right) \end{aligned} $$

现在，记 $\omega_i=1/\theta^{2i/d}$, 我们再推导一下线性相关性，就得到：

$$ \begin{aligned} \begin{bmatrix} PE(p+k, 2i)\\ PE(p+k, 2i+1)\\ \end{bmatrix}&=\begin{bmatrix} \sin \omega_i(p+k)\\ \cos \omega_i(p+k) \end{bmatrix}\\ &=\begin{bmatrix} \sin \omega_i(\omega_ip)\cos(\omega_ik)+\cos w_i(\omega_ip)\sin(\omega_ik)\\ \cos \omega_i(\omega_ip)\cos(\omega_ik)-\sin w_i(\omega_ip)\sin(\omega_ik) \end{bmatrix}\\ &= \begin{bmatrix} \cos(\omega_ik)& \sin(\omega_ik)\\ -\sin(\omega_ik)& \cos(\omega_ik) \end{bmatrix}\begin{bmatrix} PE(p, 2i)\\ PE(p, 2i+1)\\ \end{bmatrix} \end{aligned} $$

也就是说，Sinusoidal 位置编码满足线性相关性。对于 Sinusoidal 位置编码我们也可以进行可视化：

相对位置编码

前面介绍了绝对位置编码，每个位置的位置编码是固定的。但是绝对位置编码的问题是，模型比较难以学习相对位置关系。

举个例子，我们提到上下文时，通常会使用“上一节”，“上一章”这些表示相对位置关系的词。

因此，我们希望让模型学习相对位置关系而不是绝对位置关系，因为相对关系更符合我们的认知。

RoPE

RoPE 由苏剑林老师提出，最早应用于 LLaMA 架构（没有确认），后续被大多数模型所采用。

之前的 PE 大多数关注于加性位置编码，也就是假设位置编码的形式为 $\bm{x}+\bm{p}$, 基于这种假设，已有的工作基本都集中于优化下面的 Q 和 K 的内积

$$ \langle f_q(\bm{x}_q, m), f_k(\bm{x}_k, n) \rangle $$

这里 $f_q(\bm{x}_q, m)=W_q(\bm{x}_q+\bm{p}_m)$, $f_k(\bm{x}_k, n)=W_k(\bm{x}_k+ \bm{p}_n)$.

而 RoPE 里面，作者使用了一个不同的假设： 假设内积应该仅包含两者的相对信息，也就是

$$ \langle f_q(\bm{x}_q, m), f_q(\bm{x}_k, n)\rangle := g(\bm{x}_q,\bm{x}_k, m-n) $$

这里的 $f$ 和 $g$ 都是未知函数。我们的目标就是从这个公式中推导出一个合适的位置编码出来。

不失一般性，我们可以假设

$$ f_q(\bm{x}_m,0) = \bm{x}_q,\quad f_q(\bm{x}_n, 0) = \bm{x}_k $$

这个假设代表初始条件下，我们不对输入做任何改变，也就是不增加位置信息。

2D 推导

与 RoPE 一样，我们直接使用复平面来进行推导。

我们假设 $d=2$, 注意到二维平面上的每个点都可以表示为如下形式

$$ \bm{z} = (x,y) = re^{i\theta} $$

其中 ($\mathrm{atan2}$ 定义参考维基百科)

$$ r = \|\bm{z}\|_2 = \sqrt{x^2+y^2}\in\mathbb{R},\quad \theta = \mathrm{atan2}(y, x)\in\mathbb{R}, $$

现在，对于三个向量 $f_q(\bm{x}_q, m)$, $f_q(\bm{x}_k, n)$, $g(\bm{x}_q,\bm{x}_k, m-n)$ 我们可以写出其极坐标形式：

$$ \begin{aligned} f_q(\bm{x}_q,m) &:= r_q(\bm{x}_q,m)e^{i\theta_q(\bm{x}_q,m)}\\ f_k(\bm{x}_k, n) &:= r_k(\bm{x}_k, n)e^{i\theta_k(\bm{x}_k, n)}\\ g(\bm{x}_q,\bm{x}_k, m-n) &:= r_g(\bm{x}_q,\bm{x}_k, m-n)e^{i\theta_g(\bm{x}_q,\bm{x}_k, m-n)} \end{aligned} $$

我们计算内积并比较同类项得到：

$$ \begin{aligned} r_g(\bm{x}_q,\bm{x}_k, m-n) &:= r_q(\bm{x}_q,m)r_k(\bm{x}_k, n)\\ \theta_g(\bm{x}_q,\bm{x}_k, m-n) &:= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, n) \end{aligned}\tag{3} $$

我们接下来分别推导 $r_g(\bm{x}_q,\bm{x}_k, m-n)$ 和 $\theta_g(\bm{x}_q,\bm{x}_k, m-n)$ 的形式

$r_g(\bm{x}_q,\bm{x}_k, m-n)$

我们令 $m=n=0$ 可以得到初始条件

$$ r_g(\bm{x}_q,\bm{x}_k, 0) = r_q(\bm{x}_q,0)r_k(\bm{x}_k, 0)=\|\bm{q}\|_2\|\bm{k}\|_2 $$

我们再令 $n=0$,得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_q(\bm{x}_q,m)r_k(\bm{x}_k, 0)=r_q(\bm{x}_q,m)\|\bm{k}\|_2=\frac{r_g(\bm{x}_q,\bm{x}_k, m-n)}{r_k(\bm{x}_k, n)}\|\bm{k}\|_2 $$

这里最后一个等式带入了原始等式 (3)，注意到左侧与 $n$ 无关，因此右侧我们选取 $n=1$, 得到

$$ r_g(\bm{x}_q,\bm{x}_k, m) = \frac{r_g(\bm{x}_q,\bm{x}_k, m-1)}{r_k(\bm{x}_k, 1)}\|\bm{k}\|_2 =\cdots= r_g(\bm{x}_q,\bm{x}_k, 0)\left(\frac{\|\bm{k}\|_2 }{r_k(\bm{x}_k, 1)}\right)^{m+1} $$

令 $m=0$ 我们有

$$ r_k(\bm{x}_k, 1) = \|\bm{k}\|_2. $$

因此我们最终的表达式为：

$$ r_g(\bm{x}_q,\bm{x}_k, m) = r_g(\bm{x}_q,\bm{x}_k, 0) = \|\bm{q}\|_2\|\bm{k}\|_2. $$

并且，通过分别设置 $m=0$ 以及 $n=0$ 我们还可以得到

$$ r_q(\bm{x}_q,m) = \|\bm{q}\|_2,\quad r_k(\bm{x}_k, n) = \|\bm{k}\|_2 $$

$\theta_g(\bm{x}_q,\bm{x}_k, m-n)$

令 $m=n=0$, 我们得到初始条件

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_q(\bm{x}_q,0)-\theta_k(\bm{x}_k, 0)=\theta_q-\theta_k $$

令 $n=1$, 我们有

$$ \begin{aligned} \theta_g(\bm{x}_q,\bm{x}_k, m-1) &= \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 1)\\ &=\theta_g(\bm{x}_q,\bm{x}_k, m-n) + \theta_k(\bm{x}_k, n)-\theta_k(\bm{x}_k, 1) \end{aligned} $$

这里我们带入了公式 (3)，注意到公式左边与 $n$ 无关，因此在公式右侧我们令 $n=0$, 得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m-1) = \theta_g(\bm{x}_q,\bm{x}_k, m)+ \theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1) $$

分别令 $m=1,2,\dots$ 并相加这些等式，我们得到

$$ \theta_g(\bm{x}_q,\bm{x}_k, 0) = \theta_g(\bm{x}_q,\bm{x}_k, m) + m(\theta_k(\bm{x}_k, 0)-\theta_k(\bm{x}_k, 1)) $$

即

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+(\theta_q-\theta_k)\tag{4} $$

注意到

$$ \theta_g(\bm{x}_q,\bm{x}_k, m) = \theta_q(\bm{x}_q,m)-\theta_k(\bm{x}_k, 0)=\theta_q(\bm{x}_q,m)-\theta_k $$

带入上式我们就得到

$$ \theta_q(\bm{x}_q,m) = m(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_q $$

在 (4) 式中再令 $m=m-n$，并带入 $\theta_q(\bm{x}_q,m)$ 就有

$$ \theta_k(\bm{x}_k,n) = n(\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0))+\theta_k $$

汇总

最后，我们将以上结果放在一起，就得到

$$ f_q(\bm{x}_q,m) = \bm{q}e^{im\theta}, f_v(\bm{x}_k,m) = \bm{k}e^{in\theta} $$

这里 $\theta=\theta_k(\bm{x}_k, 1)-\theta_k(\bm{x}_k, 0)$ 是一个超参数，用于控制频率。

我们记

$$ R_{\theta,m} = \begin{bmatrix} \cos m\theta & -\sin m\theta\\ \sin m\theta & \cos m\theta \end{bmatrix} $$

则我们有：

$$ f_q(\bm{x}_q,m) = R_{\theta,m}\bm{q}, f_v(\bm{x}_k,m) = R_{\theta,n}\bm{k}. $$

并且

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} \tag{5} $$

多维扩展

上面是 2D 的情况，对于多维情况，苏剑林老师通过将两个元素组对，然后分别进行处理，得到了多维的情形：

$$ R_{\theta,m}^d = \begin{bmatrix} R_{\theta_1,m} & & && & \\ & & R_{\theta_2,m} & & & \\ &&&& \ddots & \\ &&&& & R_{\theta_{d/2},m} \end{bmatrix}\in\mathbb{R}^{d\times d} $$

我们可以验证公式 (5) 仍然是成立的。

RoPE 的远程衰减性质

我们接下来看一下结果与相对距离 $m-n$ 之间的关系，注意到

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle = \bm{q}^TR_{\theta,m-n}\bm{k} = \sum_{i=1}^{d/2} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i $$

这里 $\bm{q}_i=[q_{2i},q_{2i+1}]^T$, $\bm{k}_i=[k_{2i},k_{2i+1}]^T$ 分别是对应的 pair，我们考虑其中一个分量，不妨假设 $\|\bm{q}\|_2=\|\bm{k}\|_2=1$, 我们有

$$ \begin{aligned} \bm{q}_i^TR_{\theta, m-n}\bm{k}_i &\leq \bm{q}_i^TR_{\theta, m-n}\bm{q}_i\\ &= \bm{q}_i^T\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right)\bm{q}_i\\ &\leq \lambda_{\max}\left(\frac{R_{\theta, m-n}+R_{\theta, m-n}^T}{2}\right) \\ &= \cos (m-n)\theta_i \end{aligned} $$

其中，第一个不等式是因为两个向量相等时其内积最大，第二个不等式是由于二次型最大值为矩阵的特征值。

这样我们就有

$$ \langle f_q(\bm{x}_q,m), f_v(\bm{x}_k,m)\rangle \leq \sum_{i=1}^{d/2}\cos (m-n)\theta_i. $$

我们可以简单画出对应的曲线：

这里针对不同的配置，结果也会略微不同，具体分析可以参加知乎回答 RoPE的远距离衰减

RoPE 代码实现与理解

Naive 实现

我们接下来看一下如何实现 RoPE

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0 & -\sin m\theta_0 & &&\cdots &\cdots &\cdots \\ \sin m\theta_0 & \cos m\theta_0 & &&\cdots &\cdots &\cdots\\ & & \cos m\theta_1 & -\sin m\theta_1 & \cdots &\cdots &\cdots\\ & & \sin m\theta_1 & \cos m\theta_1 & \cdots &\cdots &\cdots \\ &&&& \ddots &\vdots &\vdots\\ &&&& & \cos m\theta_{d/2} & -\sin m\theta_{d/2}\\ &&&& & \sin m\theta_{d/2} & \cos m\theta_{d/2} \end{bmatrix}\begin{bmatrix} x_1\\ x_2\\ \vdots \\ x_d \end{bmatrix} $$

在实现的时候，我们一般根据 $\sin$ 和 $\cos$ 进行分组，也就是

$$ \Theta_m\bm{x}=\begin{bmatrix} \cos m\theta_0\\ \cos m\theta_0\\ \vdots\\ \cos m\theta_{d/2}\\ \cos m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} x1\\ x2\\ \vdots\\ x_{d-1}\\ x_d\\ \end{bmatrix} + \begin{bmatrix} \sin m\theta_0\\ \sin m\theta_0\\ \vdots\\ \sin m\theta_{d/2}\\ \sin m\theta_{d/2}\\ \end{bmatrix}\odot \begin{bmatrix} -\ x_2\\ x_1\\ \vdots\\ -x_d\\ x_{d-1}\\ \end{bmatrix} $$

我们通常按照奇偶 index 来分别计算，然后通过重排序来得到最终的结果，实现代码如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


def apply_rotary_pos_emb(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
 x_even = x[..., ::2] # (seq_len, d_k_half)
 x_odd = x[..., 1::2] # (seq_len, d_k_half)
 odds = cos * x_even - sin * x_odd # (...,seq_len, d_k_half)
 evens = sin * x_even + cos * x_odd # (...,seq_len, d_k_half)
 stacked = torch.stack((odds, evens), -2) # (...,seq_len, 2, d_k_half)
 stacked_trans = rearrange(
 stacked, "... seq_len double d_k_half -> ... seq_len d_k_half double"
 ) # (...,seq_len, d_k_half, 2)
 out = rearrange(
 stacked_trans, "... seq_len d_k_half double -> ... seq_len (d_k_half double)"
 ) # (..., seq_len, d_k)
 return out

LLaMA 实现

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26


def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
 t = torch.arange(end, device=freqs.device) # type: ignore
 freqs = torch.outer(t, freqs).float() # type: ignore
 freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
 return freqs_cis

def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
 ndim = x.ndim
 assert 0 <= 1 < ndim
 assert freqs_cis.shape == (x.shape[1], x.shape[-1])
 shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
 return freqs_cis.view(*shape)


def apply_rotary_emb(
 xq: torch.Tensor,
 xk: torch.Tensor,
 freqs_cis: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
 xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
 xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
 freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
 xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
 xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
 return xq_out.type_as(xq), xk_out.type_as(xk)

在 LLaMA 中，我们首先还是计算 $\theta_i$, 然后在计算的过程中，我们将 $(x_i,x_{i+1})$ 视作一个复数，然后乘以 $\exp(im\theta)$, 最后再取实部得到最终的结果

通用实现

实际上，naive 版本的实现与现在大语言模型所采用的实现并不一致，我们先看一下现有的大语言模型的 RoPE 实现，这里我们将 LLaMA的transformer代码放在下面，

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41


def rotate_half(x):
 """Rotates half the hidden dims of the input."""
 x1 = x[..., : x.shape[-1] // 2]
 x2 = x[..., x.shape[-1] // 2 :]
 return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
 cos = cos.unsqueeze(unsqueeze_dim)
 sin = sin.unsqueeze(unsqueeze_dim)
 q_embed = (q * cos) + (rotate_half(q) * sin)
 k_embed = (k * cos) + (rotate_half(k) * sin)
 return q_embed, k_embed

class LlamaRotaryEmbedding(torch.nn.Module):
 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
 super().__init__()

 self.dim = dim
 self.max_position_embeddings = max_position_embeddings
 self.base = base
 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
 self.register_buffer("inv_freq", inv_freq)

 self.max_seq_len_cached = seq_len
 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)

 freqs = torch.einsum("i,j->ij", t, self.inv_freq)
 # Different from paper, but it uses a different permutation in order to obtain the same calculation
 emb = torch.cat((freqs, freqs), dim=-1)
 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)
 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)

 def forward(self, x, seq_len=None):
 # x: [bs, num_attention_heads, seq_len, head_size]
 if seq_len > self.max_seq_len_cached:
 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

 return (
 self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
 )

我们将上述代码翻译成公式，现在我们的 $\Theta$ 变成了 (对应 emb = torch.cat((freqs, freqs), dim=-1))

$$ \Theta = [\theta_0,\dots,\theta_{d/2},\theta_0,\dots,\theta_{d/2}]^T $$

实际上 $\sin$ 部分对应的向量现在变成了

$$ [-x_{d/2+1}, -x_{d/2+2}, \dots, -x_{d}, x_1, \dots, x_{d/2}]^T $$

我们带回到原始公式，可以得到对应的 RoPE 操作变成了

$$ R_{\theta,m}^d=\begin{bmatrix} \cos m\theta_0 & & & -\sin m\theta_0 & \cdots &\cdots &\cdots \\ & & \cos m\theta_1 & &-\sin m\theta_1 & \cdots &\cdots \\ & & & \cos m\theta_2 & &-\sin m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ & & & &\cos m\theta_{d/2 - 1} & & -\sin m\theta_{d/2 - 1} \\ \sin m\theta_0 && & \cos m\theta_0 &&\cdots &\cdots \\ & & \sin m\theta_1 & &\cos m\theta_1 & \cdots &\cdots \\ & & & \sin m\theta_2 & &\cos m\theta_2 & \cdots \\ \vdots&\vdots&\vdots&\vdots& \vdots &\vdots &\vdots\\ &&&& \sin m\theta_{d/2 - 1}& & \cos m\theta_{d/2 - 1} \end{bmatrix} $$

这列每一行的 $\cos$ 和 $\sin$ 都相差了 $d/2$ 列.

因此，这里的区别在于，原始 RoPE 计算的 pair 为 $(x_{i}, x_{i+1})$, 而 LLaMA 里的 RoPE 计算的 pair 为 $(x_{i}, x_{i+d/2})$. transformers library 使用这种方式，可以减少计算量，提高整体的计算效率。

为了适应使用 LLaMA 中实现的 RoPE 的，Huggingface 对权重进行了转换，使得基于原始 RoPE 实现的模型也可以获得加速.

假设 $d=8$，原始 RoPE 的 pair 为 [(q_0, q_1), (q_2, q_3), (q_4, q_5), (q_6, q_7)], 新的 pair 为 [(q_0, q_4), (q_1, q_5), (q_2, q_6), (q_3, q_7)]. 我们希望对 index 进行 remap，我们发现一个满足条件的 permutation 为 [0, 2, 4, 6, 1, 3, 5, 7], 也就是 q_0->q_0, q_2->q_1, …, q_7->q_7.

但是，如果我们在推理时这样做，就会降低整体速度，因此 Huggingface 的做法是改变 $W_Q$ 和 $W_K$ 的权重，具体来说，就是 $\Pi q=(\Pi W_Q)x$，左边是在线转换，右侧离线转换。转换好 $W_Q$ 之后，正常计算就可以了。具体代码为

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13


# permute for sliced rotary
def permute(w, n_heads=n_heads, dim1=dim, dim2=dim):
 return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2)

state_dict = {
 f"model.layers.{layer_i}.self_attn.q_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wq.weight"]
 ),
 f"model.layers.{layer_i}.self_attn.k_proj.weight": permute(
 loaded[f"layers.{layer_i}.attention.wk.weight"]
 ),
 ...
}

结论

本文中，我们回顾了位置编码，包括绝对位置编码和相对位置编码，我们着重介绍了 RoPE 的原理，推导以及代码实现。

Position Encoding on Jiang Yi(姜祎)'s Homepage

Notes on NoPE

Introduction

Method

Conclusion

References

Notes on ALiBi

Introduction

Method

Experiments

Conclusion

References

Notes on T5

Introduction

Method

Model

Data

Downstream Tasks

Input and Output Format

Experiments

Overall

Conclusion

References

Base of RoPE Bounds Context Length

Introduction

Method

Experiments

Conclusion

References

Notes on ViT

Introduction

Method

Experiments

Conclusion

References

Notes on CoMP

Introduction

Method

C-RoPE

Text-supervised Generative Pre-training

Vision-language Representation Alignment

Training

Experiments

Conclusion

References

Notes on YaRN

Introduction

Preliminary

Unified Perspective on Related Work

Position Interpolation

NTK-aware Interpolation

NTK-by-parts Interpolation

Dynamic NTK Interpolation

YaRN

Evaluation

Code

Conclusion

References

Notes on Position encoding

Position encoding总结

位置编码

绝对位置编码

整数位置编码

二进制位置编码

Sinusoidal

相对位置编码

RoPE

2D 推导

$r_g(\bm{x}_q,\bm{x}_k, m-n)$

$\theta_g(\bm{x}_q,\bm{x}_k, m-n)$

汇总

多维扩展

RoPE 的远程衰减性质

RoPE 代码实现与理解

Naive 实现

LLaMA 实现

通用实现

结论

参考文献