从位置走到信息：AI位置编码笔记

位置编码 RoPE ALiBi Transformer 深度学习
作者: 醉月思 发布日期: 2024-08-12 分类: 学习笔记

之前用绝对位置编码，生成到1000+token时效果明显变差。

位置编码的本质问题

Transformer的核心思想是并行计算，时间维度被直接摊平了。注意力机制只知道"这个词和那个词关系近"，却不知道"这个词在前面、那个词在后面"。

最早的解决方案也很简单：给每个位置加一个偏移。最原始的绝对位置编码就是这样：

def get_positional_encoding(max_seq_len, d_model):
    """
    经典的sinusoidal位置编码
    max_seq_len: 最大序列长度
    d_model: 模型维度
    """
    position = np.arange(max_seq_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * -(np.log(10000.0) / d_model))

    pe = np.zeros(max_seq_len, d_model)
    pe[:, 0::2] = np.sin(position * div_term)
    pe[:, 1::2] = np.cos(position * div_term)

    return pe

这段代码在最早的Transformer论文里就能找到，理论上很美：不同频率的sin/cos组合能编码相对位置，模型可以通过学习来推断"位置3和位置5的距离"。

但实际用起来问题不少。训练时max_seq_len设为512，推理时来了个长度1024的句子怎么办？截断还是padding？截断会丢失信息，padding会增加计算量。

更尴尬的是，这种编码方式对模型来说像是黑盒。我试过可视化位置编码的注意力模式，发现模型根本没学会"相对位置"这个概念，只是死记硬背了每个位置的pattern。

RoPE的旋转思路

RoPE（Rotary Positional Encoding）的出现解决了很多实际问题。它的核心想法是把位置编码嵌入到attention的计算过程里，而不是简单加在输入上。

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class RoPE(nn.Module):
    """
    旋转位置编码实现
    """
    def __init__(self, dim, max_seq_len=2048, base=10000):
        super().__init__()
        self.dim = dim
        self.max_seq_len = max_seq_len
        self.base = base

        # 预计算旋转矩阵
        self._build_rotation_matrix()

    def _build_rotation_matrix(self):
        """构建旋转矩阵"""
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim))
        position = torch.arange(self.max_seq_len).float()

        freqs = torch.outer(position, inv_freq)
        self.register_buffer('freqs', torch.polar(torch.ones_like(freqs), freqs))

    def forward(self, x, seq_len):
        """
        x: [batch_size, seq_len, num_heads, head_dim]
        """
        # 只取需要的长度
        freqs = self.freqs[:seq_len]

        # 应用旋转
        x_complex = torch.view_as_complex(x.float())
        rotated = x_complex * freqs[:, None, None, :]

        return torch.view_as_real(rotated).type_as(x)

我第一次用RoPE是在一个长文本生成的项目里。之前用绝对位置编码，生成到1000+token时效果明显变差。换成RoPE后，这个问题基本解决了。

RoPE的优势在于相对位置是显式的：两个位置的编码相除（复数域里就是相减的相位）就能得到相对距离。模型不需要再学习"如何从绝对位置推断相对位置"，这个信息直接给了它。

但也别神话RoPE。我试过在很短的序列上用它，效果反而不如简单加法编码。太复杂的机制在简单问题上可能过度工程化。

ALiBi的线性思路

ALiBi（Attention with Linear Biases）走了另一条路：不在输入上做手脚，直接修改attention score。

class ALiBi(nn.Module):
    """
    ALiBi位置编码实现
    通过在attention score上添加线性偏移来实现位置感知
    """
    def __init__(self, num_heads, max_seq_len=2048, slope_scale=0.25):
        super().__init__()
        self.num_heads = num_heads
        self.max_seq_len = max_seq_len
        self.slope_scale = slope_scale

        # 为每个head分配不同的slope
        self._build_slopes()

    def _build_slopes(self):
        """
        为不同的head分配不同的slope
        这可以让不同head关注不同尺度的相对位置
        """
        slopes = torch.pow(2, -torch.arange(self.num_heads).float() * self.slope_scale)
        self.register_buffer('slopes', slopes)

    def forward(self, attention_score, seq_len):
        """
        attention_score: [batch_size, num_heads, seq_len, seq_len]
        """
        # 构建相对位置矩阵
        positions = torch.arange(seq_len).unsqueeze(1) - torch.arange(seq_len).unsqueeze(0)

        # 应用ALiBi偏移
        alibi_bias = self.slopes[:, None, None] * positions[None, :, :, :].abs()

        # 添加到attention score
        masked_score = attention_score - alibi_bias.to(attention_score.device)

        return masked_score

ALiBi有个很实用的特性：外推性好。训练时只用长度512的序列，推理时直接用到1024甚至更长，效果下降很小。

我在一个文档摘要任务上验证过这一点。训练数据平均长度只有300多字，但实际应用时经常遇到800+字的长文档。用ALiBi可以直接外推，而RoPE在超出训练长度时衰减很快。

但ALiBi也不是银弹。它对相对位置的编码是线性的，这意味着模型可能难以捕捉非线性的位置关系。而且每个head都分不同的slope，这个参数设置需要仔细调优，否则某些head会失去位置敏感度。

实践中的选择与取舍

我整理了一个简单的对比表格，都是实际踩坑后的总结：

方案训练长度推理长度调试难度适用场景
绝对位置编码512≈512简单短文本分类、固定格式
RoPE10241.5-2倍中等长文本生成、对话
ALiBi5122-3倍较复杂文档摘要、长文本理解
无位置编码不适用不适用最简单非序列任务

真实项目里的选择往往不是"哪个最好"，而是"哪个能解决当前问题且不会被技术债拖死"。

我现在的做法是：默认用RoPE，它的社区支持最好，出问题容易找资料。只有明确需要强外推能力时才考虑ALiBi。绝对位置编码就留着老项目用，不主动选了新用。

踩过的几个坑

位置编码看似简单，实际调起来坑不少。记录几个典型的：

坑1：长度对齐问题

有一次推理时用了RoPE，但没注意到max_seq_len参数不一致。训练时设为2048，推理脚本里写成了1024。结果注意力矩阵被截断了，模型根本看不到后面的token。

# 错误示例
rope = RoPE(dim=64, max_seq_len=2048)  # 训练时

# 推理时错误设置
rope_inference = RoPE(dim=64, max_seq_len=1024)  # 完全不同的编码！

解决方法是确保所有地方都用统一的max_seq_len，或者直接在推理时取min(需要的长度, 训练长度)。

坑2：head_dim不对齐

RoPE和ALiBi都要求head_dim是偶数，而且通常需要和编码维度匹配。我见过有人直接把head_dim=64传给RoPE，但实际配置里用的是head_dim=65，导致奇数维度处理错误。

# 正确做法：确保head_dim是偶数
class Config:
    d_model = 768
    num_heads = 12
    head_dim = d_model // num_heads  # = 64，偶数，没问题

坑3：缓存问题

KV cache是生成任务的常用优化，但位置编码的缓存更新容易出错。我试过一个bug：生成到第50步时，attention用的位置编码还是第1步的缓存，导致后面的token都基于错误的位置信息。

# 错误示例
def generate_with_cache(model, input_ids, max_new_tokens):
    cache = None
    for i in range(max_new_tokens):
        # 错误：position没更新
        outputs = model(input_ids, past_key_values=cache)
        cache = outputs.past_key_values
        # ...

正确做法是每次都传入正确的position或者使用支持位置编码的缓存版本。

坑4：混合精度下的数值问题

RoPE涉及复数运算，在FP16模式下容易出现数值不稳定。特别是当位置很大时，sin/cos的值可能超出FP16表示范围。

# 解决方案：关键计算保持FP32
with torch.autocast(device_type='cuda', dtype=torch.float16):
    # 前向计算
    x = self.qkv(x)

    # RoPE计算强制使用FP32
    with torch.cuda.amp.autocast(enabled=False):
        x = self.rope(x, seq_len)

    # 其他计算继续FP16
    attention = self.attention(x)

这些问题都不是理论上的"高级错误"，就是实际代码里疏忽了。位置编码作为基础设施，一旦出错很难直接从loss上看出来。

未来可能的方向

目前的位置编码方案都不是完美解。RoPE和ALiBi各自解决了一部分问题，但又引入了新的复杂度。

我个人比较期待的是"可学习位置编码"和"自适应位置编码"的结合方向。让模型在学习任务的同时，也学习最适合自己的位置编码方式。但这方面的实践还不多，而且训练成本会增加很多。

另一个方向是针对特定任务的定制化编码。比如时间序列任务可能需要周期性的位置编码，代码生成任务可能需要层级结构的位置信息。通用的位置编码很难在所有任务上都达到最优。

收尾

位置编码这个问题，表面上是技术选型，本质上是信息表示。如何在有限的参数里把"位置"这个概念表达清楚，既不过度复杂，又足够灵活，这确实需要权衡。

我现在的判断是：没有完美的位置编码，只有适合当前任务的方案。重要的不是记住所有位置编码的公式，而是理解它们要解决什么问题，以及在什么情况下会失效。

就像文章开头引用的那句诗一样，位置编码解决的就是"你在哪里、我在哪里、我们之间是什么关系"的问题。模型需要知道这些，才能做出合理的判断。

至于用什么编码，看你手里的数据有多长、算力有多足、以及你想让模型完成什么任务。实战里很少有标准答案，多是取舍。

版权声明: 本文首发于 指尖魔法屋-从位置走到信息：AI位置编码笔记（https://blog.thinkmoon.cn/post/244-ai-positional-encoding-practice/） 转载或引用必须申明原指尖魔法屋来源及源地址！

相关阅读

根据标签、分类和主题为你挑选的相近文章