上下文窗口扩展技术：从 4K 到 1M

1. 为什么长上下文重要？

4K 上下文只能容纳约 3000 字的文档。要处理整本书、代码仓库或长对话，需要 32K-1M 的上下文窗口。但直接训练长上下文模型面临三大挑战：

训练成本：Attention 复杂度 $O(n^2)$，32K 上下文的计算量是 4K 的 64 倍
位置编码外推：训练时未见过的位置编码效果不可控
显存占用：KV Cache 随序列长度线性增长

2. 位置编码外推

2.1 Position Interpolation (PI)

最简单的方法：将位置索引线性缩放到训练范围内。

$$ m’ = m \cdot \frac{L_{train}}{L_{target}}


4K 训练 → 8K 推理时，位置 $m$ 缩放为 $m \times 0.5$。

**问题：** 高频分量被等比缩放，局部注意力模式被破坏。短距离关系（如相邻 Token）的编码精度下降。

### 2.2 NTK-aware Interpolation

NTK（Neural Tangent Kernel）理论指出，高频分量需要保留，低频分量可以缩放。修改 RoPE 的 base：

$$
b' = b \cdot s^{d/(d-2)}

其中 $s = L_{target} / L_{train}$ 是扩展倍数，$d$ 是 Head 维度。

效果： 高频部分（局部注意力）几乎不变，低频部分（长距离）被扩展。4K→16K 无需微调即可使用。

2.3 NTK-by-parts

将频率维度分为三段：

频率高于 $\lambda_{high}$（波长远小于扩展后长度）：不插值
频率低于 $\lambda_{low}$（波长远大于扩展后长度）：线性插值
中间：平滑过渡

def ntk_by_parts_rope(head_dim, orig_len, target_len, alpha=1, beta=32):
    """NTK-by-parts 位置编码"""
    s = target_len / orig_len
    base = 10000.0
    
    freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2).float() / head_dim))
    
    # 计算波长
    wavelengths = 2 * math.pi / freqs
    
    # 临界波长
    lambda_high = orig_len
    lambda_low = orig_len / (2 * math.pi) * beta
    
    # 分段缩放
    scale = torch.ones_like(freqs)
    for i in range(len(freqs)):
        wl = wavelengths[i]
        if wl < lambda_low:
            scale[i] = 1.0  # 高频不缩放
        elif wl > lambda_high:
            scale[i] = 1.0 / s  # 低频完全插值
        else:
            # 平滑过渡
            t = (lambda_high - wl) / (lambda_high - lambda_low)
            scale[i] = 1.0 - (1 - 1.0/s) * t
    
    freqs = freqs * scale
    return freqs

2.4 YaRN

YaRN 在 NTK-by-parts 基础上增加 Attention 温度缩放：

$$ t = 0.1 \cdot \ln(s) + 1, \quad \text{logits}’ = \frac{\text{logits}}{t}


温度缩放补偿了长上下文中 Attention 分布变得过尖的问题。YaRN 是目前最有效的 RoPE 外推方案之一。

### 2.5 LongRoPE

LongRoPE（Microsoft, 2024）将 RoPE 的频率缩放从均匀分组推广为**逐维度优化**：

- 不再按频率高低分区，而是为每个 2D 旋转平面独立搜索缩放因子
- 使用进化算法搜索最优缩放方案
- 可扩展到 2M 上下文

## 3. 动态 NTK

推理时根据当前序列长度动态调整 RoPE base：

$$
b'(l) = b \cdot \left(\frac{l}{L_{train}}\right)^{d/(d-2)}
$$

其中 $l$ 是当前序列长度。在短序列时保持原始编码，在长序列时自动缩放。

```python
class DynamicNTKScaling:
    def __init__(self, base=10000.0, train_len=4096, head_dim=128):
        self.base = base
        self.train_len = train_len
        self.head_dim = head_dim
    
    def get_scaling_factor(self, seq_len):
        if seq_len <= self.train_len:
            return 1.0
        s = seq_len / self.train_len
        return s ** (self.head_dim / (self.head_dim - 2))
    
    def get_base(self, seq_len):
        return self.base * self.get_scaling_factor(seq_len)

4. 分块注意力

4.1 Sliding Window Attention (SWA)

每个 Token 只关注固定大小的局部窗口：

$$ \text{Attention}(q_i, K, V) = \text{softmax}(q_i K_{[i-w, i]}^T) V_{[i-w, i]} $$

窗口大小 $w$ 通常为 4096。复杂度从 $O(n^2)$ 降到 $O(nw)$。

Mistral 使用 SWA + 全局注意力交替：每隔几层放一个全局注意力层，其余层用 SWA。

4.2 StreamingLLM

StreamingLLM 的关键发现：Attention 的"注意力 sinks"现象——第一个和最后几个 Token 获得不合理的超高注意力分数。

方法：保留初始几个 Token + 滑动窗口，丢弃中间部分：

[Attention Sinks] ... [Sliding Window]
  Token 0-3           Token (i-w) to i

效果：可以处理无限长度序列（流式），但丢失了中间信息。

4.3 Landmark Attention

将序列分为固定大小的块，每块用一个"地标"Token 代表。Attention 在块级别进行：

块内：标准 Self-Attention
块间：通过地标 Token 传递信息

复杂度从 $O(n^2)$ 降为 $O(n \cdot \sqrt{n})$。

5. Ring Attention

5.1 核心思想

Ring Attention（Liu et al., 2023）将长序列分布到多个 GPU 上，GPU 组成环形拓扑：

每个 GPU 持有一段序列的 Q, K, V
K, V 在 GPU 间环形传递
每经过一个 GPU，计算当前 Q 与该 GPU 的 K, V 的部分 Attention

5.2 计算-通信重叠

时间线:
GPU 0: 计算 Q0×K0 | 发送 K0,V0 → GPU1 | 计算 Q0×K1 | 发送 K1,V1 → GPU2 | ...
GPU 1: 计算 Q1×K1 | 发送 K1,V1 → GPU2 | 计算 Q1×K2 | 发送 K2,V2 → GPU3 | ...
GPU 2: 计算 Q2×K2 | 发送 K2,V2 → GPU3 | 计算 Q2×K3 | 发送 K3,V3 → GPU0 | ...
GPU 3: 计算 Q3×K3 | 发送 K3,V3 → GPU0 | 计算 Q3×K0 | 发送 K0,V0 → GPU1 | ...

通信（K, V 传输）与计算（部分 Attention）并行，理论上可以扩展到无限长度（只要 GPU 足够）。

5.3 实现要点

def ring_attention_forward(Q, K, V, comm):
    """简化版 Ring Attention"""
    rank = comm.get_rank()
    world_size = comm.get_size()
    
    # 本地 Q/K/V
    local_q, local_k, local_v = Q, K, V
    
    # 存储 K/V 的缓冲区
    k_buffer, v_buffer = local_k.clone(), local_v.clone()
    
    output = torch.zeros_like(local_q)
    
    for step in range(world_size):
        # 计算当前 Q 对 K_buffer, V_buffer 的部分 Attention
        block_output = compute_block_attention(local_q, k_buffer, v_buffer)
        output += block_output
        
        if step < world_size - 1:
            # 异步发送 K,V 到下一个 GPU
            send_req = comm.isend(k_buffer, dst=(rank+1) % world_size)
            # 接收前一个 GPU 的 K,V
            recv_k = torch.empty_like(k_buffer)
            comm.recv(recv_k, src=(rank-1) % world_size)
            send_req.wait()
            k_buffer = recv_k
            # V 同理
    
    return output

5.4 Flash Ring Attention

结合 Flash Attention 的 tiling 技巧，Ring Attention 可以在不实例化完整 $n \times n$ Attention 矩阵的情况下完成长序列计算。实测在 8×A100 上可以处理 1M 上下文。

6. KV Cache 优化

长上下文的另一个瓶颈是 KV Cache 显存。以 LLaMA-2 70B 为例：

上下文长度	KV Cache 大小 (FP16)
4K	10.5 GB
32K	84 GB
128K	336 GB
1M	2.6 TB

6.1 量化 KV Cache

将 KV Cache 从 FP16 量化为 INT8 或 INT4：

INT8: 显存减半，精度损失 < 1%
INT4: 显存减 75%，需要特殊处理精度

6.2 PagedAttention (vLLM)

借鉴操作系统的虚拟内存管理，将 KV Cache 分为固定大小的 Page：

不需要为最大长度预分配连续内存
动态分配/释放 Page，减少碎片
支持共享前缀（多个请求共享同一 system prompt 的 KV Cache）

7. 效果对比

方法	最大长度	需要微调	质量保持	工程复杂度
直接外推	1.5x	否	差	低
PI	4x	少量	中	低
NTK-aware	8x	否	良	低
YaRN	32x	少量	优	中
LongRoPE	500x	少量	优	中
Ring Attention	无限*	否	优	高
SWA + 全局	无限	否	中	中

*受 GPU 数量限制

8. 实践建议

4K→32K：YaRN 微调是最性价比方案
32K→128K：YaRN + Ring Attention（多 GPU）
128K+：LongRoPE + PagedAttention + 量化 KV Cache
流式场景：StreamingLLM + SWA
微调数据：长文本微调需要包含长文档的 QA 数据，否则模型有"长文本窗口"但不会利用
评估：除了 perplexity，用 Needle in a Haystack 测试验证模型是否真的能检索到长上下文中的信息

加入讨论

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

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

1. 为什么长上下文重要？#

2. 位置编码外推#

2.1 Position Interpolation (PI)#

2.3 NTK-by-parts#

2.4 YaRN#

4. 分块注意力#

4.1 Sliding Window Attention (SWA)#

4.2 StreamingLLM#

4.3 Landmark Attention#

5. Ring Attention#

5.1 核心思想#

5.2 计算-通信重叠#

5.3 实现要点#

5.4 Flash Ring Attention#

6. KV Cache 优化#

6.1 量化 KV Cache#

6.2 PagedAttention (vLLM)#

7. 效果对比#

8. 实践建议#

加入讨论#

📌 相关推荐