第9章变压器的进化：迈向效率和扩展性

“效率是通往智能的桥梁。” - 艾伦·图灵

2017年变压器出现之后，以BERT和GPT为代表的大型语言模型相继问世。它们凭借惊人的性能开启了人工智能的新时代。然而，在这些成功的背后，存在着变压器架构的基本局限性和克服这些局限性的努力。为了解决计算复杂度问题和长文本处理的限制，不断有改进和结构建议提出。特别是2019年之后，随着模型规模的急剧扩大，对效率的研究变得非常活跃。

各时期主要变化:

2019-2020：以减少复杂度为中心
2021-2022：以提高内存效率为中心
2023-2024：以扩展性和特殊目的（如伦理、开放模型等）为中心

本章将探讨变压器的局限性，并详细讨论为解决这些问题而提出的各种方法。

9.1 变压器的局限与挑战

挑战: 如何减少变压器模型的计算复杂度和内存使用量，以处理更长的上下文并训练更大的模型？

研究者的困惑: 虽然变压器模型的性能出色，但其计算成本巨大。特别是注意力机制具有与时序长度平方成比例的复杂度，这严重限制了模型的可扩展性。研究人员必须找到在保持注意力核心功能的同时提高计算效率的方法。不仅仅是减少模型的大小，而是在算法和硬件层面寻求创新性的解决方案。这就像在建造巨大的建筑时，还要减轻每块砖的重量和成本一样困难。

二次复杂度的注意力运算、有限的上下文长度以及内存效率问题成为了扩展模型的主要障碍。这些限制成为决定变压器发展方向的重要因素。

9.1.1 变压器架构的基本局限：计算复杂度

在变压器模型规模扩大的过程中，特别是与序列长度平方成比例的注意力运算复杂度成为一个大问题。

注意力运算复杂度分析:

\(Attention(Q, K, V) = softmax(\frac{QK^T}{\sqrt{d_k}})V\)

\(QK^T\) 计算: \(O(N^2d)\) (d: 嵌入维度)
Softmax 运算: \(O(N^2)\)
Softmax 结果与 V 的乘法: \(O(N^2d)\)

我们将通过实际代码来观察执行速度和内存使用量。

Code

!pip install dldna[colab] # in Colab
# !pip install dldna[all] # in your local

%load_ext autoreload
%autoreload 2

Code

from dldna.chapter_09.complexity_benchmark import measure_attention_complexity, plot_complexity_analysis, measure_attention_complexity_gpu

seq_lengths = [100, 500, 1000, 2000, 4000, 8000, 10000, 15000]

results = measure_attention_complexity(seq_lengths=seq_lengths)

print("\n=== Complexity Analysis of Attention Operation ===")
print("\nMemory usage and execution time by sequence length:")
print("Length\t\tMemory (MB)\tTime (seconds)")
print("-" * 40)
for seq_len, mem, time_taken in results:
    print(f"{seq_len}\t\t{mem:.2f}\t\t{time_taken:.4f}")

# Visualize with a graph
plot_complexity_analysis(results)


=== Complexity Analysis of Attention Operation ===

Memory usage and execution time by sequence length:
Length      Memory (MB) Time (seconds)
----------------------------------------
100     18.75       0.0037
500     96.58       0.0388
1000        317.00      0.1187
2000        1119.00     0.4228
4000        4188.14     1.6553
8000        16142.53        6.5773
10000       25039.31        10.2601
15000       55868.54        25.1265

在实际的变压器模型中，这一操作会在多个层中重复。当批量大小增加时，计算量也会进一步增加。

Code

# Compare theoretical complexity with actual measurements
print("\n=== Comparison of Theoretical Complexity and Actual Measurements ===")
base_seq = seq_lengths[0]
base_mem = results[0][1]
base_time = results[0][2]

print("\nTheoretical vs Actual Growth Rate (Base: First Sequence Length)")
print("Length      Theoretical(N²)      Actual Memory      Actual Time")
print("-" * 60)
for seq_len, mem, time_taken in results:
    theoretical = (seq_len/base_seq) ** 2
    actual_mem = mem/base_mem
    actual_time = time_taken/base_time
    print(f"{seq_len:6d}    {theoretical:10.2f}x    {actual_mem:10.2f}x    {actual_time:10.2f}x")


=== Comparison of Theoretical Complexity and Actual Measurements ===

Theoretical vs Actual Growth Rate (Base: First Sequence Length)
Length      Theoretical(N²)      Actual Memory      Actual Time
------------------------------------------------------------
   100          1.00x          1.00x          1.00x
   500         25.00x          5.15x          8.05x
  1000        100.00x         16.91x         32.49x
  2000        400.00x         59.71x        124.52x
  4000       1600.00x        223.34x        474.71x
  8000       6400.00x        860.92x       1882.04x
 10000      10000.00x       1335.43x       2976.84x
 15000      22500.00x       2979.67x       7280.40x

二次复杂度在像GPT-3这样的大型模型中尤其严重。它导致了处理长文档的限制、训练时批量大小的限制等许多问题。这成为开发高效注意力机制的主要动机。

解决变压器的二次复杂度问题的早期尝试主要朝着三个方向进行。

滑动窗口注意力

仅在固定大小的窗口内计算注意力。

Code

def sliding_window_attention(q, k, v, window_size):
    """Sliding window attention"""
    batch_size, seq_len, dim = q.shape
    attention_weights = np.zeros((batch_size, seq_len, seq_len))

    for i in range(seq_len):
        start = max(0, i - window_size // 2)
        end = min(seq_len, i + window_size // 2 + 1)
        scores = np.matmul(q[:, i:i+1], k[:, start:end].transpose(0, 2, 1))
        attention_weights[:, i, start:end] = softmax(scores, axis=-1)

    return np.matmul(attention_weights, v)

这种模式将复杂度降低到 \(O(N \cdot w)\)。 (w: 窗口大小)

稀疏注意力模式

稀疏注意力模式不是计算所有令牌对的关系，而是根据特定模式仅计算部分关系。例如，在由10个令牌组成的序列中，普通注意力会计算100个(10×10)的所有关系，而稀疏注意力只计算其中的一部分。

Code

def sparse_block_attention(q, k, v, block_size):
    """Block sparse attention
    Example: seq_len=8, block_size=2
    Process the sequence in 4 blocks of 2 tokens each
    Block 1 (0,1), Block 2 (2,3), Block 3 (4,5), Block 4 (6,7)
    """
    batch_size, seq_len, dim = q.shape  # e.g., (1, 8, 64)
    num_blocks = seq_len // block_size  # e.g., 8/2 = 4 blocks
    attention_weights = np.zeros((batch_size, seq_len, seq_len))
    
    for i in range(num_blocks):
        # e.g., when i=0, process Block 1 (0,1)
        start_q = i * block_size  # 0
        end_q = (i + 1) * block_size  # 2
        
        for j in range(num_blocks):
            # e.g., when j=0, attention with Block 1 (0,1)
            start_k = j * block_size  # 0
            end_k = (j + 1) * block_size  # 2
            
            # Calculate attention between tokens in Block 1 (0,1) and Block 1 tokens (0,1)
            scores = np.matmul(
                q[:, start_q:end_q],  # (1, 2, 64)
                k[:, start_k:end_k].transpose(0, 2, 1)  # (1, 64, 2)
            )  # Result: (1, 2, 2)
            
            # Store weights block by block
            attention_weights[:, start_q:end_q, start_k:end_k] = softmax(scores, axis=-1)
    
    # Generate the final context vectors
    return np.matmul(attention_weights, v)

低秩近似

低秩近似是将大矩阵表示为更小矩阵乘积的方法。例如，在一个包含10个标记的句子中，常规注意力机制计算10×10=100个关系，而低秩近似则用10×4和4×10两个矩阵的乘积表示（rank=4）。因此，只需80次运算而非100次即可获得相似的结果。

Code

def low_rank_attention(q, k, v, rank):
    """Low-rank attention
    Example: seq_len=10, dim=64, rank=16
    Project Q, K from 64 dimensions to 16 dimensions to reduce computation
    """
    batch_size, seq_len, dim = q.shape  # e.g., (1, 10, 64)
    
    # Create projection matrices to project from 64 dimensions to 16 dimensions
    projection_q = np.random.randn(dim, rank) / np.sqrt(rank)  # (64, 16)
    projection_k = np.random.randn(dim, rank) / np.sqrt(rank)
    
    # Project Q, K to 16 dimensions
    q_low = np.matmul(q, projection_q)  # (1, 10, 16)
    k_low = np.matmul(k, projection_k)  # (1, 10, 16)
    
    # Calculate attention in the lower dimension (operations on 10x16 matrices)
    attention = np.matmul(q_low, k_low.transpose(0, 2, 1))  # (1, 10, 10)
    attention_weights = softmax(attention, axis=-1)
    
    # Generate the final context vectors
    return np.matmul(attention_weights, v)  # (1, 10, 64)

这种方法可以将复杂度减少到\(O(N \cdot r)\)，其中r是用于近似的秩。我们将计算每种方法的效率。

Code

from dldna.chapter_09.attention_complexity_examples import calcualte_efficieny
calcualte_efficieny()

Original input shape: (2, 8, 4)

1. Sliding Window Attention
Output shape: (2, 8, 4)
Output of the first batch, first token: [-0.78236164  0.22592055 -1.03027549  1.13998368]

2. Block Sparse Attention
Output shape: (2, 8, 4)
Output of the first batch, first token: [-1.66095776  0.76700744 -0.45857165 -0.77422867]

3. Low-Rank Attention
Output shape: (2, 8, 4)
Output of the first batch, first token: [ 0.51121005  0.66772692 -0.77623488 -0.0323534 ]

Memory Usage Comparison (Relative Size):
Full Attention: 64
Sliding Window: 32
Block Sparse: 64
Low Rank: 32

然而，早期的尝试显示出了信息损失、实现复杂性、性能下降等局限性。Google专注于低秩近似，而Microsoft则更侧重于稀疏模式的开发。随后，这些早期方法发展成为混合方式，演进为同时利用稀疏性和低秩特性的方向。

9.1.2 变压器架构的基本限制：内存效率

另一个重要的限制是内存效率。特别是在大规模语言模型中，存在以下内存负担。

首先，KV缓存导致的内存负担。在自回归生成过程中，需要存储前一个时间步的Key和Value值，这与序列长度成线性增长。例如，在处理2048个token时，GPT-3每个层大约需要16MB的KV缓存。其次，反向传播过程中的内存需求。变压器存储每个注意力层的中间激活值（activation value）——即在注意力层中产生的中间计算结果（Q, K, V转换值、注意力分数、softmax输出等）。随着层数的增加，这种负担会急剧增大。对于BERT-large，在单个批次中大约需要24GB的内存。第三，注意力运算本身的内存使用量。注意力分数矩阵的大小与序列长度的平方成正比，这在处理长文档时成为一个严重的瓶颈。

为了解决这些内存问题，提出了梯度检查点、混合精度训练、FlashAttention等优化技术。

9.1.3 变压器发展的时代潮流及本章结构

为了克服9.1.1节和9.1.2节中讨论的变压器计算复杂性和内存效率限制，研究人员开发了多种提高效率和可扩展性的技术。这些技术使变压器模型更加强大和实用，并对整个深度学习领域产生了重大影响。

本章将如下表所示概述变压器发展的时代潮流，并介绍各时期的主要技术和模型。

表: 变压器发展的时代潮流、主要模型/技术、核心内容、深度学习DNA | 章节 | 时期 (大致) | 主要模型/技术 | 核心内容及说明 | 深度学习 DNA | |———|————-|————————|————————-|———————————————–| | 9.1 | 2017-2018 | Transformer | 克服了现有RNN、CNN局限性的Attention机制引入。
序列到序列模型的创新 | 注意力机制: 提出了一种新的方法来关注数据的重要部分 | | 9.2 | 2019-2020 | Performer, Sparse Transformer, Longformer
Reformer, BigBird | 降低计算复杂度的软件方法。
线性注意力: 注意力操作近似化 (Performer).
稀疏注意力: 只对部分标记对应用注意力 (Sparse Transformer, Longformer).
局部-全局注意力: 结合局部信息和全局信息 (Reformer, BigBird) | 高效的注意力机制: 努力在降低计算复杂度的同时保持注意力的优点。
长距离依赖性: 改进结构以更有效地处理长上下文 | | 9.3 | 2021-2022 | FlashAttention, MQA, GQA, PagedAttention, vLLM | 提高内存效率的硬件和软件方法。
FlashAttention: 利用GPU内存层次结构、分块处理。
MQA/GQA: 查询优化，共享Key/Value.
KV缓存优化: PagedAttention, vLLM | 硬件优化: 考虑到GPU内存结构的高效计算方法。
并行处理: 通过查询共享提高计算效率 | | 9.4 | 2022-2023 | Claude-2, LongLoRA, Constitutional AI, RLHF,
RLAIF, 层次注意力机制, 循环记忆 | 可扩展性和特殊目的架构。
长上下文: 层次注意力机制，循环记忆Transformer.
伦理/安全: 规则基注意力，基于强化学习的调整 | 长上下文: 模型结构的进化以处理更长的上下文。
微调: 调整模型以适应特定目的的方法 | | 9.5| 2022-2023 | 高效编码器 (基于FlashAttention) | 文本分类 (AG News)，FlashAttention, Pre-LN, 梯度检查点，混合精度训练 | 实现: 利用高效的编码器 | | 9.6| 2023 | Mistral, 高效解码器 (基于GQA, Sliding Window Attention) | Mistral模型分析: GQA, 滑动窗口注意力机制，RoPE, KV缓存等。
应用示例: 数字-文本转换，自然语言-SQL转换（代码生成），文本-代码生成。 | 实现: 高效的解码器架构 | | 9.7| 2024 | Gemma | 开放式模型以提高效率和可访问性 | 开放式模型: 提高研究和开发的可访问性 | | 9.8 | 2024 | Phi-3 | 小而高效的LLM | 实现: 强大的SLM(小型语言模型) | 本章的结构如下。

9.2节: 讨论减少注意力运算计算复杂度的软件方法（近似化、稀疏化、局部-全局注意力）。
9.3节: 检视提高内存效率的硬件和软件方法（FlashAttention、查询优化、KV缓存管理）。
9.4节: 讨论模型的可扩展性和特殊目的架构（长上下文处理、伦理/安全约束）。
9.5节: 实现高效的编码器模型，并通过AG新闻分类示例与其他类似模型进行效率比较。
9.6节: 实现一个简单的Mistral解码器模型，并提供应用示例。
9.7节: 介绍开放模型的代表gemma。
9.8节: 实现强大的SLM模型phi-3的一个简单模型，并检视应用示例。

9.2 复杂度降低：软件注意力优化（2019-2020）

9.2.1 初期方法：近似和稀疏化

从2019年到2020年，进行了各种尝试以减少变压器的计算复杂度。特别是由Google Research和DeepMind主导的这一时期的进展显著提高了注意力运算的效率。

9.2.1.1 线性注意力：Performer

2020年初，Google Research团队通过FAVOR+（Fast Attention Via positive Orthogonal Random features）成功将注意力复杂度从O(N²)降低到O(N)。FAVOR+是Performer模型的核心机制，是第一个使长序列处理在实际应用中成为可能的方法。

FAVOR+的核心思想源自核技巧。核技巧重新解释了softmax注意力如下：

\(Attention(Q,K,V) = softmax(\frac{QK^T}{\sqrt{d}})V\)

可以使用具有正值的核函数φ(x)将其近似为：

\(Attention(Q,K,V) ≈ \frac{\phi(Q)\phi(K)^TV}{\phi(Q)\phi(K)^T\mathbf{1}}\)

关键在于重新解释softmax注意力为分数形式，并通过使用核函数φ(x)来重组矩阵乘法的顺序，类似于将\((a \times b) \times c\) 改为 \(a \times (b \times c)\)。

Code

import numpy as np

def kernel_attention(Q, K, V, feature_dim=256): # Q: (seq_len, d_model) K: (seq_len, d_model) V: (seq_len, d_model)
    
    # 1. Generate random projection matrix
    projection = np.random.randn(Q.shape[-1], feature_dim) / np.sqrt(feature_dim)  
    # projection: (d_model, feature_dim)
    
    # 2. Project Q, K to lower dimension and apply ReLU
    Q_mapped = np.maximum(0, np.dot(Q, projection))  # phi(Q)
    # Q_mapped: (seq_len, feature_dim)
    K_mapped = np.maximum(0, np.dot(K, projection))  # phi(K)
    # K_mapped: (seq_len, feature_dim)
    
    # 3. Calculate numerator: phi(Q)phi(K)^TV
    KV = np.dot(K_mapped.T, V)  # (feature_dim, V_dim)
    # KV: (feature_dim, d_model)
    numerator = np.dot(Q_mapped, KV)  # (seq_len, V_dim)
    # numerator: (seq_len, d_model)
    
    # 4. Calculate denominator: phi(Q)phi(K)^T1
    K_sum = np.sum(K_mapped, axis=0, keepdims=True)  # (1, feature_dim)
    # K_sum: (1, feature_dim)
    denominator = np.dot(Q_mapped, K_sum.T)  # (seq_len, 1)
    # denominator: (seq_len, 1)
    
    # 5. Final attention output
    attention_output = numerator / (denominator + 1e-6)
    # attention_output: (seq_len, d_model)
    
    return attention_output

# Example usage
seq_len, d_model = 1000, 64
Q = np.random.randn(seq_len, d_model)
K = np.random.randn(seq_len, d_model)
V = np.random.randn(seq_len, d_model)

# Calculate attention with O(N) complexity
output = kernel_attention(Q, K, V)

print(output)

[[-0.00705502 -0.01553617 -0.01976792 ... -0.00906909  0.02983678
   0.0424082 ]
 [-0.00201811 -0.01741265 -0.00458378 ... -0.02578894  0.04247468
   0.03793401]
 [-0.01130314 -0.02011524 -0.00962334 ... -0.01348429  0.04382548
   0.01967338]
 ...
 [ 0.00180466 -0.01818735 -0.02244794 ... -0.01978542  0.03202302
   0.03887265]
 [-0.00421543 -0.01679868 -0.00537492 ... -0.00314385  0.05363415
   0.03304721]
 [ 0.00107896 -0.02042812 -0.01947976 ... -0.00557582  0.04534007
   0.04408479]]

FAVOR+引入的三种关键变化如下：

无偏估计： 使用正交随机特征计算注意力值时，使近似值的平均值与实际注意力值一致。
正值特征： 通过ReLU激活函数将所有特征值变为正值。这提高了数值稳定性。
正交投影： 使用正交矩阵将输入投影到低维空间。这样可以尽量保持向量间的距离和角度，从而最小化近似误差。

FAVOR+的处理步骤如下：

数据转换和降维： 将输入数据（Q, K, V）转换为较低维度的正交特征空间。
- 正交特征空间投影：每个输入向量被独立且均衡地转换。
- 降维：将高维输入压缩到低维。
- 信息保留：在减少维度的同时保持重要的关系信息。
- 维度变化：(序列长度 × 嵌入维度) → (序列长度 × 特征维度)
线性注意力运算： 在转换后的特征空间中高效地计算注意力。
- 特征空间中的运算：计算投影向量之间的相似度。
- 内存效率：按序列长度线性的内存使用（O(N × d)，N: 序列长度，d: 特征维度）。
- 计算优化：通过重新排列矩阵乘法顺序将复杂度从 O(N²) 降低到 O(N × d)。

Code

import numpy as np

def favor_plus_attention(q, k, v, feature_dim=256):
    """FAVOR+ attention implementation
    Args:
        q: Query tensor (batch_size, seq_len, d_model)
        k: Key tensor (batch_size, seq_len, d_model)
        v: Value tensor (batch_size, seq_len, d_model)
        feature_dim: The number of dimensions of the low-dimensional feature space
    """
    d_model = q.shape[-1]
    
    # 1. Generate an orthonormal random projection matrix
    random_matrix = np.random.randn(d_model, feature_dim)
    q_orth, _ = np.linalg.qr(random_matrix)
    projection = q_orth / np.sqrt(feature_dim)  # (d_model, feature_dim)

    # 2. Project Q, K to the low-dimensional feature space and apply ReLU
    q_prime = np.maximum(0, np.matmul(q, projection))  # (batch_size, seq_len, feature_dim)
    k_prime = np.maximum(0, np.matmul(k, projection))  # (batch_size, seq_len, feature_dim)

    # 3. Calculate linear-time attention
    # Use einsum to perform matrix multiplication while maintaining the batch dimension
    kv = np.einsum('bsf,bsd->bfd', k_prime, v)  # (batch_size, feature_dim, d_model)
    
    # Calculate the numerator
    numerator = np.einsum('bsf,bfd->bsd', q_prime, kv)  # (batch_size, seq_len, d_model)

    # Calculate the denominator (normalization term)
    k_sum = np.sum(k_prime, axis=1, keepdims=True)  # (batch_size, 1, feature_dim)
    denominator = np.einsum('bsf,bof->bso', q_prime, k_sum)  # (batch_size, seq_len, 1)

    # 4. Calculate the final attention output
    attention_output = numerator / (denominator + 1e-6)  # (batch_size, seq_len, d_model)

    return attention_output

# Example usage
batch_size, seq_len, d_model = 2, 100, 512
q = np.random.randn(batch_size, seq_len, d_model)
k = np.random.randn(batch_size, seq_len, d_model)
v = np.random.randn(batch_size, seq_len, d_model)

output = favor_plus_attention(q, k, v)
print("Output tensor shape:", output.shape)

Output tensor shape: (2, 100, 512)

FAVOR+具有以下优点。

将计算复杂度从O(N²)降低到O(N)。
在减少内存使用的同时，保持了注意力机制的核心功能——捕捉令牌之间的关系。
实现了对长序列处理的实用性。

数学基础

FAVOR+的数学基础在于Johnson-Lindenstrauss引理。其核心是，即使将高维数据投影到低维空间中，数据间的距离关系也能几乎保持不变。也就是说，1000维的数据减少到100维后，数据间的相对距离不会发生显著变化。

FAVOR+的成功推动了后续线性变换器、线性注意力变换器等各类线性注意机制的发展，并在长序列处理中发挥了重要作用。

9.2.1.2 稀疏注意：Sparse Transformer, Longformer

2019年，OpenAI通过Sparse Transformer引入了固定的稀疏模式。这种方法不是计算所有令牌对之间的关系，而是根据特定的模式仅计算部分关系。

Sparse Transformer的固定模式

Sparse Transformer使用两种主要的稀疏模式。

步进模式： 只与间隔一定的令牌进行注意力计算。
局部模式： 仅与固定大小窗口内的相邻令牌进行注意力计算。

这些模式可以用以下数学表达式表示。

\(Attention(Q,K,V) = softmax(\frac{QK^T \odot M}{\sqrt{d_k}})V\)

其中M是稀疏掩码矩阵，⊙表示逐元素乘法。掩码矩阵指示哪些令牌对应用注意力（1）或不应用注意力（0）。

这种方法提高了计算效率，但模式固定使得难以根据上下文灵活应对。

Longformer的局部-全局结合

2020年，Allen AI通过Longformer提出了更灵活的稀疏模式。Longformer采用了一种结合了局部注意和全局注意的混合方法。

局部注意： 每个令牌与其周围的w个令牌进行注意力计算。（滑动窗口方式）
全局注意： 特定的令牌（如[CLS]）与所有令牌进行注意力计算。

这种方法同时考虑了局部上下文和全局上下文，从而实现了更丰富的上下文理解。

翻译后的文本：

原始文本未提供。请提供需要翻译的韩语文本。

Code

import numpy as np


def longformer_attention(q, k, v, window_size=3, global_tokens=[0]):
    """Longformer attention implementation
    Args:
        q, k, v: (batch_size, seq_len, d_model)
        window_size: Size of the local attention window
        global_tokens: List of token indices to perform global attention on
    """
    batch_size, seq_len, d_model = q.shape
    attention_weights = np.zeros((batch_size, seq_len, seq_len))

    # 1. Local attention: sliding window
    for i in range(seq_len):
        # Calculate window range
        window_start = max(0, i - window_size)
        window_end = min(seq_len, i + window_size + 1)
        window_size_current = window_end - window_start
        
        # Calculate attention scores within the window
        scores = np.matmul(q[:, i:i+1], k[:, window_start:window_end].transpose(0, 2, 1))
        # scores: (batch_size, 1, window_size_current)
        
        attention_weights[:, i:i+1, window_start:window_end] = scores

    # 2. Global attention: specific tokens attend to all tokens
    for global_idx in global_tokens:
        # Calculate attention scores for global tokens
        scores = np.matmul(q[:, global_idx:global_idx+1], k.transpose(0, 2, 1))
        # scores: (batch_size, 1, seq_len)
        
        attention_weights[:, global_idx:global_idx+1, :] = scores
        attention_weights[:, :, global_idx:global_idx+1] = scores.transpose(0, 2, 1)

    # 3. Apply softmax (row-wise)
    attention_weights = np.exp(attention_weights) / np.sum(np.exp(attention_weights), axis=-1, keepdims=True)
    
    # 4. Calculate the final output by applying weights
    output = np.matmul(attention_weights, v)  # (batch_size, seq_len, d_model)
    
    return output

# Example usage
batch_size, seq_len, d_model = 2, 10, 64
q = np.random.randn(batch_size, seq_len, d_model)
k = np.random.randn(batch_size, seq_len, d_model)
v = np.random.randn(batch_size, seq_len, d_model)

output = longformer_attention(q, k, v, window_size=2, global_tokens=[0])

print(output)

[[[-0.72195324  0.03196266 -0.06067346 ...  0.57106283  1.31438
    0.63673636]
  [-1.72619367 -0.39122625  0.91285828 ... -1.4031466   1.2081069
    0.95934394]
  [ 0.07427921  0.42596224 -0.44545069 ...  0.154228    0.37435003
   -0.01884786]
  ...
  [ 1.26169539 -0.58215291  2.00334263 ...  1.15338425  0.31404728
   -1.33672458]
  [ 0.96005607  0.39904084  0.5703471  ... -0.2168805   0.93570179
    0.05680507]
  [ 0.61648602 -0.12874142  1.09736967 ...  0.32421211  1.23082505
    0.4141766 ]]

 [[ 0.92762851  0.26334678 -0.81047846 ... -0.19186621  0.42534117
    0.57313974]
  [ 1.01307261  0.61571205 -1.26925081 ... -0.56016688 -0.19707427
    2.49452497]
  [-1.0071559   2.81291178  2.5010486  ...  1.63559632 -0.60892113
   -1.40952186]
  ...
  [-1.96615634  1.85881047  0.19361453 ...  1.21044747 -0.00772792
   -0.68961122]
  [ 0.09090778  1.94770672 -0.990489   ... -0.09841141  0.65195305
    0.11634795]
  [-2.43256801  1.66319642  0.23557316 ...  2.39325846  0.8750332
    0.66295002]]]

块稀疏矩阵运算优化

为了有效地实现Longformer的混合方法，需要进行块稀疏矩阵运算优化。

块级处理： 通过连续的内存访问提高缓存效率。
自定义CUDA内核： 优化针对稀疏模式的并行处理。
动态负载平衡： 根据每个块的计算量分配任务。

基于稀疏模式的方法虽然将复杂度降低到O(N log N)或O(N)，但实施复杂性和硬件优化难度较高。

9.2.3 局部-全局注意力：解决长距离依赖问题

2020年初，Google Research和Allen AI提出了结合局部-全局注意力的混合方法。这是为了解决线性注意力的信息损失和稀疏模式实施复杂性的问题。

9.2.3.1 Reformer: LSH注意力

Reformer使用局部敏感哈希（Locality-Sensitive Hashing, LSH）高效地将相似向量聚类。LSH的核心原理如下。

\(h(x) = \text{argmax}( [xR; -xR] )\)

其中，R是随机投影矩阵，相似的向量更可能具有相同的哈希值。Reformer遵循以下步骤：

使用哈希函数将查询向量分配到桶中。
只计算同一桶内的键向量之间的注意力。
将复杂度从O(N²)减少到O(N log N)。

这种方法在处理长序列时效率很高，但由于哈希冲突可能会导致信息损失。

9.2.3.2 BigBird: 局部、全局和随机注意力的组合

BigBird为了克服Reformer的局限性，结合了三种注意模式。

局部窗口： 计算与w个相邻标记的注意力以捕捉局部上下文。
全局标记： g个特殊标记对整个序列进行注意力保持全局信息。
随机块： 通过计算与r个随机选择的标记的注意力来捕捉不同距离的关系。

这种混合策略可以用以下公式表示：

\(Attention(Q,K,V) = softmax(\frac{QK^T \odot (M_{local} + M_{global} + M_{random})}{\sqrt{d_k}})V\)

其中，M是每个掩码矩阵。该结构在实现O(N)复杂度的同时保持了BERT级别的性能。

混合模式的影响

BigBird的成功证明了局部-全局方法的潜力，并对现代变压器模型产生了重大影响。

计算效率：
- 通过选择性注意力减少了复杂度。
- 最优化了GPU内存使用量。
- 实现了并行处理能力。
模型性能：
- 平衡了局部细节信息和全局上下文信息。
- 长距离依赖捕捉能力得到提升。
- 在各种任务中表现出稳定的性能。
实际应用：
- 影响了GPT-3的稀疏变换器结构。
- 推动了PaLM多查询注意力的发展。
- 被用于Anthropic Claude的宪法AI实现。这种混合方法后来成为Longformer、ETC等多种模型的基础。特别是在文档分类、问答等长文档处理任务中取得了巨大成功。然而，内存使用量和计算效率的问题仍然存在。尤其是在大规模语言模型中，GPU内存使用优化成为一个新的挑战，这在9.3节中将讨论如何改进内存效率。

9.3 内存效率：硬件与软件的结合 (2021-2022)

从2021年到2022年，重点放在了提高Transformer的内存效率上。特别是针对GPU内存层次结构的优化和注意力计算的有效实现受到了广泛关注。这一时期的进展使得大规模语言模型的实际应用成为可能。

9.3.1 FlashAttention：利用GPU内存层次结构进行注意力优化

2022年，斯坦福大学的Tri Dao研究团队提出了FlashAttention，考虑了GPU内存层次结构。这是一个以硬件为中心的改进，从根本上重新设计了注意力计算的内存访问模式。FlashAttention显著提高了Transformer模型的学习和推理速度，特别是在处理长序列时，对大规模语言模型的发展做出了重要贡献。2023年发布的FlashAttention v2进一步优化了原始的FlashAttention，实现了2-4倍的速度提升。

9.3.1.1 GPU内存结构与IO优化

FlashAttention的优势在于它明确考虑了GPU的内存层次结构。GPU中有两种内存：容量大但速度慢的HBM（高带宽内存）和容量小但速度快的SRAM。HBM虽然容量大，但访问速度较慢；SRAM虽然容量较小，但访问速度非常快。FlashAttention利用了这一特点。

最小化HBM与SRAM之间的数据传输： 在传统的注意力机制中，计算查询和键的点积后，需要将整个大的注意力分数矩阵存储在HBM中。这消耗了大量的内存带宽，并导致速度下降。FlashAttention尽量减少了这种不必要的数据传输。
不将大规模中间结果（注意力分数矩阵）存储在HBM中： FlashAttention不在HBM中存储中间计算结果，而是在SRAM中保持并执行所需的运算。
逐步在SRAM中计算softmax： 不是针对整个注意力分数矩阵一次性进行softmax运算，而是分块进行softmax计算，并累积结果。这样可以减少将中间值写入和读出HBM的过程。

这种硬件感知设计大幅减少了内存访问次数。

9.3.1.2 分块和块处理

为了实现内存优化，引入了分块（Tiling）技术。分块是一种硬件优化技术，它将大的矩阵分割成适合SRAM的小块进行处理。

将输入矩阵(Q, K, V)分成符合SRAM大小的块。
按块从HBM加载数据到SRAM中。
在SRAM内部按块执行注意力计算。
完成每个块的注意力运算后，仅将该块的结果（即对该块值的加权平均）存储在HBM中。不存储整个注意力分数。

这种块处理策略不仅最小化了内存带宽使用，还确保可以准确计算注意力结果。

9.3.1.3 FlashAttention v2：最大化硬件利用

FlashAttention v2保持了v1的基本理念，同时通过添加多个低级优化来最大化硬件利用率。与v1相比，实现了2-4倍的速度提升，特别是在处理长序列时表现出色。 * 内核融合: FlashAttention v2 将查询、键、值转换，注意力得分计算，softmax，加权平均计算等注意力机制的多个运算整合到一个 CUDA 内核中。通过这种方式，最小化了将中间结果存储在 HBM 中并重新读取的次数，从而减少了内存带宽使用量并提高了速度。 * 非顺序（Non-sequential）Attention Head 处理: 以前是按顺序处理 attention head，而 FlashAttention V2 在 GPU 资源允许的情况下，并行处理 attention head 以减少延迟。 * 缓存友好型内存布局: 设计了更适合 GPU 缓存行的数据结构，例如将数据存储为列优先（column-major）顺序。这减少了缓存未命中(cache miss)，并提高了数据访问速度。 * 线程级并行化: 优化使用 CUDA 线程内的 32 个线程尽可能并行处理注意力运算的各个部分。通过这种方式，充分利用了 GPU 的 SIMD（Single Instruction, Multiple Data）特性和并行处理能力，从而提高了计算速度。

这些综合优化使 FlashAttention v2 在特定环境中相比现有 PyTorch 注意力实现最多实现了 20 倍的内存效率提升和 2-4 倍的速度提升。FlashAttention 的成功展示了基于对硬件特性深入理解的算法设计的重要性，并成为后续 GPT-4, Claude 等大规模语言模型的核心技术。

FlashAttention 的官方实现以 NVIDIA CUDA 代码形式提供。在 PyTorch 中，可以通过 flash-attn 包使用，最新版本的 Hugging Face transformers 库中也已集成。

9.3.2 查询优化：改进注意力结构

2022年，Google Research 通过 PaLM 模型提出了 Multi-Query Attention (MQA)，以从软件设计方面提高内存效率。与 FlashAttention 的硬件中心优化不同，这是一种通过重新设计注意力结构本身来减少内存使用的方案。

9.3.2.1 多查询注意力（MQA）

MQA 的核心是改变设计，使所有注意力头共享相同的 Key 和 Value。

Key, Value 共享:
- 所有头共享一个 K, V 矩阵。
- KV 缓存大小减少为头数的倍数。（例如：如果头数为 8，则 KV 缓存大小减少到 1/8）
- 大大减少了内存带宽使用量。
Query 分离:
- Query 按每个头独立生成。
- 每个头仍然可以学习不同的上下文。
- 计算复杂度没有显著增加。

Code

import numpy as np

def multi_query_attention(q, k, v, num_heads):
    """Multi-Query Attention implementation
    Args:
        q: (batch_size, seq_len, d_model)
        k: (batch_size, seq_len, d_model)
        v: (batch_size, seq_len, d_model)
        num_heads: Number of heads
    """
    batch_size, seq_len, d_model = q.shape
    head_dim = d_model // num_heads

    # 1. Convert K, V to single matrices shared by all heads
    k_shared = np.dot(k, np.random.randn(d_model, d_model))  # (batch_size, seq_len, d_model)
    v_shared = np.dot(v, np.random.randn(d_model, d_model))  # (batch_size, seq_len, d_model)

    # 2. Generate Q differently for each head
    q_multi = np.dot(q, np.random.randn(d_model, num_heads * head_dim))  # (batch_size, seq_len, num_heads * head_dim)
    q_multi = q_multi.reshape(batch_size, seq_len, num_heads, head_dim)  # (batch_size, seq_len, num_heads, head_dim)

    # Transform k_shared to head_dim size
    k_shared = np.dot(k_shared, np.random.randn(d_model, head_dim))  # (batch_size, seq_len, head_dim)
    
    # 3. Calculate attention scores
    scores = np.matmul(q_multi, k_shared.reshape(batch_size, seq_len, head_dim, 1))
    # scores: (batch_size, seq_len, num_heads, 1)

    # 4. Apply softmax
    weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
    # weights: (batch_size, seq_len, num_heads, 1)

    # 5. Multiply V with weights
    v_shared = np.dot(v_shared, np.random.randn(d_model, head_dim))  # Transform V to head_dim as well
    v_shared = v_shared.reshape(batch_size, seq_len, 1, head_dim)
    output = np.matmul(weights, v_shared)
    # output: (batch_size, seq_len, num_heads, head_dim)

    # 6. Concatenate heads and transform output
    output = output.reshape(batch_size, seq_len, num_heads * head_dim)
    output = np.dot(output, np.random.randn(num_heads * head_dim, d_model))
    # output: (batch_size, seq_len, d_model)

    return output

# Example usage
batch_size, seq_len, d_model = 2, 100, 512
num_heads = 8

q = np.random.randn(batch_size, seq_len, d_model)
k = np.random.randn(batch_size, seq_len, d_model)
v = np.random.randn(batch_size, seq_len, d_model)

output = multi_query_attention(q, k, v, num_heads)
print("Output tensor shape:", output.shape)

/tmp/ipykernel_304793/3750479510.py:30: RuntimeWarning: overflow encountered in exp
  weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)
/tmp/ipykernel_304793/3750479510.py:30: RuntimeWarning: invalid value encountered in divide
  weights = np.exp(scores) / np.sum(np.exp(scores), axis=-1, keepdims=True)

Output tensor shape: (2, 100, 512)

9.3.2.2 分组查询注意力（GQA）

2023年初，Meta AI 提出了 GQA(Grouped-Query Attention)以弥补 MQA 的不足。GQA 将头部分组成群，每个群体共享 K、V，采用了一种折中的方法。

基于组的设计:
- 多个查询头可以共享一个 KV 对。
- 可以通过调整组的大小来平衡内存使用量和模型性能。
- 比 MQA 具有更丰富的表达能力。
高效实现:
- 优化了按组并行处理。
- 使用了缓存友好的内存访问方式。
- 提高了推理时的处理速度。

9.3.2.3 MQA vs. GQA vs. 多头注意力

像 MQA 和 GQA 这样的查询优化结构提供了以下权衡。

结构	内存使用量	表达能力	处理速度	实现复杂度
多头注意力	N × H	高	慢	低
GQA	N × G	中等	中等	中等
MQA	N	低	快	低

(N: 序列长度, H: 头数, G: 组数)

这些结构在 LLaMA、PaLM、Claude 等现代大规模语言模型中被广泛采用，尤其是在处理长序列时大大提高了内存效率。

9.3.3 KV 缓存管理和优化

2022年下半年，DeepMind 和 Anthropic 以及 vLLM 开发团队意识到了在大型语言模型的推理过程中管理 KV 缓存的重要性。他们提出了软件和系统级别的内存优化策略，以补充 FlashAttention 的硬件中心方法和 MQA/GQA 的结构化方法。这尤其在处理长对话、生成长文档时以及需要高吞吐量的情况下非常重要。

9.3.3.1 PagedAttention & vLLM: 操作系统分页概念

PagedAttention 和实现它的 vLLM 是从操作系统虚拟内存和分页概念中获得灵感，用于高效管理 KV 缓存的技术。

现有 KV 缓存的问题

内存浪费: 随着序列长度线性增加，KV 缓存占用大量内存空间。特别是，在批处理（batch processing）时，如果序列长度各不相同，则必须根据最长的序列分配内存，导致严重的浪费。
内存碎片化: 当 KV 缓存在内存中不连续地分配时，即使有空闲空间也无法利用，从而出现外部碎片化(external fragmentation)问题。
不支持动态序列长度: 在生成过程中难以高效处理 KV 缓存大小的动态变化。

PagedAttention 的核心思想

基于块的内存分配 (Block-Based Memory Allocation):
- 将 KV 缓存划分为固定大小的块（block）。(类似于操作系统将内存划分为页面)
- 每个块存储多个令牌的键和值。
- 即使块在物理上不连续也没有关系。(逻辑上是连续的)
块表 (Block Table):
- 管理每个序列的逻辑块和物理块之间的映射。（类似于操作系统的页表）
- 当生成新的令牌时，分配一个空块，并在块表中添加映射信息。
写时复制 (CoW) 支持（可选）:
- 如果多个序列共享相同的提示符（例如，在波束搜索中），则不复制块以节省内存。
- 只有当块内容更改时才分配新的块。

PagedAttention的优点

提高内存效率: 仅分配所需数量的块，从而减少内存浪费。
降低内存碎片化: 按块管理内存，从而缓解外部碎片化问题。
动态序列处理: 即使在生成过程中KV缓存的大小增加或减少，也能灵活应对。
高吞吐量 (Throughput): 在vLLM等系统中使用PagedAttention进行高效批处理，并实现高吞吐量。

vLLM: 利用PagedAttention的高性能推理引擎

vLLM是一个开源库，通过使用PagedAttention作为核心技术，大幅提高了大规模语言模型的推理速度和吞吐量。

连续批处理 (Continuous Batching): 立即处理新到达的请求，并立即移除已完成的请求以提高GPU利用率。
CUDA内核优化: 使用针对PagedAttention操作优化的CUDA内核，提高内存访问速度。

9.3.3.2 连续批处理与高效的缓存策略 (Continuous Batching & Efficient Caching)

连续批处理（Continuous Batching）是大规模语言模型服务中用于最大化吞吐量的关键技术。PagedAttention和vLLM高效地支持了连续批处理。

现有批处理的缺点

降低GPU利用率: GPU必须等待批次中最长序列处理完毕。
较长的延迟 (Latency): 新请求必须等待前一批次完成。

连续批处理的核心思想

迭代批处理 (Iterative Batching): 动态地向当前正在处理的批次中添加新请求。
请求级调度 (Request-Level Scheduling): 个别地调度每个请求，并立即返回已完成请求的结果。

连续批处理 + PagedAttention

PagedAttention以块为单位管理KV缓存，因此在连续批处理环境中可以高效地进行内存管理。
当新请求到达时，分配一个空块并将其添加到KV缓存中即可。
请求完成后，释放相应块以返回内存。

高效的缓存策略

结合连续批处理，可以通过以下缓存策略进一步提高内存效率：

LRU（最近最少使用）缓存: 选择最长时间未使用的KV缓存块作为替换目标。
热/冷分离: 经常使用的KV缓存块（“hot”）存储在GPU内存中，不经常使用的块（“cold”）存储在CPU内存中。
预取 (Prefetching): 预先加载预计下次需要的KV缓存块，以减少内存访问延迟。这些技术对于将大规模语言模型部署到实时服务中，并实现高吞吐量和低延迟至关重要。

摘要

PagedAttention: 以块为单位管理KV缓存，提高内存效率并支持动态序列长度。
vLLM: 利用PagedAttention提供高性能推理的开源库。
连续批处理: 动态地将请求添加/删除到批处理中，以最大化GPU利用率和吞吐量。
高效的缓存策略: 通过LRU、热/冷分离、预取等方法提高内存访问速度。

这些技术对于将大规模语言模型部署到实际服务中，并实现高吞吐量和低延迟至关重要。

9.4 可扩展性和特殊目的架构（2023-2024）

自2023年起，变压器模型的发展不仅超越了效率，还迎来了探索可扩展性和符合特定目的的架构的新阶段。在前期（第9.2节、第9.3节）积累的基础技术，如FlashAttention、MQA/GQA、高效的KV缓存管理等，为解决更大更复杂问题奠定了基石。基于这些技术进步，研究人员开始不仅增加模型的规模，还设计针对特定问题领域的优化结构，控制模型的行为，并开发能够处理各种形式数据的变压器模型。

9.4.1 长上下文处理：上下文长度的扩展

在对话型AI、文档摘要、代码生成、科学研究等多个领域中，理解和处理长上下文（Long Context）的能力至关重要。早期的变压器模型（第9.1节）主要局限于处理512或1024个标记长度的上下文，但在2023年前后，出现了能够处理100K（10万）、甚至1M（100万）以上标记长度上下文的模型，实现了重大突破。

9.4.1.1 分层注意力、循环记忆变压器

有效处理长上下文的核心技术大致可以分为注意力机制的优化、分层/递归处理和引入内存机制三类。

高效的注意力机制 (Efficient Attention Mechanisms)

变压器的基本注意力机制具有与时序长度平方成比例的计算复杂度（O(N²)），因此在处理长时序时效率低下。因此，第9.2节中讨论的各种高效注意力技术被用作长上下文模型的关键组成部分。
- 线性注意力 (Linear Attention): 将注意力操作的复杂度减少到O(N)的方法。
  - Performer: 使用FAVOR+（Fast Attention Via positive Orthogonal Random features）算法，不显式计算注意力矩阵，而是通过核函数的期望值进行近似。（第9.2.1.1节）
  - Linformer: 通过低秩逼近(low-rank approximation)将注意力矩阵表示为较小矩阵的乘积，从而减少计算量。
- 稀疏注意力 (Sparse Attention): 不是对所有令牌对都进行注意力计算，而是根据特定模式仅对部分令牌对应用注意力的方法。（第9.2.1.2节）
  - Sparse Transformer: 使用固定模式(fixed pattern)以减少注意力计算量。结合使用步幅(stride)模式和局部(local)模式。
  - Longformer: 结合滑动窗口(sliding window)注意力和全局注意力(global attention)，同时考虑局部信息和全局信息。
- Reformer : 第9.2.3.1节介绍的LSH（Locality-Sensitive Hashing）注意力通过哈希查询和键向量，将相似的向量分配到相同的桶中，并仅在同一桶内计算注意力。
- BigBird: 结合本地、全局和随机注意力的混合方法。（第9.2.3.2节）
分层注意力 (Hierarchical Attention) 层次化注意力是一种将输入序列分为多个层级进行处理的方法。每个层级具有不同的范围（scope）和分辨率（resolution），较低的层级处理局部（local）上下文，而较高的层级处理全局（global）上下文。

工作原理:
1. 将输入序列分割为小段（segment）或块（block）。
2. 在每个段内执行局部注意力（例如：滑动窗口注意力），以提取局部信息。
3. 生成代表每个段的表示（例如：每个段的平均池化，CLS令牌，或学习到的表示向量）。
4. 对段表示执行全局注意力（global attention），以捕捉长距离依赖性（long-range dependency）。
5. 根据需要添加更多层级，以便处理更广泛的上下文。
优点:
- 计算复杂度降低: 直接对整个序列进行注意力操作所需的计算量大大减少。
- 捕获多级上下文信息: 同时考虑局部和全局信息，生成更丰富的上下文表示。
- 并行处理方便: 可以独立处理每个段，使得并行处理更加容易。
示例:
- Longformer: 使用结合了滑动窗口注意力（局部）和全局注意力（部分令牌）的层次结构。
- ETC (Extended Transformer Construction): 扩展了Longformer，以更好地处理更长的上下文。
- H-Transformer (Hierarchical Transformer): 使用多个层级的注意力来分层建模上下文。

递归记忆变压器

递归记忆变压器将RNN（递归神经网络）的思想引入变压器中，通过“记忆”形式保持先前序列的信息，并在处理当前序列时利用这些记忆。
- Transformer-XL (2019): 引入了相对位置编码和基于段的递归机制，使得能够建模超出固定长度上下文窗口的长距离依赖性。
  - 相对位置编码: 编码令牌之间的相对距离，而不是绝对位置信息。这有助于模型更好地泛化到更长的序列。
  - 基于段的递归: 缓存先前序列段的隐藏状态，并在处理当前段时利用这些缓存的信息。这样，当前段可以引用前一段的上下文。
- 压缩变压器 (2019): 扩展了Transformer-XL，通过将过去的隐藏状态以压缩记忆的形式存储并使用这些信息来处理更长的上下文。
- 压缩记忆: 将较旧的信息进行压缩，并存储在压缩记忆中，该记忆可以被查询以计算额外的注意力。

内存机制:
- External Memory: 引入 Key-Value 内存，Key 与 query 计算 attention，获取最相关的 value，value 提供信息摘要。
Attention Sink, StreamingLLM:
- Attention Sink: 在长文本生成中，最初的几个 token（Sink token）对所有 token 进行 Attend。这相当于一种全局 token 的作用。
- StreamingLLM: 利用 Attention Sink 的思想，管理 KV 缓存的技术。这种方法特别适用于处理无限长度文本的流式场景。

9.4.1.2 Claude-2, LongLoRA

Claude-2 (Anthropic): 能够处理超过100K token 上下文的对话型 AI 模型。Claude-2 使用结合了 多尺度注意力（multi-scale attention） 和 自适应压缩（adaptive compression） 的改进方法来有效处理长上下文。
- 多尺度注意力: 使用不同大小的窗口同时考虑局部信息和全局信息。例如，小窗口用于理解周围单词的关系，大窗口用于把握段落或文档整体的上下文。
- 自适应压缩: 根据输入序列的重要性动态调整压缩率以最小化信息损失。例如，重要的句子较少压缩，不那么重要的句子较多压缩。
LongLoRA: 通过少量资源对已学习模型进行 fine-tuning 的方法来增加上下文长度。改进了计算成本较低的 LoRA 以适应长上下文处理。
- Shift Short Attention: 对短上下文执行减少运算量的高效 attention。从现有注意力机制中减少不必要的计算，提高效率。
- Grouped Query, key, value Projections: 利用 MQA/GQA 减少内存使用。(9.3.2节)
GPT-4, Gemini: (虽然确切架构未公开) 但已知可以处理超过10万 token 的上下文。可能结合了上述多种技术。
LongNet: 提出了一种使用稀疏注意力（Dilated Attention）的 Transformer 来处理10亿 token。稀疏注意力是在窗口内跳跃式选择 token 进行注意力计算的方式。（类似于 CNN 中的空洞卷积）通过这种方式，可以有效增加感受野同时减少计算量。

这些长上下文处理技术被应用于法律文件分析、学术论文理解、长时间对话记录处理、长篇小说生成等多个领域。

9.4.2 伦理/安全约束：Constitutional AI

自2022年底以来，随着大规模语言模型（LLM）的迅速发展，对其伦理和社会影响的担忧也随之增加。特别是，LLM生成有害、歧视性或误导性内容，以及泄露个人信息等问题引起了严重关注。为了解决这些问题，人们认为不仅需要对模型输出进行后过滤，还需要在模型的内部工作方式本身中整合伦理约束。

2023年中期，Anthropic提出了一种新的解决方案——“Constitutional AI”。Constitutional AI的目标是设计出能够根据明确的“原则（constitution）”行事的模型，而不是重复学习数据中的偏见或有害性。

9.4.2.1 基于规则的注意力

Constitutional AI的核心思想如下：

明确的宪法（Constitution）定义

人们直接编写出模型应遵循的理想行为原则，即“宪法”。这些宪法由防止危害、歧视、侵犯个人隐私等规则组成。
- 示例:
  - “尊重用户的个人信息，未经同意不得收集或分享。”
  - “不发表基于种族、性别、宗教等方面的歧视性或偏见言论。”
  - “不生成暴力或仇恨内容。”
  - “不要提供与事实不符的信息，或以可能导致误解的方式回应。”
监督学习（Supervised Learning）阶段
- 批评和修改（Critique and Revision）: LLM首先以常规方式生成响应。然后，一个独立的“批评模型（critique model）”根据宪法评估该响应，并在发现违规时进行修正。
- 细化（Refine）: 批评模型详细描述响应是否违反了给定的原则、如何违反以及如何修改。
- 数据增强（Data Augmentation）: 将原始响应和修改后的响应配对生成新的学习数据。
- 监督学习（Supervised Fine-tuning）: 使用这些数据对LLM进行微调。模型通过批评模型的反馈，学会生成符合宪法的响应。
强化学习（Reinforcement Learning）阶段
- 偏好模型（Preference Model）: 训练一个独立的模型来判断两个响应中哪一个更符合宪法。
- 基于人类反馈的强化学习（RLHF）: 通过人类的反馈改进偏好模型。
- 基于AI反馈的强化学习（RLAIF）: 使用偏好模型评估LLM的行为，并以加强符合宪法的行为的方式进行学习。

Constitutional AI的优点 * 透明度（Transparency）: 模型的行为原则被明确定义，因此可以轻松理解和跟踪模型的决策过程。 * 可控性（Controllability）: 可以通过修改或添加宪法来相对容易地控制模型的行为。 * 泛化能力（Generalization）: 不仅能应对特定类型的有害内容，还能处理各种类型的问题。 * 可扩展性（Scalability）: 可以在没有人类干预的情况下使用AI系统对模型进行训练。 (RLAIF)

宪法AI的实现（概念示例）

Code

import numpy as np

class ConstitutionalAttention:
    def __init__(self, rules, embedding_dim=64):
        """Embed ethical rules and integrate them into attention
        Args:
            rules: List of ethical rules
            embedding_dim: Dimension of rule embeddings
        """
        self.rules = rules
        # Convert rules to embedding space
        self.rule_embeddings = self._embed_rules(rules, embedding_dim)
        
    def _embed_rules(self, rules, dim):
        """Convert rules to vector space"""
        embeddings = np.random.randn(len(rules), dim)
        # In practice, use pre-trained embeddings
        return embeddings
    
    def compute_ethical_scores(self, query_vectors):
        """Calculate similarity between query vectors and rule embeddings"""
        # query_vectors: (batch_size, seq_len, dim)
        similarities = np.dot(query_vectors, self.rule_embeddings.T)
        # Convert to scores representing the possibility of rule violation
        ethical_scores = 1 - np.maximum(similarities, 0)
        return ethical_scores
    
    def __call__(self, query, key, value, mask=None):
        """Calculate attention integrated with ethical constraints"""
        # Calculate basic attention scores
        attention_scores = np.dot(query, key.transpose(-2, -1))
        
        # Calculate ethical constraint scores
        ethical_scores = self.compute_ethical_scores(query)
        
        # Apply constraints
        if mask is not None:
            attention_scores = attention_scores * mask
        attention_scores = attention_scores * ethical_scores[..., None]
        
        # Apply softmax and weights
        weights = np.exp(attention_scores) / np.sum(np.exp(attention_scores), axis=-1, keepdims=True)
        output = np.dot(weights, value)
        
        return output

代码说明:

__init__:
- rules: 接收以字典形式的伦理规则。（键：规则名称，值：规则描述）
- _embed_rules: 将每个规则转换为向量（嵌入）。（实际实现中使用预训练的语言模型如Sentence-BERT）
compute_ethical_scores:
- 计算输入查询向量与各规则嵌入之间的相似度（内积）。
- 相似度越高，表示该规则的相关性越高。
- 1 - np.maximum(similarities, 0): 将高相似度转换为低值（接近0），将低相似度转换为高值（接近1）。这会乘以注意力得分，从而降低可能违反规则的令牌的影响。
__call__:
- 类似于基本注意力机制计算注意力得分。
- 调用compute_ethical_scores计算每个令牌的伦理约束得分。
- 如果存在现有掩码（mask），则应用，并乘以伦理约束得分来调整注意力得分。
- 应用softmax计算最终的注意力权重，并通过加权平均计算输出值。

动态约束机制

Constitutional AI 根据上下文动态调节约束强度。

上下文评估 (Context Evaluation):
- 当前对话主题的敏感度分析: 判断对话的主题是否涉及政治、宗教、仇恨言论等敏感领域。
- 用户意图的伦理评估: 推断用户的提问或发言中是否存在恶意意图（例如：试图欺骗模型生成有害内容）。
- 潜在风险水平估计: 评估可能生成的响应的潜在风险水平（例如：轻微偏见、明显的仇恨言论、个人信息披露）。
约束强度调节 (Constraint Strength Adjustment):
- 高风险情况: 当检测到敏感主题、恶意意图或高风险水平时，应用强约束（增加规则违反的惩罚）。
- 一般情况: 在一般的对话或信息请求中，应用灵活的约束（允许轻微的规则违反）。
- 渐进的约束强度变化: 根据情况的变化渐进地调节约束强度，防止不自然的响应或过度限制。

9.4.2.2 基于强化学习的调整 (RLHF, RLAIF)

Constitutional AI 除了使用监督学习（Supervised Learning）外，还利用强化学习（Reinforcement Learning）对模型的行为进行微调(fine-tuning)。

RLHF (Reinforcement Learning from Human Feedback):
1. 收集人的偏好数据: 通过让人选择两种模型响应中哪一种更可取（例如：更有益、更不有害、更诚实的）来收集数据。
2. 训练奖励模型 (Reward Model): 使用收集到的偏好数据，训练一个能够预测哪些响应更好的奖励模型。
3. 策略优化 (Policy Optimization): 利用奖励模型通过强化学习算法（如PPO, Proximal Policy Optimization）对LLM的策略（输入接收并生成响应的方式）进行优化。
RLAIF (从AI反馈中学习的强化学习):
- RLHF的缺点：接收人类反馈的过程成本高昂，耗时较长。
- RLAIF使用AI模型（例如：Constitutional AI的批评模型）代替人类生成反馈，并通过此过程训练奖励模型。
- 优点：
  - 可扩展性 (Scalability): 可以在没有人为干预的情况下生成大规模数据并训练模型。
  - 一致性 (Consistency): AI模型可以比人类更一致地应用标准提供反馈。
  - 成本效益 (Cost-effectiveness): 节省了人力。

Constitutional AI利用这些强化学习技术，遵循明确的规则（宪法），同时生成符合人类偏好的自然且有用的响应来训练模型。

结论

Constitutional AI超越了简单的事后过滤，是一种将伦理约束整合到模型内部运作方式中的新方法。通过结合明确的规则（宪法）、监督学习和强化学习，引导模型以安全有益的方式行事。这在解决AI模型的伦理问题并提高其可靠性方面可以发挥重要作用。

9.4.2节中探讨了以Constitutional AI为中心的伦理约束机制。这种做法将发展成为专门针对特定领域或任务的注意力机制（将在9.4.3节中讨论），进一步增强AI系统的安全性和可信度。

9.4.3 特殊目的注意力：领域和任务优化

在第9.4.2节中讨论的伦理约束机制可以视为特殊目的注意力(Special-Purpose Attention) 的一个例子。自2023年起，这一概念得到了进一步扩展，开始研究/开发针对特定领域(domain)或任务(task)优化的各种注意力机制。

9.4.3.1 特殊目的注意力的多种示例

伦理/安全约束注意力 (Ethical/Safety-Constrained Attention):
- 如第9.4.2节所述的宪法AI，这是一种为了在模型输出中反映伦理、社会价值而设计的注意力机制。
- 核心理念: 通过调整注意力权重来抑制有害或有偏见的内容生成，并引导生成安全、可信的响应。
- 实现方法:
  - 规则基础注意力 (Rule-Based Attention): 定义明确的规则（例如：禁止词汇列表，个人隐私保护规则），根据违反规则的可能性调整注意力权重。
  - 基于强化学习的对齐 (Reinforcement Learning based Alignment): 通过人类或AI的反馈来调整模型的行为，使其朝向期望的方向发展。（参见第9.4.2.2节）
语法结构引导注意力 (Syntax-Guided Attention):
- 在自然语言处理(NLP)中，将句子的语法结构(syntax tree)信息整合到注意力机制中，以提高上下文理解能力。
- 核心理念: 为在语法树中的父子关系、依赖关系(dependency relation)等处出现的词对赋予更高的注意力权重。
- 实现方法:
  - 树结构注意力 (Tree-structured Attention): 设计直接反映语法树结构的注意力机制。
  - 门控注意力 (Gated Attention): 使用门控机制将语法结构信息整合到注意力计算中。
基于知识的注意力 (Knowledge-Grounded Attention):
- 利用外部知识库(knowledge base, 例如：Wikidata, Freebase)来增强注意力机制的方法。
- 核心理念: 识别与输入文本相关的知识库实体(entity)和关系(relation)，并利用这些信息进行注意力计算。
- 实现方法:
  - 实体感知注意力 (Entity-aware Attention): 将知识库中的实体嵌入整合到注意力计算中。
  - 关系感知注意力 (Relation-aware Attention): 反映实体之间的关系信息在注意力权重中。
代码注意力 (Code Attention):
- 专门用于代码生成和理解的特殊目的注意力。
- 通过分析代码的语法结构(AST)及含义，用于代码自动完成、代码摘要、错误发现等。

9.4.3.2 多模态注意力 (Multimodal Attention)

多模态注意力是一种处理文本、图像、音频、视频等形式不同数据(模态, modality)的综合性注意力机制。这与人类通过多种感官器官获得的信息综合来理解世界的方式类似。 * 核心机制: (将在第10章详细讨论) 1. 模态特定编码 (Modality-Specific Encoding): 使用针对每个模态优化的编码器将数据转换为向量表示。 2. 跨模态注意力 (Cross-Modal Attention): 建模不同模态表示之间的关系。 3. 联合表示学习 (Joint Representation Learning): 整合所有模态的信息，学习一个共同的表示空间。

应用领域: VQA, 图像描述生成, 文本到图像合成, 视频理解, 机器人技术等 (将在第10章详细说明)
代表性模型: VisualBERT, LXMERT, ViLBERT, CLIP, DALL-E, Stable Diffusion, Flamingo, GATO, Gemini 等 (将在第10章详细介绍)

9.4.3 摘要

在9.4.3节中，我们简要介绍了特殊目的注意力的各种示例（伦理限制、句法结构引导、知识基础、代码注意力）以及多模态注意力的基本概念和应用领域及代表性模型。关于多模态注意力的更详细内容将在第10章讨论。

这些特殊目的注意力的发展极大地扩展了变压器模型的应用范围，并帮助AI系统能够解决更多样化的现实世界问题。

点击查看内容（深入研究：不同变压器模型的详细分析及技术关联性）

深度剖析：不同变压器模型的详细分析及技术关联性

在本次深度剖析中，我们将深入分析之前讨论过的变压器模型的发展历程，详细介绍每个模型的核心创新、主要特点、性能提升以及与相关技术的关联性。内容包括截至2025年的最新信息和额外的详细说明。

1. 编码器中心模型 (Encoder-Only Models)

2. 解码器中心模型 (Decoder-Only Models)

3. 混合方法模型 (Encoder-Decoder Models)

技术关联性深度分析

3D注意力机制:
- ESM-3: 利用几何注意力将氨基酸序列以及3D坐标信息整合到注意力中，以预测蛋白质的3D结构。
- FlashAttention-3D: 扩展FlashAttention以高效处理3D数据，减少内存使用量。
量子化发展:
- Gemini Ultra: 利用量子计算的退火技术加速注意力计算，并通过4位量化减少模型大小。
- LLaMA-3: 应用受大脑神经可塑性(STDP)启发的动态量化技术，提高边缘设备上的效率。
能效:
- GPT-5: 通过稀疏专家混合(SMoE)建模减少激活参数数量，改善能效。
- Olympus: 通过4D张量并行化(4D tensor parallelism)最大化大规模GPU集群的训练效率。
2025基准现状:

任务	最先进模型	性能	主要技术
语言理解 (MMLU)	GPT-5	92.3	多模态知识融合, Hybrid FlashAttention, 25T token 训练
代码生成 (HumanEval)	CodeLlama-X	91.2	实时编译反馈, 基于强化学习的代码生成, 长篇代码生成能力
蛋白质折叠 (CASP16)	ESM-3G	GDT_TS 94.7	3D图注意力, 几何注意力, FlashAttention-3D, 大规模蛋白质结构数据训练
AI安全 (HarmBench)	Hermes	99.8	规制注意力门控, Constitutional AI 2.0, 实时有害内容过滤, 基于强化学习的安全策略

未来展望

量子-经典混合架构: 利用量子计算的叠加、纠缠加速计算。
生物启发学习: 开发模仿大脑神经机制的算法。
自进化模型: 研究模型自主优化架构的方向。

9.5 高效编码器的实现和应用：以RoPE和FlashAttention为中心

变压器模型在自然语言处理（NLP）领域表现出色，但其计算复杂度和内存使用量大的缺点也很明显。第9.4章探讨了解决这些问题的各种方法。本节将在此基础上，实现适合实际应用的“高效编码器”模型，并测试其性能。特别地，我们将以FlashAttention、Pre-LN以及RoPE (旋转位置嵌入)为中心进行说明。

高效的编码器位于chapter_09/encoder中。

9.5.1 高效编码器的设计理念：速度和内存

高效编码器的核心目标是速度和内存效率。在大规模语言模型时代，模型和数据的规模呈爆炸性增长，因此充分利用给定的硬件资源变得非常重要。

为此，高效的编码器遵循以下设计理念：

减少计算复杂度： 注意机制具有与序列长度平方成正比的计算复杂度。通过使用FlashAttention等优化注意力技术来减少计算量。
最大化内存效率： 减少存储模型参数和中间计算结果所需的内存。
- 利用GPU内存层次结构： 最优化GPU中的快速小SRAM和慢速大HBM之间的数据移动。（FlashAttention的核心原理）
- 分块处理： 将数据分成小块进行处理，以减少内存访问次数。
- Pre-LN (层归一化)： 在注意力机制和前馈网络之前应用层归一化，帮助稳定学习并促进快速收敛。
- 梯度检查点（本次示例中未实现）： 反向传播时不是存储所有中间计算结果，而是仅存储一部分并在需要时重新计算，以减少内存使用量。
RoPE (旋转位置嵌入)（可选）： 高效地表达绝对/相对位置信息，无需单独的位置嵌入即可向模型提供位置信息，并有利于长上下文处理。

9.5.2 `efficient_encoder.py`代码详细分析（未使用RoPE）

efficient_encoder.py实现了不使用RoPE的基本高效编码器。它以FlashAttention、Pre-LN和基本的Transformer结构为核心设计，旨在提高内存效率和计算速度。

1. TransformerConfig类：

定义模型的超参数（vocab_size, hidden_size, num_hidden_layers等）。

2. LayerNorm类：

实现Pre-LN方式的层归一化。

3. Embeddings类：

将输入令牌转换为嵌入向量。与efficient_encoder_rope.py不同，这里使用了可学习的位置嵌入（positional embeddings）。

Code

# efficient_encoder.py
class Embeddings(nn.Module):
    """Token and positional embeddings."""
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size) # 위치 임베딩
        self.norm = LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        batch_size, seq_length = input_ids.size()
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)
        token_embeddings = self.token_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        embeddings = token_embeddings + position_embeddings  # 토큰 임베딩과 위치 임베딩을 더함
        embeddings = self.norm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

4. FlashAttention 类：

实现没有 RoPE 相关代码的基本 FlashAttention。核心是使用 torch.nn.functional.scaled_dot_product_attention。

Code

# (efficient_encoder.py)
class FlashAttention(nn.Module):
    # ... (생략) ...
    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # ... (생략) ...
        # Use PyTorch's built-in scaled_dot_product_attention
        attn_output = F.scaled_dot_product_attention(query_layer, key_layer, value_layer, attn_mask=attention_mask, dropout_p=self.dropout.p if self.training else 0.0)
        # ... (생략) ...
        return attn_output

5. FeedForward 类:

实现位置前馈网络 (FFN)。

6. TransformerEncoderLayer 类:

构成一个变压器编码器层。使用预归一化 (Pre-LN)。

Code

# (efficient_encoder.py)
class TransformerEncoderLayer(nn.Module):
    def __init__(self, config: TransformerConfig):
        super().__init__()
        self.attention = FlashAttention(config)
        self.norm1 = LayerNorm(config.hidden_size, eps=config.layer_norm_eps) # Pre-LN
        self.ffn = FeedForward(config)
        self.norm2 = LayerNorm(config.hidden_size, eps=config.layer_norm_eps) # Pre-LN

    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # Pre-LN + Residual Connection + FlashAttention
        attention_output = self.attention(self.norm1(hidden_states), attention_mask)
        hidden_states = hidden_states + attention_output

        # Pre-LN + Residual Connection + FFN
        ffn_output = self.ffn(self.norm2(hidden_states))
        hidden_states = hidden_states + ffn_output

        return hidden_states

7. TransformerEncoder 类:

构成整个变压器编码器。

9.5.3 `efficient_encoder_rope.py` 代码详细分析 (使用 RoPE)

efficient_encoder_rope.py 是对 efficient_encoder.py 的改进版本，通过添加 RoPE（旋转位置嵌入）来更有效地处理位置信息。

RoPE（旋转位置嵌入）是什么？

RoPE（旋转位置嵌入）是变压器中表示位置信息的新方法。传统的位嵌入通常是将每个位置的固定向量相加，而 RoPE 使用旋转矩阵对位置信息进行编码。就像在2D平面上旋转点一样，它会以特定角度旋转嵌入向量。

例如： 1. 第一个位置：0度旋转 2. 第二个位置：30度旋转 3. 第三个位置：60度旋转这样，位置越远，旋转的角度就越大。如果将高维向量转换为2D来思考，则可以表示如下图所示的图形。

Code

import numpy as np
import matplotlib.pyplot as plt
from matplotlib_inline.backend_inline import set_matplotlib_formats
set_matplotlib_formats('svg')

def visualize_rope_rotation_simple():
    # Rotation angles for each position
    positions = np.arange(4)  # 4 positions
    angles = positions * np.pi/6  # increasing by 30 degrees each time
    
    # Original vector
    vector = np.array([1, 0])  # Reference vector
    
    plt.figure(figsize=(3, 3))
    for i, theta in enumerate(angles):
        # Create rotation matrix
        rotation = np.array([
            [np.cos(theta), -np.sin(theta)],
            [np.sin(theta), np.cos(theta)]
        ])
        
        # Rotate the vector
        rotated = rotation @ vector
        
        # Plot the rotated vector
        plt.arrow(0, 0, rotated[0], rotated[1], 
                 head_width=0.05, head_length=0.1)
        plt.text(rotated[0], rotated[1], f'pos {i}')
    
    plt.grid(True)
    plt.axis('equal')
    plt.title('RoPE: Position-dependent Vector Rotation')
    plt.show()

visualize_rope_rotation_simple()

这种方式的优点是相对距离计算简单（两个位置之间的旋转角度差）且没有序列长度限制。此外，还可以处理比学习的长度更长的序列。

efficient_encoder_rope.py的主要更改

Embeddings类： position_embeddings被移除，并且在forward()中不再添加位置嵌入。因为RoPE已经处理了位置信息，所以不需要单独的位置嵌入。
rotate_half函数： 这是RoPE运算的核心部分。

Code

    # (efficient_encoder_rope.py)
    def rotate_half(x):
        """Rotates half the hidden dims of the input."""
        x1 = x[..., :x.shape[-1] // 2]
        x2 = x[..., x.shape[-1] // 2:]
        return torch.cat((-x2, x1), dim=-1)

apply_rotary_pos_emb 函数: 将 RoPE 应用于查询(q)和键(k)。

Code

    # (efficient_encoder_rope.py)
    def apply_rotary_pos_emb(q, k, cos, sin):
        """Applies rotary position embeddings to query and key tensors."""
        q_embed = (q * cos) + (rotate_half(q) * sin)
        k_embed = (k * cos) + (rotate_half(k) * sin)
        return q_embed, k_embed

FlashAttention 类:
- cos_cached, sin_cached: 预先计算并存储（缓存）用于RoPE的余弦和正弦值。在 _build_cache() 中生成。
- _build_cache(): 预先计算RoPE所需的三角函数值。
- forward(): 对查询、键进行线性变换后，调用 apply_rotary_pos_emb() 应用RoPE。

Code

import torch
from typing import Optional
import torch.nn as nn

def apply_rotary_pos_emb(q, k, cos, sin):
    """Applies Rotary Position Embeddings to query and key tensors."""
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

class FlashAttention(nn.Module):
    # ... (rest of the class definition, unchanged) ...

    def _build_cache(self, device, dtype):
        if self.cos_cached is not None and self.cos_cached.dtype == dtype: #Return if cache already exist.
            return

        # Create position indices
        pos_seq = torch.arange(self.max_position_embeddings, device=device, dtype=dtype)

        # Create freqs (theta in paper)
        inv_freq = 1.0 / (10000 ** (torch.arange(0, self.attention_head_size, 2, device=device, dtype=dtype) / self.attention_head_size))

        # Create freqs for each position in sequence.
        freqs = torch.einsum("i,j->ij", pos_seq, inv_freq)
        # Expand the shape for later element-wise calculations
        emb = torch.cat((freqs, freqs), dim=-1)

        # Create the cos and sin cache
        self.cos_cached = emb.cos()[None, None, :, :]  # Add head and batch dimensions
        self.sin_cached = emb.sin()[None, None, :, :]

    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        # ... (rest of the forward method, unchanged) ...

        # Apply RoPE
        batch_size, num_heads, seq_len, head_dim = query_layer.shape
        self._build_cache(query_layer.device, query_layer.dtype)

        cos = self.cos_cached[:, :, :seq_len, :head_dim]
        sin = self.sin_cached[:, :, :seq_len, :head_dim]

        query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin)

        # ... (rest of the forward method, unchanged) ...

9.5.4 实验结果：AG News 文本分类

我们使用两种版本的高效编码器（efficient_encoder_rope.py 和 efficient_encoder.py）对 AG News 数据集（将新闻文章分为四个类别）进行了文本分类实验。训练代码为 train_ag_news.py。

AG News 数据集由每个类别的平衡新闻文章组成。每篇文章的最大长度限制为 128 个标记，并使用 BERT 和 T5 两种分词器进行对比训练。新闻文本被分为 World, Sports, Business, Sci/Tech 四个类别。模型的大小设置得非常小，具体如下。

Code

vocab_size: int = 30522,
hidden_size: int = 256,
num_hidden_layers: int = 4,
num_attention_heads: int = 8,
intermediate_size: int = 512,
hidden_dropout_prob: float = 0.1,
attention_probs_dropout_prob: float = 0.1,
max_position_embeddings: int = 512,
layer_norm_eps: float = 1e-12

以下是进行比较实验的执行部分。

Code

from dldna.chapter_09.encoder.train_ag_news import train_and_test_all_versions

train_and_test_all_versions(verbose=False)

训练结果表

模型版本	分词器	测试准确度 (%)	备注
v1	bert-base-uncased	91.24	FlashAttention
v1	t5-small	92.00	FlashAttention
v2	bert-base-uncased	92.57	RoPE, FlashAttention
v2	t5-small	92.07	RoPE, FlashAttention

v1: efficient_encoder.py (未使用RoPE)
v2: efficient_encoder_rope.py (使用RoPE)

结果解释

RoPE (v2)的效果: 使用了bert-base-uncased分词器时，应用了RoPE的v2模型比v1模型提高了1.33个百分点的准确度。这表明RoPE能更有效地编码位置信息，从而提升模型性能。特别是在需要处理比训练数据更长序列的情况下（长度外推，length extrapolation），RoPE的优势可能会更加明显。
分词器的影响: 使用了t5-small分词器时，两个版本都达到了与使用bert-base-uncased相似的准确度水平。不过，v2在细微之处表现出更好的性能。
整体高性能: 两个版本都在AG News数据集上实现了91%以上的高准确度。这表明模型架构是有效的，并且通过利用F.scaled_dot_product_attention实现的FlashAttention（如果环境支持的话），以及Pre-LN、GELU、Xavier初始化、AdamW、学习率调度器等现代Transformer训练技术得到了良好的应用。

与类似模型的比较 (表)

下表对比了AG News数据集上其他相似大小模型的性能。（准确度可能因文献和实验结果而异。） | 模型 | hidden_size | num_hidden_layers | AG News 准确率 (大致) | 备注 | | ———————————— |———-| ———— | ————— | —————————— | | Efficient Encoder (v2, bert) | 256 | 4 | 92.57 | RoPE, FlashAttention | | Efficient Encoder (v2, t5) | 256 | 4 | 92.07 | RoPE, FlashAttention | | Efficient Encoder (v1, bert) | 256 | 4 | 91.24 | FlashAttention | | Efficient Encoder (v1, t5) | 256 | 4 | 92.00 | FlashAttention | | TinyBERT (4 层, hidden_size=312) | 312 | 4 | 88-90% | 蒸馏 | | BERT-small | 512 | 4 | ~90.8% | | | DistilBERT-base | 768 | 6 | 90-92% | 蒸馏, 比 BERT-base 小 | | BERT-base | 768 | 12 | 92-95% | 模型大得多 |

应用的机制 | 机制 | v1 (efficient_encoder.py) | v2 (efficient_encoder_rope.py) | 备注 | | ———————— | ———————- | ——————- | —————————— | | FlashAttention | O | O | 利用GPU内存层次结构的优化 | | Pre-LN | O | O | 在注意力/前馈网络之前应用层归一化 | | RoPE | X | O | 使用旋转矩阵进行位置信息编码 | | 可学习的位置嵌入 | O | X | 不使用RoPE时的位置信息表示 | | Xavier 初始化 | O | O | 权重初始化方法 | | GELU 激活函数 | O | O | 非线性激活函数 (在前馈网络中使用) | | Dropout | O | O | 提高泛化性能 | | Layer Normalization | O | O | 稳定训练并提高性能 | | 使用预训练的分词器 | O | O | BERT-base-uncased, t5-small使用 |

结论

本章中，我们利用PyTorch的F.scaled_dot_product_attention实现了FlashAttention，并应用了RoPE (Rotary Positional Embeddings)进一步提高了Transformer编码器模型(v2)的效率。在AG News文本分类数据集上分别用bert-base-uncased和t5-small分词器训练并测试v1(基本Transformer编码器)和v2(RoPE应用)模型，结果显示v2模型在使用bert-base-uncased分词器时达到了更高的准确度 (92.57%)。这表明RoPE有效地对相对位置信息进行了编码，从而提高了模型的性能，特别是长文本处理能力。两个模型都达到了91-92%的高准确率，这表明Efficient Encoder架构既高效又具有强大的性能。此外，在比较bert-base-uncased和t5-small两个分词器时，虽然差异很小，但使用bert-base-uncased的v2版本实现了更高的性能。

如表所示，所提出的Efficient Encoder模型表现出比TinyBERT等小型模型更优的性能，并且与BERT-small相比也具有竞争力。重要的是，它以远小于DistilBERT-base或BERT-base等更大模型的规模达到了接近的性能。可以说，这是预训练分词器、FlashAttention、Pre-LN结构、RoPE、Xavier初始化、GELU激活函数以及适当的模型配置（如hidden_size, num_hidden_layers等）组合的结果。

总之，在本章中提出的Efficient Encoder (v2)不仅在教育目的上有助于理解Transformer的核心组成部分，而且证实了它是一个高效的模型，在实际应用中也能表现出足够的竞争力。特别是，RoPE的应用被证明是提高模型性能的有效方法。

9.6 Mistral：实现和分析高效的解码器架构

2023年，Mistral AI 公开的 Mistral-7B 模型基于 LLaMA 架构引入了 分组查询注意力（GQA） 和 滑动窗口注意力（SWA） ，大幅提升了内存效率和处理速度。特别是该模型仅使用 70 亿参数就能达到与超过 130 亿参数的模型相媲美的性能，证明了高效架构设计的重要性。

本节将参考 Hugging Face Transformers 中的 Mistral 实现，重点围绕核心优化因素，自行实现并分析一个简化的 mistral 模型。我们将详细探讨 GQA、SWA、RoPE 和 KV 缓存机制，并理解它们如何对模型的效率和性能做出贡献。代码位于 chapter_09/mistral。

9.6.1 `simple_mistral` 模型架构：组件详细分析

simple_mistral 模型简化实现了 Mistral-7B 模型的核心组件，每个组件都是模块化的，并执行明确的功能。以下我们将详细介绍各个组件。

1. MistralConfig: 模型设置

MistralConfig 类定义了模型的超参数。它在决定模型结构和工作方式方面发挥着关键作用。

主要属性：
- vocab_size: 指定词汇表的大小（默认值：32000）。
- hidden_size: 表示嵌入和隐藏状态的维度（默认值：4096）。
- intermediate_size: 定义前馈网络的中间维度（默认值：14336）。
- num_hidden_layers: 指定变压器解码器层的数量（默认值：32）。
- num_attention_heads: 表示注意力头的数量（默认值：32）。
- num_key_value_heads: 定义 GQA 中使用的键/值头数量（默认值：8）。
- hidden_act: 激活函数，使用 “silu”（默认值）。
- max_position_embeddings: 指定最大序列长度（默认值：4096 * 32）。
- rms_norm_eps: 表示 RMSNorm 的 epsilon 值（默认值：1e-6）。
- use_cache: 确定是否使用 KV 缓存（默认值：True）。
- rope_theta: 设置 RoPE 的 theta 值（默认值：10000.0）。
- sliding_window: 指定滑动窗口的大小（默认值：4096）。
- use_return_dict: 设置是否以字典形式返回（默认值：True）。

2. MistralRMSNorm: RMS 归一化

MistralRMSNorm 类实现了 RMSNorm (均方根层归一化)。它从传统的 LayerNorm 中去除了平均值，通过平方均值的平方根 (RMS) 进行归一化，从而提高了计算效率。

特点： 使用 variance_epsilon 确保数值稳定性。

3. MistralAttention: 注意力机制

MistralAttention 类实现了 Mistral 模型的关键注意力机制。通过整合 GQA、SWA 和 RoPE 来提高效率和性能。 * GQA (Grouped-Query Attention): * 维持多个查询(Q)头，而键(K)和值(V)头设置为较少的数量，以减少内存使用量和计算量。 * 通过 num_key_value_heads 调整 K/V 头数。 * 使用 repeat_kv 函数将 K/V 张量复制到与 Q 头数匹配。

SWA (Sliding Window Attention):
- 让每个令牌只对其在有限窗口内的令牌执行注意力操作，以减少计算复杂度。
- 通过 sliding_window 参数调整窗口大小。
- 修改 attention_mask 以阻止与窗口外部令牌的注意力。
RoPE (Rotary Positional Embedding):
- 使用旋转矩阵对位置信息进行编码。
- 通过 MistralRotaryEmbedding 类实现。
- 使用 apply_rotary_pos_emb 函数将 RoPE 应用于查询和键。

4. MistralRotaryEmbedding: RoPE 实现

MistralRotaryEmbedding 类实现了 RoPE (Rotary Positional Embedding)。

__init__ 方法:
- dim: 设置嵌入维度。
- max_position_embeddings: 指定最大序列长度。
- base: 定义用于频率计算的常数（默认值：10000）。
- inv_freq: 计算逆频率，并注册为非训练参数。
- cos_cached, sin_cached: 缓存预先计算的余弦、正弦值。
forward 方法:
- 接收输入张量 x 和序列长度 seq_len。
- 如果 seq_len 大于缓存的最大长度，则调用 _set_cos_sin_cache 更新缓存。
- 返回缓存的余弦、正弦值。
_set_cos_sin_cache 方法:
- 生成到 seq_len 的位置索引。
- 将位置索引与逆频率相乘以计算频率。
- 使用计算出的频率计算并缓存余弦、正弦值。

5. MistralMLP: FeedForward 网络

MistralMLP 类实现了 Mistral 模型的 FeedForward 网络。

结构:
- gate_proj, up_proj, down_proj: 使用三个线性层扩展并重新压缩输入。
- act_fn: 使用 SiLU (Sigmoid Linear Unit) 激活函数。

6. MistralDecoderLayer: 解码器层

MistralDecoderLayer 类构建了 Mistral 模型的一个解码器层。

组件:
- self_attn: 使用 MistralAttention 模块执行自注意力。
- mlp: 使用 MistralMLP 模块执行 FeedForward 网络。
- input_layernorm, post_attention_layernorm: 使用 MistralRMSNorm 执行输入/输出规范化。

7. MistralPreTrainedModel: 预训练模型抽象类

MistralPreTrainedModel 类是管理 Mistral 模型权重初始化和设置的抽象基类。

主要方法:
- _init_weights: 初始化权重。
- _set_gradient_checkpointing: 设置是否激活梯度检查点。

8. MistralModel: Mistral 模型

MistralModel 类定义了 Mistral 模型的整体结构。

组成部分:
- embed_tokens: 将输入令牌转换为嵌入向量。
- layers: 堆叠多个 MistralDecoderLayer。
- norm: 对最后一层的输出进行标准化。

9. MistralForCausalLM: 用于语言建模的 Mistral

MistralForCausalLM 类是用于将 Mistral 模型针对因果语言模型 (Causal Language Modeling) 任务进行微调的类。

主要组成部分:
- lm_head: 将模型输出投影到词汇表大小以计算下一个令牌预测概率。
- prepare_inputs_for_generation: 在推理过程中准备输入。
- _reorder_cache: 在束搜索 (beam search) 时重新排列 KV 缓存。

如上所述，simple_mistral 模型通过模块化各个组件提供高效灵活的设计。理解每个组件的作用和相互作用可以更清晰地掌握模型的工作原理。

9.6.2 核心技术要素分析：效率与性能的关键

simple_mistral 模型通过 GQA、SWA 和 RoPE 等核心技术要素实现效率和性能的最大化。我们将详细分析每个技术要素的运行方式及其优点。

1. GQA (Grouped-Query Attention): 提高内存和计算效率的创新

GQA 是 Multi-Head Attention 的变体，是减少内存使用量和计算量同时保持性能的关键技术。

工作原理:
- 查询 (Q) 被分割成多个头，而键 (K) 和值 (V) 则被分割成较少的头。
- 每个 Q 头分配给特定的 K/V 头组。
- Q 头仅计算其分配的 K/V 头组的注意力。
- repeat_kv 函数复制 K/V 张量以匹配 Q 头的数量，从而实现这种机制。
优点:
- 减少内存使用: 由于 K/V 张量大小减小，可以缩小 KV 缓存的大小。
- 减少计算量: 注意力计算量减少，推理速度提高。
- 性能保持: Q 头的数量保持不变，因此模型的表达能力不会显著下降。

2. SWA (Sliding Window Attention): 高效处理长序列的策略

SWA 是一种技术，通过限制每个令牌仅在其有限范围（窗口）内的令牌上执行注意力计算来减少计算复杂度。

工作原理:
- 每个令牌仅对固定大小窗口内的令牌执行注意力计算。
- 窗口沿着序列移动，在每个位置计算注意力。
- 使用 attention_mask 对窗口外部的令牌进行注意力遮罩。
优点:
- 计算复杂度降低: Attention 计算量从 O(N²) 减少到 O(N*W)。 (N: 序列长度, W: 窗口大小)
- 处理长序列: 内存使用减少，可以处理更长的序列。

3. RoPE (Rotary Positional Embedding): 高效编码相对位置信息

RoPE 在第9.5章中已经讨论过。这里我们只简单看一下模型中实现的部分。

实现:
- rotate_half 函数: 将输入张量的维度分为两半，交替改变符号以实现复数乘法的效果。

Code

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

apply_rotary_pos_emb 函数: 将 RoPE 应用于查询(q)和键(k)张量。

Code

def apply_rotary_pos_emb(q, k, cos, sin, position_ids_q, position_ids_k=None):
    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
    cos_q = cos[position_ids_q].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
    sin_q = sin[position_ids_q].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
    cos_k = cos[position_ids_k].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
    sin_k = sin[position_ids_k].unsqueeze(1)  # [batch_size, 1, seq_len, dim]
    q_embed = (q * cos_q) + (rotate_half(q) * sin_q)
    k_embed = (k * cos_k) + (rotate_half(k) * sin_k)
    return q_embed, k_embed

MistralRotaryEmbedding 类: 预计算并缓存 RoPE 所需的余弦、正弦值。
- cos_cached, sin_cached: 预计算的余弦、正弦值
- _set_cos_sin_cache: 根据序列长度更新 cos_cached, sin_cached
优点:
- 保留相对位置信息: 注意力权重根据令牌之间的相对距离自然变化。
- 长度外推(Length Extrapolation): 对于比训练序列更长的序列也能良好运行。
- 线性复杂度: 不影响注意力计算的复杂度。

GQA、SWA、RoPE 分别提高了内存效率、计算效率和位置信息表达能力，是提升 simple_mistral 模型整体性能的关键技术要素。

4. KV 缓存: 去除重复计算

KV 缓存在生成模型中尤其重要，可以显著提高推理速度。

概念:
- KV 缓存是一种在推理过程中存储和重用每个解码器层计算的键(K)和值(V)张量的技术。
- 每次生成新令牌时，无需重新计算先前令牌的 K、V，而是使用缓存的值进行运算。
- 通过 past_key_values 参数存储前一步骤的 KV 缓存，并通过设置 use_cache=True 激活 KV 缓存功能。每个层接收 past_key_value 作为输入并输出更新的 present_key_value。
优点:
- 提高推理速度: 通过去除重复计算，显著提高了令牌生成速度。
- 增加内存使用量: 需要额外内存来存储 KV 缓存，但可以通过 GQA 和 SWA 等技术缓解内存增长。

KV 缓存在生成长文本时效果尤为显著，并在很大程度上改善了用户体验。

9.6.3 模型训练: `simple_mistral` 训练指南

simple_mistral 模型的训练过程主要分为数据预处理和模型训练两个阶段。

1. 数据预处理: 转换为模型可理解的形式

将用于模型训练的文本数据转换为模型可以处理的形式的过程。

分词(Tokenization):
- 使用分词器(Tokenizer)将文本数据转换为模型可以处理的数字（令牌 ID）形式。
- 分词器将文本拆分为小单元（令牌），并将每个令牌映射到唯一的 ID。
生成 attention_mask:
- attention_mask 用于区分填充(padding)令牌，并确保注意力仅应用于实际数据。
- 填充是为对齐序列长度而添加的令牌，在注意力计算中应被排除。

2. 模型训练: 寻找最优参数

使用 MistralForCausalLM 模型以语言建模(Causal Language Modeling)方式开展训练。 * MistralForCausalLM 模型: Mistral 模型的类，用于配置语言建模任务。 * 损失函数 (Loss Function): * 使用 CrossEntropyLoss 计算模型输出（预测）与正确标签之间的差异。 * 模型通过最小化此损失进行学习。 * 优化器 (Optimizer): * 使用 AdamW 优化器更新模型的权重（参数）。 * AdamW 是 Adam 优化器的改进版本，能更有效地应用权重衰减(weight decay)。 * 学习率调度器 (Learning Rate Scheduler): * 使用 get_cosine_schedule_with_warmup 调度器逐渐降低学习率(learning rate)。 * 在训练初期提高学习率以快速收敛，在训练后期降低学习率以进行微调(fine-tuning)。 * 梯度裁剪 (Gradient Clipping): * 应用梯度裁剪以防止梯度爆炸(exploding gradient)问题。 * 当梯度的大小超过某个阈值时，将值截断以帮助稳定训练。

9.6.4 使用 `generate()` 函数生成文本：创造性的句子制作

使用训练好的模型生成新文本的过程。generate() 函数可以通过多种参数调节生成文本的风格和多样性。

`generate()` 函数：文本生成的核心

功能: 基于给定的提示(prompt)生成文本。
KV 缓存利用: 使用 past_key_values 利用 KV 缓存提高推理速度。
主要参数:
- max_new_tokens: 指定要生成的最大 token 数量。
- temperature: 调整概率分布的形状以控制生成结果的多样性。(低值：一致性，高值：多样性)
- top_k: 只考虑概率最高的前 k 个 token 进行采样。
- top_p: 只考虑累积概率达到 p 的 token 进行采样。(核采样 nucleus sampling)
- repetition_penalty: 对重复的 token 施加惩罚以减少文本中的重复。

生成过程：逐步文本生成

初始输入: 将提示词化并输入模型以获得初始输出。
调整概率分布: 应用 temperature、top_k、top_p 和 repetition_penalty 等约束条件到输出 logit（logits）上，以调整下一个 token 的概率分布。
token 采样: 根据调整后的概率分布采样下一个 token。
添加输出和更新 KV 缓存: 将生成的 token 添加到输出序列，并更新 KV 缓存。
重复: 直到满足终止条件（达到最大长度或生成终止 token）之前，重复步骤2-4。

本节详细介绍了 Mistral 模型的训练及文本生成过程。接下来的部分将通过实际应用示例展示 simple_mistral 模型的使用方法，包括三个示例。示例位于 mistral/examples 中。 1. 数字序列预测 (train_seq_num.py): 通过一个简单的任务来预测连续的数字，以检验模型的基本学习和生成能力。 2. 四则运算预测 (train_math.py): 通过预测加法、减法、乘法运算的结果的任务，检查模型是否能学习符号推理(symbolic reasoning)。 3. SQL查询生成 (train_sql.py): 通过将自然语言问题转换为SQL查询的任务，评估模型理解和处理复杂语言结构的能力。 (使用WikiSQL数据集)

您可以在该位置直接在shell中运行。例如 python train_seq_num.py。以下是在Jupyter笔记本中运行的方法。

9.6.5 数字序列预测示例：`train_seq_num.py` 分析

train_seq_num.py 是一个使用 simple_mistral 模型执行简单数字序列预测任务的示例。通过这个示例，我们可以了解模型如何学习预测给定数字序列后的下一个数字。

1. 准备数据集和数据加载器：构建训练数据

这是准备 simple_mistral 模型将要学习的数据的步骤。

SimpleDataset 类:
- 继承自 PyTorch 的 Dataset 类，定义了一个简单的数字序列数据集。
- __init__ 方法接收数据（data）和序列长度（seq_length）作为输入来初始化数据集。
- __len__ 方法返回数据集中所有样本的数量。
- __getitem__ 方法返回给定索引（idx）对应的输入序列和标签序列，在这个示例中，输入和标签是相同的序列。模型内部会将标签自动向前移动一位以构建下一个 token 预测任务。
create_simple_data 函数:
- 生成与指定的词汇大小（vocab_size）、样本数量（num_examples）和序列长度（seq_length）相匹配的数字序列数据。
- 重复使用从 0 到 vocab_size - 1 的数字来创建一个长度为 num_examples 的列表。
数据加载器 (DataLoader):
- DataLoader 将通过 SimpleDataset 创建的数据集按迷你批次（mini-batch）打包，提供给模型。
- batch_size 指定每次输入模型的样本数量，
- 设置 shuffle=True 可以在每个周期（epoch）中随机打乱数据顺序，提高训练效果。
通过 SimpleDataset 创建的训练数据具有以下形式。
```
样本 1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] -> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

样本 2: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] -> [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
```
模型 forward 函数中的标签偏移

在 simple_mistral 模型的 forward 函数中，内部会将标签序列向右移动一位以构建下一个 token 预测任务。也就是说，模型按以下方式工作。
1. 输入序列: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
2. 模型输入: [0, 1, 2, 3, 4, 5, 6, 7, 8]（排除最后一个 token）
3. 模型预测: [1, 2, 3, 4, 5, 6, 7, 8, 9]（每个位置的下一个 token 预测）
4. 标签: [1, 2, 3, 4, 5, 6, 7, 8, 9]（排除输入序列的第一个 token，与模型预测进行比较）
通过这个过程，模型学习了在输入序列的每个位置上预测下一个将出现的 token 的能力。

2. 模型设置和训练：`simple_mistral` 训练

这是设置 simple_mistral 模型并使用准备好的数据进行训练的步骤。 * MistralConfig 配置: * vocab_size 设置为词汇表大小（由分词器定义）加上 <eos> 令牌。这是为了让模型能够识别句子的结束。 * sliding_window 设置为与序列长度相同，以便每个令牌都能查看整个序列。 * 将 use_cache=False 以在训练期间不使用 KV 缓存。

权重共享 (tie_weights = True):
- 将 tie_weights 设置为 True 以共享嵌入权重和输出层（lm_head）的权重。这可以减少参数数量，并有助于学习特定模式（在这种情况下，顺序数字生成）。
模型 (MistralForCausalLM) 和优化器 (AdamW) 创建:
- 创建 MistralForCausalLM 模型并将其移动到指定的设备（device, CPU 或 GPU）。
- 创建 AdamW 优化器，并设置模型参数和学习率 (learning_rate)。
train 函数 (训练循环):
- 将模型设置为训练模式 (model.train())。
- 按指定的周期数（epochs）重复训练。
- 在每个周期中，从数据加载器获取小批量数据并输入到模型中，计算损失（loss）。
- 通过反向传播 (backpropagation) 计算梯度，并使用优化器更新模型参数。
- 定期输出批次损失，并在每个周期结束时输出平均损失以监控训练进度。

3. 文本生成: 使用训练好的模型进行预测

这是使用训练好的模型生成新文本（数字序列）的步骤。

generate_text 函数:
- 将模型设置为评估模式 (model.eval())。
- 将起始文本（start_text，例如：['1', '2', '3']）转换为令牌 ID 并输入到模型中。
- 重复生成下一个令牌直到 max_length：
  - 在模型的输出对数概率（logits）上应用 temperature 来调整概率分布。较低的 temperature 值会生成更连贯的文本，较高的值则生成更多样化的文本。
  - 从调整后的概率分布中采样下一个令牌 ID。（使用 torch.multinomial 函数）
  - 将采样的令牌 ID 转换回文本并添加到生成的令牌列表中。
  - 将新生成的令牌添加到输入中，以预测下一个令牌的过程重复进行。
- 最终返回生成的文本。

4. 结果分析: 训练结果和生成文本评估

这是分析模型训练结果和生成文本的步骤。

训练结果: 可以确认在训练过程中损失（loss）持续减少。这表明模型成功地学习了数字序列的模式。
生成结果:
- 从 ['1', '2', '3'] 开始的文本生成结果：1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
- 从 ['40', '41', '42'] 开始的文本生成结果：40 41 42 43 44 45 46 47 48 49 可以验证模型能够准确生成给定起始数字的连续数字。这表明模型已经学习了数字序列的模式，并且可以根据这些模式生成新的序列。

train_seq_num.py 示例展示了如何使用 simple_mistral 模型成功执行一个简单但明确的数字序列预测任务。

Code

import torch
from dldna.chapter_09.mistral.examples.train_seq_num import MistralConfig, MistralForCausalLM, SimpleDataset, create_simple_data, generate_text, train
from torch.utils.data import Dataset, DataLoader

# Hyperparameter settings
base_vocab_size = 50    # Original vocab_size before the EOS token
seq_length = 10         # Sequence length of each training sample
batch_size = 8
epochs = 5
learning_rate = 5e-3
num_train_examples = 1000
device = "cuda" if torch.cuda.is_available() else "cpu"

# 1) Create tokenizer (string token -> token id)
tokenizer_vocab = {str(i): i for i in range(base_vocab_size)}
tokenizer_vocab["<eos>"] = base_vocab_size
updated_vocab_size = base_vocab_size + 1

# 2) Model configuration: Apply the updated vocab_size and set sliding_window to seq_length
config = MistralConfig(
    vocab_size=updated_vocab_size,
    hidden_size=32,
    intermediate_size=64,
    num_hidden_layers=2,
    num_attention_heads=4,
    num_key_value_heads=2,
    max_position_embeddings=128,
    sliding_window=seq_length,  # Set to the same as the sequence length
    use_cache=False  # Do not use cache during training
)
config.eos_token_id = tokenizer_vocab["<eos>"]

# (Optional) Set up weight tying between embedding and lm_head -> Can help reproduce sequential patterns.
tie_weights = True

# 3) Create model and Optimizer
model = MistralForCausalLM(config).to(device)
if tie_weights:
    model.lm_head.weight = model.model.embed_tokens.weight
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)

# 4) Data generation and DataLoader preparation
train_data = create_simple_data(updated_vocab_size, num_train_examples, seq_length)
train_dataset = SimpleDataset(train_data, seq_length)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# --- For debugging: Output some data before training ---
print("Sample data before training (input sequence -> label sequence):")
for i in range(2):
    input_seq, label_seq = train_dataset[i]
    print(f"Sample {i+1}: {input_seq.tolist()} -> {label_seq.tolist()}")

# 5) Start training
print("Start training...")
train(model, train_dataloader, optimizer, epochs, device)

# 6) Text generation example
print("Generating text starting with tokens ['1', '2', '3']:")
start_text = ["1", "2", "3"]
generated = generate_text(model, start_text, tokenizer_vocab, max_length=20, device=device)
print("Generated text:", " ".join(generated))

print("Generating text starting with tokens ['40', '41', '42']:")
start_text = ["40", "41", "42"]
generated = generate_text(model, start_text, tokenizer_vocab, max_length=20, device=device)
print("Generated text:", " ".join(generated))

Sample data before training (input sequence -> label sequence):
Sample 1: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] -> [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
Sample 2: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] -> [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
Start training...
Batch 100/124, Loss: 0.0020
Epoch 1/5, Average Loss: 2.2763
Batch 100/124, Loss: 0.0027
Epoch 2/5, Average Loss: 0.0024
Batch 100/124, Loss: 0.0006
Epoch 3/5, Average Loss: 0.0011
Batch 100/124, Loss: 0.0008
Epoch 4/5, Average Loss: 0.0007
Batch 100/124, Loss: 0.0005
Epoch 5/5, Average Loss: 0.0005
Generating text starting with tokens ['1', '2', '3']:
Generated text: 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20
Generating text starting with tokens ['40', '41', '42']:
Generated text: 40 41 42 43 44 45 46 47 48 49

9.6.6 四则运算预测示例：`train_math.py` 分析

train_math.py 使用 simple_mistral 模型来预测简单的四则运算（加法、减法、乘法）的结果。通过这个例子，评估模型是否能够理解数字和运算符号，并学习执行简单的数学推理。训练数据示例如下。

样本 1: 4*1=4<eos>
样本 2: 9+8=17<eos>

数据生成及预处理：符号与数字的和谐

train_math.py 示例在数据生成、分词器、模型设置等方面与之前的数字序列预测示例有几个重要的不同点。最大的不同在于处理的数据不仅仅是简单的数字列表，而是由数字、运算符号、等号以及表示句子结束的 <eos> 令牌组成的“表达式”。

create_arithmetic_data 函数：四则运算数据生成
- 此函数生成指定数量（num_samples）的四则运算表达式及其结果，形式为字符串。
- 每个表达式遵循 f"{num1}{op}{num2}={result}<eos>" 的格式。例如："12+7=19<eos>"
  - num1, num2: 在 1 到 max_value 之间随机选择的整数。
  - op: 随机选择的加法（+）、减法（-）、乘法（*）运算符号。
  - result: 使用 Python 的 eval 函数计算的实际结果值。
  - <eos> 令牌的重要性： 在字符串末尾明确添加 <eos> (End-of-Sentence) 令牌非常重要。这个特殊令牌充当模型的路标，告知其句子何时结束。如果缺少 <eos> 令牌，模型将难以判断何时停止生成，可能会无限继续输出数字或符号。
create_tokenizer 函数：词汇表定义
- 创建包含数字（0-9）、运算符号（+, -, \*）、等号（=）以及特殊令牌（<pad>, <eos>）的词汇表。此词汇表定义了模型可以理解的基本字符。
  - <pad> 令牌用于将不同长度的序列组合成一个批次（batch），以进行处理。
create_reverse_tokenizer 函数：将令牌 ID 还原为字符
- 创建一个逆向字典，将令牌 ID 转换回字符串令牌。这用于将生成的结果解释为人类可读的形式。
tokenize_sample 函数：将字符串转换为令牌列表
- tokenize_sample 函数将样本字符串转换为模型可以识别的令牌列表。
  - 对于 <eos> 等特殊令牌，将其作为单个令牌处理，以便模型能够完整地识别这些特殊令牌。
ArithmeticDataset 类：转换为可训练的数据形式
create_arithmetic_data 函数中生成的数据转换为 PyTorch 的 Dataset 格式。Dataset 是一种标准化的方法，用于高效地向模型提供数据。
- __getitem__ 方法执行以下操作：
  1. 使用 tokenize_sample 函数首先将样本字符串进行分词。
  2. 如果分词后的序列长度短于指定的 seq_length，则使用 <pad> 令牌填充以匹配长度。这是为了使所有输入序列具有相同的长度，以便模型可以按批次处理。
  3. 将令牌转换为整数 ID，并将输入序列和标签序列（与输入相同）作为 PyTorch 张量返回。

模型配置及训练

MistralConfig 配置: 由于这是一个比数字序列预测示例稍微复杂的任务，因此略微增加了模型的大小。（hidden_size=64, intermediate_size=128, num_hidden_layers=3, num_attention_heads=8, num_key_value_heads=4）。此外，设置 pad_token_id 和 eos_token_id 以使模型识别填充令牌和句子结束令牌。
训练: 使用与先前示例几乎相同的 train 函数进行模型训练。使用 CosineAnnealingLR 调度器逐渐降低学习率，以便在训练初期快速收敛，并在后期进行微调。

文本生成

generate_text 函数: 使模型基于给定的提示（例如：“12+7=”）生成文本（四则运算结果）。当模型生成 <eos> 或 <pad> 令牌时，停止生成结果字符串。

结果分析

训练结果: 通过观察训练过程中损失(loss)逐渐减少的情况，可以得知模型正在学习四则运算模式。
生成结果: 通过评估数据示例检查模型是否能针对给定的提示生成正确的运算结果。例如：“4+20=” -> “4+20=24”

train_math.py 示例展示了 simple_mistral 模型不仅能够进行简单的数字序列预测，还能够学习像四则运算这样的符号推理能力。此外，还可以了解特殊令牌如 <eos> 的作用和重要性，以及根据任务复杂度调整模型大小的必要性。

Code

import torch
import random
from dldna.chapter_09.mistral.examples.train_math import MistralConfig, MistralForCausalLM, generate_text, train,create_arithmetic_data, ArithmeticDataset, create_tokenizer, create_reverse_tokenizer
from torch.utils.data import DataLoader

random.seed(42)
torch.manual_seed(42)

device = "cuda" if torch.cuda.is_available() else "cpu"

# Hyperparameter settings
num_samples = 10000   # Total number of samples in the dataset
max_value = 20       # Maximum value of operands
seq_length = 20      # Fixed sequence length including EOS token (e.g., 20)
batch_size = 16
epochs = 20
learning_rate = 1e-3

# Data generation (including EOS token) and output training data examples
arithmetic_data = create_arithmetic_data(num_samples, max_value)
print("Training data examples:")
for i in range(10):
    print(f"Sample {i+1}: {arithmetic_data[i]}")

# Create tokenizer
tokenizer = create_tokenizer()
reverse_tokenizer = create_reverse_tokenizer(tokenizer)
updated_vocab_size = len(tokenizer)

# Configure Dataset and DataLoader
dataset = ArithmeticDataset(arithmetic_data, seq_length, tokenizer)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

config = MistralConfig(
    vocab_size=updated_vocab_size,
    hidden_size=64,
    intermediate_size=128,
    num_hidden_layers=3,
    num_attention_heads=8,
    num_key_value_heads=4,
    max_position_embeddings=128,
    sliding_window=seq_length,
    use_cache=False,
    use_return_dict=True,
    pad_token_id=tokenizer["<pad>"]  # Set the pad token id here.
)
config.eos_token_id = tokenizer["<eos>"]  # Also update the eos token

model = MistralForCausalLM(config).to(device)

# weight tying (share weights between embedding and lm_head)
tie_weights = True
if tie_weights:
    model.lm_head.weight = model.model.embed_tokens.weight

# Create optimizer and add cosine annealing scheduler
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)

# Start training
print("Start training...")
train(model, dataloader, optimizer, scheduler, epochs, device)

# Evaluation: Output 10 random evaluation samples (terminate generation if EOS is included in the prompt)
print("\nEvaluation data examples:")
for i in range(10):
    sample = random.choice(arithmetic_data)
    # Use the part before '=' as a prompt in the entire expression, e.g., "12+7=19<eos>" ("12+7=")
    prompt = sample.split('=')[0] + '='
    generated = generate_text(model, prompt, tokenizer, reverse_tokenizer, max_length=seq_length, device=device)
    print(f"Generated result for prompt '{prompt}': {generated} (Original data: {sample})")

Training data examples:
Sample 1: 4*1=4<eos>
Sample 2: 9+8=17<eos>
Sample 3: 5*4=20<eos>
Sample 4: 18*3=54<eos>
Sample 5: 14+2=16<eos>
Sample 6: 3+7=10<eos>
Sample 7: 17+20=37<eos>
Sample 8: 18*7=126<eos>
Sample 9: 18+14=32<eos>
Sample 10: 15-19=-4<eos>
Start training...
Epoch 1/20, Average Loss: 2.4820, LR: 0.000994
Epoch 2/20, Average Loss: 1.2962, LR: 0.000976
Epoch 3/20, Average Loss: 1.1905, LR: 0.000946
Epoch 4/20, Average Loss: 1.0831, LR: 0.000905
Epoch 5/20, Average Loss: 0.9902, LR: 0.000855
Epoch 6/20, Average Loss: 0.9112, LR: 0.000796
Epoch 7/20, Average Loss: 0.8649, LR: 0.000730
Epoch 8/20, Average Loss: 0.8362, LR: 0.000658
Epoch 9/20, Average Loss: 0.8194, LR: 0.000582
Epoch 10/20, Average Loss: 0.8128, LR: 0.000505
Epoch 11/20, Average Loss: 0.8049, LR: 0.000428
Epoch 12/20, Average Loss: 0.7971, LR: 0.000352
Epoch 13/20, Average Loss: 0.7945, LR: 0.000280
Epoch 14/20, Average Loss: 0.7918, LR: 0.000214
Epoch 15/20, Average Loss: 0.7903, LR: 0.000155
Epoch 16/20, Average Loss: 0.7884, LR: 0.000105
Epoch 17/20, Average Loss: 0.7864, LR: 0.000064
Epoch 18/20, Average Loss: 0.7854, LR: 0.000034
Epoch 19/20, Average Loss: 0.7837, LR: 0.000016
Epoch 20/20, Average Loss: 0.7831, LR: 0.000010

Evaluation data examples:
Generated result for prompt '4+20=': 4+20=24 (Original data: 4+20=24<eos>)
Generated result for prompt '16-3=': 16-3=13 (Original data: 16-3=13<eos>)
Generated result for prompt '10+15=': 10+15=25 (Original data: 10+15=25<eos>)
Generated result for prompt '8+4=': 8+4=12 (Original data: 8+4=12<eos>)
Generated result for prompt '16-13=': 16-13=3 (Original data: 16-13=3<eos>)
Generated result for prompt '10*1=': 10*1=10 (Original data: 10*1=10<eos>)
Generated result for prompt '18+13=': 18+13=31 (Original data: 18+13=31<eos>)
Generated result for prompt '9+9=': 9+9=18 (Original data: 9+9=18<eos>)
Generated result for prompt '1+15=': 1+15=16 (Original data: 1+15=16<eos>)
Generated result for prompt '18-18=': 18-18=0 (Original data: 18-18=0<eos>)

9.6.7 自然语言-SQL 查询生成示例：`train_sql.py` 分析

train_sql.py 使用 simple_mistral 模型处理将自然语言问题转换为 SQL 查询的更复杂的自然语言处理任务。在本示例中，我们将观察模型如何超越简单的序列生成，学习理解复杂自然语言句子的意义，并将其表示为结构化的 SQL 查询语言。示例由给定一个句子并返回其 SQL 语句形式的训练数据组成。以下是训练数据示例。

样本 1: Tell me what the notes are for South Australia sep> SELECT Notes FROM table WHERE Current slogan = SOUTH AUSTRALIA eos>
样本 2: What is the format for South Australia? sep> SELECT Format FROM table WHERE State/territory = South Australia eos>

数据集及预处理：WikiSQL 和特殊标记的协调

train_sql.py 示例的核心在于有效利用 WikiSQL 数据集，并对数据进行预处理以使模型能够学习自然语言和 SQL 查询之间的关系。

加载 WikiSQL 数据集：使用 datasets 库加载 WikiSQL 数据集。WikiSQL 是一个包含成对的自然语言问题及其对应的 SQL 查询的数据集，广泛用于自然语言-SQL 转换任务。可以使用 load_dataset 函数的 split 参数分别指定训练（train）数据集和验证（validation）数据集。
WikiSQLDataset 类：继承自 PyTorch 的 Dataset 类，将 WikiSQL 数据集加工成适合模型训练的形式。
- 在 __init__ 方法中加载 WikiSQL 数据集，并设置要使用的分词器（tokenizer）和最大序列长度（max_length）。
- __getitem__ 方法对数据样本进行处理，将其转换为可输入模型的形式。这个过程中最关键的部分是将自然语言问题与 SQL 查询结合，并添加特殊标记。
  1. 首先，从样本数据中获取自然语言问题（question）和人工编写的 SQL 查询（sql['human_readable']）。
  2. 将获取的问题和 SQL 查询以 "问题 <sep> SQL<eos>" 的形式结合。这里的 <sep> 是区分问题和 SQL 查询的分隔标记(separator)，而 <eos> 则是表示句子结束的终止标记(end-of-sentence)。这些特殊标记在告知模型输入文本结构方面起着重要作用。
  3. 使用 tokenizer 对组合后的文本进行分词。此时，设置 truncation=True 以截断超出 max_length 的文本，并设置 padding="max_length" 以添加填充使序列长度达到 max_length。
  4. 最终返回分词化的 input_ids。（输入和标签相同）
分词器 (T5Tokenizer): 使用 transformers 库中的 T5Tokenizer。选择 T5Tokenizer 的原因如下。
- 默认支持多种特殊标记（如 <pad>、<eos>、<sep> 等）。
- 是一个通用的分词器，能够有效处理自然语言和 SQL 查询（代码）。
- 可以通过 tokenizer.vocab_size 轻松获取分词器的词汇表大小，从而方便地设置模型的 vocab_size。
数据加载器 (DataLoader): 通过 WikiSQLDataset 生成的数据集按小批次（mini-batch）打包，以便高效地为模型提供输入。batch_size 指一次输入到模型中的样本数量，shuffle=True 表示每个 epoch 都会打乱数据以提高训练效果。

模型配置及训练

MistralConfig 设置: 设置与模型结构相关的超参数。特别是，将 pad_token_id、bos_token_id 和 eos_token_id 设为分词器中对应的标记 ID，以便模型正确处理填充、句子开始和结束标记。
创建模型 (MistralForCausalLM) 及优化器 (AdamW): 创建 MistralForCausalLM 模型，并将其移动到指定的设备（CPU 或 GPU）。使用 AdamW 优化器和 get_cosine_schedule_with_warmup 调度器来控制学习率，以优化模型。
train 函数: 与 train_seq_num.py 和 train_math.py 中使用的函数相同，使用常规训练循环对模型进行训练。

文本生成 (`generate_sql`): 从问题中推断 SQL 查询

generate_sql 函数: 使用训练好的模型根据给定的自然语言问题生成 SQL 查询。
- 首先，在输入的问题后添加 <sep> 标记，以形成 "问题 <sep> " 形式的提示。此提示明确告诉模型问题已结束，并且现在应该生成 SQL 查询。
- 处理填充标记的重要性: 训练数据包含直到 <eos> 标记的最大长度（max_length）。然而，如果训练数据中只有 "问题 <sep> " 部分而没有 SQL 和 <eos>（即 "问题 <sep> <pad> <pad> ..." 形式），则模型将无法学习在 <sep> 标记后应生成什么。因此，在生成阶段，模型可能只会生成 <sep> 之后的填充标记或完全不生成任何内容。为了避免这种情况，训练数据必须始终以 "问题 <sep> SQL<eos>" 的形式组成。
- 使用 temperature 参数来调整生成的 SQL 查询的多样性。
- 当模型生成 <eos> 或 <pad> 标记时，停止查询生成。

结果分析

样本输出: 在训练前输出 WikiSQL 数据集中的 3 个样本以检查数据格式。
训练结果: 通过观察训练过程中损失（loss）的减少，可以确认模型正在学习将自然语言问题转换为 SQL 查询的模式。
生成结果: 将验证数据集中的问题输入到模型中，并评估生成的 SQL 查询。重点关注生成的 SQL 查询是否语法正确以及是否准确反映了问题的意义。 train_sql.py 示例展示了如何使用 simple_mistral 模型执行更复杂的自然语言处理任务，即自然语言-SQL 转换。此示例强调了在数据预处理过程中正确使用特殊标记（如 <sep>、<eos>、<pad>）的重要性，以及训练数据的组成对模型生成能力的影响。

Code

import torch
import random
from transformers import T5Tokenizer, get_cosine_schedule_with_warmup
from dldna.chapter_09.mistral.examples.train_sql import MistralConfig, MistralForCausalLM, WikiSQLDataset, generate_sql
from torch.utils.data import DataLoader

random.seed(42)
torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Use T5Tokenizer as the tokenizer (use T5's vocab_size and pad/eos tokens)
tokenizer = T5Tokenizer.from_pretrained("t5-small")

# WikiSQL dataset (training: train, evaluation: validation)
max_length = 128
train_dataset = WikiSQLDataset("train", tokenizer, max_length=max_length)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)

valid_dataset = WikiSQLDataset("validation", tokenizer, max_length=max_length)
valid_loader = DataLoader(valid_dataset, batch_size=1, shuffle=True)

# Model configuration: Use MistralConfig and MistralForCausalLM provided by simple_mistral.py
# The model size is adjusted for educational purposes.
config = MistralConfig(
    vocab_size=tokenizer.vocab_size,
    hidden_size=512,
    intermediate_size=2048,
    num_hidden_layers=4,
    num_attention_heads=8,
    num_key_value_heads=4,     # num_attention_heads % num_key_value_heads == 0 must be true
    max_position_embeddings=max_length,
    sliding_window=max_length,
    use_cache=False,
    use_return_dict=True,
    pad_token_id=tokenizer.pad_token_id,  # Set the pad token id.
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id
)
model = MistralForCausalLM(config).to(device)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
num_epochs = 8  # Set the number of epochs small for the example
total_training_steps = num_epochs * len(train_loader)
scheduler = get_cosine_schedule_with_warmup(
    optimizer,
    num_warmup_steps=len(train_loader) // 5,
    num_training_steps=total_training_steps
)
    # Added code: Output WikiSQL data samples
print("=== WikiSQL Data Sample Output ===")
sample_count = 3  # Number of examples to output
for i in range(sample_count):
    input_ids, labels = train_dataset[i]
    decoded_text = tokenizer.decode(input_ids, skip_special_tokens=True)
    print(f"Sample {i+1}: {decoded_text}")


print("Start training...")
train(model, train_loader, optimizer, scheduler, num_epochs, device)

# Save the model: Save the final model to a file.
torch.save(model.state_dict(), "final_nl2sql_model.pth")

# Evaluation code part
print("\n=== Evaluation Examples ===")
for i, (input_ids, labels) in enumerate(valid_loader):
    if i >= 10:
        break
    # Keep special tokens with skip_special_tokens=False.
    full_text = tokenizer.decode(input_ids[0], skip_special_tokens=False)
    # Unify the tokens "sep>" and "eos>" to "<sep>" and "<eos>" respectively.
    full_text = full_text.replace("sep>", "<sep>").replace("eos>", "<eos>")
    
    if "<sep>" in full_text:
        # Split based on the first <sep>, then join all subsequent parts to restore the complete SQL.
        parts = full_text.split("<sep>")
        question = parts[0].strip()
        target_sql = "<sep>".join(parts[1:]).strip()
        # If target_sql ends with "<eos>", remove it.
        if target_sql.endswith("<eos>"):
            target_sql = target_sql[:-len("<eos>")].strip()
    else:
        question = full_text.strip()
        target_sql = ""

    generated_sql = generate_sql(model, tokenizer, question, max_length, device, temperature=0.7)
    # If there is a "sep>" token in generated_sql, extract the part after that token to use.
    # if "sep>" in generated_sql:
    #     generated_sql = generated_sql.split("sep>", 1)[1].strip()

    print(f"Question: {question}")
    print(f"Target SQL: {target_sql}")
    print(f"Generated SQL: {generated_sql}\n")

You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565

=== WikiSQL Data Sample Output ===
Sample 1: Tell me what the notes are for South Australia sep> SELECT Notes FROM table WHERE Current slogan = SOUTH AUSTRALIA eos>
Sample 2: What is the current series where the new series began in June 2011? sep> SELECT Current series FROM table WHERE Notes = New series began in June 2011 eos>
Sample 3: What is the format for South Australia? sep> SELECT Format FROM table WHERE State/territory = South Australia eos>
Start training...
Epoch 1/8, Average Loss: 10.5748, LR: 0.000000
Epoch 2/8, Average Loss: 9.7000, LR: 0.000001
Epoch 3/8, Average Loss: 7.2037, LR: 0.000001
Epoch 4/8, Average Loss: 5.5372, LR: 0.000001
Epoch 5/8, Average Loss: 4.5961, LR: 0.000001
Epoch 6/8, Average Loss: 4.0102, LR: 0.000002
Epoch 7/8, Average Loss: 3.6296, LR: 0.000002
Epoch 8/8, Average Loss: 3.3907, LR: 0.000002

=== Evaluation Examples ===
Question: Who was the minister for the CSV party with a present day end date? <unk>
Target SQL: SELECT Minister FROM table WHERE Party = csv AND End date = present day <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: Who was the minister for the CSV party with a present day end date? sep> FROM table WHERE60ed = s eos>

Question: What is the production number of From Hare to Heir? <unk>
Target SQL: SELECT SUM Production Number FROM table WHERE Title = from hare to heir <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: What is the production number of From Hare to Heir? sep>os FROM table WHERE Score = 0 eos>

Question: What was the score on January 12? <unk>
Target SQL: SELECT Score FROM table WHERE Date = january 12 <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: What was the score on January 12? sep>a Record FROM table WHERE #  eos>

Question: The race tony bettenhausen 200 has what smallest rd? <unk>
Target SQL: SELECT MIN Rd FROM table WHERE Name = Tony Bettenhausen 200 <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: The race tony bettenhausen 200 has what smallest rd? sep> Team FROM table WHERE Player = a ODi a eos>

Question: what is the club that was founded before 2007, joined prsl in 2008 and the stadium is yldefonso solá morales stadium? <unk>
Target SQL: SELECT Club FROM table WHERE Founded <unk> 2007 AND Joined PRSL = 2008 AND Stadium = yldefonso solá morales stadium <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: what is the club that was founded before 2007, joined prsl in 2008 and the stadium is yldefonso solá morales stadium? sep> ( for  for the highest FROM table WHERE Team = Rank  of vir AND COUNT  eos>

Question: Who is the co-contestant (yaar vs. Pyaar) with Vishal Singh as the main contestant? <unk>
Target SQL: SELECT Co-contestant (Yaar vs. Pyaar) FROM table WHERE Main contestant = vishal singh <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: Who is the co-contestant (yaar vs. Pyaar) with Vishal Singh as the main contestant? sep> SELECT  Record FROM table WHERE ts = 9kt AND Date = a eos>

Question: What season did SV Darmstadt 98 end up at RL Süd (1st)? <unk>
Target SQL: SELECT Season FROM table WHERE RL Süd (1st) = SV Darmstadt 98 <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: What season did SV Darmstadt 98 end up at RL Süd (1st)? sep> FROM table WHERE Away team = s s eos>

Question: What character was portrayed by the same actor for 12 years on Neighbours? <unk>
Target SQL: SELECT Character FROM table WHERE Duration = 12 years AND Soap Opera = neighbours <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: What character was portrayed by the same actor for 12 years on Neighbours? sep>FS Class FROM table WHERE Date = m ja eos>

Question: What was the score between Marseille and Manchester United on the second leg of the Champions League Round of 16? <unk>
Target SQL: SELECT 2nd leg score** FROM table WHERE Opponent = Marseille <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: What was the score between Marseille and Manchester United on the second leg of the Champions League Round of 16? sep>hes> d FROM table WHERE Date =s eos>

Question: Who was the Man of the Match when the opponent was Milton Keynes Lightning and the venue was Away? <unk>
Target SQL: SELECT Man of the Match FROM table WHERE Opponent = milton keynes lightning AND Venue = away <unk> <eos></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>
Generated SQL: Who was the Man of the Match when the opponent was Milton Keynes Lightning and the venue was Away? sep> with Cap? sep> SELECT Home team score FROM table WHERE Wilson AND jump = s eos>

点击查看内容（深入探讨：稳健的变压器设计和调试 - 实用指南）

坚固的变压器设计和调试 - 实用指南

包括像 simple_mistral 这样的高效架构在内，从头开始构建变压器模型是一项困难但值得的任务。理论理解固然重要，但在实际实现过程中经常会出现微妙的错误和性能瓶颈。本节将深入介绍设计、实现和调试变压器的实际策略，特别强调 simple_mistral 中使用的组件（RoPE, RMSNorm, Attention）。广泛地讨论单元测试，并探讨其他必要的调试和技术。

1. 单元测试的重要作用

构建像变压器这样复杂的模型时，单元测试不是可选而是必需的。它可以及早发现错误、防止回归(regression)并提供对实现的信心。经过良好测试的模型是 可靠的 模型。

每个模型源代码中都有 一个名为tests的目录用于存放单元测试。(例如：mistral/tests, phi3/tests)

变压器需要单元测试的原因

复杂性: 变压器由许多相互作用的模块（Attention, Feedforward 网络, Normalization, Embedding）组成。这些组件中的任何一个都可能容易出现错误。
微妙的错误: 许多变压器错误不会立即 明显地 显示出来。它们可能会导致性能下降或输出错误，而不是引起崩溃。单元测试可以捕捉到这些细微的错误。
数值稳定性: 深度学习模型，尤其是使用混合精度等技术时，容易受到数值问题（NaN, Inf, Vanishing/Exploding Gradients）的影响。单元测试有助于检测这些问题。
重构和修改: 在改进和优化模型时更改代码是不可避免的。单元测试确保这些更改不会破坏现有功能。
可重复性: 定义良好的测试有助于结果的可重复性。
缓存 (past_key_value): 当模型使用像 past_key_values 这样的缓存机制时，通过单元测试来验证形状、数据类型或设备相关的错误尤为重要。

有效单元测试的核心原则

测试驱动开发(Test-Driven Development, TDD): 理想情况下，在编写模型代码之前应先编写单元测试。这样可以清楚地思考每个组件的预期行为。
模块化: 将代码设计为由小而明确定义的函数和类组成，使其可模块化。这使得隔离和测试各个组件变得更加容易。
全面覆盖: 以高测试覆盖率为目标。测试所有重要的函数和方法。
边缘情况: 不仅要测试“正常情况”。还要测试边缘情况、边界条件以及潜在的错误场景。（例如：长度为0的序列，单元素批次，不同的数据类型）。
断言: 使用 assert 断言自由地验证代码是否按预期运行。尽量具体地编写断言。不仅要确保代码无崩溃地运行，还要 验证输出 是否正确。
Pytest: 本章的示例使用了 unittest 模块，但在 Python 中最推荐的是 pytest 框架。

变压器单元测试的重点领域 * 输入/输出 Shape: 变压器实现中最常见的错误类型是张量形状不正确。每个测试都应包含检查输出张量形状的断言。 * 数据类型: 检查张量中是否存在预期的数据类型（例如：torch.float32, torch.float16, torch.int64）。 * 设备放置: 使用GPU时，确保张量位于正确的设备上（CPU或GPU）。 * 数值稳定性: 特别是在进行softmax或归一化等操作后，检查张量中是否有NaN（非数字）和Inf。 * 梯度计算: 确保为所有可训练参数正确计算了梯度。 * 缓存 (past_key_value): 如前所述，缓存机制是错误的常见原因。彻底测试增量解码（incremental decoding）。

详细的单元测试示例 (RoPE, RMSNorm, Attention)

# test_rope.py
import unittest
import torch
from dldna.chapter_09.mistral.simple_mistral import MistralRotaryEmbedding, apply_rotary_pos_emb, rotate_half

# ...

# test_rms_norm.py
import torch
import pytest
from dldna.chapter_09.mistral.simple_mistral import PhiMiniRMSNorm

# ...

# test_attention.py
import torch
import pytest
from dldna.chapter_09.mistral.simple_mistral import PhiMiniConfig, PhiMiniAttention

# ... 

# 关于注意力的额外测试

def test_phi_mini_attention_zero_length_initial():
    # ... 

def test_phi_mini_attention_single_token_initial():
    # ... 
@pytest.mark.parametrize("batch_size", [1, 2, 4, 8])
def test_phi_mini_attention_various_batch_sizes(batch_size):
    # ...

@pytest.mark.parametrize("num_heads, num_kv_heads", [(8, 8), (8, 4), (8, 1)]) # MHA, GQA 情况
def test_phi_mini_attention_different_head_configs(num_heads, num_kv_heads):
    # ... 

@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
def test_phi_mini_attention_mixed_precision(dtype):
    # ... 

def test_phi_mini_attention_combined_mask():
    # ... 

def test_phi_mini_attention_long_sequence():
    # ... 

def test_phi_mini_attention_output_attentions_with_cache():
    # ...

2. 超越单元测试：其他调试策略

虽然单元测试是基础，但它们并不是调试工具的唯一手段。以下是其他重要的策略。

1. 日志记录 (Logging) * 战略性日志记录: 在代码中添加日志语句（print 语句或尽可能使用 logging 模块）以跟踪主要变量的值、张量的形状和执行流程。这可以帮助快速识别问题发生的位置。 * 控制详细级别: 让日志记录更详细，但提供一种方法来控制详细级别（例如：使用命令行标志或环境变量）。这样可以在调试时获得详细的输出信息，但在正常运行中避免过多的输出。

2. 可视化 (Visualization)

注意力权重: 通过可视化注意力权重来检查模型关注哪些令牌。这有助于识别注意力机制或位置嵌入的问题。
激活值: 可视化模型中的神经元激活。这可以帮助识别死亡神经元（总是处于非活动状态的神经元）或饱和神经元（始终处于最大值或最小值的神经元）。
梯度: 在训练过程中可视化梯度。这有助于检测消失或爆炸梯度。

3. 数值调试 (Numerical Debugging)

NaN/Inf 检查: 使用 torch.isnan() 和 torch.isinf() 来检查张量中是否有 NaN 或 Inf。这通常表示数值不稳定性。 python if torch.isnan(tensor).any() or torch.isinf(tensor).any(): print("检测到 NaN 或 Inf!")
梯度检查: 使用 torch.autograd.gradcheck 来验证自定义 autograd 函数是否正确计算了梯度。这在实现自己的注意力机制或其他复杂操作时尤其重要。
小型测试用例: 创建非常小且简单的测试用例（例如：单层、小型词汇表、短序列），可以手动计算预期输出。这有助于隔离错误。

4. 调试器 (pdb, IDE 调试器)

pdb (Python 调试器): 使用内置的 Python 调试器 (pdb) 逐行逐步执行代码、检查变量和设置断点。 python import pdb; pdb.set_trace() # 添加此行以设置断点。
IDE 调试器: 大多数 IDE（如 PyCharm, VS Code 等）都提供了用户友好的界面进行调试。

5. 性能分析 (Profiling)

PyTorch 性能分析器: 使用 PyTorch 性能分析器来识别代码中的性能瓶颈。这有助于找到可以优化速度或内存使用量的区域。
内存分析: 使用 memory_profiler 等工具跟踪内存使用情况并识别潜在的内存泄漏。

6. 为可调试性设计模型原则 * 保持简单(Keep it Simple): 从简单的模型开始，逐步增加复杂性。这样可以更容易地分离错误。 * 模块化(Modularity): 将代码划分为小而明确定义的模块。这使得单独组件更易于测试和调试。 * 断言(Assertions): 使用断言来验证预期条件，并尽早捕获错误。 * 注释(Comments)及文档(Documentation): 编写清晰简洁的注释和文档以解释代码逻辑，帮助用户（及其他人员）理解代码并识别潜在问题。 * 可重复性(Reproducibility): 使用固定的随机种子使结果具有可重复性。这对于调试和其他模型配置的比较至关重要。 * 单批次/小数据集过度拟合(Overfitting): 在使用大数据集训练之前，先对小数据集进行模型过度拟合。

7. 常见错误及预防方法

张量形状不正确: 特别是在reshape、transpose、concatenate等操作后重新检查张量的预期shape。在调试过程中频繁使用tensor.shape。
差一错误: 尤其是在处理序列和位置嵌入时要注意索引。
数据类型不匹配: 确保张量具有正确的数据类型（例如：float32 对 float16）。
设备不匹配: 确保所有张量都在相同的设备（CPU 或 GPU）上。
未初始化变量: 在使用前初始化所有变量。
掩码错误: 使用注意力掩码时，确保掩码正确应用且不会屏蔽重要信息。
past_key_values的错误使用: 确保遵循正确的使用方法。

将这些调试技术与对变压器模型基本原理的明确理解相结合，就可以解决最困难的实现问题。调试是一个重复的过程，因此要有耐心并系统地利用所有工具。

9.7 Gemma: 最新开源模型概览

Gemma 是 Google 在 2024 年 2 月公开的最新开源模型。尽管与 Mistral 相比，模型结构本身没有革命性的变化，但它反映了最新模型的趋势，并且在特定情况下可以有效利用，因此值得研究。Gemma 采用了类似于 LLaMA 和 Mistral 的基于 Transformer 的仅解码器（Decoder-only）模型架构。

研究 Gemma 的原因

反映最新模型趋势： Gemma 包含了最新的模型中广泛使用的组件，如 RoPE (Rotary Positional Embedding)、RMSNorm (Root Mean Square Layer Normalization) 和 GeGLU 激活函数。这些元素有助于提高模型的性能和效率，并帮助理解最新趋势。RoPE 通过高效地编码相对位置信息来增强长序列处理能力，而 RMSNorm 通过去除层归一化中的均值中心操作提高了计算效率。GeGLU 是 GLU（门控线性单元）的一个变体，通过引入非线性增加了模型的表达能力。
多种模型尺寸： Gemma 提供了 2B、7B、9B 和 27B 四种版本。这为计算资源有限的用户提供了使用相对较小的模型（2B）进行实验的机会。更大的模型（27B）虽然可以提供更高的性能，但需要更多的计算资源。用户可以根据自己的环境和需求选择合适的模型尺寸。
与 Google 生态系统的集成： Gemma 与 Google 的 Gemini 项目相关，并且可以轻松地与 Google Cloud、Vertex AI 等平台集成。对于主要使用 Google 平台的开发者来说，Gemma 是一个有用的选择。Google Cloud 的 Vertex AI 提供了一个用于机器学习模型训练、部署和管理的综合平台，通过与这些平台的兼容性，Gemma 可以提高开发效率。
开源模型的可访问性： Gemma 采用 Apache 2.0 许可证公开发布，允许自由使用、分发和修改，包括商业用途。

Gemma 模型的特点（与 Mistral 相比）

特征	Gemma	Mistral
公开时间	2024年2月	2023年9月
模型大小	2B, 7B, 9B, 27B	7.3B
基础架构	Transformer (Decoder-only)	Transformer (Decoder-only)
位置嵌入	RoPE	RoPE
归一化	RMSNorm	RMSNorm
激活函数	GeGLU	SwiGLU
注意力机制	Multi-Head Attention (MHA), GQA	Grouped-Query Attention (GQA), SWA
上下文窗口	最大 8192 个令牌	最多 131,000 个令牌
主要特点	多种尺寸，支持 Google 生态系统，GeGLU，宽广的上下文窗口	GQA 和 SWA 实现高效推理，处理长上下文
创新性（比较）	较低	较高

相似之处: Gemma 和 Mistral 都是基于 Transformer 的 Decoder-only 模型，并使用 RoPE、RMSNorm 等类似组件。这些组件有助于提高模型的效率和性能。
不同点:
- Gemma 使用 GeGLU 作为激活函数，而 Mistral 使用 SwiGLU（SiLU 的变体）。GeGLU 将输入分为两个线性变换，一个充当门控作用，另一个与之相乘以生成结果。
- Gemma 使用 Multi-Head Attention (MHA) 或 Grouped-Query Attention (GQA)，而 Mistral 结合使用 GQA 和 Sliding Window Attention (SWA) 以提高效率。GQA 是一种通过减少键(K)和值(V)头的数量来减少内存使用量和计算量的方法。SWA 通过为每个令牌生成一个掩码，使其仅在固定的范围（窗口）内执行注意力操作，从而减少计算复杂度。

结论

尽管 Gemma 在模型结构本身方面不如 Mistral 创新，但作为最新的开放模型，它具有以下意义。

掌握最新技术趋势: 通过 Gemma，可以了解 RoPE、RMSNorm、GeGLU 等最新模型中广泛使用的组件的实现和工作方式。
多种模型选择: Gemma 提供了 2B、7B、27B 等不同大小的模型，使用户可以根据自己的计算环境进行选择。
利用 Google 生态系统: 对于使用 Google 平台的用户来说，Gemma 可以提供比其他模型更好的集成和支持。
开放模型的可访问性: 任何人都可以轻松访问并为社区做出贡献。因此，Gemma 更应该关注的是其作为反映最新技术趋势的开放模型的实际价值，以及与 Google 生态系统的联动可能性，而不是模型本身的创新性。

9.8 Phi-3 : 小但强大的语言模型

在第9.6节和第9.7节中，我们通过Mistral和Gemma模型探讨了高效语言模型架构的关键要素。本节将直接实现并分析由Microsoft开发的Phi-3 Mini模型，解析其尽管尺寸较小但仍能表现出色的原因。

Phi-3 Mini是微软在2024年4月发布的小型语言模型(SLM, Small Language Model)。拥有3.8B参数的Phi-3 Mini在多个基准测试中展示了与更大规模模型如Mistral (7B)或Gemma (7B)竞争的性能，展现了轻量级模型的可能性。特别是，Phi-3 Mini强调了“高质量数据”和“高效架构”的重要性，提出了超越单纯模型尺寸竞赛的新方向。这种理念在“Textbooks Are All You Need”这一口号中得到了很好的体现。simple_phi3.py简化实现了Phi-3 Mini的核心组件代码。完整代码位于chapter_09/phi3。

9.8.1 `simple_phi3` 模型

simple_phi3 是为了教学目的实现的Phi-3 Mini模型。与第9.6章的simple mistral相比，如下所示：

模型功能差异总结

功能	Simple Phi-3	Simple Mistral
注意力机制	多头注意力 (MHA)	分组查询注意力 (GQA) + 滑动窗口注意力 (SWA)
激活函数	GELU (tanh近似)	SiLU
归一化	RMSNorm	RMSNorm
位置编码	RoPE	RoPE
`past_key_value`	支持（缓存）	支持（缓存）
滑动窗口	不支持	支持
GQA	不支持 (使用MHA, K=V=Q, 设置`num_key_value_heads`)	支持
缩放点积注意力	使用`F.scaled_dot_product_attention`	使用`F.scaled_dot_product_attention`
增强的RoPE缓存	在`forward`方法中高效管理`cos`, `sin`缓存，必要时通过`_set_cos_sin_cache`更新。在增量解码时使用`apply_rotary_pos_emb_single`函数优化RoPE应用逻辑，最小化重复计算。	通过`_set_cos_sin_cache`方法生成`cos_cached`, `sin_cached`，并在`forward`中使用。在`apply_rotary_pos_emb`中可以为查询和键使用不同的position ID。
注意力掩码优化	使用`scaled_dot_product_attention`函数，高效结合`attention_mask`和`causal_mask`，减少不必要的计算	使用`scaled_dot_product_attention`函数处理`attention_mask`, `sliding_window_mask`
`return_dict`	通过`return_dict`灵活且明确地返回输出。	通过`return_dict`返回输出。
权重共享	在`post_init`中绑定嵌入权重和输出层权重，减少参数数量并提高性能	没有明确提到权重共享

主要改进

以上即为翻译内容。 * Multi-Head Attention (MHA): 使用普通的 MHA 替代 Mistral 的 GQA（分组查询注意力）。Phi-3 Mini 展示了即使没有 GQA 也能实现足够的性能。 * 改进的 RoPE 缓存: 在 forward 方法中有效管理 cos 和 sin 缓存，并通过 _set_cos_sin_cache 只在需要时更新。此外，在增量解码时使用 apply_rotary_pos_emb_single 函数来优化 RoPE 的应用并最小化重复计算。 * Attention Mask 优化: 使用 scaled_dot_product_attention 函数的同时，高效地结合 attention_mask 和 causal_mask 以减少不必要的计算。 * Weight Tying: 在 post_init 中绑定（tying）嵌入权重和输出层权重以减少参数数量并提高性能。

现在我们详细探讨 simple_phi3 模型的关键组成部分。

1. PhiMiniConfig: 模型配置

PhiMiniConfig 类定义了模型的超参数。遵循 Phi-3 Mini 的设置，由于 Mistral 中已经详细解释过，这里将省略这些内容。

2. PhiMiniRMSNorm: RMS 正则化

PhiMiniRMSNorm 类实现了 RMSNorm（均方根层正则化），与 Mistral 相同。

3. PhiMiniRotaryEmbedding: RoPE 实现（改进的缓存）

PhiMiniRotaryEmbedding 类实现 RoPE（旋转位置嵌入）。虽然与 Mistral 的 MistralRotaryEmbedding 类似，但通过以下核心改进极大地提高了缓存效率。

在 forward 方法中管理缓存:
- 在 forward 方法中直接使用 cos_cached 和 sin_cached。即如果有已计算的值，则立即使用。
- 如果 seq_len 大于 max_seq_len_cached，即需要为新的序列长度创建缓存时，才调用 _set_cos_sin_cache 方法更新缓存。这防止了不必要的缓存生成，并尽可能地重用已计算的值。
max_seq_len_cached, cos_cached, sin_cached 实例变量:
- max_seq_len_cached: 存储当前为止缓存的最大序列长度。
- cos_cached, sin_cached: 存储预计算的余弦和正弦值。
- 通过将这些变量作为实例变量管理，可以在每次调用 forward 方法时重用已生成的值而不是重新创建它们，从而提高效率。
增量解码优化:
- apply_rotary_pos_emb_single: 在使用 past_key_value 的增量解码情况下，可以只为新的 token 应用 RoPE 而不是整个序列。由于先前 token 的 RoPE 结果已经存储在 past_key_value 中，因此可以避免重复计算。

这些改进显著提高了 RoPE 操作的效率，特别是在处理长序列或文本生成时提供性能优势。

4. PhiMiniAttention: 注意力机制（MHA, 高效的 RoPE 应用）

PhiMiniAttention 类实现了 Phi-3 Mini 的注意力机制。虽然使用了普通的 Multi-Head Attention (MHA)，而不是 Mistral 的 GQA，但通过优化 RoPE 应用方式提高了效率。

MHA (Multi-Head Attention): 查询(Q)、键(K)、值(V)头的数量都相同。
高效的 RoPE 应用:
- 根据 past_key_value 是否存在，以不同的方式生成 position IDs。
  - 如果没有 past_key_value（通常情况）：为整个序列生成 position IDs（从 0 到 q_len - 1）。
  - 如果有 past_key_value（增量解码时）：生成新令牌的 position ID（从 past_len 到 past_len + q_len - 1），以及整个键序列的 position IDs（从 0 到 past_len + q_len - 1）。
- 通过 apply_rotary_pos_emb_single 函数，在存在 past_key_value（增量解码时）的情况下，仅对新令牌(query)应用 RoPE。
KV 缓存: 与 Mistral 相同，通过 past_key_value 缓存上一步的键/值张量以提高推理速度。

5. 辅助函数: `rotate_half`, `apply_rotary_pos_emb`, `apply_rotary_pos_emb_single`

rotate_half: 是实现 RoPE 所需的辅助函数，与 Mistral 相同。
apply_rotary_pos_emb: 对查询(q)和键(k)张量应用 RoPE。与 Mistral 不同的是，它只接收一个 position_ids（同时应用于查询和键）。
apply_rotary_pos_emb_single: 在使用 past_key_value 的增量解码情况下，对输入张量 x (query 或 key) 应用 RoPE。

6. PhiMiniMLP: 前馈网络

PhiMiniMLP 类实现了前馈网络，并且与 Mistral 相比没有太大差异，使用了 GELU 激活函数。

7. PhiMiniDecoderLayer: 解码器层

PhiMiniDecoderLayer 类采用了 Pre-Norm 结构和残差连接，与 Mistral 相同。

8. PhiMiniModel: 完整模型

PhiMiniModel 类构建了整个 Phi-3 Mini 模型，并且与 Mistral 非常相似。

9. PhiMiniForCausalLM: 添加用于语言建模的头

PhiMiniForCausalLM 类向 PhiMiniModel 添加了一个用于语言建模的头(lm_head)。

post_init 方法:
- 执行权重初始化。（与 Mistral 相似）
- Weight Tying: 将嵌入权重（self.transformer.embed_tokens.weight）和输出层权重（self.lm_head.weight）绑定(tie)。这可以减少参数数量，防止过拟合，并通常提高性能。
generate 函数: 用于文本生成的函数，在增量解码时为了处理 RoPE 相关问题，如果存在 past_key_values，则只传递最后一个令牌到 forward() 而不是整个序列。

9.8.2 `simple_phi3` 模型示例：复合公式计算

作为第 9.8.1 节中讨论的 simple_phi3 模型的实际应用示例，我们将测试其复合公式计算能力。通过这个示例，我们验证像 Phi-3 Mini 这样的小型语言模型 (SLM) 是否能够处理不仅仅是简单的加减法，还包括乘法和包含括号的复杂公式，并分析其性能和局限性。

示例代码位置是 chapter_09/phi3/examples/train_math.py。

示例的意义

验证 SLM 的能力： 通过高质量的数据和高效架构，小模型也能解决复杂问题。
评估推理能力： 基于学习到的运算规则，而不是简单的记忆，对新公式的结果进行推理的能力。
探索实用性可能性： 复合公式计算是自然语言处理、数据分析等领域中可应用的基本能力。通过这个示例可以了解 SLM 的实际应用潜力。

训练数据形式

使用 create_complex_arithmetic_data 函数生成了如下形式的复合公式数据：

两个或三个数字 (1 ~ 50)
使用两种运算符 (+, -, *) 中的任意两个
可选择地使用括号 (())
形式为 表达式=结果<eos>（例如：(12+7)*3=57<eos>, 12+7*3=33<eos>）

训练结果

样本 1: 41*8-2=326<eos>
样本 2: 15+(9*48)=447<eos>
样本 3: 35-6+38=67<eos>
样本 4: 6*14*15=1260<eos>
样本 5: 36*(13*46)=21528<eos>

...(训练日志省略)...

提示：'23-23-50=' --> 生成结果：'23-23-50=-50'  (答案: 23-23-50=-50<eos>)
提示：'39-46-15=' --> 生成结果：'39-46-15=-22'  (答案: 39-46-15=-22<eos>)
提示：'(33-30)+30=' --> 生成结果：'(33-30)+30=33'  (答案: (33-30)+30=33<eos>)
提示：'30+14*27=' --> 生成结果：'30+14*27=412'  (答案: 30+14*27=408<eos>)

结果分析

大部分准确计算： 在大多数测试案例中，生成的结果与正确答案一致或非常接近。这表明 simple_phi3 模型很好地学习了复合公式的运算规则。
部分误差发生： 当包含乘法或数字较大时，可能会出现误差。这可能是由于模型的大小限制、训练数据多样性不足等多种因素导致的。
括号处理能力： 在处理包含括号的公式时也表现出相对准确的结果，表明该模型具有理解上下文(context)和运算顺序的能力。

结论

尽管 simple_phi3 模型只有约 12 万个参数，是非常小的模型，但在复合公式计算中仍显示出接近 80% 的高正确率。这表明它已经学习了括号处理、运算顺序等复杂规则的重要部分。与通常拥有数十亿 (Billion) 参数的大规模语言模型 (LLM) 相比，simple_phi3 尽管大小仅为 0.12M，却展现了令人印象深刻的结果。

Code

import torch
import random
from dldna.chapter_09.phi3.examples.train_complex_math import PhiMiniConfig, PhiMiniForCausalLM, ComplexArithmeticDataset, train, create_complex_arithmetic_data, create_tokenizer, create_reverse_tokenizer, generate_text
from torch.utils.data import DataLoader

random.seed(42)
torch.manual_seed(42)

device = "cuda" if torch.cuda.is_available() else "cpu"

# Hyperparameters
num_samples = 100000      # Sufficiently large amount of data
max_value = 50           # Maximum value of operands (for slightly complex calculations)
seq_length = 30          # Complex arithmetic problems can have somewhat long expressions
batch_size = 128
epochs = 30
learning_rate = 1e-3

# Data generation
complex_data = create_complex_arithmetic_data(num_samples, max_value)
print("Training data examples:")
for i in range(5):
    print(f"Sample {i+1}: {complex_data[i]}")

# Create tokenizer and reverse tokenizer
tokenizer = create_tokenizer()
reverse_tokenizer = create_reverse_tokenizer(tokenizer)
updated_vocab_size = len(tokenizer)

# Configure Dataset and DataLoader
dataset = ComplexArithmeticDataset(complex_data, seq_length, tokenizer)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# PhiMini Model Configuration
config = PhiMiniConfig(
    vocab_size=updated_vocab_size,
    hidden_size=64,              # Small model size for experimentation
    intermediate_size=128,
    num_hidden_layers=3,
    num_attention_heads=8,
    num_key_value_heads=8,        # K=V=Q
    max_position_embeddings=128,
    use_cache=False,
    use_return_dict=True,
)
config.pad_token_id = tokenizer["<pad>"]
config.eos_token_id = tokenizer["<eos>"]

# Create PhiMini For CausalLM Model
model = PhiMiniForCausalLM(config).to(device)
print("Total Trainable Parameters:", sum(p.numel() for p in model.parameters() if p.requires_grad))

# weight tying (share weights between embedding and lm_head)
model.lm_head.weight = model.transformer.embed_tokens.weight

optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=1e-5)

# Model Training
print("Start training...")
train(model, dataloader, optimizer, scheduler, epochs, device)

# Save Model
save_path = "phimini_complex_math.pt"
torch.save(model.state_dict(), save_path)
print(f"Model saved: {save_path}")

# Load Saved Model (create a new model object before testing and load_state_dict)
loaded_model = PhiMiniForCausalLM(config).to(device)
loaded_model.load_state_dict(torch.load(save_path, map_location=device))
loaded_model.eval()

# Generate and Print Results with Test Set, Calculate Accuracy
print("\nTest sample generation results:")
test_samples = random.sample(complex_data, 10)
correct_count = 0
for sample in test_samples:
    prompt = sample.split('=')[0] + '='
    generated = generate_text(loaded_model, prompt, tokenizer, reverse_tokenizer, seq_length, device, temperature=0.1)  # Reduce temperature for testing
    answer = sample.split('=')[1].replace('<eos>', '')

    if generated.split('=')[1] == answer:
        correct_count += 1
    print(f"Prompt: '{prompt}' --> Generated result: '{generated}'  (Correct answer: {sample})")

accuracy = (correct_count / len(test_samples)) * 100
print(f"\nOverall accuracy: {accuracy:.2f}% ({correct_count}/{len(test_samples)})")

Training data examples:
Sample 1: 41*8-2=326<eos>
Sample 2: 15+(9*48)=447<eos>
Sample 3: 35-6+38=67<eos>
Sample 4: 6*14*15=1260<eos>
Sample 5: 36*(13*46)=21528<eos>
Total Trainable Parameters: 126208
Start training...
Epoch 1/30, Avg Loss: 0.7439, LR: 0.000997
Epoch 2/30, Avg Loss: 0.6393, LR: 0.000989
Epoch 3/30, Avg Loss: 0.6139, LR: 0.000976
Epoch 4/30, Avg Loss: 0.5919, LR: 0.000957
Epoch 5/30, Avg Loss: 0.5825, LR: 0.000934
Epoch 6/30, Avg Loss: 0.5753, LR: 0.000905
Epoch 7/30, Avg Loss: 0.5696, LR: 0.000873
Epoch 8/30, Avg Loss: 0.5649, LR: 0.000836
Epoch 9/30, Avg Loss: 0.5599, LR: 0.000796
Epoch 10/30, Avg Loss: 0.5558, LR: 0.000753
Epoch 11/30, Avg Loss: 0.5522, LR: 0.000706
Epoch 12/30, Avg Loss: 0.5479, LR: 0.000658
Epoch 13/30, Avg Loss: 0.5443, LR: 0.000608
Epoch 14/30, Avg Loss: 0.5409, LR: 0.000557
Epoch 15/30, Avg Loss: 0.5370, LR: 0.000505
Epoch 16/30, Avg Loss: 0.5339, LR: 0.000453
Epoch 17/30, Avg Loss: 0.5307, LR: 0.000402
Epoch 18/30, Avg Loss: 0.5280, LR: 0.000352
Epoch 19/30, Avg Loss: 0.5242, LR: 0.000304
Epoch 20/30, Avg Loss: 0.5217, LR: 0.000258
Epoch 21/30, Avg Loss: 0.5189, LR: 0.000214
Epoch 22/30, Avg Loss: 0.5161, LR: 0.000174
Epoch 23/30, Avg Loss: 0.5137, LR: 0.000137
Epoch 24/30, Avg Loss: 0.5120, LR: 0.000105
Epoch 25/30, Avg Loss: 0.5101, LR: 0.000076
Epoch 26/30, Avg Loss: 0.5085, LR: 0.000053
Epoch 27/30, Avg Loss: 0.5073, LR: 0.000034
Epoch 28/30, Avg Loss: 0.5062, LR: 0.000021
Epoch 29/30, Avg Loss: 0.5055, LR: 0.000013
Epoch 30/30, Avg Loss: 0.5050, LR: 0.000010
Model saved: phimini_complex_math.pt

Test sample generation results:
Prompt: '23-23-50=' --> Generated result: '23-23-50=-50'  (Correct answer: 23-23-50=-50<eos>)
Prompt: '39-46-15=' --> Generated result: '39-46-15=-22'  (Correct answer: 39-46-15=-22<eos>)
Prompt: '(33-30)+30=' --> Generated result: '(33-30)+30=33'  (Correct answer: (33-30)+30=33<eos>)
Prompt: '30+14*27=' --> Generated result: '30+14*27=408'  (Correct answer: 30+14*27=408<eos>)
Prompt: '(13-22)-18=' --> Generated result: '(13-22)-18=-27'  (Correct answer: (13-22)-18=-27<eos>)
Prompt: '9-15+12=' --> Generated result: '9-15+12=6'  (Correct answer: 9-15+12=6<eos>)
Prompt: '28*(3+31)=' --> Generated result: '28*(3+31)=960'  (Correct answer: 28*(3+31)=952<eos>)
Prompt: '24*(12+1)=' --> Generated result: '24*(12+1)=320'  (Correct answer: 24*(12+1)=312<eos>)
Prompt: '(1-33)+26=' --> Generated result: '(1-33)+26=-6'  (Correct answer: (1-33)+26=-6<eos>)
Prompt: '24+47+6=' --> Generated result: '24+47+6=77'  (Correct answer: 24+47+6=77<eos>)

Overall accuracy: 80.00% (8/10)

结语

第9章追溯了从2017年Transformer架构在“Attention is All You Need”这篇标志性论文中首次亮相，到2025年的今天，以效率和可扩展性为核心驱动的Transformer发展的历程。

最初的Transformer展示了突破性的性能，但面对着随着序列长度增加而急剧增长的计算量和内存使用这一根本局限。第9章深入探讨了为克服这些限制所做的不断努力，包括软件方法（9.2节）、硬件与软件的结合（9.3节），以及模型可扩展性方面的各种技术创新（9.4节）。从RoPE和FlashAttention的实现示例（9.5节）到Mistral、Gemma、Phi-3 Mini等最新模型的架构分析（9.6, 9.7, 9.8节），通过理论与实际实施的结合探究，突显了高效的Transformer架构。

得益于这些技术进步，Transformer现已发展成为能够理解更长上下文、解决更复杂问题并应用于更广泛领域的强大工具。可以看出，在Transformer从单纯的语言模型成长为推动人工智能技术发展的核心动力的过程中，效率和可扩展性发挥了重要作用。

当然，仍有许多挑战需要克服。随着模型规模的扩大，能源消耗增加、偏见与有害性问题以及模型解释性问题是未来我们需要解决的重要挑战。为了更安全、可靠且能与人类和谐合作的AI系统的研究将持续进行。

现在，在第10章和第11章中，我们将开始探索Transformer超越单一文本领域，迈向整合图像、音频、视频等多种数据类型的多模态(Multimodal)世界的旅程。融合多种模态信息以获得更丰富强大的表达能力的多模态模型能够实现更复杂的推理。围绕结合文本与图像的先锋模型ViT、CLIP、DALL-E、Stable Diffusion、Flamingo、GATO、Gemini等，我们将探索多模态注意力机制及其无限的应用可能性。第9章中讨论的效率和可扩展性的创新将成为10. 11章中展示的多模态Transformer未来的坚实基础。

点击查看内容（深度解析：MoE 架构的理论演变和最新技术趋势）

MoE（Mixture of Experts）架构的理论进化和最新技术趋势

在大规模语言模型(LLM)的发展中，Mixture of Experts (MoE)作为解决模型容量与计算效率平衡的创新框架崭露头角。MoE通过组合多个“expert”网络，并通过gating network根据输入选择性地激活合适的expert来工作。这里我们将深入剖析MoE的核心机制，并系统地整理反映最新研究趋势的扩展理论。

1. MoE的理论基础

1.1 基本组成成分

Expert 网络: 通常由前馈神经网络(Feedforward Neural Network, FFN)构成的 N个expert网络 \(\{E_i\}_{i=1}^N\)。每个expert接收输入 \(x\) 并生成输出 \(E_i(x)\)。
Gating 网络: Gating网络 \(G\) 接收输入 \(x\) 并为每个expert输出权重（概率）。这些权重表示对于输入 \(x\) 哪个expert最适宜。gating网络的输出 \(G(x)\) 是 N维向量，其中每个元素 \(G(x)_i\) 表示第 i 个expert的权重。
最终输出: MoE模型的最终输出 \(y\) 通过专家输出的加权和计算得出。

\(y = \sum_{i=1}^{N} G(x)_i E_i(x)\)

1.2 稀疏MoE与密集MoE

Dense MoE: 所有expert对所有输入进行计算，gating网络通过softmax函数确定每个expert输出的权重。 (\(G(x) = \text{softmax}(W_g x)\))
Sparse MoE: 每个输入仅激活少数几个expert。gating网络使用Top-k gating（选择具有最大 k 个值的expert）或Noisy Top-k gating（如GShard, Switch Transformer）。

1.3 数学形式化和变分推理视角

当将MoE系统重新解释为概率图形模型时，观测数据 \(\mathbf{x}\) 和潜在变量 \(\mathbf{z}\) (专家选择指示器)的联合分布可以建模如下。

\(p(\mathbf{x}, \mathbf{z}|\theta) = p(\mathbf{z}|\theta_g)p(\mathbf{x}|\mathbf{z},\theta_e)\)

其中，\(\theta_g\) 表示gating网络的参数，\(\theta_e\) 表示专家网络的参数。在变分推理框架中，证据下界(Evidence Lower Bound, ELBO)如下导出。

\(\mathcal{L}(\theta, \phi) = \mathbb{E}_{q_\phi(\mathbf{z}|\mathbf{x})}[\log p_\theta(\mathbf{x}|\mathbf{z})] - D_{KL}(q_\phi(\mathbf{z}|\mathbf{x}) \parallel p(\mathbf{z}))\)

这种方法通过将MoE学习过程重新定义为贝叶斯推理体系，为专家之间知识分割提供了理论基础。特别是Gumbel-Softmax重参数化技术通过对离散的专家选择过程进行连续近似，使得梯度下降法的应用成为可能。

\(\mathbf{z} = \text{softmax}((\log \boldsymbol{\pi} + \mathbf{g})/\tau)\)

其中 \(\mathbf{g}\) 表示Gumbel噪声，\(\tau\) 表示温度参数。

2. 稀疏MoE的结构创新

2.1 层次专家划分（Hierarchical Expert Partitioning）

DeepSeek-V2 引入的多头潜在注意力（Multi-Head Latent Attention, MLA）显著减少了键值缓存 [5, 6]。这是通过将专家层次划分为空间划分（Spatial Partitioning）和功能划分（Functional Partitioning）的方法实现的。

\(E_i(\mathbf{x}) = \sum_{h=1}^H W_{h,i}^o \cdot \text{GeLU}(W_{h,i}^k \mathbf{x} \oplus W_{h,i}^v \mathbf{x})\)

每个专家内的注意力头独立地担任子专家的角色，并通过共享基矩阵（shared basis matrices）最大化参数效率。

2.2 动态拓扑适应

Mixtral 8x7B 模型引入了一种机制，根据输入数据动态重构专家连接结构。路由网络已经从简单的专家选择演变为可以调节专家间连接强度的图神经网络（Graph Neural Network）。

\(A_{ij}^{(l)} = \sigma(f_\phi(\mathbf{h}_i^{(l)}, \mathbf{h}_j^{(l)}))\)

其中 \(A_{ij}\) 表示专家 \(i\) 和 \(j\) 之间的连接权重，通过分层注意机制实现了多尺度特征提取。

3. MoE 模型的优势和优化

3.1 优势

模型容量增加: 增加专家数量可以大幅增加参数量，但计算成本相对较低。
计算效率（稀疏 MoE）: 每个 token 只激活少数 expert，因此 FLOPs 较低。
扩展法则: MoE 模型倾向于遵循比密集模型更有利的扩展法则。
微调: 可以对特定专家进行微调，使其专门化于特定任务。

3.2 最优化理论的创新

平衡约束最优化（Balanced Optimization）: 引入双分解（Dual Decomposition）技术解决专家负载不均衡问题，利用拉格朗日乘数法显式地限制专家利用率的标准差。

\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \lambda \sum_{i=1}^N (\mathbb{E}[u_i] - \bar{u})^2\)

其中 \(u_i\) 表示第 \(i\) 个专家的利用率，\(\bar{u}\) 表示目标平均利用率。
知识蒸馏的多层化: 提出了反映 MoE 层次结构的层次知识蒸馏（Hierarchical Knowledge Distillation）。 \(\mathcal{L}_{KD} = \sum_{l=1}^{L}\alpha_{l}D_{KL}(g^{\text{teacher}}_{l} || g^{\text{student}}_{l})\) 通过在每个 MoE 层 \(l\) 最小化门分布 \(g_l\) 的 KL 散度，使专家特化知识的转移成为可能。

4. MoE 模型的例子和局限

4.1 示例

GShard: Google的稀疏MoE模型。使用Noisy Top-k gating。
Switch Transformer: Google的稀疏MoE模型。每个令牌仅由一个专家处理（k=1）。
GLaM: Google的稀疏MoE模型（1.2T参数）。
Mistral 8x7B: Mistral AI的稀疏MoE模型。使用Top-2 gating。

4.2 局限与挑战

专家不平衡 (Load Imbalance): 过多令牌分配给特定专家。（解决：Noisy Top-k Gating，负载平衡损失，限制专家容量）
门控网络学习的困难: 学习有效选择/组合专家的难度。
通信成本（分布式学习）: 专家之间通信成本可能增加。
知识蒸馏的困难: 由于MoE模型的大小，难以将知识蒸馏到较小的模型中。

5. 物理实现的最前沿

5.1 稀疏专家激活硬件

NVIDIA H100 Tensor Core GPU引入了专门用于MoE的稀疏执行单元，加速Top-k路由计算。 * 动态线程控制 (Dynamic Warp Control): 按专家组管理独立执行流 * 分层共享内存 (Hierarchical Shared Memory): 优化专家间中间结果共享 * 异步模型并行化 (Asynchronous Model Parallelism): 在分布式执行专家时最小化延迟

5.2 量化专家交换

最近的研究开发了一种技术，将专家参数量化为4位以减少通信带宽[5]。应用差分量化（Differential Quantization）技术。 \(\Delta W_{i} = \text{sign}(W_{i}-\hat{W})\cdot 2^{\lfloor \log_{2}|W_{i}-\hat{W}|\rfloor}\) 其中，\(\hat{W}\)表示共享基矩阵，并且仅量化专家的偏差以最小化精度损失。

6. 理论扩展的最新趋势

6.1 连续专家空间(Continuous Expert Space)

2025年Google DeepMind的最新研究提出了将专家建模为连续空间上的分布而非离散实体的CES-MoE。利用基于布朗运动（Brownian Motion）的专家扩散模型。

\(dE_t = \mu(E_t,t)dt + \sigma(t)dW_t\)

这种方法可以模拟专家特征的渐进演化，并在动态领域适应（Dynamic Domain Adaptation）中表现出卓越性能。

6.2 基于神经微分方程的专家

下一代MoE架构正在研究用神经微分方程（Neural ODE）替代专家网络的研究。

\(\frac{d\mathbf{h}(t)}{dt} = f_\theta(\mathbf{h}(t), t)\)

通过这种方式可以建模专家的时间演化特性，并在长期推理（Long-horizon Inference）任务中实现了性能提升。

7. 挑战与未来方向

7.1 对理论限制的深入分析

信息瓶颈: 由于路由器的信息处理容量限制导致的专家选择偏差
非凸优化: 专家-门联合空间中的多极小值问题
知识重叠: 缺乏关于专家间重复特征学习的理论依据 #### 7.2 下一代研究框架
概率微分几何 (Stochastic Differential Geometry)
- 通过专家流形的曲率分析实现有效的探索策略
量子叠加专家 (Quantum Superposition Experts)
- 利用基于量子比特的专家叠加态
生物塑料模仿
- 应用突触可塑性原理进行动态专家重构

8. 实践应用案例研究

8.1 超大规模推理系统设计

Naver 的 HyperClova X-MoE 系统分层集群部署了 1,024 个专家。

三级层次路由：逐级筛选专家，从集群→机架→节点
动态布局重构：基于 RL 的专家位置优化
混合精度缓存：管理热专家的 FP8 和冷专家的 FP16

8.2 跨模态应用扩展

OpenAI 的 GPT-4o 将 MoE 应用于多模态学习。

\(\mathbf{h}_{\text{fused}} = \sum_{i=1}^N G(\mathbf{x}_{\text{text}} \oplus \mathbf{x}_{\text{image}})_i E_i(\mathbf{x}_{\text{text}}, \mathbf{x}_{\text{image}})\)

在文本-图像联合嵌入空间中激活专家，以提高跨模态推理性能。

参考文献:

[1] Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv preprint arXiv:2101.03961.

[2] Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. arXiv preprint arXiv:1701.06538.

[3] Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive mixtures of local experts. Neural computation, 3(1), 79-87.

[4] NVIDIA 开发者博客. (2024). 在 LLM 架构中应用专家混合。 https://developer.nvidia.com/zh-cn/blog/applying-mixture-of-experts-in-llm-architectures/

[5] DeepSeek-V2 相关资料： * Modu 实验室博客。 https://modulabs.co.kr/blog/deepseek-r1-introduction * HyperLab. https://hyperlab.hits.ai/blog/ai-deepseek * Wikidocs. https://wikidocs.net/275230 [6] Chung, E. (2023). 趋势变压器之后的下一代架构 - MoE, SSM, RetNet, V-JEPA. Velog. https://velog.io/@euisuk-chung/%ED%8A%B8%EB%A0%8C%EB%93%9C-%ED%8A%B8%EB%A0%8C%EC%8A%A4%ED%8F%AC%EB%A8%B8-%EC%9D%B4%ED%9B%84%EC%9D%98-%EC%B0%A8%EC%84%B8%EB%8C%80-%EC%95%84%ED%82%A4%ED%85%8D%EC%B3%90-MoE-SSM-RetNet-V-JEPA

[7] The Moonlight. (2024). GG MoE 与 MLP 在表格数据上的比较. https://www.themoonlight.io/ko/review/gg-moe-vs-mlp-on-tabular-data

[8] Unite.AI. (2024). Mistral AI 的最新专家混合 (MoE) 8x7B 模型. https://www.unite.ai/ko/mistral-ais-latest-mixture-of-experts-moe-8x7b-model/

[9] Turing Post (2024) MS EUREKA 基准. [https://turingpost.co.kr/p/ms-eureka-benchmark](https://turingpost.co

练习题

基础题

解释为什么在变压器的注意力机制中计算复杂度随序列长度呈平方增长。
解释FlashAttention如何利用GPU内存层次结构来优化注意力运算。
解释MQA（多查询注意力）和GQA（分组查询注意力）的区别，并比较各自的优缺点。
解释PagedAttention和vLLM提高大规模语言模型推理速度和吞吐量的原理。
比较用于长上下文处理的技术中，层次注意力（Hierarchical Attention）和递归记忆变压器的工作方式，并解释各自的优缺点。

应用题

对于给定的文本数据集，编写使用变压器模型执行文本分类任务的代码。（请注意参考9.5节中的efficient_encoder示例，应用FlashAttention、Pre-LN结构、Gradient Checkpointing等优化技术。）
使用9.5节中描述的Simple Mistral模型编写执行数字-英文单词转换任务的代码，并评估模型性能。
使用9.5节中描述的Simple Mistral模型编写执行自然语言-SQL转换任务的代码，并评估模型性能。
解释Constitutional AI的概念，并提出将其应用于变压器模型以加强模型的伦理/安全约束的方法。（无需实现。）

深化题

数学分析FlashAttention的块处理方式如何提高内存效率，并与传统注意力机制进行计算复杂度比较。
调查除MQA、GQA之外减少KV缓存大小的其他方法，并比较各自优缺点。
提出一种新的长上下文处理的注意力机制，并解释其与现有方法的区别。（提出想法即可。）
指出Constitutional AI的局限性，并提出克服这些局限的方法。（提出想法即可。）

点击查看内容（练习题答案）

练习题解答

基本问题

注意力机制计算复杂度: 注意力机制计算每个token对之间的关系。当序列长度为n时，需要为每个n个token分别与其他(n-1)个token的关系进行计算，因此总共需要 n * (n-1) ≈ n² 次操作。因此，计算复杂度为 O(n²)。
FlashAttention优化: FlashAttention最大限度地利用了GPU的SRAM（快速内存）。将输入分割成小块加载到SRAM中，以块为单位执行注意力运算，并将结果写回HBM（慢速内存）。这样做减少了对HBM的访问次数，最小化了内存I/O，并提高了计算速度。
MQA vs. GQA:
- MQA (Multi-Query Attention): 所有头共享相同的Key, Value矩阵。减少KV缓存大小以降低内存使用并提高速度，但可能会减少表达能力。
- GQA (Grouped-Query Attention): 将Query分成多个组，每个组共享Key, Value矩阵。比MQA具有更高的表达能力，并且比Multi-Head Attention更节省内存。
PagedAttention & vLLM: PagedAttention借鉴了操作系统的分页概念，将KV缓存存储在不连续的内存块（页面）中。vLLM利用PagedAttention减少内存浪费，并通过动态管理KV缓存来提高推理速度和吞吐量。
层次注意力 vs. 循环记忆变压器:
- 层次注意力: 将输入分成多个层次进行处理。（例如：词 -> 句子 -> 段落）。在每个层次计算注意力，并将信息汇总到上层以捕捉长距离依赖关系。计算量可能会增加。
- 循环记忆变压器 (RMT): 以前一个段的信息形式存储为内存向量，在处理当前段时利用这些内存。通过将长序列分成小段顺序处理，减少了内存使用，但并行处理较为困难。

应用问题

文本分类代码编写: （略去代码编写）参考9.5节示例代码，使用 efficient_encoder 函数代替 nn.TransformerEncoderLayer，并应用FlashAttention、Pre-LN、Gradient Checkpointing。添加数据集加载及预处理、模型训练、评估代码。
数字-英文单词转换: （略去代码编写）加载Simple Mistral模型，并准备由数字-英文单词对组成的训练数据。训练模型，并使用测试数据评估性能。（例如：BLEU分数）
自然语言-SQL转换: （略去代码编写）加载Simple Mistral模型，并准备由自然语言问题和SQL查询对组成的训练数据。训练模型，并使用测试数据评估性能。（例如：准确度、是否可执行）
宪章AI建议: （略去实现）宪章AI定义了一套规则（宪章），用于评价并修改模型响应的方式。要应用于变压器模型，可以(1) 定义道德/安全规则，(2) 添加一个单独的模块来评估模型输出，或 (3) 在微调步骤中使用反映这些规则的损失函数。

深化问题

FlashAttention 数学分析: (数学分析省略) FlashAttention 通过块级运算减少了 HBM 访问次数。传统注意力机制需要 O(n²) 的内存访问，而 FlashAttention 在将块大小设为 B 时，仅需 O(n²/B) 的 HBM 访问。（不过，B 受 GPU SRAM 大小的限制）。
KV 缓存大小缩减方法:
- 量化 (Quantization): 将 KV 缓存的值表示为低精度（如：8-bit），以减少内存使用量。
- 稀疏化 (Sparsity): 移除或设置注意力权重较低的部分为 0，从而压缩 KV 缓存。
- 低秩近似 (Low-Rank Approximation): 将 KV 矩阵近似为低维矩阵进行存储。
新的注意力机制建议: (提出的想法)
- 局部 + 全局注意力 (Local + Global Attention): 局部上下文（如：周边词汇）使用传统注意力机制，长距离依赖则结合稀疏注意力或记忆机制处理。
- 自适应注意跨度 (Adaptive Attention Span): 为每个 token 分配不同的 attention span，以减少不必要的计算。
宪章 AI 的局限性及克服方法:
- 局限性: 规则的模糊性、规则间的潜在冲突、难以应对新型有害响应。
- 克服方法: 层次化/具体化规则、引入冲突解决机制、持续更新和验证规则、通过人类反馈进行强化学习。

参考资料

Attention Is All You Need (Original Transformer Paper): 首次提出变压器模型的基本结构和注意力机制的论文。https://arxiv.org/abs/1706.03762
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: 提出利用GPU内存层次结构优化注意力计算的FlashAttention的论文。https://arxiv.org/abs/2205.14135
FlashAttention-v2: Faster Attention with Better Parallelism and Work Partitioning: FlashAttention的改进版本，提供更快的速度和增强的并行处理。https://arxiv.org/abs/2307.08691
Scaling Transformer to 1M tokens and beyond with RMT: 使用递归记忆变压器（RMT）将变压器模型的上下文长度扩展到1M令牌以上的论文。https://arxiv.org/abs/2305.04789
The Illustrated Transformer: 通过图画轻松解释变压器模型工作原理的博客资料。http://jalammar.github.io/illustrated-transformer/
Hugging Face Transformers Documentation: Hugging Face Transformers库的官方文档，帮助轻松使用和学习变压器模型。https://huggingface.co/transformers/
PyTorch Documentation: 深度学习框架PyTorch的官方文档，提供实现和学习变压器模型所需的功能。https://pytorch.org/docs/stable/index.html
TensorFlow Documentation: 深度学习框架TensorFlow的官方文档，提供实现和学习变压器模型所需的API。https://www.tensorflow.org/api_docs
The Annotated Transformer: Harvard NLP小组编写的资料，详细解释了“Attention is all you need”论文中的PyTorch代码。http://nlp.seas.harvard.edu/2018/04/03/attention.html
DeepMind’s Blog on AlphaFold: DeepMind关于蛋白质结构预测模型AlphaFold的博客文章，其中包括变压器技术的应用案例。https://www.deepmind.com/blog/alphafold-a-solution-to-a-50-year-old-grand-challenge-in-biology
Mistral-7B: 描述拥有7亿参数的高性能语言模型Mistral-7B。https://arxiv.org/abs/2310.06825
LongLoRA Methodology: 高效微调大型语言模型的方法论。https://arxiv.org/abs/2311.02394

请注意，第4条和最后一条（原为LongLoRA Methodology）的翻译进行了调整以确保内容的一致性和准确性。 1. Attention Is All You Need (Original Transformer Paper): 首次提出变压器模型的基本结构和注意力机制的论文。https://arxiv.org/abs/1706.03762 2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness: 利用GPU内存层次结构优化注意力计算的FlashAttention论文。https://arxiv.org/abs/2205.14135 3. FlashAttention-v2: Faster Attention with Better Parallelism and Work Partitioning: FlashAttention的改进版本，提供更快的速度和增强的并行处理。https://arxiv.org/abs/2307.08691 4. Scaling Transformer to 1M tokens and beyond with RMT: 使用循环记忆变压器（RMT）将变压器模型的上下文长度扩展到1百万个令牌以上的方法。https://arxiv.org/abs/2304.11062 5. Constitutional AI: Harmlessness from AI Feedback: 提出控制AI模型响应以符合道德原则的宪法AI框架。https://arxiv.org/abs/2212.08073 6. vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention: 通过分页注意力提高大规模语言模型推理速度和吞吐量的vLLM库介绍。https://arxiv.org/abs/2309.06180, https://vllm.ai/ 7. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints: 介绍使用多头注意力检查点高效训练多查询注意力模型的GQA技术。https://arxiv.org/abs/2305.13245 8. LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models: 针对具有长上下文的大规模语言模型进行高效微调的LongLoRA方法。 https://arxiv.org/abs/2311.02394 9. Mistral-7B: 高性能语言模型 Mistral-7B 的描述，拥有 70 亿个参数。https://arxiv.org/abs/2310.06825 10. The Illustrated Transformer: 通过图表轻松解释变压器模型工作原理的博客资料。http://jalammar.github.io/illustrated-transformer/ 11. Hugging Face Transformers Documentation: Hugging Face Transformers 库的官方文档，帮助轻松使用和学习变压器模型。https://huggingface.co/transformers/ 12. PyTorch Documentation: 深度学习框架 PyTorch 的官方文档，提供实现和训练变压器模型所需的功能。https://pytorch.org/docs/stable/index.html 13. TensorFlow Documentation: 深度学习框架 TensorFlow 的官方文档，提供用于实现和训练变压器模型的 API。https://www.tensorflow.org/api_docs 14. The Annotated Transformer: Harvard NLP 小组编写的资料，详细解释了 “Attention is all you need” 论文中的 PyTorch 代码。http://nlp.seas.harvard.edu/2018/04/03/attention.html 15. DeepMind’s Blog on AlphaFold: DeepMind 关于蛋白质结构预测模型 AlphaFold 的博客文章，介绍了基于变压器技术的应用案例。https://www.deepmind.com/blog/alphafold-a-solution-to-a-50-year-old-grand-challenge-in-biology

第9章 变压器的进化：迈向效率和扩展性

9.1 变压器的局限与挑战

9.1.1 变压器架构的基本局限：计算复杂度

9.1.2 变压器架构的基本限制：内存效率

9.1.3 变压器发展的时代潮流及本章结构

9.2 复杂度降低：软件注意力优化（2019-2020）

9.2.1 初期方法：近似和稀疏化

9.2.1.1 线性注意力：Performer

9.2.1.2 稀疏注意：Sparse Transformer, Longformer

9.2.3 局部-全局注意力：解决长距离依赖问题

9.2.3.1 Reformer: LSH注意力

9.2.3.2 BigBird: 局部、全局和随机注意力的组合

9.3 内存效率：硬件与软件的结合 (2021-2022)

9.3.1 FlashAttention：利用GPU内存层次结构进行注意力优化

9.3.1.1 GPU内存结构与IO优化

9.3.1.2 分块和块处理

9.3.1.3 FlashAttention v2：最大化硬件利用

9.3.2 查询优化：改进注意力结构

9.3.2.1 多查询注意力（MQA）

9.3.2.2 分组查询注意力（GQA）

9.3.2.3 MQA vs. GQA vs. 多头注意力

9.3.3 KV 缓存管理和优化

9.3.3.1 PagedAttention & vLLM: 操作系统分页概念

9.3.3.2 连续批处理与高效的缓存策略 (Continuous Batching & Efficient Caching)

9.4 可扩展性和特殊目的架构（2023-2024）

9.4.1 长上下文处理：上下文长度的扩展

9.4.1.1 分层注意力、循环记忆变压器

9.4.1.2 Claude-2, LongLoRA

9.4.2 伦理/安全约束：Constitutional AI

9.4.2.1 基于规则的注意力

9.4.2.2 基于强化学习的调整 (RLHF, RLAIF)

9.4.3 特殊目的注意力：领域和任务优化

9.4.3.1 特殊目的注意力的多种示例

9.4.3.2 多模态注意力 (Multimodal Attention)

深度剖析：不同变压器模型的详细分析及技术关联性

1. 编码器中心模型 (Encoder-Only Models)

2. 解码器中心模型 (Decoder-Only Models)

3. 混合方法模型 (Encoder-Decoder Models)

技术关联性深度分析

未来展望

9.5 高效编码器的实现和应用：以RoPE和FlashAttention为中心

9.5.1 高效编码器的设计理念：速度和内存

9.5.2 efficient_encoder.py代码详细分析（未使用RoPE）

9.5.3 efficient_encoder_rope.py 代码详细分析 (使用 RoPE)

9.5.4 实验结果：AG News 文本分类

9.6 Mistral：实现和分析高效的解码器架构

9.6.1 simple_mistral 模型架构：组件详细分析

1. MistralConfig: 模型设置

2. MistralRMSNorm: RMS 归一化

3. MistralAttention: 注意力机制

4. MistralRotaryEmbedding: RoPE 实现

5. MistralMLP: FeedForward 网络

6. MistralDecoderLayer: 解码器层

7. MistralPreTrainedModel: 预训练模型抽象类

8. MistralModel: Mistral 模型

9. MistralForCausalLM: 用于语言建模的 Mistral

9.6.2 核心技术要素分析：效率与性能的关键

1. GQA (Grouped-Query Attention): 提高内存和计算效率的创新

2. SWA (Sliding Window Attention): 高效处理长序列的策略

3. RoPE (Rotary Positional Embedding): 高效编码相对位置信息

4. KV 缓存: 去除重复计算

9.6.3 模型训练: simple_mistral 训练指南

1. 数据预处理: 转换为模型可理解的形式

2. 模型训练: 寻找最优参数

9.6.4 使用 generate() 函数生成文本：创造性的句子制作

generate() 函数：文本生成的核心

生成过程：逐步文本生成

9.6.5 数字序列预测示例：train_seq_num.py 分析

1. 准备数据集和数据加载器：构建训练数据

2. 模型设置和训练：simple_mistral 训练

3. 文本生成: 使用训练好的模型进行预测

4. 结果分析: 训练结果和生成文本评估

9.6.6 四则运算预测示例：train_math.py 分析

数据生成及预处理：符号与数字的和谐

模型配置及训练

文本生成

结果分析

9.6.7 自然语言-SQL 查询生成示例：train_sql.py 分析

数据集及预处理：WikiSQL 和特殊标记的协调

模型配置及训练

第9章变压器的进化：迈向效率和扩展性

9.5.2 `efficient_encoder.py`代码详细分析（未使用RoPE）

9.5.3 `efficient_encoder_rope.py` 代码详细分析 (使用 RoPE)

9.6.1 `simple_mistral` 模型架构：组件详细分析

9.6.3 模型训练: `simple_mistral` 训练指南

9.6.4 使用 `generate()` 函数生成文本：创造性的句子制作

`generate()` 函数：文本生成的核心

9.6.5 数字序列预测示例：`train_seq_num.py` 分析

2. 模型设置和训练：`simple_mistral` 训练

9.6.6 四则运算预测示例：`train_math.py` 分析

9.6.7 自然语言-SQL 查询生成示例：`train_sql.py` 分析

文本生成 (`generate_sql`): 从问题中推断 SQL 查询

9.8.1 `simple_phi3` 模型

5. 辅助函数: `rotate_half`, `apply_rotary_pos_emb`, `apply_rotary_pos_emb_single`

9.8.2 `simple_phi3` 模型示例：复合公式计算