YaRN：利用温度缩放修正熵变实现长上下文微调
大家好，今天我们来深入探讨一个在长上下文语言模型微调中非常重要的技术——YaRN（Yet Another RoPE extension），它通过温度缩放来修正因 RoPE (Rotary Position Embedding) 外推而导致的熵变，从而提升长序列模型的性能。
1. 长上下文语言模型的挑战
近年来，大型语言模型（LLMs）在各种自然语言处理任务中表现出色。然而，大多数LLMs的训练数据都限制在相对较短的上下文长度（例如，2048或4096个token）。当模型应用于超出训练范围的长序列时，性能往往会显著下降。这主要是因为：
- 位置编码的外推问题： 现有的位置编码方法，如绝对位置编码、相对位置编码，在超出训练长度时，要么失去意义，要么导致性能下降。RoPE作为一种流行的相对位置编码，在一定程度上缓解了这个问题，但外推到远大于训练长度的序列时，仍然面临性能衰退。
- 注意力机制的复杂性： 注意力机制的计算复杂度与序列长度呈平方关系，导致长序列推理的计算成本显著增加。
- 信息丢失： 当序列过长时，模型可能会丢失早期token的信息，导致长程依赖关系难以捕捉。
2. RoPE (Rotary Position Embedding) 的原理与局限性
RoPE是一种相对位置编码，它通过旋转操作将位置信息嵌入到query和key向量中。具体来说，对于一个长度为L的序列，RoPE将位置信息表示为旋转矩阵，将query和key向量进行旋转，使得它们之间的点积能够反映相对位置关系。
RoPE的数学表达:
假设query向量为 q，key向量为 k，位置索引为 m 和 n。RoPE的目标是设计一个函数 f(q, m) 和 f(k, n)，使得：
q^T k 表示没有位置信息
f(q, m)^T f(k, n) 能够编码位置信息 (m - n)RoPE通过以下方式实现：
f(q, m) = R_m q
f(k, n) = R_n k其中 R_m 和 R_n 是旋转矩阵。对于二维向量 (x, y)，旋转操作定义为：
R_θ (x, y) = (x cos θ - y sin θ, x sin θ + y cos θ)对于更高维度的向量，RoPE将向量分成多个二维子向量，并对每个子向量应用不同的旋转角度。旋转角度的计算公式如下：
θ_i = m / (10000^(2i/d)) (对于偶数维度 i)
θ_i = n / (10000^(2i/d)) (对于奇数维度 i)其中 d 是向量的维度。
RoPE的优点：
- 相对位置编码： RoPE直接编码相对位置信息，对绝对位置不敏感，更符合语言建模的需求。
- 旋转不变性： RoPE具有旋转不变性，使得模型能够更好地泛化到不同的序列长度。
- 理论上的外推性： RoPE理论上可以外推到任意长度，但实际应用中性能会下降。
RoPE的局限性：
尽管RoPE具有上述优点，但在外推到远大于训练长度的序列时，仍然存在性能问题。主要原因是：
- 频率拥挤： 当序列长度远大于训练长度时，RoPE的旋转角度会变得非常小，导致不同位置的向量之间的区分度降低，出现频率拥挤现象。
- 熵变： 外推时，RoPE引入的旋转操作可能会改变query和key向量的分布，导致熵值发生变化，从而影响模型的性能。
3. YaRN：温度缩放修正熵变
YaRN的核心思想是通过温度缩放来修正因 RoPE 外推而导致的熵变。 具体来说，YaRN 对 RoPE 的旋转角度进行缩放，从而调整频率，减少频率拥挤，并减小熵变。
YaRN的实现方法：
YaRN 在 RoPE 的旋转角度计算公式中引入一个温度参数 τ：
θ_i = m / (B * 10000^(2i/d)) (对于偶数维度 i)
θ_i = n / (B * 10000^(2i/d)) (对于奇数维度 i)其中 B 是温度参数，它是一个大于1的常数。通过调整 B 的值，可以控制旋转角度的大小，从而影响 RoPE 的频率和熵值。实际上，原始的 RoPE 可以看作是 B=1 的情况。
温度参数 B 的选择：
选择合适的温度参数 B 至关重要。如果 B 过大，会导致旋转角度过小，使得不同位置的向量之间的区分度降低。如果 B 过小，则无法有效缓解频率拥挤和熵变问题。YaRN 的作者通过实验发现，选择合适的 B 值可以显著提升长序列模型的性能。
YaRN 的优点：
- 缓解频率拥挤： 通过调整旋转角度，YaRN 可以缓解频率拥挤问题，提高不同位置向量之间的区分度。
- 修正熵变： YaRN 可以减小 RoPE 外推引起的熵变，使得模型的输出分布更加稳定。
- 易于实现： YaRN 的实现非常简单，只需要在 RoPE 的旋转角度计算公式中引入一个温度参数即可。
4. YaRN 的代码实现 (PyTorch)
import torch
import torch.nn as nn
import math
class RoPE(nn.Module):
def __init__(self, dim, base=10000.0, device=None):
super().__init__()
self.dim = dim
self.base = base
self.inv_freq = 1. / (self.base ** (torch.arange(0, dim, 2).float() / dim))
self.device = device
def forward(self, q, k, seq_len):
t = torch.arange(seq_len, device=self.device).type_as(self.inv_freq)
freqs = torch.outer(t, self.inv_freq)
emb = torch.cat((freqs, freqs), dim=-1)
cos_emb = emb.cos()
sin_emb = emb.sin()
return self._rotate_half(q) * cos_emb + self._rotate_half(q) * sin_emb,
self._rotate_half(k) * cos_emb + self._rotate_half(k) * sin_emb
def _rotate_half(self, x):
x = x.float()
b, n, _, d = x.shape
x = x.reshape(b, n, -1, 2, d // 2)
x1 = x[..., 0, :]
x2 = x[..., 1, :]
x = torch.stack((-x2, x1), dim=-2)
return x.reshape(b, n, -1, d).type_as(x)
class YaRN(nn.Module):
def __init__(self, dim, base=10000.0, yarn_scale = 1.0, device=None):
super().__init__()
self.dim = dim
self.base = base
self.yarn_scale = yarn_scale
self.inv_freq = 1. / ((self.base * self.yarn_scale) ** (torch.arange(0, dim, 2).float() / dim))
self.device = device
def forward(self, q, k, seq_len):
t = torch.arange(seq_len, device=self.device).type_as(self.inv_freq)
freqs = torch.outer(t, self.inv_freq)
emb = torch.cat((freqs, freqs), dim=-1)
cos_emb = emb.cos()
sin_emb = emb.sin()
return self._rotate_half(q) * cos_emb + self._rotate_half(q) * sin_emb,
self._rotate_half(k) * cos_emb + self._rotate_half(k) * sin_emb
def _rotate_half(self, x):
x = x.float()
b, n, _, d = x.shape
x = x.reshape(b, n, -1, 2, d // 2)
x1 = x[..., 0, :]
x2 = x[..., 1, :]
x = torch.stack((-x2, x1), dim=-2)
return x.reshape(b, n, -1, d).type_as(x)
# 示例用法
if __name__ == '__main__':
batch_size = 2
seq_len = 2048 # 长序列长度
dim = 128 # 向量维度
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 创建随机的query和key向量
q = torch.randn(batch_size, seq_len, 1, dim).to(device)
k = torch.randn(batch_size, seq_len, 1, dim).to(device)
# 初始化RoPE和YaRN
rope = RoPE(dim=dim, device=device)
yarn = YaRN(dim=dim, yarn_scale = 2.0, device=device)
# 使用RoPE和YaRN进行位置编码
q_rope, k_rope = rope(q, k, seq_len)
q_yarn, k_yarn = yarn(q, k, seq_len)
print("RoPE q shape:", q_rope.shape)
print("YaRN q shape:", q_yarn.shape)
# 计算相似度矩阵 (示例)
attention_scores_rope = torch.matmul(q_rope.transpose(-2, -1), k_rope)
attention_scores_yarn = torch.matmul(q_yarn.transpose(-2, -1), k_yarn)
print("RoPE Attention Scores Shape:", attention_scores_rope.shape)
print("YaRN Attention Scores Shape:", attention_scores_yarn.shape)代码解释：
- RoPE类: 实现了原始的 RoPE 位置编码。
forward函数接受 query (q)、key (k) 和序列长度 (seq_len) 作为输入，并返回经过 RoPE 编码后的 query 和 key 向量。 - YaRN类: 实现了带有温度缩放的 RoPE 位置编码。 与 RoPE 类相比，YaRN 类在初始化时接受一个额外的参数
yarn_scale，用于控制温度参数 B。在forward函数中，YaRN 使用缩放后的频率计算旋转角度，从而实现温度缩放。 - 示例用法: 展示了如何使用 RoPE 和 YaRN 进行位置编码，以及如何计算注意力分数。
5. 实验结果与分析
YaRN 在多个长序列语言模型任务中取得了显著的性能提升。例如，在长文本摘要、长文本分类等任务中，YaRN 能够显著提高模型的准确率和召回率。
以下是一个示例性的实验结果表格：
| 模型 | 上下文长度 | 指标 (例如：Perplexity) |
|---|---|---|
| RoPE (B=1) | 4096 | 15.2 |
| RoPE (B=1) | 8192 | 22.5 |
| YaRN (B=2) | 8192 | 18.7 |
| YaRN (B=4) | 8192 | 17.5 |
| YaRN (B=8) | 8192 | 18.2 |
从表格中可以看出，当序列长度增加到8192时，RoPE的性能显著下降。而使用YaRN，通过调整温度参数B，可以有效地提升模型的性能。
实验分析：
- 温度参数的影响： 实验结果表明，选择合适的温度参数 B 对于 YaRN 的性能至关重要。过大或过小的 B 值都可能导致性能下降。
- 长序列性能： YaRN 在长序列上的性能明显优于原始的 RoPE，表明 YaRN 能够有效地缓解频率拥挤和熵变问题。
- 任务相关性： 最佳的温度参数 B 可能因任务而异，需要根据具体任务进行调整。
6. YaRN 的进一步讨论与优化
虽然 YaRN 在长上下文语言模型微调中表现出色，但仍有一些可以进一步讨论和优化的方向：
- 自适应温度缩放： 当前的 YaRN 使用固定的温度参数 B，未来可以考虑使用自适应的温度缩放方法，根据序列长度和任务特点动态调整 B 值。
- 与其他位置编码方法结合： 可以将 YaRN 与其他位置编码方法（如 ALiBi, xPos）结合使用，以进一步提升长序列模型的性能。
- 计算效率优化： 可以对 YaRN 的计算过程进行优化，以降低计算成本，提高推理速度。例如，可以预先计算旋转矩阵，并使用高效的矩阵运算库进行计算。
7. 总结：YaRN 的核心价值
YaRN 通过引入温度缩放机制，有效地修正了 RoPE 在长序列外推中产生的熵变问题，从而显著提升了长上下文语言模型的性能。它实现简单，效果显著，是长序列建模领域的一个重要进展。通过调整温度参数，我们可以更好地控制 RoPE 的频率特性，使其更好地适应不同的序列长度和任务需求。