核心摘要

注意力机制(Attention)是一类用于动态加权信息的可微计算。本指南从直觉出发,讲解缩放点积、掩码、Query-Key-Value 投影、多头设计、Transformer 集成,以及实现和解释的边界,并提供 NumPy 与 PyTorch 示例。

引言

人类注意力类比有助于建立直觉,但不是神经机制断言。神经网络注意力经过多条研究路线发展,包括约 2014 年神经机器翻译研究,将信息加权做成可微计算。

从机器翻译到 ChatGPT、图像识别和语音处理,注意力机制已成为现代 AI 系统的重要组件。2017 年论文《Attention Is All You Need》提出了 Transformer 架构;它除了注意力,还包含前馈网络、归一化、残差连接和位置信息。

在本指南中,你将学到:

  • 注意力机制的直觉理解和设计动机
  • 自注意力(Self-Attention)的数学原理
  • Query、Key、Value的计算过程
  • 多头注意力(Multi-Head Attention)的工作方式
  • 注意力分数的可视化与解释
  • 注意力机制在Transformer中的应用
  • 完整的Python代码实现

什么是注意力机制

注意力机制是一种让神经网络能够动态聚焦于输入中最相关部分的技术。与传统方法对所有输入一视同仁不同,注意力机制会为每个输入元素分配不同的权重,让模型"关注"最重要的信息。

graph LR subgraph "传统方法" I1[输入1] --> E1[等权重] I2[输入2] --> E2[等权重] I3[输入3] --> E3[等权重] E1 --> O1[输出] E2 --> O1 E3 --> O1 end subgraph "注意力机制" A1[输入1] --> W1[权重 0.7] A2[输入2] --> W2[权重 0.2] A3[输入3] --> W3[权重 0.1] W1 --> O2[输出] W2 --> O2 W3 --> O2 end

为什么需要注意力机制

在注意力机制出现之前,序列模型(如RNN、LSTM)面临几个关键问题:

  1. 信息瓶颈:编码器必须将整个输入序列压缩成固定长度的向量,长序列信息容易丢失
  2. 长距离依赖:相距较远的元素难以建立有效联系
  3. 计算效率:必须按顺序处理,无法并行化

注意力机制让一个层内的位置可以直接访问允许范围内的其他位置,但稠密分数会带来二次复杂度和额外显存;因果掩码、稀疏模式和硬件内核会改变具体行为与成本。

注意力机制的直觉理解

想象你在图书馆查找资料:

  • Query(查询):你心中的问题——"我想找关于机器学习的书"
  • Key(键):每本书的标签或摘要——帮助你判断相关性
  • Value(值):书的实际内容——你最终要获取的信息

注意力机制的工作方式类似:用Query去匹配所有Key,找到最相关的,然后提取对应的Value。

自注意力机制详解

自注意力从序列表示生成 Query、Key 和 Value,并根据掩码让每个位置关注允许的位置;它是 Transformer 的核心组件之一,但不是整个架构。

Query、Key、Value的计算

自注意力的核心是将输入转换为三个向量:Query、Key和Value。

python
import numpy as np

class SelfAttention:
    def __init__(self, d_model, d_k):
        """
        初始化自注意力层
        d_model: 输入维度
        d_k: Query/Key/Value的维度
        """
        self.d_k = d_k
        self.W_q = np.random.randn(d_model, d_k) * 0.1
        self.W_k = np.random.randn(d_model, d_k) * 0.1
        self.W_v = np.random.randn(d_model, d_k) * 0.1
    
    def compute_qkv(self, X):
        """
        计算Query、Key、Value
        X: 输入矩阵 (seq_len, d_model)
        """
        Q = np.matmul(X, self.W_q)  # (seq_len, d_k)
        K = np.matmul(X, self.W_k)  # (seq_len, d_k)
        V = np.matmul(X, self.W_v)  # (seq_len, d_k)
        return Q, K, V

每个输入token通过三个不同的线性变换,分别得到:

  • Query:表示"我在寻找什么"
  • Key:表示"我包含什么信息"
  • Value:表示"我要传递什么内容"

缩放点积注意力

有了Q、K、V,接下来计算注意力分数:

python
def scaled_dot_product_attention(Q, K, V, mask=None):
    """
    缩放点积注意力
    Q: 查询矩阵 (seq_len, d_k)
    K: 键矩阵 (seq_len, d_k)
    V: 值矩阵 (seq_len, d_v)
    mask: 可选的掩码矩阵
    """
    d_k = K.shape[-1]
    
    scores = np.matmul(Q, K.T) / np.sqrt(d_k)
    
    if mask is not None:
        scores = np.where(mask == 0, -1e9, scores)
    
    attention_weights = softmax(scores, axis=-1)
    
    output = np.matmul(attention_weights, V)
    
    return output, attention_weights

def softmax(x, axis=-1):
    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)

注意力计算的数学公式:

code
Attention(Q, K, V) = softmax(QK^T / √d_k) V

为什么要缩放

除以√d_k是为了防止点积值过大。当d_k较大时,点积的方差也会变大,导致softmax输出趋近于one-hot分布,梯度变得极小。缩放操作保持了梯度的稳定性。

graph TB subgraph "注意力计算流程" Q[Query] --> MM1[矩阵乘法] K[Key] --> MM1 MM1 --> Scale[缩放 ÷√d_k] Scale --> Mask["掩码 可选"] Mask --> SM[Softmax] SM --> MM2[矩阵乘法] V[Value] --> MM2 MM2 --> Out[输出] end

多头注意力机制

单个注意力头也可能表示多种关系,但容量受投影维度限制。多头注意力在多个学习到的子空间中并行计算,再通过输出投影混合;每个头的语义角色并不是预先指定的。

python
class MultiHeadAttention:
    def __init__(self, d_model, num_heads):
        """
        多头注意力
        d_model: 模型维度
        num_heads: 注意力头数量
        """
        assert d_model % num_heads == 0
        
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.d_model = d_model
        
        self.W_q = np.random.randn(d_model, d_model) * 0.1
        self.W_k = np.random.randn(d_model, d_model) * 0.1
        self.W_v = np.random.randn(d_model, d_model) * 0.1
        self.W_o = np.random.randn(d_model, d_model) * 0.1
    
    def split_heads(self, x):
        """将输入分割成多个头"""
        seq_len = x.shape[0]
        x = x.reshape(seq_len, self.num_heads, self.d_k)
        return x.transpose(1, 0, 2)  # (num_heads, seq_len, d_k)
    
    def forward(self, X):
        """
        前向传播
        X: 输入 (seq_len, d_model)
        """
        Q = np.matmul(X, self.W_q)
        K = np.matmul(X, self.W_k)
        V = np.matmul(X, self.W_v)
        
        Q = self.split_heads(Q)
        K = self.split_heads(K)
        V = self.split_heads(V)
        
        heads_output = []
        for i in range(self.num_heads):
            head_out, _ = scaled_dot_product_attention(Q[i], K[i], V[i])
            heads_output.append(head_out)
        
        concat = np.concatenate(heads_output, axis=-1)
        
        output = np.matmul(concat, self.W_o)
        
        return output

多头注意力的优势

graph TB Input[输入序列] --> H1["头1: 语法关系"] Input --> H2["头2: 语义关系"] Input --> H3["头3: 位置关系"] Input --> H4["头4: 指代关系"] H1 --> Concat[拼接] H2 --> Concat H3 --> Concat H4 --> Concat Concat --> Linear[线性变换] Linear --> Output[输出]

不同的注意力头可能学习到:

  • 语法结构:主语-谓语-宾语关系
  • 语义相似性:同义词、近义词
  • 位置模式:相邻词、固定距离词
  • 指代关系:代词与其指代对象

注意力分数的可视化

注意力权重可以可视化,帮助我们理解模型在"看"什么:

python
import matplotlib.pyplot as plt

def visualize_attention(attention_weights, tokens):
    """
    可视化注意力权重
    attention_weights: 注意力权重矩阵 (seq_len, seq_len)
    tokens: token列表
    """
    fig, ax = plt.subplots(figsize=(10, 10))
    
    im = ax.imshow(attention_weights, cmap='Blues')
    
    ax.set_xticks(range(len(tokens)))
    ax.set_yticks(range(len(tokens)))
    ax.set_xticklabels(tokens, rotation=45, ha='right')
    ax.set_yticklabels(tokens)
    
    for i in range(len(tokens)):
        for j in range(len(tokens)):
            text = ax.text(j, i, f'{attention_weights[i, j]:.2f}',
                          ha='center', va='center', fontsize=8)
    
    ax.set_xlabel('Key')
    ax.set_ylabel('Query')
    ax.set_title('Attention Weights')
    
    plt.colorbar(im)
    plt.tight_layout()
    plt.show()

tokens = ['我', '喜欢', '机器', '学习']
attention = np.array([
    [0.4, 0.3, 0.2, 0.1],
    [0.2, 0.3, 0.3, 0.2],
    [0.1, 0.2, 0.4, 0.3],
    [0.1, 0.2, 0.3, 0.4]
])

热力图只显示特定层、特定头、特定输入和某次运行的路由权重。即使出现对角线或词对集中,也不能自动解释最终预测的因果原因。

注意力机制在Transformer中的应用

Transformer架构中有三种不同的注意力应用:

编码器自注意力

编码器中的自注意力让每个位置都能关注输入序列的所有位置:

python
class EncoderLayer:
    def __init__(self, d_model, num_heads, d_ff):
        self.self_attention = MultiHeadAttention(d_model, num_heads)
        self.feed_forward = FeedForward(d_model, d_ff)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
    
    def forward(self, x):
        attn_output = self.self_attention.forward(x)
        x = self.norm1.forward(x + attn_output)
        
        ff_output = self.feed_forward.forward(x)
        x = self.norm2.forward(x + ff_output)
        
        return x

解码器掩码自注意力

解码器中使用掩码防止关注未来位置:

python
def create_causal_mask(seq_len):
    """创建因果掩码,防止看到未来信息"""
    mask = np.triu(np.ones((seq_len, seq_len)), k=1)
    return mask == 0  # True表示可以关注,False表示需要掩盖

def masked_self_attention(Q, K, V):
    """带掩码的自注意力"""
    seq_len = Q.shape[0]
    mask = create_causal_mask(seq_len)
    return scaled_dot_product_attention(Q, K, V, mask)

交叉注意力

解码器通过交叉注意力关注编码器的输出:

下面只是接口示意,并非可直接运行的实现。生产级交叉注意力需要为解码器设置独立 Query 投影,为编码器设置 Key/Value 投影,并处理掩码和形状检查。

python
class CrossAttention:
    def __init__(self, d_model, num_heads):
        self.attention = MultiHeadAttention(d_model, num_heads)
    
    def forward(self, decoder_input, encoder_output):
        """
        交叉注意力
        decoder_input: 解码器输入,用于生成Query
        encoder_output: 编码器输出,用于生成Key和Value
        """
        pass
graph TB subgraph "Transformer中的三种注意力" subgraph "编码器" EI[输入] --> ESA["自注意力 全部可见"] end subgraph "解码器" DI[输出历史] --> DSA["掩码自注意力 只看过去"] DSA --> CA[交叉注意力] ESA --> CA end end

完整代码实现

以下是一个完整的自注意力层实现:

python
import numpy as np

class CompleteAttentionLayer:
    def __init__(self, d_model=512, num_heads=8, dropout_rate=0.1):
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        scale = np.sqrt(2.0 / (d_model + self.d_k))
        self.W_q = np.random.randn(d_model, d_model) * scale
        self.W_k = np.random.randn(d_model, d_model) * scale
        self.W_v = np.random.randn(d_model, d_model) * scale
        self.W_o = np.random.randn(d_model, d_model) * scale
        
        self.dropout_rate = dropout_rate
    
    def softmax(self, x, axis=-1):
        exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
        return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
    
    def dropout(self, x, training=True):
        if not training or self.dropout_rate == 0:
            return x
        mask = np.random.binomial(1, 1 - self.dropout_rate, x.shape)
        return x * mask / (1 - self.dropout_rate)
    
    def split_heads(self, x, batch_size):
        x = x.reshape(batch_size, -1, self.num_heads, self.d_k)
        return x.transpose(0, 2, 1, 3)
    
    def forward(self, x, mask=None, training=True):
        batch_size = x.shape[0] if len(x.shape) == 3 else 1
        if len(x.shape) == 2:
            x = x[np.newaxis, :, :]
        
        Q = np.matmul(x, self.W_q)
        K = np.matmul(x, self.W_k)
        V = np.matmul(x, self.W_v)
        
        Q = self.split_heads(Q, batch_size)
        K = self.split_heads(K, batch_size)
        V = self.split_heads(V, batch_size)
        
        scores = np.matmul(Q, K.transpose(0, 1, 3, 2)) / np.sqrt(self.d_k)
        
        if mask is not None:
            scores = np.where(mask == 0, -1e9, scores)
        
        attention_weights = self.softmax(scores)
        attention_weights = self.dropout(attention_weights, training)
        
        context = np.matmul(attention_weights, V)
        
        context = context.transpose(0, 2, 1, 3)
        context = context.reshape(batch_size, -1, self.d_model)
        
        output = np.matmul(context, self.W_o)
        
        if batch_size == 1:
            output = output.squeeze(0)
        
        return output, attention_weights


if __name__ == "__main__":
    d_model = 512
    num_heads = 8
    seq_len = 10
    
    attention = CompleteAttentionLayer(d_model, num_heads)
    
    x = np.random.randn(seq_len, d_model)
    
    output, weights = attention.forward(x)
    
    print(f"输入形状: {x.shape}")
    print(f"输出形状: {output.shape}")
    print(f"注意力权重形状: {weights.shape}")

实践指南

使用PyTorch实现

在实际项目中,推荐使用深度学习框架:

python
import torch
import torch.nn as nn

class AttentionLayer(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        self.attention = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            dropout=dropout,
            batch_first=True
        )
    
    def forward(self, x, mask=None):
        output, weights = self.attention(x, x, x, attn_mask=mask)
        return output, weights

d_model = 512
num_heads = 8
seq_len = 20
batch_size = 4

layer = AttentionLayer(d_model, num_heads)
x = torch.randn(batch_size, seq_len, d_model)
output, weights = layer(x)

注意力机制的调优技巧

  1. 头数选择:该实现要求 d_model 能被 num_heads 整除;应比较配置并测量质量、显存、吞吐和算子效率
  2. 缩放因子:标准做法是除以√d_k,某些变体使用可学习的缩放
  3. Dropout:在注意力权重上应用dropout可以防止过拟合
  4. 位置编码:注意力机制本身不包含位置信息,需要额外添加

总结

注意力机制的核心要点:

  1. 动态权重分配:根据输入内容动态决定关注哪些部分
  2. Query-Key-Value:通过三个向量实现信息的查询、匹配和提取
  3. 缩放点积:使用√d_k缩放保持梯度稳定
  4. 多头并行:多个注意力头关注不同类型的信息
  5. 诊断而非证明:注意力权重可以帮助检查计算,但不能单独解释模型决策

注意力机制是理解Transformer、GPT、BERT等现代大语言模型的基础。掌握这些原理,将帮助你更好地使用和开发AI应用。

常见问题

注意力机制和人类注意力有什么区别?

注意力机制是受人类注意力启发的数学模型,但两者有本质区别。人类注意力是生物神经系统的复杂过程,涉及意识、情感等因素;而注意力机制是纯粹的数学计算,通过点积和softmax实现权重分配。模型的"注意力"只是一种比喻,表示不同输入元素对输出的贡献程度。

为什么Transformer只用注意力不用RNN?

RNN 的递归结构限制了训练并行性,而自注意力可以在一个层内同时计算所有位置。两者在内存、延迟和流式处理上各有权衡,性能结论取决于任务、序列长度、硬件和架构。Transformer 也不是“纯注意力”模型,因为它还包含前馈网络和归一化层。

多头注意力的头数如何选择?

该实现的主要约束是 d_model 必须能被 num_heads 整除。在固定 d_model 时,改变头数主要改变每个头的维度,并不会自动增加投影参数。应比较少量配置,并测量质量、显存、吞吐和算子效率。

注意力机制的计算复杂度为什么是O(n²)?

朴素稠密自注意力会形成 n×n 的分数矩阵,因此分数计算和分数存储项随序列长度呈二次增长,此外还存在投影和算子开销。优化内核可以减少显式物化的内存,但不改变精确算法;稀疏或近似变体会引入质量和硬件权衡,端到端并不普遍是 O(n)。

如何理解注意力权重的可视化结果?

注意力权重热力图显示了每个Query位置对每个Key位置的关注程度。高权重(深色)表示强关联。常见模式包括:对角线高亮(自身关注)、特定词对高亮(语义关联)、句首/句尾高亮(特殊token)。但要注意,注意力权重不等于因果解释,高权重不一定意味着该位置对最终预测最重要。