Transformer实战：用Python手写Cross-Attention机制（附完整代码）

Transformer实战：用Python手写Cross-Attention机制（附完整代码）

当我们在处理序列数据时，如何让模型学会关注不同序列之间的相关性？Cross-Attention机制正是解决这一问题的关键。与Self-Attention不同，Cross-Attention允许一个序列（如翻译中的目标语言）动态关注另一个序列（如源语言）中的重要部分。这种机制在机器翻译、文本摘要等任务中表现出色。

本文将带你从零实现Cross-Attention的核心逻辑，通过PyTorch代码展示Q/K/V矩阵计算、注意力分数归一化等关键步骤。我们不仅会剖析数学原理，还会提供可直接运行的Colab示例，帮助你快速将这一技术应用到自己的NLP项目中。

1. Cross-Attention核心原理拆解

Cross-Attention的核心思想是建立两个不同序列之间的动态关联。想象你正在翻译一句话：当生成目标语言的每个词时，模型需要知道源语言中哪些部分最值得关注。这种"跨序列关注"正是Cross-Attention的魔力所在。

1.1 三个关键矩阵：Q, K, V

在Cross-Attention中，三个矩阵扮演着核心角色：

• 查询矩阵(Query): 来自目标序列，表示"我想知道什么"
• 键矩阵(Key): 来自源序列，表示"我能提供什么信息"
• 值矩阵(Value): 来自源序列，包含实际的特征表示

它们的计算过程可以表示为：

Q = target_sequence @ W_q  # [batch_size, tgt_len, d_k]
K = source_sequence @ W_k  # [batch_size, src_len, d_k] 
V = source_sequence @ W_v  # [batch_size, src_len, d_v]

其中W_q, W_k, W_v是可学习的参数矩阵，d_k和d_v是维度大小。

1.2 注意力分数计算

模型通过计算Q和K的点积来评估目标序列每个位置与源序列所有位置的相关性：

scores = Q @ K.transpose(-2, -1) / sqrt(d_k)  # [batch_size, tgt_len, src_len]

这里除以√d_k是为了防止点积结果过大导致梯度消失问题。接着用softmax将分数转换为概率分布：

attn_weights = torch.softmax(scores, dim=-1)

1.3 加权求和输出

最后，用注意力权重对Value矩阵加权求和，得到融合了两个序列信息的输出：

output = attn_weights @ V  # [batch_size, tgt_len, d_v]

这个过程可以用一个简单的表格来总结：

2. PyTorch完整实现

现在让我们用PyTorch实现一个完整的Cross-Attention层。这个实现包含多头注意力机制，这是Transformer中的标准配置。

2.1 基础实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class CrossAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        
        # 定义Q,K,V的投影矩阵
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        
    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)
        
        # 1. 投影到Q,K,V空间
        Q = self.W_q(query)  # [batch_size, tgt_len, d_model]
        K = self.W_k(key)    # [batch_size, src_len, d_model]
        V = self.W_v(value)  # [batch_size, src_len, d_model]
        
        # 2. 分割多头
        Q = Q.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
        
        # 3. 计算注意力分数
        scores = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.d_k, dtype=torch.float32))
        
        # 4. 应用mask（可选）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)
        
        # 5. softmax归一化
        attn_weights = F.softmax(scores, dim=-1)
        
        # 6. 加权求和
        output = torch.matmul(attn_weights, V)
        
        # 7. 合并多头
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
        
        # 8. 最终投影
        output = self.W_o(output)
        
        return output, attn_weights

2.2 使用示例

让我们看一个简单的使用例子，模拟机器翻译中的编码器-解码器交互：

# 假设编码器输出（源语言表示）
encoder_output = torch.randn(2, 10, 512)  # [batch_size, src_len, d_model]

# 解码器输入（目标语言表示）
decoder_input = torch.randn(2, 8, 512)    # [batch_size, tgt_len, d_model]

# 创建Cross-Attention层
cross_attn = CrossAttention(d_model=512, num_heads=8)

# 前向计算
output, attn_weights = cross_attn(
    query=decoder_input,  # Q来自解码器
    key=encoder_output,   # K来自编码器
    value=encoder_output  # V来自编码器
)

print(f"输出维度: {output.shape}")      # [2, 8, 512]
print(f"注意力权重维度: {attn_weights.shape}")  # [2, 8, 10, 10]

3. 注意力可视化技巧

理解模型关注什么是调试和解释模型行为的重要手段。下面我们实现一个简单的注意力可视化函数：

import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention(attention_weights, source_tokens=None, target_tokens=None):
    """
    绘制注意力权重热力图
    :param attention_weights: [tgt_len, src_len]的注意力权重矩阵
    :param source_tokens: 源语言token列表（可选）
    :param target_tokens: 目标语言token列表（可选）
    """
    plt.figure(figsize=(10, 8))
    ax = sns.heatmap(attention_weights, cmap="YlGnBu")
    
    if source_tokens is not None:
        ax.set_xticks(range(len(source_tokens)))
        ax.set_xticklabels(source_tokens, rotation=45)
    
    if target_tokens is not None:
        ax.set_yticks(range(len(target_tokens)))
        ax.set_yticklabels(target_tokens, rotation=0)
    
    plt.xlabel("Source Sequence")
    plt.ylabel("Target Sequence")
    plt.title("Cross-Attention Weights")
    plt.show()

使用示例：

# 模拟的注意力权重（假设batch中第一个样本）
sample_attn = attn_weights[0].mean(dim=0)  # 平均多头注意力

# 模拟的token（实际使用时替换为真实token）
src_tokens = ["The", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"]
tgt_tokens = ["Le", "rapide", "renard", "marron", "saute", "par-dessus", "le", "chien"]

plot_attention(sample_attn.detach().numpy(), src_tokens, tgt_tokens)

这将生成一个热力图，显示目标语言每个词（纵轴）与源语言词（横轴）之间的注意力强度。

4. 性能优化技巧

在实际应用中，Cross-Attention可能会成为计算瓶颈，特别是处理长序列时。以下是几个优化建议：

4.1 内存高效的注意力计算

原始实现需要存储中间注意力矩阵（大小为batch×heads×tgt_len×src_len），对于长序列这会消耗大量内存。我们可以使用以下技巧：

# 内存高效的注意力计算
with torch.backends.cuda.sdp_kernel(enable_flash=True):
    output = F.scaled_dot_product_attention(
        Q, K, V, 
        attn_mask=mask,
        dropout_p=0.1 if self.training else 0
    )

提示：PyTorch 2.0+提供了优化的scaled_dot_product_attention实现，自动选择最快的可用内核（flash attention、memory-efficient attention等）。

4.2 关键参数选择

4.3 批处理优化

当处理不同长度的序列时，使用pad_sequence和attention_mask：

from torch.nn.utils.rnn import pad_sequence

# 假设我们有不同长度的源序列
src_sequences = [torch.randn(5, 512), torch.randn(7, 512), torch.randn(4, 512)]
padded_src = pad_sequence(src_sequences, batch_first=True)  # [3, 7, 512]

# 创建attention mask
src_lengths = torch.tensor([5, 7, 4])
src_mask = torch.arange(padded_src.size(1))[None, :] < src_lengths[:, None]  # [3, 7]
src_mask = src_mask.unsqueeze(1).unsqueeze(2)  # [3, 1, 1, 7]

# 在计算注意力分数时应用mask
scores = scores.masked_fill(src_mask == 0, -1e9)

5. 实际应用案例

让我们看一个简化的机器翻译任务中如何使用Cross-Attention的例子。假设我们已经有了编码器和解码器的基本框架：

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model, num_heads, ff_dim, dropout=0.1):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)  # 自注意力
        self.cross_attn = CrossAttention(d_model, num_heads)     # 交叉注意力
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, d_model)
        )
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout = nn.Dropout(dropout)
        
    def forward(self, x, encoder_output, src_mask, tgt_mask):
        # 自注意力
        attn_output, _ = self.self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + self.dropout(attn_output))
        
        # 交叉注意力
        cross_output, attn_weights = self.cross_attn(
            query=x,
            key=encoder_output,
            value=encoder_output,
            mask=src_mask
        )
        x = self.norm2(x + self.dropout(cross_output))
        
        # 前馈网络
        ffn_output = self.ffn(x)
        x = self.norm3(x + self.dropout(ffn_output))
        
        return x, attn_weights

在这个实现中，每个解码器层依次执行：

1. 自注意力：关注已生成的目标序列部分
2. 交叉注意力：关注编码器输出的源序列
3. 前馈网络：进一步变换特征

实际训练时，我们会观察到随着模型学习，注意力权重会逐渐变得更加集中和有意义。例如在翻译任务中，模型会学会将目标语言的形容词与源语言对应的形容词对齐。