Transformer注意力机制可视化工具：annotated-transformer详解

引言：破解注意力黑箱困境

你是否曾困惑于Transformer模型如何"思考"？当输入"我爱机器学习"时，模型究竟关注哪个词来生成"机器学习"的翻译？注意力权重（Attention Weights）作为Transformer的核心，却常常被视作难以捉摸的黑箱。annotated-transformer项目通过完整实现《Attention is All You Need》论文，提供了从模型训练到注意力可视化的全流程工具链，让你直观理解注意力机制的工作原理。

读完本文你将获得：

• 掌握3种注意力可视化技术（编码器自注意力/解码器自注意力/编码器-解码器交叉注意力）
• 学会使用Altair构建交互式注意力热力图
• 理解注意力掩码（Mask）对模型行为的影响
• 获取可复现的注意力可视化代码模板

核心功能解析：从模型结构到可视化实现

1. 注意力机制的工程实现

annotated-transformer通过MultiHeadedAttention类实现了多头注意力机制，关键代码如下：

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super().__init__()
        assert d_model % h == 0
        self.d_k = d_model // h  # 每个头的维度：512/8=64
        self.h = h               # 头数：8
        self.linears = clones(nn.Linear(d_model, d_model), 4)  # 4个线性层
        self.attn = None         # 存储注意力权重
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        if mask is not None:
            mask = mask.unsqueeze(1)  # 扩展维度以适配多头
        nbatches = query.size(0)

        # 1) 线性投影并分拆为h个头
        query, key, value = [
            lin(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
            for lin, x in zip(self.linears, (query, key, value))
        ]

        # 2) 计算注意力权重
        x, self.attn = attention(
            query, key, value, mask=mask, dropout=self.dropout
        )

        # 3) 拼接多头结果并应用最终线性层
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

关键设计亮点：

• 通过clones函数创建4个共享维度的线性层（Q/K/V投影+输出投影）
• 使用张量变形（view/transpose）实现多头并行计算
• 将注意力权重存储在self.attn属性中，为后续可视化提供数据来源

2. 三种注意力类型及其可视化价值

项目实现了Transformer中的全部三种注意力机制，每种类型提供独特的语言学洞察：

注意力类型	输入来源	核心作用	可视化价值
编码器自注意力	编码器前一层输出	捕捉输入序列内部依赖关系	展示句法结构（如主谓宾关系）
解码器自注意力	解码器前一层输出	建模输出序列的时序依赖	揭示语言模型的上下文依赖
编码器-解码器交叉注意力	解码器查询,编码器输出	建立源语言到目标语言的对齐	可视化翻译过程中的词对齐关系

代码定位：在DecoderLayer类中清晰区分了三种注意力的调用逻辑：

class DecoderLayer(nn.Module):
    def forward(self, x, memory, src_mask, tgt_mask):
        m = memory
        # 解码器自注意力（带掩码）
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        # 编码器-解码器交叉注意力
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

3. 注意力掩码：控制信息流的关键

为防止模型关注未来信息，annotated-transformer实现了两种核心掩码：

3.1 后续位置掩码（Subsequent Mask）

def subsequent_mask(size):
    "掩盖后续位置，确保预测仅依赖已知输出"
    attn_shape = (1, size, size)
    subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8)
    return subsequent_mask == 0  # True表示允许注意力

可视化效果：生成下三角矩阵，禁止当前位置关注后续位置：

# 生成掩码可视化
def example_mask():
    LS_data = pd.concat([
        pd.DataFrame({
            "Subsequent Mask": subsequent_mask(20)[0][x,y].flatten(),
            "目标位置": y,
            "当前位置": x,
        }) for y in range(20) for x in range(20)
    ])
    
    return alt.Chart(LS_data).mark_rect().encode(
        x="目标位置:O", y="当前位置:O",
        color="Subsequent Mask:Q"  # 蓝色表示允许注意力
    ).properties(width=400, height=400)

3.2 填充掩码（Padding Mask）

用于屏蔽输入序列中的填充符号（如<PAD>），确保模型不关注无意义的填充 token。两种掩码在训练时会合并使用：

def make_std_mask(tgt, pad):
    "创建目标序列的掩码：结合填充掩码和后续位置掩码"
    tgt_mask = (tgt != pad).unsqueeze(-2)
    tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
    return tgt_mask

实战教程：构建你的第一个注意力可视化工具

1. 环境准备与模型训练

# 克隆项目仓库
git clone https://gitcode.com/gh_mirrors/an/annotated-transformer
cd annotated-transformer

# 安装依赖
pip install -r requirements.txt  # 包含torch/altair/pandas等

# 启动训练（使用合成数据快速验证）
python the_annotated_transformer.py --epochs 10 --batch_size 32

2. 提取注意力权重

训练完成后，通过修改forward方法保存注意力权重：

# 在MultiHeadedAttention类中添加
def get_attention_weights(self):
    return self.attn  # 返回最后一次前向传播的注意力权重

# 推理时提取
model.eval()
with torch.no_grad():
    output, attn_weights = model.inference(src)  # 需要修改inference方法返回权重

3. 使用Altair创建交互式热力图

import altair as alt
import pandas as pd

def visualize_attention(attn_weights, src_tokens, tgt_tokens, layer=0, head=0):
    """
    可视化指定层和头的注意力权重
    attn_weights: 注意力权重张量 (batch, heads, tgt_len, src_len)
    src_tokens: 源序列token列表
    tgt_tokens: 目标序列token列表
    """
    # 提取指定层、头的权重 (tgt_len, src_len)
    attn = attn_weights[layer][head].cpu().numpy()
    
    # 转换为DataFrame
    df = pd.DataFrame(attn, index=tgt_tokens, columns=src_tokens)
    df = df.stack().reset_index()
    df.columns = ["目标token", "源token", "注意力权重"]
    
    # 创建热力图
    return alt.Chart(df).mark_rect().encode(
        x=alt.X('源token:O', axis=alt.Axis(labelAngle=-45)),
        y=alt.Y('目标token:O'),
        color=alt.Color('注意力权重:Q', scale=alt.Scale(scheme='blueorange')),
        tooltip=['源token', '目标token', '注意力权重']
    ).properties(title=f'Layer {layer+1}, Head {head+1}', width=600, height=400)

4. 多维度注意力分析

通过对比不同头和层的注意力模式，揭示模型的分工机制：

def compare_attention_heads(attn_weights, src_tokens, tgt_tokens, layer=0):
    """比较同一层不同头的注意力模式"""
    charts = []
    for head in range(8):  # 遍历8个注意力头
        chart = visualize_attention(attn_weights, src_tokens, tgt_tokens, layer, head)
        charts.append(chart)
    
    # 水平拼接8个头的可视化结果
    return alt.hconcat(*charts).resolve_scale(color='independent')

典型模式：不同头会展现出不同的注意力偏好：

• 语法头：关注句法结构（如冠词与名词的关系）
• 语义头：关注语义关联（如同义词或上下位词）
• 位置头：主要关注相邻位置的token

高级应用：注意力模式分析与模型诊断

1. 跨层注意力演化分析

通过可视化不同层的注意力变化，观察模型如何逐步构建抽象表示：

def layer_attention_evolution(attn_weights, src_tokens, tgt_tokens, head=0):
    """展示同一头在不同层的注意力演化"""
    charts = []
    for layer in range(6):  # Transformer通常有6层
        chart = visualize_attention(attn_weights, src_tokens, tgt_tokens, layer, head)
        charts.append(chart)
    
    return alt.vconcat(*charts).resolve_scale(color='independent')

观察结论：低层通常关注局部词序，中层关注短语结构，高层关注长距离语义依赖。

2. 注意力权重统计分析

通过量化分析注意力分布，揭示模型行为特征：

def analyze_attention_distribution(attn_weights):
    """统计注意力权重的分布特征"""
    stats = {
        '平均注意力熵': [],
        '最大注意力值': [],
        '注意力分散度': []  # 1 - 最大权重占比
    }
    
    for layer in range(6):
        for head in range(8):
            attn = attn_weights[layer][head]
            entropy = -torch.sum(attn * torch.log(attn + 1e-10), dim=-1).mean().item()
            max_attn = attn.max(dim=-1)[0].mean().item()
            stats['平均注意力熵'].append(entropy)
            stats['最大注意力值'].append(max_attn)
            stats['注意力分散度'].append(1 - max_attn)
    
    # 转换为DataFrame并可视化
    df = pd.DataFrame(stats, index=[f'L{l}H{h}' for l in range(6) for h in range(8)])
    return alt.Chart(df.reset_index()).mark_bar().encode(
        x='index:N', y='平均注意力熵:Q', color='index:N'
    ).properties(width=800)

诊断价值：

• 低熵高最大权重：模型对某些token有强烈依赖（可能过拟合）
• 高熵低最大权重：注意力分散（可能欠拟合或任务难度高）
• 跨层波动大：模型训练不稳定的信号

常见问题与解决方案

1. 可视化结果噪声过大

可能原因：模型未充分训练或学习率设置不当。

解决方案：

# 调整优化器参数（源自论文的预热策略）
def get_std_opt(model):
    """使用论文推荐的优化器配置"""
    return NoamOpt(
        model.src_embed[0].d_model, 2, 4000,
        torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)
    )

2. 注意力权重过于稀疏

可能原因：多头注意力头数过多或隐藏层维度不足。

解决方案：调整模型超参数：

# 创建模型时调整头数和维度
model = make_model(
    src_vocab=10000, tgt_vocab=10000,
    N=6, d_model=512, h=8,  # 8头注意力
    d_ff=2048, dropout=0.3  # 增加dropout减轻过拟合
)

3. 中文注意力可视化乱码

解决方案：配置Altair支持中文字体：

alt.renderers.set_embed_options(
    font='SimHei',  # 设置中文字体
    scaleFactor=2.0  # 提高分辨率
)

总结与扩展

annotated-transformer项目不仅是Transformer的忠实实现，更是理解注意力机制的强大工具。通过本文介绍的可视化技术，你可以：

1. 教学演示：直观展示注意力机制原理
2. 模型诊断：通过注意力模式发现模型缺陷
3. 语言学研究：分析模型对句法/语义结构的学习
4. 改进优化：基于注意力分布设计更好的模型结构

扩展方向：

• 结合梯度归因方法（如Grad-CAM）定位关键注意力头
• 开发注意力权重编辑工具，干预模型预测
• 构建注意力模式库，实现模型行为的标准化评估

希望本文提供的工具和方法，能帮助你揭开注意力机制的神秘面纱，构建更可解释、更可靠的Transformer模型。立即克隆项目，开始你的注意力探索之旅吧！

附录：关键代码索引

功能	代码位置	核心类/函数
多头注意力实现	第350-400行	MultiHeadedAttention
注意力掩码	第280-300行	subsequent_mask
位置编码	第580-620行	PositionalEncoding
模型构建	第700-750行	make_model
可视化基础	第1200-1250行	example_mask