长上下文技术演进：从位置编码到环形注意力

引言

早期Transformer模型受限于 $O(n^2)$ 的注意力复杂度，上下文长度被限制在512或2048 tokens。到2026年，主流模型已普遍支持128K-1M tokens的上下文窗口。这一进步并非单一技术的产物，而是位置编码、注意力机制、架构设计等多个方向协同演进的结果。本文将系统梳理这一技术演进脉络。

位置编码：长上下文的基石

Absolute Position Embedding

原始Transformer使用learned absolute position embedding：

$$ h_i = x_i + \text{PE}[i], \quad \text{PE}[i] \in \mathbb{R}^{d_{\text{model}}} $$

问题：无法外推到训练长度之外的位置，泛化性差。

Sinusoidal Position Encoding

$$ \text{PE}(pos, 2i) = \sin\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) $$ $$ \text{PE}(pos, 2i+1) = \cos\left(\frac{pos}{10000^{2i/d_{\text{model}}}}\right) $$

Sinusoidal编码具有相对位置信息，但外推能力仍然有限。

RoPE：旋转位置编码

RoPE（Rotary Position Embedding）是当前主流的位置编码方案，通过旋转矩阵实现相对位置编码：

$$ f_q(x_m, m) = R_\Theta^m q(x_m), \quad f_k(x_n, n) = R_\Theta^n k(x_n) $$

其中：

$$ R_\Theta^m = \begin{pmatrix} \cos m\theta_1 & -\sin m\theta_1 & 0 & 0 \ \sin m\theta_1 & \cos m\theta_1 & 0 & 0 \ 0 & 0 & \cos m\theta_2 & -\sin m\theta_2 \ 0 & 0 & \sin m\theta_2 & \cos m\theta_2 \ \vdots & \vdots & \vdots & \vdots \end{pmatrix} $$

内积仅依赖相对位置：

$$ \langle f_q(x_m, m), f_k(x_n, n) \rangle = \langle q(x_m), k(x_n) \rangle \cdot g(m - n) $$

其中 $g(\cdot)$ 是只与相对位置差有关的函数。

RoPE外推与插值

RoPE原生支持相对位置编码，但外推到训练长度之外仍然困难。2023-2024年发展出几种方案：

1. Position Interpolation（位置插值）

$$ \text{PE}{\text{interp}}(pos) = \text{PE}\left(\frac{pos \cdot L{\text{train}}}{L_{\text{target}}}\right) $$

将目标长度线性映射到训练长度，无需重新训练即可支持更长上下文。

2. NTK-Aware Interpolation

通过扩展RoPE的基底频率而非简单插值，更好地保持高频信息：

def ntk_aware_rope(base, max_pos, scale_factor):
    """NTK-Aware的RoPE频率扩展"""
    # 计算所需的最小alpha
    alpha = scale_factor if scale_factor > 1 else 1
    # 对每个频率分量应用NTK感知的缩放
    scaled_base = base * alpha
    inv_freq = 1.0 / (scaled_base ** (torch.arange(0, dim, 2).float() / dim))
    return inv_freq

3. YaRN（Yet another RoPE extensioN）

YaRN结合NTK-Aware和部分维度的缩放，是目前效果最好的RoPE扩展方案：

$$ \text{YaRN}(pos, i) = \text{NTK}\alpha(\min(pos, L{\text{train}}), i) \times w(i) $$

其中 $w(i)$ 是维度相关的加权函数，对高频维度（小i）采用更小缩放。

ALiBi：无位置编码的注意力偏置

ALiBi（Attention with Linear Biases）完全不使用位置编码，而是直接在注意力分数上添加线性偏置：

$$ A_{ij} = \frac{Q_i K_j^T}{\sqrt{d}} - m \times (j - i) $$

其中 $m$ 是head特定的斜率。ALiBi在推理时天然支持任意长度外推，但训练速度较慢。

注意力机制优化：突破二次复杂度

Sparse Attention

Sparse Attention通过限制每个token注意力的范围，将复杂度降至 $O(n\sqrt{n})$ 或 $O(n \log n)$。

Sliding Window + Global Attention（Longformer）：

def longformer_attention(Q, K, V, window_size=512, global_indices=None):
    B, H, N, D = Q.shape
    attn_mask = torch.full((N, N), float('-inf'))
    
    # 滑动窗口
    for i in range(N):
        start = max(0, i - window_size)
        end = min(N, i + window_size + 1)
        attn_mask[i, start:end] = 0
    
    # 全局token（如[CLS]）可看到所有token
    if global_indices is not None:
        attn_mask[global_indices, :] = 0
        attn_mask[:, global_indices] = 0
    
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (D ** 0.5)
    scores += attn_mask
    attn = F.softmax(scores, dim=-1)
    return torch.matmul(attn, V)

Ring Attention：环形分布式注意力

Ring Attention（环形注意力）是处理超长上下文的关键技术，通过序列维度上的分布式计算，使上下文长度可以随GPU数量线性扩展。

核心思想：将序列维度切分到多个设备，通过Ring形式的All-Gather通信，每个设备只需存储和计算部分KV。

def ring_attention(Q_shard, K_shard, V_shard, ring_size):
    """
    Q_shard: [B, H, N/ring_size, D] 本设备上的Q
    K_shard, V_shard: 同上
    """
    # 初始化输出累加器
    O_local = torch.zeros_like(Q_shard)
    denominator = torch.zeros(Q_shard.shape[:-1])
    
    # Ring通信：每个设备轮流接收上一设备的K/V
    for step in range(ring_size):
        # 计算本地注意力分数
        scores = torch.matmul(Q_shard, K_shard.transpose(-2, -1)) / (D ** 0.5)
        attn = F.softmax(scores, dim=-1)
        
        # 累加到输出
        O_local += torch.matmul(attn, V_shard)
        denominator += attn.sum(dim=-1)
        
        # Ring传递KV
        K_shard = all_gather_ring(K_shard)
        V_shard = all_gather_ring(V_shard)
    
    return O_local / denominator.unsqueeze(-1)

Ring Attention使Gemini 1.5 Pro能处理10M tokens的上下文，GPT-5也采用了类似的分布式注意力机制。

Flash Attention：IO感知优化

Flash Attention v2/v3虽然不是直接降低理论复杂度，但通过分块计算和IO优化，将长上下文的实际计算效率提升了3-5倍。

关键技术：

Tiling：将QKV分块，使每个块能放入SRAM
Recomputation：反向传播时重新计算注意力，而非存储 $O(n^2)$ 的注意力矩阵
内核融合：将softmax、dropout等操作融合到单个CUDA kernel

上下文扩展的工程实践

渐进式扩展

从短上下文到长上下文的扩展通常采用渐进式策略：

class ProgressiveContextExtension:
    def __init__(self, model, max_length_schedule):
        self.model = model
        self.schedule = max_length_schedule  # e.g., [2048, 8192, 32768, 131072]
    
    def train_step(self, step, data):
        current_max_len = self.schedule[min(step // 1000, len(self.schedule) - 1)]
        data = truncate_or_pad(data, max_len=current_max_len)
        
        # 位置编码调整
        if hasattr(self.model, 'adjust_rope'):
            self.model.adjust_rope(current_max_len)
        
        loss = self.model(data)
        return loss.backward()

长度外推（Length Extrapolation）

长度外推指模型在短上下文上训练，在长上下文上推理的能力。主流策略：

Position Interpolation：如前所述，插值位置编码
Blockwise Attention：将长序列分块处理，块间通过压缩表示交互
Streaming Attention：维护一个固定大小的KV Cache，超出部分压缩或丢弃

长上下文评测

评测基准	长度范围	评测维度
RULER	4K-128K	检索、多跳推理
L-Eval	3K-200K	长文档QA
LongBench	5K-100K	多任务长文本
Needle-in-a-Haystack	2K-128K	信息检索
InfinityBench	10K-1M	超长上下文

各方案性能对比

方案	理论复杂度	训练支持	外推能力	代表模型
标准Attention	$O(n^2)$	✅	❌	GPT-3
Flash Attention	$O(n^2)$	✅	需插值	Llama 3
Sparse Attention	$O(n \log n)$	✅	部分	Longformer
Ring Attention	$O(n^2)$ 分布式	✅	✅	Gemini 1.5
Mamba/SSM	$O(n)$	✅	✅	Mamba2
RetNet	$O(n)$	✅	✅	RetNet
ALiBi	$O(n^2)$ 但可外推	✅	✅	BLOOM

超长上下文的挑战

注意力稀释

上下文越长，注意力分数越分散，模型难以聚焦于关键信息。解决方案：

Attention Sink：保留初始token的高注意力权重
压缩记忆：使用可学习的压缩token压缩历史信息

位置信息退化

超长上下文中，RoPE的旋转频率接近周期边界，位置信息模糊。NTK-Aware和YaRN通过扩展频率基底缓解这一问题。

训练稳定性

超长序列训练时梯度容易出现数值不稳定。技术：

梯度检查点：降低显存，支持更长序列
混合精度训练：BF16在长序列训练上更稳定
Loss Masking：忽略过长序列的部分loss

结语

长上下文技术的发展是2023-2026年间大模型领域最显著的进步之一。从RoPE的逐步完善到Ring Attention的分布式突破，从Flash Attention的工程优化到Mamba的架构创新，这些技术共同将上下文窗口从最初的512 tokens扩展到了百万token级别。未来，随着MoE架构和线性注意力技术的成熟，我们或许能看到更长、更高效的长上下文处理方案。

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

引言#

位置编码：长上下文的基石#

Absolute Position Embedding#

Sinusoidal Position Encoding#

RoPE：旋转位置编码#

RoPE外推与插值#

ALiBi：无位置编码的注意力偏置#

注意力机制优化：突破二次复杂度#

Sparse Attention#

Ring Attention：环形分布式注意力#

Flash Attention：IO感知优化#

上下文扩展的工程实践#

渐进式扩展#

长度外推（Length Extrapolation）#

长上下文评测#

各方案性能对比#

超长上下文的挑战#

注意力稀释#

位置信息退化#

训练稳定性#

结语#

加入讨论#

📌 相关推荐

引言