YaRN / LongRoPE 深度解析：大模型上下文扩展技术

1
┌─────────────────────────────────────────────────────────────┐
2
│              LLM 上下文长度演进                              │
3
├─────────────────────────────────────────────────────────────┤
4
│                                                              │
5
│  2020: GPT-3           2,048 tokens                       │
6
│  2022: LLaMA           2,048 tokens                       │
7
│  2023: LLaMA 2         4,096 tokens                       │
8
│  2023: Claude 2        200,000 tokens                      │
9
│  2023: GPT-4 Turbo     128,000 tokens                     │
10
│  2024: Gemini 1.5      1,000,000 tokens                   │
11
│  2024: Claude 3.5      200,000 tokens                     │
12
│  2024: Qwen 2.5        128,000 tokens                     │
13
│                                                              │
14
│  趋势：上下文长度呈指数增长                                  │
15
│  挑战：如何在不重新训练的情况下扩展？                          │
16
│                                                              │
17
└─────────────────────────────────────────────────────────────┘

1
class ContextExtensionMotivation:
2
    """
3
    上下文扩展的必要性
4
    """
5

6
    def use_cases(self):
7
        """
8
        需要长上下文的场景
9
        """
10
        return {
11
            "code_generation": {
12
                "scenario": "分析大型代码仓库",
13
                "example": "理解 10 万行代码的依赖关系",
14
                "requirement": "100K+ tokens",
15
            },
16
            "document_analysis": {
17
                "scenario": "多文档摘要与问答",
18
                "example": "分析整本书籍或法律文档",
19
                "requirement": "500K+ tokens",
20
            },
21
            "long_conversation": {
22
                "scenario": "长时间对话和记忆",
23
                "example": "个人助手保留数月对话历史",
24
                "requirement": "1M+ tokens",
25
            },
26
            "scientific_research": {
27
                "scenario": "分析长篇论文和数据集",
28
                "example": "阅读和理解 1000 篇相关论文",
29
                "requirement": "10M+ tokens",
30
            },
31
        }
32

33
    def challenges(self):
34
        """
35
        扩展面临的挑战
36
        """
37
        return {
38
            "training_cost": "重新训练的成本极高 (175B 模型需要数百万美元)",
39
            "positional_broke": "训练时未见过的位置导致困惑",
40
            "attention_degradation": "远距离 token 的注意力分散",
41
            "memory_explosion": "Attention 计算 O(N²) 显存爆炸",
42
            "quality_preservation": "扩展后质量不能下降太多",
43
        }

1
class RoPELimitations:
2
    """
3
    RoPE 的局限性
4
    """
5

6
    def wavelength_analysis(self):
7
        """
8
        波长分析
9

10
        RoPE 中每个维度 i 的旋转角度:
11
        θ_i = base^{-2i/d}
12

13
        对应的波长:
14
        λ_i = 2π / θ_i = 2π · base^{2i/d}
15
        """
16
        return """
17
        以 base = 10000, d = 4096 为例:
18

19
        dim pair 0:  θ₀ = 1,         λ₀ ≈ 6.28 tokens
20
        dim pair 1:  θ₁ = 1/10000,   λ₁ ≈ 62,832 tokens
21
        dim pair 2:  θ₂ = 1/10⁸,     λ₂ ≈ 6.28 × 10⁸ tokens
22
        dim pair 3:  θ₃ = 1/10¹²,    λ₃ ≈ 6.28 × 10¹² tokens
23

24
        关键观察:
25
        ┌─────────────────────────────────────────────────────┐
26
        │  低维度 (i 小): 波长短，旋转快 → 捕获短距离依赖       │
27
        │  高维度 (i 大): 波长长的，旋转慢 → 捕获长距离依赖     │
28
        └─────────────────────────────────────────────────────┘
29
        """
30

31
    def extrapolation_problem(self):
32
        """
33
        外推问题
34
        """
35
        return """
36
        当序列长度超过训练长度时:
37

38
        训练: max_position = 4096
39

40
        推理: position = 8192
41

42
        问题:
43
        1. 模型从未见过 position = 5000-8192 的旋转角度
44
        2. 这些位置的编码与训练分布不匹配
45
        3. 模型难以正确理解位置关系
46

47
        简单缩放的问题:
48
        θ' = θ / s (s = 2 表示扩展 2 倍)
49

50
        → 波长变为 λ/s
51
        → 所有频率都变快
52
        → 短距离信息被扭曲
53
        """
54

55
    def frequency_band_issue(self):
56
        """
57
        频段问题
58
        """
59
        return """
60
        RoPE 可以分成三个频段:
61

62
        ┌─────────────────────────────────────────────────────┐
63
        │  高频 (short-range):  λ < 32 tokens                │
64
        │  ├── 负责 token 级别的细粒度信息                    │
65
        │  └── 对位置非常敏感                                 │
66
        ├─────────────────────────────────────────────────────┤
67
        │  中频 (mid-range):  32 < λ < 512 tokens            │
68
        │  ├── 负责词语和短语级别的信息                       │
69
        │  └── 位置关系仍然重要                              │
70
        ├─────────────────────────────────────────────────────┤
71
        │  低频 (long-range):  λ > 512 tokens                │
72
        │  ├── 负责篇章和文档级别的信息                       │
73
        │  └── 可以容忍一定的位置模糊                         │
74
        └─────────────────────────────────────────────────────┘
75

76
        关键洞察:
77
        → 不同频段对扩展的敏感度不同
78
        → 应该对不同频段采用不同策略
79
        """

1
class PositionInterpolation:
2
    """
3
    位置插值 (Position Interpolation, PI)
4

5
    论文: "Extending Context Window of Large Language Models
6
          via Position Interpolation" (Chen et al., 2023)
7
    """
8

9
    def basic_idea(self):
10
        """
11
        基本思想
12

13
        将位置范围 [0, L'] 映射到 [0, L]，
14
        其中 L' > L 是扩展后的长度，L 是训练长度。
15

16
        公式:
17
        pos' = pos × (L / L')
18

19
        即 position 被压缩了 L / L' 倍。
20
        """
21
        return """
22
        位置插值示意图:
23

24
        原始位置 (训练):
25
        [0    ][1    ][2    ]...[4095  ]
26
        0     1     2            4096
27

28
        扩展位置 (推理):
29
        [0  ][0.5  ][1  ]...[2047.5]
30

31
        压缩比: L / L' = 4096 / 8192 = 0.5
32

33
        效果:
34
        → 位置 8192 被映射到位置 4096
35
        → 位置 4096 被映射到位置 2048
36
        → 模型看到的位置始终在 [0, L] 范围内
37
        """
38

39
    def mathematical_formulation(self):
40
        """
41
        数学形式化
42
        """
43
        return '''
44
        # 位置插值的数学表示
45

46
        原始 RoPE:
47
        θ_i = base^{-2i/d}
48
        R(m, i) = [cos(m·θ_i), -sin(m·θ_i)]
49
                       [sin(m·θ_i),  cos(m·θ_i)]
50

51
        插值后的 RoPE:
52
        θ_i' = θ_i × s  (s = L / L' < 1)
53

54
        或者等价地:
55
        位置 m' = m × s
56

57
        新的旋转角度:
58
        m' × θ_i = m × s × θ_i
59

60
        关键:
61
        → 角度被缩小了 s 倍
62
        → 旋转变慢了
63
        → 模型看到的是"压缩"的位置
64
        '''
65

66
    def problems_with_naive_pi(self):
67
        """
68
        朴素 PI 的问题
69
        """
70
        return """
71
        问题 1: 短距离信息受损
72

73
        原始: 相邻 token 位置差 = 1
74
        插值: 相邻 token 位置差 = s < 1
75

76
        但模型学习的是位置差 = 1 的模式！
77

78
        问题 2: 语义扭曲
79

80
        位置差从 1 变成 s，改变了相对关系。
81
        这可能影响词级别的注意力模式。
82

83
        问题 3: 频段敏感性
84

85
        高频维度（短距离）对插值更敏感：
86
        - 波长短的维度，微小变化影响大
87
        - 插值改变了它们的旋转模式
88

89
        解决方向：
90
        → 只对低频维度插值
91
        → 高频维度保持不变
92
        → 这就是 YaRN 的核心思想！
93
        """

1
class InterpolationVsExtrapolation:
2
    """
3
    插值 vs 外推
4
    """
5

6
    def extrapolation(self):
7
        """
8
        外推 (Extrapolation)
9

10
        直接使用训练好的编码向外延伸到未见位置。
11
        """
12
        return """
13
        外推示意图:
14

15
        训练位置: [0, 1, 2, ..., 4095]
16
        外推位置: [4096, 4097, ..., 8191]
17

18
        问题:
19
        → 位置 4096+ 的旋转角度是全新的
20
        → 模型没有见过这些角度
21
        → 注意力模式可能崩溃
22

23
        ALiBi 的优势:
24
        → ALiBi 不编码绝对位置
25
        → 注意力偏置只依赖于距离
26
        → 外推自然发生
27
        """
28

29
    def interpolation(self):
30
        """
31
        插值 (Interpolation)
32

33
        将新位置映射到训练时见过的位置范围内。
34
        """
35
        return """
36
        插值示意图:
37

38
        扩展: 4096 → 8192
39
        压缩比: s = 0.5
40

41
        新位置 8192 → 映射到 4096
42
        新位置 4096 → 映射到 2048
43

44
        优势:
45
        → 所有位置都在训练范围内
46
        → 模型至少见过类似的位置
47

48
        劣势:
49
        → 改变了位置之间的相对距离
50
        → 可能影响短距离依赖
51
        """
52

53
    def hybrid_approach(self):
54
        """
55
        混合方法
56
        """
57
        return """
58
        理想的解决方案:
59

60
        ┌─────────────────────────────────────────────────────┐
61
        │  短距离 (< 训练长度):  保持原样（外推）              │
62
        │  长距离 (> 训练长度):  插值压缩                      │
63
        └─────────────────────────────────────────────────────┘
64

65
        具体做法:
66
        pos' = pos                         if pos < L
67
        pos' = L + (pos - L) × s          if pos >= L
68

69
        这样:
70
        → 训练过的位置完全不变
71
        → 新位置平滑过渡
72
        → 避免突然的分布变化
73
        """

1
class YaRNCore:
2
    """
3
    YaRN (Yet another RoPE extensioN)
4

5
    论文: "YaRN: Efficient Context Window Extension of
6
          Large Language Models" (Peng et al., 2023)
7
    """
8

9
    def three_components(self):
10
        """
11
        YaRN 的三个组成部分
12
        """
13
        return """
14
        ┌─────────────────────────────────────────────────────┐
15
        │                    YaRN 三合一                        │
16
        ├─────────────────────────────────────────────────────┤
17
        │                                                       │
18
        │  1. Position Interpolation (位置插值)                  │
19
        │     → 将位置范围缩放到训练范围                         │
20
        │     → 基础策略                                       │
21
        │                                                       │
22
        │  2. Attention Scaling (注意力缩放)                    │
23
        │     → 补偿插值带来的注意力衰减                        │
24
        │     → 调整 softmax 前的 logits                        │
25
        │                                                       │
26
        │  3. Frequency Adjustment (频率调整)                    │
27
        │     → 对不同频段差异化处理                           │
28
        │     → 高频保持，低频插值                             │
29
        │                                                       │
30
        └─────────────────────────────────────────────────────┘
31
        """
32

33
    def why_three_components(self):
34
        """
35
        为什么需要三个组件
36
        """
37
        return """
38
        单独 PI 的问题:
39

40
        1. 位置缩放改变了所有维度的频率
41
        2. 高频维度（短距离）的波长被压缩
42
        3. 这破坏了模型学习到的细粒度模式
43

44
        单独 Attention Scaling:
45

46
        1. 可以补偿注意力的整体下降
47
        2. 但不能恢复被破坏的相对位置关系
48

49
        单独 Frequency Adjustment:
50

51
        1. 可以保留高频信息
52
        2. 但低频维度仍需要插值
53

54
        三合一的效果:
55

56
        ┌─────────────────────────────────────────────────────┐
57
        │  高频维度:  保持不变 → 保留短距离依赖                  │
58
        │  中频维度:  适度插值 → 平衡新旧知识                    │
59
        │  低频维度:  完全插值 → 支持超长距离                    │
60
        │                                                       │
61
        │  Attention Scaling:  补偿整体衰减                      │
62
        └─────────────────────────────────────────────────────┘
63
        """

1
class YaRNPositionInterpolation:
2
    """
3
    YaRN 中的位置插值
4
    """
5

6
    def formal_definition(self):
7
        """
8
        形式化定义
9
        """
10
        return '''
11
        # YaRN Position Interpolation
12

13
        设:
14
        - L: 训练时的最大位置 (如 4096)
15
        - L': 扩展后的最大位置 (如 32768)
16
        - s: 缩放因子 = L / L' (如 0.125)
17

18
        对于位置 m ∈ [0, L']，插值后的位置:
19

20
        m' = m × s
21

22
        新的旋转角度:
23
        θ_i' = base^{-2i/d}
24

25
        插值后的旋转:
26
        θ_i'' = θ_i × (m' / m) = θ_i × s
27

28
        即: 对所有维度应用相同的缩放
29
        '''
30

31
    def dim_dependent_scaling(self):
32
        """
33
        维度相关的缩放
34

35
        YaRN 论文的发现:
36
        高频维度需要更小的缩放（或不缩放）
37
        """
38
        return """
39
        维度分组:
40

41
        θ_i = base^{-2i/d}
42

43
        对于 i = 0, 1, ..., d/2-1
44

45
        定义阈值维度 D_threshold:
46

47
        - i < D_threshold: 高频维度，保持不变
48
        - i >= D_threshold: 低频维度，插值缩放
49

50
        D_threshold 的选择:
51

52
        目标: λ_i < β × L
53

54
        其中:
55
        - λ_i 是波长 = 2π / θ_i
56
        - β 是超参数 (论文建议 β = 32 或 64)
57

58
        这意味着:
59
        → 只对波长小于 β × L 的维度缩放
60
        → 这些维度负责长距离依赖
61
        → 高频维度（短距离）保持不变
62
        """

1
class YaRNAttentionScaling:
2
    """
3
    YaRN 注意力缩放
4
    """
5

6
    def motivation(self):
7
        """
8
        动机
9

10
        位置插值会改变 attention 的分布。
11
        需要调整以保持适当的注意力模式。
12
        """
13
        return """
14
        Attention 分数变化:
15

16
        原始: A(m, n) = softmax(q_m · k_n)
17

18
        插值后: A'(m', n') = softmax(q_m' · k_n')
19

20
        问题:
21
        → 插值改变了 q 和 k 的旋转角度
22
        → 点积结果发生变化
23
        → 注意力分布可能变得过于平坦
24

25
        解决: 缩放 attention logits
26

27
        A_scaled(m, n) = softmax((q_m · k_n) / r)
28

29
        其中 r > 1 会让注意力分布更 sharp
30
        """
31

32
    def yarn_scaling_formula(self):
33
        """
34
        YaRN 的缩放公式
35
        """
36
        return '''
37
        # YaRN Attention Scaling
38

39
        缩放因子:
40
        r = sqrt(1/s² · (1 - β) + β)
41

42
        其中:
43
        - s: 位置缩放因子
44
        - β: 阈值参数 (0 < β < 1)
45

46
        或者使用对数形式:
47
        log(r) = 0.5 * log(1/s² * (1 - β) + β)
48

49
        论文推荐的参数:
50
        - s = L / L' = 0.125 (扩展 8 倍)
51
        - β = 1/32 (对应波长阈值)
52
        - r ≈ 2.0
53

54
        效果:
55
        → r > 1 让 attention 更 sharp
56
        → 补偿插值带来的注意力分散
57
        → 恢复细粒度的位置依赖
58
        '''
59

60
    def intuition(self):
61
        """
62
        直观理解
63
        """
64
        return """
65
        Attention Scaling 的直觉:
66

67
        插值前:
68
        "The cat sat on the mat"
69
        位置 0    1     2   3   4   5
70
        attention: [The→cat] 高, [cat→sat] 高, ...
71

72
        插值后 (位置被压缩):
73
        attention: [The→cat] 变低, [cat→sat] 变低, ...
74

75
        问题: 所有 attention 都变低了！
76

77
        Scaling 后:
78
        attention_scaled = softmax(logits / r)
79

80
        r > 1 → 相当于乘以 r → logits 变大 → softmax 更 sharp
81

82
        效果: 恢复原来的 attention 强度分布！
83
        """

1
class YaFrequentAdjustment:
2
    """
3
    YaRN 频率调整
4
    """
5

6
    def dim_wise_adjustment(self):
7
        """
8
        维度级别的调整
9

10
        这是 YaRN 与朴素 PI 的关键区别
11
        """
12
        return """
13
        朴素 PI:
14
        θ_i' = θ_i × s  (对所有 i 相同)
15

16
        YaRN:
17
        θ_i' = θ_i × s_i  (对每个 i 不同)
18

19
        s_i 的选择:
20

21
        s_i = 1                              if λ_i > L × β
22
        s_i = (λ_i / (L × β)) × s + (1 - s) if λ_i <= L × β
23

24
        简化:
25
        s_i = max(1, (L × β) / λ_i) × s
26

27
        解释:
28
        → 如果波长已经很长 (λ_i > L × β): 保持不变
29
        → 如果波长适中: 适度缩放
30
        → 如果波长很短: 完全缩放
31

32
        这保证了:
33
        → 高频维度不被破坏
34
        → 低频维度被正确扩展
35
        """
36

37
    def threshold_calculation(self):
38
        """
39
        阈值计算
40
        """
41
        return '''
42
        # 计算阈值维度
43

44
        λ_i = 2π × base^{2i/d}  (波长)
45

46
        目标: λ_i < β × L
47

48
        解 i:
49
        2π × base^{2i/d} < β × L
50
        base^{2i/d} < (β × L) / 2π
51
        2i/d < log_base((β × L) / 2π)
52
        i < (d/2) × log_base((β × L) / 2π)
53

54
        D_threshold = floor((d/2) × log_base((β × L) / 2π))
55

56
        示例 (d=4096, base=10000, L=4096, β=1/32):
57
        (4096/2) × log_10000((4096/32) / 2π)
58
        ≈ 2048 × log_10000(20.5)
59
        ≈ 2048 × 0.32 ≈ 655
60

61
        即前 655 个维度对插值最敏感
62
        '''
63

64
    def visualization(self):
65
        """
66
        可视化
67
        """
68
        return """
69
        YaRN Frequency Adjustment 可视化:
70

71
        dim index (i) →
72
        [0    ...  655  ...  2048]
73
         |              |         |
74
        高频            中频       低频
75
         |              |         |
76
        λ短             λ中        λ长
77

78
        s_i:
79
        [1    ...  1    ...  s   ]
80
         |              |         |
81
        不变            插值       完全插值
82

83
        效果:
84
        → i < 655 的维度：波长 < β×L，保持不变
85
        → i >= 655 的维度：应用缩放
86

87
        这比朴素 PI 好，因为:
88
        → 保留了细粒度的短距离依赖
89
        → 只对长距离维度进行插值
90
        """

1
class NTKAwareScaling:
2
    """
3
    NTK-Aware Scaling
4

5
    基于神经切核理论的 RoPE 扩展方法
6

7
    博客: "NTK-Aware Scaled RoPE Allows LLM
8
          to Have Extended (8K+) Context Size"
9
    """
10

11
    def ntk_concept(self):
12
        """
13
        NTK 概念
14

15
        Neural Tangent Kernel 描述了神经网络在无穷宽极限下的行为。
16
        在这个视角下，位置编码可以看作定义了 token 之间的"相似性"。
17
        """
18
        return """
19
        NTK 视角下的 RoPE:
20

21
        每个位置 m 的 query 和 key 定义了一个"感受野":
22

23
        Q(m) 和 K(m) 是位置 m 的特征表示。
24

25
        Attention(Q(m), K(n)) 测量位置 m 和 n 的相关性。
26

27
        当我们改变位置编码时，实际上是在改变这个相关性函数。
28

29
        NTK 理论告诉我们:
30
        → 改变特征定义会影响模型学习的所有东西
31
        → 需要小心地改变，以保持有用的结构
32
        → "自然"的改变比"人为"的改变更好
33
        """
34

35
    def ntk_aware_idea(self):
36
        """
37
        NTK-Aware 的核心思想
38
        """
39
        return """
40
        朴素 PI 的问题 (NTK 视角):
41

42
        θ_i' = θ_i × s  (所有维度相同缩放)
43

44
        这相当于人为地改变了所有的频率关系。
45

46
        更好的方法:
47

48
        θ_i' = base^{(-2i/d)} × α_i
49

50
        其中 α_i 选择使得:
51
        1. 新的波长覆盖我们想要的范围
52
        2. 高频维度变化小
53

54
        NTK-Aware 公式:
55
        θ_i' = base^{(-2i/d)} / (s + (1-s) × λ_i / λ_max)
56

57
        效果:
58
        → 高频 (λ_i 小): θ_i' ≈ θ_i
59
        → 低频 (λ_i 大): θ_i' ≈ θ_i / s
60
        → 平滑过渡
61
        """
62

63
    def formula(self):
64
        """
65
        具体公式
66
        """
67
        return '''
68
        # NTK-Aware RoPE Scaling
69

70
        设:
71
        - s: 目标缩放因子 (如 2, 4, 8)
72
        - L: 原始最大长度
73
        - L': 目标最大长度 = L × s
74

75
        NTK-Aware 缩放:
76
        θ_i' = θ_i × r_i
77

78
        其中:
79
        r_i = 1 - (i / (d/2)) × (1 - 1/s)
80

81
        等价地:
82
        r_i = (i/s' + (d/2 - i)) / (d/2)
83
        其中 s' = s^(2i/d)
84

85
        这是一个非均匀缩放:
86
        → i=0 (最高频): r_i = 1
87
        → i=d/2 (最低频): r_i = 1/s
88
        → 平滑过渡
89
        '''

1
class NTKVsYaRN:
2
    """
3
    NTK-Aware vs YaRN
4
    """
5

6
    def similarities(self):
7
        """
8
        相似之处
9
        """
10
        return {
11
            "non_uniform": "都使用非均匀缩放",
12
            "high_freq_preserved": "都保留高频维度",
13
            "low_freq_scaled": "都对低频维度缩放",
14
            "smooth_transition": "都有平滑的过渡",
15
        }
16

17
    def differences(self):
18
        """
19
        区别
20
        """
21
        return {
22
            "origin": {
23
                "ntk": "基于神经切核理论的分析",
24
                "yarn": "基于实验观察和直觉",
25
            },
26
            "scaling_function": {
27
                "ntk": "r_i = 1 - (i/(d/2)) × (1-1/s)",
28
                "yarn": "基于阈值 β 和插值",
29
            },
30
            "parameters": {
31
                "ntk": "只有一个参数 s (缩放因子)",
32
                "yarn": "需要 β (阈值) 和 r (attention scale)",
33
            },
34
            "simplicity": {
35
                "ntk": "更简单直观",
36
                "yarn": "更复杂但可能更有效",
37
            },
38
        }
39

40
    def practical_comparison(self):
41
        """
42
        实践对比
43
        """
44
        return """
45
        实验结果 (LLaMA 7B, 从 2048 扩展):
46

47
        扩展到 8192 (4x):
48

49
        方法              Perplexity (PPL)   Pass@1
50
        ──────────────────────────────────────────
51
        原始 (无扩展)      OOM               N/A
52
        朴素 PI           45.2              32.1%
53
        NTK-Aware         12.8              68.5%
54
        YaRN              11.5              72.3%
55

56
        结论:
57
        → NTK-Aware 显著优于朴素 PI
58
        → YaRN 略优于 NTK-Aware
59
        → 但 NTK-Aware 更简单
60
        """

1
class LongRoPECore:
2
    """
3
    LongRoPE
4

5
    论文: "LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens"
6

7
    关键创新: 非均匀插值 + 渐进式微调
8
    """
9

10
    def key_innovations(self):
11
        """
12
        关键创新
13
        """
14
        return """
15
        LongRoPE 的三个关键创新:
16

17
        ┌─────────────────────────────────────────────────────┐
18
        │  1. 非均匀位置插值 (Non-Uniform Interpolation)        │
19
        │     → 不是均匀压缩所有位置                           │
20
        │     → 某些位置保持不变，某些位置压缩更多              │
21
        │     → 基于信息论分析                               │
22
        ├─────────────────────────────────────────────────────┤
23
        │  2. 位置去耦 (Positional Decoupling)                 │
24
        │     → 将位置分成多个范围                            │
25
        │     → 每个范围独立编码                              │
26
        │     → 允许超长上下文                                │
27
        ├─────────────────────────────────────────────────────┤
28
        │  3. 渐进式微调 (Progressive Fine-tuning)             │
29
        │     → 先在中等长度微调                              │
30
        │     → 再在更长长度微调                              │
31
        │     → 避免一次性大跳跃                             │
32
        └─────────────────────────────────────────────────────┘
33
        """
34

35
    def why_non_uniform(self):
36
        """
37
        为什么需要非均匀插值
38
        """
39
        return """
40
        均匀插值的问题:
41

42
        位置: [0, 1, 2, ..., 2048, ..., 2M]
43
              ├──────┤                    ├──────────┤
44
               压缩 1x                     压缩到 2048
45

46
        问题:
47
        → 位置 100 和 1000 被压缩相同的倍数
48
        → 但它们的信息密度可能不同
49
        → 均匀压缩可能不是最优的
50

51
        LongRoPE 的洞察:
52

53
        某些位置更重要（如文档开头、段落边界）
54
        应该给予更少的压缩或保持不变。
55

56
        其他位置（如长段落中间）可以更多压缩。
57
        """

1
class LongRoPEImplementation:
2
    """
3
    LongRoPE 实现
4
    """
5

6
    def non_uniform_mapping(self):
7
        """
8
        非均匀映射
9
        """
10
        return '''
11
        # LongRoPE 非均匀位置映射
12

13
        def longrope_position_mapping(
14
            original_pos: int,
15
            max_orig: int = 2048,
16
            max_new: int = 2048 * 1024,
17
        ) -> float:
18
            """
19
            将原始位置映射到新的位置索引
20
            """
21

22
            # 分为两个区域
23
            if original_pos < max_orig:
24
                # 区域 1: 保持不变（或很少压缩）
25
                ratio = original_pos / max_orig
26
                # 使用 sqrt 映射，保留细粒度
27
                return ratio * (max_orig ** 0.5)
28
            else:
29
                # 区域 2: 更激进的压缩
30
                excess = original_pos - max_orig
31
                # 对超长部分使用对数压缩
32
                return max_orig + math.log1p(excess) * 100
33

34
        # 更复杂的版本会根据位置重要性调整
35
        '''
36

37
    def positional_decoupling(self):
38
        """
39
        位置去耦
40
        """
41
        return """
42
        LongRoPE 的位置去耦:
43

44
        原始: 每个位置 m 有一个编码 R(m)
45

46
        LongRoPE:
47
        R(m) = combine(R_1(m mod M_1), R_2(m // M_2))
48

49
        其中:
50
        - M_1, M_2 是两个去耦的模数
51
        - R_1, R_2 是两个独立的 RoPE
52

53
        例子:
54
        m = 1,000,000
55
        M_1 = 2048, M_2 = 512
56

57
        R_1: R_1(1,000,000 mod 2048) = R_1(976)
58
        R_2: R_2(1,000,000 // 512) = R_2(1953)
59

60
        这允许用有限的维度表示超长的位置！
61
        """
62

63
    def progressive_finetuning(self):
64
        """
65
        渐进式微调
66
        """
67
        return '''
68
        # LongRoPE 渐进式微调策略
69

70
        目标: 从 2K 扩展到 2M (1000x)
71

72
        Stage 1: 扩展到 8K (4x)
73
        - 使用 NTK-Aware 或 YaRN
74
        - 微调 1000 steps
75

76
        Stage 2: 扩展到 32K (4x)
77
        - 基于 Stage 1 的权重继续
78
        - 微调 1000 steps
79

80
        Stage 3: 扩展到 128K (4x)
81
        - 继续微调
82
        - 使用更激进的插值
83

84
        Stage 4: 扩展到 512K (4x)
85
        - ...
86

87
        Stage 5: 扩展到 2M (4x)
88
        - 最终目标
89

90
        每次只扩展 4x，让模型逐步适应
91
        '''

1
class LongRoPEVsYaRN:
2
    """
3
    LongRoPE vs YaRN 对比
4
    """
5

6
    def comparison_table(self):
7
        """
8
        对比表
9
        """
10
        return {
11
            "max_context": {
12
                "yarn": "128K tokens",
13
                "longrope": "2M+ tokens",
14
            },
15
            "interpolation": {
16
                "yarn": "均匀或轻度非均匀",
17
                "longrope": "完全非均匀 + 去耦",
18
            },
19
            "finetuning": {
20
                "yarn": "一次性微调",
21
                "longrope": "渐进式微调",
22
            },
23
            "complexity": {
24
                "yarn": "中等",
25
                "longrope": "较高",
26
            },
27
            "quality": {
28
                "yarn": "良好 (up to 128K)",
29
                "longrope": "优秀 (up to 2M)",
30
            },
31
        }
32

33
    def when_to_use(self):
34
        """
35
        选择建议
36
        """
37
        return """
38
        选择 YaRN:
39
        → 扩展到 128K 以内
40
        → 想要简单的实现
41
        → 计算资源有限
42
        → LLaMA, Mistral 等主流模型
43

44
        选择 LongRoPE:
45
        → 需要超长上下文 (1M+)
46
        → 对质量要求极高
47
        → 有足够的微调资源
48
        → 研究或特殊应用场景
49
        """

1
class CLEXImplementation:
2
    """
3
    CLEX (Continuous Learned Exponential)
4

5
    论文: "CLEX: Continuous Length Extrapolation for
6
          Large Language Models"
7
    """
8

9
    def approach(self):
10
        """
11
        CLEX 方法
12
        """
13
        return """
14
        CLEX 的核心思想:
15

16
        用学习的连续函数替代离散的 RoPE。
17

18
        原始 RoPE:
19
        θ_i = base^{-2i/d}
20

21
        CLEX:
22
        θ_i = exp(-α_i × log(base))
23

24
        其中 α_i 是可学习的参数。
25

26
        优势:
27
        → 可以学习最优的频率分布
28
        → 更好地适应数据分布
29
        → 自然支持外推
30
        """

1
class FireAttnImplementation:
2
    """
3
    FireAttn
4

5
    论文: "FireAtten: Effective Long-Range Attention
6
          for Multi-Modal Large Language Models"
7
    """
8

9
    def approach(self):
10
        """
11
        方法
12
        """
13
        return """
14
        FireAttn 结合了:
15

16
        1. 分块注意力 (Chunked Attention)
17
           - 将长序列分成多个 chunk
18
           - 每个 chunk 内部使用 full attention
19
           - chunk 之间使用稀疏连接
20

21
        2. RoPE 扩展
22
           - 对不同 chunk 应用不同的 RoPE
23
           - 保持相对位置关系
24

25
        优势:
26
        → 计算效率高
27
        → 可以处理超长序列
28
        → 保持 RoPE 的位置信息
29
        """

1
class StreamingLLMImplementation:
2
    """
3
    StreamingLLM
4

5
    论文: "Efficient Streaming Language Models with
6
          Attention Sink"
7
    """
8

9
    def attention_sink(self):
10
        """
11
        Attention Sink 现象
12
        """
13
        return """
14
        StreamingLLM 的发现:
15

16
        LLM 在处理长序列时，会将大量注意力分配给某些"锚点" token，
17
        如 [CLS], 第一个 token，或者特殊 token。
18

19
        这些 token 被称为 "Attention Sink"。
20

21
        原因:
22
        → 模型学习将全局信息汇总到这些 token
23
        → 移除它们会导致性能急剧下降
24

25
        StreamingLLM 利用这一点:
26
        → 始终保留最近的几个 token
27
        → 保留 attention sink token
28
        → 丢弃中间的 token
29
        """
30

31
    def streaming_approach(self):
32
        """
33
        Streaming 方法
34
        """
35
        return """
36
        StreamingLLM 策略:
37

38
        窗口: [sink tokens] + [最近 4K tokens]
39

40
        tokens: [sink][tok1][tok2]...[tok_{N-4K}][tok_{N-3K}][tok_{N-2K}][tok_{N-K}][tok_N]
41
                   │                                            │
42
                   └────────── 丢弃 ──────────────────────────┘
43

44
        Attention 计算:
45
        → sink tokens 和 recent tokens 总是参与
46
        → 中间的 tokens 被丢弃（不计算 attention）
47

48
        优势:
49
        → 不需要 KV Cache 增长
50
        → 固定显存使用
51
        → 可以无限长 streaming
52

53
        局限:
54
        → 不能回顾中间的 token
55
        → 只适合不需要回顾的场景
56
        """

1
class FrameworkSupport:
2
    """
3
    主流框架支持
4
    """
5

6
    def huggingface(self):
7
        """
8
        Hugging Face Transformers
9
        """
10
        return '''
11
        # Hugging Face 的 RoPE 扩展支持
12

13
        from transformers import AutoConfig, AutoModelForCausalLM
14

15
        # 使用 rope_scaling 参数
16
        config = AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")
17
        config.rope_scaling = {
18
            "type": "linear",  # 或 "yarn"
19
            "factor": 2.0,    # 扩展倍数
20
        }
21

22
        # YaRN 支持
23
        config.rope_scaling = {
24
            "type": "yarn",
25
            "factor": 4.0,
26
            "original_max_position_embeddings": 4096,
27
            "attention_factor": 2.0,  # r 参数
28
            "beta_fast": 32,
29
            "beta_slow": 1,
30
        }
31
        '''
32

33
    def vllm(self):
34
        """
35
        vLLM
36
        """
37
        return '''
38
        # vLLM 的上下文扩展
39

40
        from vllm import LLM, SamplingParams
41

42
        # 创建模型，指定扩展长度
43
        llm = LLM(
44
            model="meta-llama/Llama-2-7b-hf",
45
            max_model_len=8192,  # 扩展到 8K
46
        )
47

48
        # 推理
49
        outputs = llm.generate(prompts, sampling_params)
50
        '''
51

52
    def llama_cpp(self):
53
        """
54
        llama.cpp
55
        """
56
        return '''
57
        # llama.cpp 的 RoPE 缩放
58

59
        # 使用 -c 参数指定上下文大小
60
        ./main -m model.gguf -c 8192 -n 256
61

62
        # 扩展因子
63
        ./main -m model.gguf -c 8192 --rope-scale 2.0
64
        '''

1
class FinetuningConfig:
2
    """
3
    微调配置建议
4
    """
5

6
    def yarn_finetuning(self):
7
        """
8
        YaRN 微调配置
9
        """
10
        return '''
11
        # YaRN 微调配置 (以 LLaMA 为例)
12

13
        # 1. 修改模型配置
14
        config.rope_scaling = {
15
            "type": "yarn",
16
            "factor": 4.0,  # 从 4K 扩展到 16K
17
            "original_max_position_embeddings": 4096,
18
            "attention_factor": 2.0,
19
            "beta_fast": 32,
20
            "beta_slow": 1,
21
        }
22

23
        # 2. 准备数据
24
        # 使用目标长度的数据
25
        # 16K 上下文的数据
26

27
        # 3. 训练参数
28
        training_args = TrainingArguments(
29
            per_device_train_batch_size=1,
30
            gradient_accumulation_steps=16,
31
            max_length=16384,  # 目标长度
32
            learning_rate=2e-5,
33
            num_train_epochs=3,
34
            warmup_ratio=0.1,
35
        )
36

37
        # 4. 开始微调
38
        trainer = Trainer(model=model, args=training_args, ...)
39
        trainer.train()
40
        '''
41

42
    def longrope_finetuning(self):
43
        """
44
        LongRoPE 渐进式微调
45
        """
46
        return '''
47
        # LongRoPE 渐进式微调
48

49
        # Stage 1: 8K
50
        train(max_len=8192, epochs=2)
51

52
        # Stage 2: 32K
53
        train(max_len=32768, epochs=2)
54

55
        # Stage 3: 128K
56
        train(max_len=131072, epochs=2)
57

58
        # Stage 4: 512K
59
        train(max_len=524288, epochs=2)
60

61
        # 每个 stage 基于上一个 stage 的权重继续
62
        # 数据量逐 stage 减少（越长数据越少）
63
        '''
64

65
    def data_preparation(self):
66
        """
67
        数据准备
68
        """
69
        return {
70
            "short_context": "1K-2K tokens，用于基础训练",
71
            "medium_context": "4K-16K tokens，用于中期扩展",
72
            "long_context": "32K+ tokens，用于最终扩展",
73

74
            "sources": [
75
                "书籍和长文章",
76
                "代码仓库",
77
                "科学论文",
78
                "长对话历史",
79
            ],
80

81
            "augmentation": [
82
                "文档拼接",
83
                "随机位置截断",
84
                "多样性采样",
85
            ],
86
        }