把Self-Attention换到Cross时踩过的坑

作者: 醉月思发布日期: 2023-05-23 分类: 技术实践

序列一长，Self-Attention 就 OOM；head 一多，权重又全挤在头尾——这是我第一次自己写 Attention 时碰到的两个坎。

从一个简单问题开始

让我真正想搞清楚 Attention 的，是序列标注里想把 BiLSTM-CRF 换成带 Attention 的版本：序列一长 BiLSTM 还是丢远距离信息，能不能让模型自己挑重点？

第一次实现的时候，直接照着论文公式写了个Self-Attention层：

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

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=8):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        # x: [batch_size, seq_len, embed_dim]
        batch_size, seq_len, _ = x.shape

        # 生成Q, K, V
        qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # [3, batch_size, num_heads, seq_len, head_dim]
        q, k, v = qkv[0], qkv[1], qkv[2]

        # 计算注意力分数
        attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)  # [batch_size, num_heads, seq_len, seq_len]

        # 应用mask
        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2) == 0, float('-inf'))

        # softmax和输出
        attn = F.softmax(attn, dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
        return self.out(out)

跑起来后第一个问题就来了：序列长度稍微长一点就报CUDA out of memory。

第一个坑：序列长度和显存的矛盾

报错信息很直接：32的batch_size、512的序列长度，在8GB显存的GPU上根本跑不动。问题出在注意力矩阵上：[batch_size, num_heads, seq_len, seq_len]这个张量的空间复杂度是O(n²)，n一上去，显存就扛不住了。

我当时试了几种解决办法：

减小batch_size：从32降到16，能跑，但训练慢了将近一倍。
截断序列：把超过256长度的序列截断。

def truncate_or_pad(seq, max_len=256):
    if len(seq) > max_len:
        return seq[:max_len]
    return seq + [0] * (max_len - len(seq))

问题暂时解决了，但留下了一个疑问：Self-Attention的计算到底有没有必要对所有位置都计算？

第二个坑：理解Multi-Head Attention

打印 attention weight 之后才发现：每个 head 独立学 Q、K、V 投影。head 太多、数据又不够时，很多 head 学不到东西，就退化成「只盯边界」。

当时我的embedding维度是512，head数量设了8，每个head_dim=64。后来把head数量降到4，每个head_dim=128，问题就缓解了很多：

class SelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=4):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"

        self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        self.out = nn.Linear(embed_dim, embed_dim)
        self.dropout = nn.Dropout(0.1)  # 加dropout防止过拟合

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        qkv = self.qkv(x).reshape(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2) == 0, float('-inf'))

        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)  # 防止过拟合

        out = (attn @ v).transpose(1, 2).reshape(batch_size, seq_len, self.embed_dim)
        return self.out(out)

改动很小：head数量减半、加了dropout、把输出投影前后的维度都收紧了一些。效果却很明显：训练收敛快了约30%，注意力分布也更合理了。

这时候才体会到，Multi-Head 得在数据量和模型复杂度之间找平衡，不是 head 越多越好。

第三个坑：什么时候该用Cross-Attention

真正让我意识到Self-Attention和Cross-Attention区别的，是一个问答系统的任务。任务是这样的：给定一段参考文本和一个问题，要模型从参考文本里提取答案片段。

一开始我直接把问题和参考文本拼成一个序列，用Self-Attention处理：

# 拼接query和context
query_tokens = tokenizer(question, return_tensors='pt')['input_ids']  # [1, query_len]
context_tokens = tokenizer(context, return_tensors='pt')['input_ids']  # [1, context_len]

combined = torch.cat([query_tokens, context_tokens], dim=1)  # [1, query_len + context_len]
output = self_attention_layer(combined)

结果很糟糕：模型要么总是输出整个context，要么就输出一些不相关的片段。打印attention权重后发现，问题在于query和context之间的注意力信号被Self-Attention的"全连接"性质冲淡了——模型把所有位置等同看待，但实际情况是：context的信息应该根据query来动态筛选。

这时候才意识到应该用Cross-Attention：

class CrossAttention(nn.Module):
    def __init__(self, embed_dim, num_heads=4):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads

        self.q_proj = nn.Linear(embed_dim, embed_dim)
        self.k_proj = nn.Linear(embed_dim, embed_dim)
        self.v_proj = nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key_value, mask=None):
        # query: [batch_size, query_len, embed_dim]
        # key_value: [batch_size, key_len, embed_dim]
        batch_size, query_len, _ = query.shape
        key_len = key_value.shape[1]

        Q = self.q_proj(query).reshape(batch_size, query_len, self.num_heads, self.head_dim)
        K = self.k_proj(key_value).reshape(batch_size, key_len, self.num_heads, self.head_dim)
        V = self.v_proj(key_value).reshape(batch_size, key_len, self.num_heads, self.head_dim)

        Q = Q.permute(0, 2, 1, 3)  # [batch_size, num_heads, query_len, head_dim]
        K = K.permute(0, 2, 1, 3)
        V = V.permute(0, 2, 1, 3)

        attn = (Q @ K.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if mask is not None:
            attn = attn.masked_fill(mask.unsqueeze(1).unsqueeze(2) == 0, float('-inf'))

        attn = F.softmax(attn, dim=-1)
        out = (attn @ V).transpose(1, 2).reshape(batch_size, query_len, self.embed_dim)
        return self.out(out)

用Cross-Attention重写后，query作为Q，context作为K和V，效果立竿见影：F1值从0.47提升到0.63。

这里的关键区别是：Self-Attention是同一个序列内部做注意力，适合序列建模；Cross-Attention是两个序列之间做注意力，适合信息检索和融合。

第四个坑：注意力权重可视化

模型训练出来后，最直观的验证方法就是看注意力权重到底学了什么。但第一次可视化的时候，我掉进了坑里。

最原始的画法是直接用matplotlib热图：

import matplotlib.pyplot as plt
import seaborn as sns

def plot_attention(attn_weights, tokens, heads=[0, 1, 2, 3]):
    # attn_weights: [batch_size, num_heads, seq_len, seq_len]
    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
    for i, head in enumerate(heads):
        ax = axes[i // 2, i % 2]
        sns.heatmap(attn_weights[0, head].detach().cpu().numpy(),
                    xticklabels=tokens,
                    yticklabels=tokens,
                    cmap='viridis',
                    ax=ax)
        ax.set_title(f'Head {head}')
    plt.tight_layout()
    plt.show()

跑出来的图有几个问题：

中文标签显示乱码：matplotlib默认字体不支持中文。
序列太长看不清：256个token在热图上挤成一团。
head之间的差异不明显：有些head学的模式很像。

后来改用seaborn的clustermap加字体配置：

plt.rcParams['font.sans-serif'] = ['SimHei']  # 中文字体
plt.rcParams['axes.unicode_minus'] = False  # 负号显示

def plot_attention_cluster(attn_weights, tokens, head_idx=0, max_tokens=50):
    # 只取前50个token避免太挤
    attn = attn_weights[0, head_idx, :max_tokens, :max_tokens].detach().cpu().numpy()
    short_tokens = tokens[:max_tokens]

    sns.clustermap(attn,
                   xticklabels=short_tokens,
                   yticklabels=short_tokens,
                   cmap='YlOrRd',
                   figsize=(10, 8),
                   row_cluster=True,
                   col_cluster=True)
    plt.show()

这样能看出来不同head确实学到了不同的模式：有的关注语义相似词，有的关注句法结构，有的关注位置关系。

但更实用的可视化方法是直接看关键位置的注意力分布：

def plot_token_attention(attn_weights, tokens, target_token, head_idx=0):
    target_idx = tokens.index(target_token)
    attn = attn_weights[0, head_idx, target_idx].detach().cpu().numpy()

    plt.figure(figsize=(12, 4))
    plt.bar(range(len(tokens)), attn)
    plt.xticks(range(len(tokens)), tokens, rotation=45, ha='right')
    plt.title(f'Attention distribution for "{target_token}" (Head {head_idx})')
    plt.tight_layout()
    plt.show()

# 查看某个词的关注点
plot_token_attention(attn_weights, tokens, '人工智能')

这种可视化更有实践意义：能直观看到模型对某个词的理解是否合理。

实际项目中的一些选择

做完这些实验后，我在实际项目中总结了一些选择标准：

Self-Attention适合这些场景：

序列内部建模，比如文本分类、机器翻译的编码器
需要捕捉长距离依赖的任务
序列长度不会太长（<512）

Cross-Attention适合这些场景：

有明显的信息检索需求，比如QA、摘要
两个序列之间有明确的主从关系
一个序列需要另一个序列的信息来增强

实用调参经验：

head数量：embedding_dim在128-256之间用2-4个head，512-768用4-8个head，再大需要考虑显存。
head_dim：至少64，最好是128，太小每个head学不到有意义的东西。
dropout：0.1-0.3之间，注意力权重dropout防过拟合效果明显。
序列长度：如果实在太长（>1024），考虑用稀疏注意力或者分块注意力。

一些没完全解决的问题

在实践中还有一些问题没想明白：

注意力权重的可解释性到底有多强？：有时候看起来合理的注意力分布，可能是模型学到的某种分布模式，而不是真正的"关注"。
如何设计更好的注意力偏置？：比如位置编码，目前的相对位置编码方法很多，但没有一个普适的最佳方案。
注意力机制的效率优化：除了显存问题，计算效率也是瓶颈，特别是在低算力设备上。

这些问题可能需要更多的实践和实验才能回答。

写在最后

Attention 对我最大的价值，是把「该看哪里」这件事显式化了。Self-Attention 管序列内部，Cross-Attention 管两个序列之间的信息筛选——QA 那趟从 F1 0.47 到 0.63，就是选错机制的典型代价。

head 数量、dropout、序列长度这些调参，没有万能公式，得结合数据量和显存慢慢试。可视化也别省，中文乱码、序列太长挤成一团，都是真实会踩的坑。

代码环境：Python 3.9, PyTorch 1.12.1, CUDA 11.3 GPU：NVIDIA RTX 3090 24GB 测试序列长度：128-512 batch_size：16-32 训练轮数：50 epochs

版权声明: 本文首发于指尖魔法屋-把Self-Attention换到Cross时踩过的坑（https://blog.thinkmoon.cn/post/243-attention-mechanism-deep-dive-self-cross-attention/）转载或引用必须申明原指尖魔法屋来源及源地址！

关于Transformer 架构的几点记录

后来硬着头皮从头实现了一遍 Transformer，过程中踩了不少坑，也算是对这个架构有了些实际体会。

📁 学习笔记

📅 2024-06-10

AI注意力即一切实践笔记

当时我在做一个机器翻译的项目，用的还是基于 LSTM 的 seq2seq 模型，那个下午训练曲线又震荡得厉害，我盯着屏幕看了二十分钟，想的是：这…

📁 学习笔记

📅 2024-08-17

从LeNet走到Transformer：神经网络架构设计笔记

后来发现 Transformer 更通用，但迁移过程踩了不少坑。

📁 技术实践

📅 2023-05-28