Transformer模型中的自注意力机制：从生物学灵感到大模型核心

你是否曾好奇，为何在阅读一段文字时，大脑能瞬间抓住核心论点，而忽略无关的细节？这种高效筛选信息的能力，正是生物进化赋予我们的“注意力”天赋。如今，这一源自神经科学的灵感，已经深刻地重塑了人工智能的底层架构。Transformer模型及其核心——自注意力机制，正是这一跨学科思想碰撞的璀璨结晶。它不仅是驱动GPT、BERT等大语言模型运转的引擎，更是理解现代深度学习为何能处理海量、复杂序列数据的关键钥匙。对于希望深入模型内部、理解其设计哲学与实现细节的研究者和工程师而言，自注意力机制是绕不开的核心课题。本文将带你从生物学的源头出发，穿越计算模型的抽象，最终抵达Transformer的工程实现，构建一个完整而深刻的理解框架。

1. 从生物注意力到计算模型：一次思想的迁徙

人类大脑处理信息的能力并非无限。面对汹涌而来的感官洪流，我们依赖一套精密的“注意力系统”来分配有限的认知资源。这套系统大致可分为两类：自下而上的刺激驱动注意和自上而下的目标导向注意。前者如突然的声响会吸引你的目光，是外界刺激的“被动”捕获；后者则像你在嘈杂咖啡馆里专注于朋友的谈话，是内在意图的“主动”聚焦。

注意：这种“主动”与“被动”的二分法，为计算模型的注意力机制设计提供了最根本的范式。计算模型中的“查询”对应自上而下的目标，“键”对应自下而上的刺激特征，而最终的“聚焦”过程，就是根据查询与键的匹配程度，对“值”（原始感官信息）进行加权聚合。

将这一生物学原理转化为数学模型，其核心挑战在于如何量化“重要性”。早期的一个经典尝试是1964年的Nadaraya-Watson核回归。这个模型虽然简单，却清晰地勾勒出注意力机制的雏形：对于一个查询点（比如你想预测的年份），模型并非平等看待所有历史数据点，而是根据它们与查询点的“距离”或“相似度”来分配权重，距离越近（相似度越高）的数据点，对预测结果的贡献越大。

用数学公式表达，对于一个查询 q 和一系列键值对 (k1, v1), (k2, v2), ...，注意力汇聚的输出是值的加权和：

输出 = Σ_i (α(q, k_i) * v_i)

其中，权重 α(q, k_i) 通过一个评分函数计算，再经softmax归一化，确保所有权重之和为1。这个框架的通用性极强，后续所有复杂的注意力变体，都建立在这个基础之上。

从生物学启示到基础数学模型，我们可以总结出注意力机制的几个关键特性：

• 动态性：权重并非固定，而是根据每次查询动态生成。
• 选择性：模型有能力忽略大部分信息，只聚焦于相关的少数。
• 可解释性（一定程度）：通过分析注意力权重，我们可以窥见模型做出决策时“看”向了哪里。

2. 自注意力：让序列自己与自己对话

在传统的循环神经网络（RNN）中，序列信息的传递是串行的。一个词的信息需要一步步“流经”整个序列才能影响到远处的另一个词，这导致了长距离依赖学习的困难。卷积神经网络（CNN）虽然能并行处理，但其感受野受限于卷积核大小，捕捉长程关系需要堆叠很多层。

自注意力机制的提出，彻底改变了这一局面。它的核心思想非常直观：序列中的每个元素，都应该直接与序列中的所有其他元素进行交互，以确定彼此的相关性。

2.1 核心计算：Query, Key, Value

自注意力为输入序列中的每个位置（例如句子中的每个词）生成三个向量：查询（Query）、键（Key）和值（Value）。这通常通过三个不同的可学习线性变换实现。

• Query (Q)：可以理解为该位置“想要寻找什么”。
• Key (K)：可以理解为该位置“拥有什么特征，可供被寻找”。
• Value (V)：可以理解为该位置“真正要贡献的信息内容”。

计算过程分为四步：

1. 计算注意力分数：对于序列中的第 i 个位置，用它的查询向量 q_i 与所有位置的键向量 k_j 进行点积，得到一组分数。这衡量了位置 i 与位置 j 的关联强度。

   分数_ij = q_i · k_j^T
   

2. 缩放与归一化：为了防止点积结果过大导致softmax梯度消失，将分数除以键向量维度的平方根 √d_k。然后对所有分数应用softmax函数，将其转化为和为1的概率分布，即注意力权重。

   权重_ij = softmax(分数_ij / √d_k)
   

3. 加权求和：用得到的注意力权重对所有的值向量 v_j 进行加权求和，得到位置 i 新的表示。

   输出_i = Σ_j (权重_ij * v_j)
   

4. 矩阵并行化：上述过程可以完美地并行化。将整个序列的Q、K、V分别堆叠成矩阵，一次矩阵运算即可完成所有位置的注意力计算。

import torch
import torch.nn.functional as F

def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    缩放点积注意力计算。
    Args:
        Q: 查询矩阵，形状为 (batch_size, seq_len_q, d_k)
        K: 键矩阵，形状为 (batch_size, seq_len_k, d_k)
        V: 值矩阵，形状为 (batch_size, seq_len_v, d_v) (通常 seq_len_k == seq_len_v)
        mask: 可选的掩码矩阵，用于在解码时屏蔽未来信息。
    Returns:
        注意力输出，形状为 (batch_size, seq_len_q, d_v)
        注意力权重矩阵
    """
    d_k = Q.size(-1)
    # 计算点积并缩放
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
    
    # 应用掩码（如果有）
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)
    
    # 应用softmax得到注意力权重
    attn_weights = F.softmax(scores, dim=-1)
    
    # 加权求和
    output = torch.matmul(attn_weights, V)
    
    return output, attn_weights

# 示例：假设我们有一个批次大小为2，序列长度为5，维度为64的输入
batch_size, seq_len, d_model = 2, 5, 64
d_k = d_v = 64
X = torch.randn(batch_size, seq_len, d_model) # 输入序列

# 假设我们有一个简单的线性层来生成Q, K, V（实际中会使用不同的参数）
W_Q = torch.nn.Linear(d_model, d_k)
W_K = torch.nn.Linear(d_model, d_k)
W_V = torch.nn.Linear(d_model, d_v)

Q = W_Q(X)
K = W_K(X)
V = W_V(X)

output, attn = scaled_dot_product_attention(Q, K, V)
print(f"输入形状: {X.shape}")
print(f"注意力输出形状: {output.shape}")
print(f"注意力权重形状: {attn.shape}") # 应为 (2, 5, 5)

2.2 为何自注意力如此强大？

自注意力机制带来了几个革命性的优势：

• 全局视野：每个词元在单层内就能直接“看到”序列中的所有其他词元，无论距离多远。这极大地缓解了长距离依赖问题。
• 极致并行：计算不依赖于前一步的隐藏状态，所有位置的注意力可以同时计算，充分利用GPU等硬件加速。
• 灵活的关联模式：模型可以学习到语法（如主谓一致）、语义（如指代消解）等多种复杂的依赖关系，这些关系通过注意力权重矩阵直观呈现。

下表对比了自注意力与RNN、CNN在处理序列数据时的关键差异：

特性	循环神经网络 (RNN/LSTM)	卷积神经网络 (CNN)	自注意力 (Self-Attention)
计算复杂度	O(n) （序列长度）	O(k*n) （k为卷积核宽度）	O(n²) （序列长度的平方）
并行化程度	低（顺序依赖）	高（局部并行）	极高（完全并行）
长程依赖	困难（梯度问题）	需要深层堆叠	优秀（单层全局交互）
路径长度	O(n)	O(log_k(n))	O(1)
可解释性	隐状态，难以解释	卷积核，部分可解释	注意力权重，直观可视

提示：自注意力的O(n²)复杂度是其主要的计算瓶颈，尤其是在处理超长序列（如长文档、高分辨率图像）时。这也是催生“稀疏注意力”、“线性注意力”等改进变体的直接原因。

3. 多头注意力与位置编码：增强与补全

基本的自注意力机制虽然强大，但仍有两个关键问题需要解决：1）如何让模型同时关注来自不同表示子空间的信息？2）如何让完全并行的模型感知序列的顺序？

3.1 多头注意力：并行化的“专家委员会”

想象一下，一个句子中的词与词之间的关系是多元的。例如在“他打开了银行账户，然后去河边的银行散步”中，“银行”与“账户”是语义关联（金融机构），与“河边”是语义关联（河流沿岸）。单一的注意力机制可能难以同时捕捉这些不同类型的关系。

多头注意力的设计灵感正在于此。它不满足于只做一次注意力计算，而是将模型容量分割成多个“头”（head），让每个头独立地在不同的子空间里学习关注不同的模式。

其实现步骤如下：

1. 线性投影：将原始的Q、K、V矩阵，分别通过 h（头数）组不同的线性变换层，投影到 d_k、d_k、d_v 维度。通常设置 d_k = d_v = d_model / h，以保持总参数量不变。
2. 并行注意力：在每个投影后的子空间上，独立执行缩放点积注意力计算，得到 h 个输出矩阵。
3. 拼接与输出：将 h 个头的输出在特征维度上拼接起来，最后通过一个线性变换层映射回 d_model 维度。

class MultiHeadAttention(torch.nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model必须能被num_heads整除"
        
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # 定义生成Q, K, V的线性层和最后的输出层
        self.W_q = torch.nn.Linear(d_model, d_model)
        self.W_k = torch.nn.Linear(d_model, d_model)
        self.W_v = torch.nn.Linear(d_model, d_model)
        self.W_o = torch.nn.Linear(d_model, d_model)
        
        self.dropout = torch.nn.Dropout(dropout)
        
    def split_heads(self, x):
        """将输入张量重塑为多头形式。"""
        batch_size, seq_len, _ = x.size()
        # 重塑为 (batch_size, seq_len, num_heads, d_k)
        x = x.view(batch_size, seq_len, self.num_heads, self.d_k)
        # 转置为 (batch_size, num_heads, seq_len, d_k) 以方便并行计算
        return x.transpose(1, 2)
    
    def combine_heads(self, x):
        """将多头输出合并回原始形状。"""
        batch_size, _, seq_len, _ = x.size()
        # 转置回来并重塑
        x = x.transpose(1, 2).contiguous()
        return x.view(batch_size, seq_len, self.d_model)
    
    def forward(self, Q, K, V, mask=None):
        batch_size = Q.size(0)
        
        # 1. 线性投影并分头
        Q = self.split_heads(self.W_q(Q))
        K = self.split_heads(self.W_k(K))
        V = self.split_heads(self.W_v(V))
        
        # 2. 在每个头上应用缩放点积注意力
        attn_output, attn_weights = scaled_dot_product_attention(Q, K, V, mask)
        
        # 3. 合并多头
        attn_output = self.combine_heads(attn_output)
        
        # 4. 最终线性投影
        output = self.W_o(attn_output)
        
        return output, attn_weights

# 使用示例
d_model = 512
num_heads = 8
mha = MultiHeadAttention(d_model, num_heads)
Q = torch.randn(4, 10, d_model) # (batch_size, seq_len, d_model)
output, weights = mha(Q, Q, Q) # 自注意力
print(f"多头注意力输出形状: {output.shape}") # 应保持 (4, 10, 512)

3.2 位置编码：为无序注入顺序

自注意力机制对输入序列的顺序是完全不敏感的。打乱输入词的顺序，只要词本身不变，计算出的注意力权重和输出在理论上就是一样的。这显然不符合语言（乃至大多数序列数据）的本质。

因此，Transformer必须显式地告诉模型每个词的位置信息。这就是位置编码的使命。原始Transformer论文采用了一种巧妙且固定的正弦余弦编码方式：

对于位置 pos 和维度 i，其编码值计算如下：

PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))

其中，d_model 是模型维度。这种编码具有几个优良性质：

• 唯一性：每个位置都有唯一的编码。
• 相对位置可学习：对于固定的偏移量 k，PE(pos+k) 可以表示为 PE(pos) 的线性函数，这使得模型能够轻松学习到相对位置关系。
• 值域有界：正弦余弦函数的值在[-1, 1]之间，与词嵌入向量的范围匹配。

import math

class PositionalEncoding(torch.nn.Module):
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        self.dropout = torch.nn.Dropout(p=dropout)
        
        # 创建位置编码矩阵
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        
        pe[:, 0::2] = torch.sin(position * div_term) # 偶数维度用sin
        pe[:, 1::2] = torch.cos(position * div_term) # 奇数维度用cos
        
        pe = pe.unsqueeze(0) # 增加批次维度: (1, max_len, d_model)
        self.register_buffer('pe', pe) # 注册为缓冲区，不参与训练但会保存
        
    def forward(self, x):
        """
        Args:
            x: 输入张量，形状为 (batch_size, seq_len, d_model)
        """
        x = x + self.pe[:, :x.size(1)] # 将位置编码加到输入上
        return self.dropout(x)

# 可视化位置编码
import matplotlib.pyplot as plt
d_model = 64
max_len = 100
pe = PositionalEncoding(d_model, max_len, dropout=0)
dummy_input = torch.zeros(1, max_len, d_model)
pos_encoded = pe(dummy_input).squeeze(0).detach().numpy()

plt.figure(figsize=(10, 6))
plt.pcolormesh(pos_encoded.T, cmap='RdBu')
plt.xlabel('位置 (Position)')
plt.ylabel('维度 (Dimension)')
plt.colorbar(label='编码值')
plt.title('正弦余弦位置编码可视化')
plt.show()

在实际应用中，位置编码与词嵌入向量直接相加，作为编码器和解码器的输入。这样，模型在计算注意力时，每个词的表征就同时包含了“我是谁”（语义）和“我在哪”（位置）的信息。

4. Transformer架构全景：编码器与解码器的交响

有了自注意力（多头注意力）和位置编码这两块基石，我们就可以搭建完整的Transformer模型了。Transformer采用经典的编码器-解码器架构，但内部完全由注意力机制和前馈神经网络堆叠而成。

4.1 编码器：信息的抽象与提炼

编码器由N个完全相同的层堆叠而成（原论文中N=6）。每一层包含两个核心子层：

1. 多头自注意力子层：让输入序列中的每个位置都能关注整个序列，整合上下文信息。
2. 前馈神经网络子层：一个简单的两层全连接网络，通常中间层的维度更大（例如4倍 d_model），并带有ReLU激活函数。它对每个位置独立进行变换，因此称为“位置式前馈网络”。

每个子层周围都包裹着两个至关重要的组件：

• 残差连接：将子层的输入直接加到其输出上，即 输出 = LayerNorm(子层(输入) + 输入)。这极大地缓解了深层网络中的梯度消失问题，使得训练非常深的网络成为可能。
• 层归一化：在残差相加之后进行层归一化，稳定激活值的分布，加速训练收敛。

class EncoderLayer(torch.nn.Module):
    """Transformer编码器的一个层。"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(d_model, d_ff),
            torch.nn.ReLU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(d_ff, d_model)
        )
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.norm2 = torch.nn.LayerNorm(d_model)
        self.dropout1 = torch.nn.Dropout(dropout)
        self.dropout2 = torch.nn.Dropout(dropout)
        
    def forward(self, x, mask=None):
        # 子层1: 多头自注意力 + 残差 & 层归一化
        attn_output, _ = self.self_attn(x, x, x, mask)
        x = x + self.dropout1(attn_output)
        x = self.norm1(x)
        
        # 子层2: 前馈网络 + 残差 & 层归一化
        ffn_output = self.ffn(x)
        x = x + self.dropout2(ffn_output)
        x = self.norm2(x)
        
        return x

4.2 解码器：自回归生成与交叉注意力

解码器同样由N个相同的层堆叠。每一层包含三个子层：

1. 掩码多头自注意力子层：这是解码器特有的。在训练时，为了确保模型在预测第 t 个词时只能看到前 t-1 个词（即保持自回归特性），需要通过一个掩码矩阵将未来位置的信息屏蔽掉（设置为负无穷大，softmax后权重为0）。
2. 编码器-解码器注意力子层：这是连接编码器和解码器的桥梁。该子层的**查询（Q）来自解码器上一层的输出，而键（K）和值（V）**则来自编码器的最终输出。这使得解码器在生成每一个词时，都能有选择地“回顾”输入序列的全部信息。
3. 前馈神经网络子层：与编码器中的相同。

解码器的每个子层同样有残差连接和层归一化。

class DecoderLayer(torch.nn.Module):
    """Transformer解码器的一个层。"""
    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.cross_attn = MultiHeadAttention(d_model, num_heads, dropout)
        self.ffn = torch.nn.Sequential(
            torch.nn.Linear(d_model, d_ff),
            torch.nn.ReLU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(d_ff, d_model)
        )
        self.norm1 = torch.nn.LayerNorm(d_model)
        self.norm2 = torch.nn.LayerNorm(d_model)
        self.norm3 = torch.nn.LayerNorm(d_model)
        self.dropout1 = torch.nn.Dropout(dropout)
        self.dropout2 = torch.nn.Dropout(dropout)
        self.dropout3 = torch.nn.Dropout(dropout)
        
    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # 子层1: 掩码多头自注意力
        attn_output, _ = self.self_attn(x, x, x, tgt_mask)
        x = x + self.dropout1(attn_output)
        x = self.norm1(x)
        
        # 子层2: 编码器-解码器注意力
        attn_output, _ = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = x + self.dropout2(attn_output)
        x = self.norm2(x)
        
        # 子层3: 前馈网络
        ffn_output = self.ffn(x)
        x = x + self.dropout3(ffn_output)
        x = self.norm3(x)
        
        return x

def generate_square_subsequent_mask(sz):
    """生成一个上三角为1，对角线及以下为0的掩码矩阵，用于屏蔽未来信息。"""
    mask = torch.triu(torch.ones(sz, sz), diagonal=1).bool()
    # 将True（未来位置）转换为 -inf，False（过去和当前位置）转换为 0
    return mask.float().masked_fill(mask == 1, float('-inf'))

4.3 从理论到实践：一个简化的Transformer

将编码器、解码器、嵌入层、位置编码和最后的线性输出层组合起来，就构成了一个完整的Transformer模型。虽然现代库（如Hugging Face Transformers）提供了高度优化和封装的实现，但理解这个基本骨架对于调试模型、进行定制化修改至关重要。

class Transformer(torch.nn.Module):
    """一个简化的Transformer模型，用于序列到序列任务。"""
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_len, dropout=0.1):
        super().__init__()
        self.encoder_embedding = torch.nn.Embedding(src_vocab_size, d_model)
        self.decoder_embedding = torch.nn.Embedding(tgt_vocab_size, d_model)
        self.positional_encoding = PositionalEncoding(d_model, max_len, dropout)
        
        self.encoder_layers = torch.nn.ModuleList([
            EncoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)
        ])
        self.decoder_layers = torch.nn.ModuleList([
            DecoderLayer(d_model, num_heads, d_ff, dropout) for _ in range(num_layers)
        ])
        
        self.final_linear = torch.nn.Linear(d_model, tgt_vocab_size)
        self.dropout = torch.nn.Dropout(dropout)
        
    def encode(self, src, src_mask):
        src_embedded = self.dropout(self.positional_encoding(self.encoder_embedding(src)))
        enc_output = src_embedded
        for layer in self.encoder_layers:
            enc_output = layer(enc_output, src_mask)
        return enc_output
    
    def decode(self, tgt, enc_output, src_mask, tgt_mask):
        tgt_embedded = self.dropout(self.positional_encoding(self.decoder_embedding(tgt)))
        dec_output = tgt_embedded
        for layer in self.decoder_layers:
            dec_output = layer(dec_output, enc_output, src_mask, tgt_mask)
        return dec_output
    
    def forward(self, src, tgt, src_mask=None, tgt_mask=None):
        enc_output = self.encode(src, src_mask)
        dec_output = self.decode(tgt, enc_output, src_mask, tgt_mask)
        output = self.final_linear(dec_output)
        return output

# 模型参数示例
src_vocab_size = 10000
tgt_vocab_size = 10000
d_model = 512
num_heads = 8
num_layers = 6
d_ff = 2048
max_len = 100
dropout = 0.1

model = Transformer(src_vocab_size, tgt_vocab_size, d_model, num_heads, num_layers, d_ff, max_len, dropout)
print(f"模型参数量: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")

5. 超越原始Transformer：演进、变体与实战工具

自2017年诞生以来，Transformer已远不止于最初的机器翻译架构。它已成为一个强大的通用序列建模框架，催生了无数变体和改进。

5.1 核心架构的演进方向

• 效率优化：原始自注意力的O(n²)复杂度是处理长序列的瓶颈。稀疏注意力（如Longformer的滑动窗口注意力、BigBird的随机+全局注意力）、线性注意力（通过核函数近似）和分块递归（如Transformer-XL）等方法被提出，旨在保持性能的同时大幅降低计算和内存开销。
• 结构简化：仅编码器模型（如BERT）专注于理解与表征，通过掩码语言建模等任务进行预训练；仅解码器模型（如GPT系列）专注于生成，通过自回归语言建模进行预训练。这两种简化架构在各自领域取得了巨大成功。
• 多模态融合：Vision Transformer将图像分割为图块序列进行处理；Audio Transformer处理音频频谱图。跨模态的Transformer（如CLIP、DALL-E）通过注意力机制对齐不同模态（文本、图像）的表示空间。

5.2 现代开发实战：拥抱Hugging Face生态

对于绝大多数开发者和研究者而言，从头实现Transformer已非必要。Hugging Face的 transformers 库提供了极其丰富的预训练模型和易用的API，是进行NLP研究和应用开发的瑞士军刀。

快速使用预训练模型进行推理：

from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM

# 使用pipeline，一行代码实现翻译
translator = pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr")
result = translator("Hello, how are you?", max_length=40)
print(result[0]['translation_text'])  # 输出: Bonjour, comment allez-vous ?

# 更细粒度的控制：文本摘要
model_name = "facebook/bart-large-cnn"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

article = """The Apollo program, also known as Project Apollo, was the third United States human spaceflight program carried out by the National Aeronautics and Space Administration (NASA), which succeeded in landing the first humans on the Moon from 1969 to 1972."""
inputs = tokenizer([article], max_length=1024, return_tensors="pt", truncation=True)
summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=50, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
print(f"摘要: {summary}")

在自己的数据集上进行微调：

from transformers import Trainer, TrainingArguments, AutoModelForSequenceClassification, AutoTokenizer
from datasets import load_dataset
import evaluate
import numpy as np

# 1. 加载数据集和模型
dataset = load_dataset("imdb")
model_name = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

# 2. 预处理数据
def preprocess_function(examples):
    return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=512)

tokenized_datasets = dataset.map(preprocess_function, batched=True)

# 3. 定义训练参数
training_args = TrainingArguments(
    output_dir="./results",
    evaluation_strategy="epoch",
    learning_rate=2e-5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    num_train_epochs=3,
    weight_decay=0.01,
    logging_dir='./logs',
)

# 4. 加载评估指标
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

# 5. 创建Trainer并开始训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"].shuffle(seed=42).select(range(1000)), # 示例，取部分数据
    eval_dataset=tokenized_datasets["test"].shuffle(seed=42).select(range(200)),
    tokenizer=tokenizer,
    compute_metrics=compute_metrics,
)

trainer.train()

从生物大脑中筛选关键信息的本能，到数学公式中优雅的加权求和，再到支撑起千亿参数大模型的工程巨构，自注意力机制完成了一次从灵感到核心的华丽蜕变。理解它，不仅仅是理解一组矩阵乘法，更是理解当代AI如何学会像我们一样“聚焦”与“关联”。当你下次与ChatGPT对话，或使用翻译软件时，不妨想象一下，在那些硅基芯片的深处，正有无数个“注意力头”在忙碌地计算着权重，编织着意义。这或许就是技术最迷人的地方：它将人类认知中最精妙的部分抽象出来，并以一种前所未有的规模与速度运行着。