可视化 Attention 头实战：它到底在关注什么？

“Attention 很重要”这句话大家都听过。

但真正训练模型时，更关键的问题是：

你的 attention 头到底学到了什么？

如果你从不看注意力图，只盯着 loss，很多问题会被隐藏很久，比如：

• 某些头完全退化（几乎均匀分布）；
• 模型过度关注最近 token，长程依赖失效；
• padding 或特殊符号被异常高频关注。

这篇文章给你一套从 0 到 1 的 attention 可视化工作流。

---

1. 先回答：为什么要看 attention 可视化

可视化不是“为了好看”，而是为了定位问题：

1. 验证机制是否生效：因果掩码是否正确；
2. 发现头部冗余：哪些头在重复做同一件事；
3. 诊断数据问题：异常符号是否污染注意力；
4. 辅助架构决策：是否需要减少/增加 head 数。

它是你在“loss 曲线之外”的第二双眼睛。

---

2. 你需要导出哪些中间结果

要做可视化，最关键的是拿到每层每头的注意力权重：

• 形状通常为 [B, n_head, T, T]；
• 其中最后两个维度表示“query 位置 -> key 位置”权重。

在 PyTorch 中，常见做法是：

• 在 attention 模块里临时保留 softmax 后的权重；
• 推理时开启开关 return_attn=True；
• 只导出小 batch，避免显存压力。

注意：训练阶段频繁保存全量注意力会明显拖慢速度，建议只在评估 step 抽样记录。

---

3. 最小可视化代码：先画出一张热力图

下面给一个简化示意（Matplotlib）：

import matplotlib.pyplot as plt

# attn: [n_head, T, T]
# tokens: 长度为 T 的 token 字符串列表

def plot_head(attn, tokens, head_idx=0):
    m = attn[head_idx].detach().cpu().numpy()
    plt.figure(figsize=(8, 6))
    plt.imshow(m, cmap="viridis")
    plt.xticks(range(len(tokens)), tokens, rotation=90)
    plt.yticks(range(len(tokens)), tokens)
    plt.title(f"Attention Head {head_idx}")
    plt.colorbar()
    plt.tight_layout()
    plt.show()

先把第一张图画出来，再谈“高级解释”。

---

4. 读图方法：四种常见模式

4.1 对角线强（local pattern）

说明模型主要看邻近 token，擅长局部依赖。

4.2 前缀聚焦（prefix sink）

很多 query 都指向开头 token（如 BOS），常见于格式控制。

4.3 远程跳跃（long-range link）

能看到非邻近区域高亮，可能在捕捉长程语义或结构对应关系。

4.4 全局均匀（head collapse）

几乎一片平均色，说明该头信息量有限，可能退化。

---

5. 推荐工具链（按复杂度递进）

方案 A：Matplotlib/Seaborn 自绘（最灵活）

适合：

• 快速调试；
• 自定义指标叠加；
• 与训练日志联动。

方案 B：BertViz / Transformer 可视化工具（上手快）

适合：

• 交互式查看多层多头；
• 演示与教学；
• 快速比较不同输入样本。

方案 C：实验平台整合（TensorBoard/W&B）

适合：

• 长期项目持续跟踪；
• 版本间对比；
• 团队协作复盘。

建议从 A 起步，稳定后再接 B/C。

---

6. 把可视化变成“可量化监控”

只看图很主观，建议配合数值指标：

1. Attention Entropy（熵）
   • 低熵：更聚焦；
   • 高熵：更分散。
2. Head Diversity（头间差异）
   • 衡量不同头是否学到不同模式。
3. Special Token Focus Ratio
   • 统计 <bos>/<pad>/<eos> 被关注比例。
4. Long-range Attention Ratio
   • 统计超过距离阈值的注意力占比。

把这些指标按 step 画曲线，你会比“人工看几张图”更早发现异常。

---

7. 三个真实可用的排障案例

案例一：模型总复读

现象：输出“因此因此因此……”。

可视化发现：

• 多个头只盯最近 1~2 个 token；
• 长程注意力几乎为 0。

处理：

• 调整训练数据中重复样本；
• 增加上下文长度；
• 在采样阶段加入 repetition penalty。

案例二：中文标点异常堆积

现象：句号、逗号比例过高。

可视化发现：

• 某些头异常偏向标点 token；
• 清洗后残留了格式化噪声文本。

处理：

• 重做语料清洗规则；
• 限制噪声模板数据比例；
• 复训并对比 attention 指标。

案例三：长文问答答非所问

现象：短问答还行，长上下文明显跑偏。

可视化发现：

• 深层头部对远端上下文关注不足；
• 实际有效感受野偏短。

处理：

• 增加 block_size 并继续训练；
• 检查 RoPE/位置编码实现；
• 对长样本做专门评估集。

---

8. 可视化时的注意事项

1. 不要只看单个样本：至少比较多种输入类型。
2. 不要只看第一层：浅层与深层职责不同。
3. 不要只看“好看图”：要结合任务指标与失败案例。
4. 注意分词影响：tokenizer 不同，attention 图会明显不同。
5. 警惕过度解读：attention 可解释性有边界，应结合 probing 与消融实验。

---

9. 一套建议执行流程（可直接落地）

• 第 1 周：为模型加 return_attn 开关，完成单样本热力图；
• 第 2 周：接入熵、头多样性、长程比例三项指标；
• 第 3 周：建立“失败样本库”并固定可视化报表；
• 第 4 周：把可视化结论反馈到 tokenizer、数据、采样策略调优。

你会得到一个闭环：

训练 -> 观测 -> 解释 -> 调优 -> 再训练。

---

10. 小结

attention 可视化的价值，不在于“证明 Transformer 很神奇”，而在于帮你把模型行为从黑箱变成灰箱。

当你开始系统地看图、看指标、看失败样本，你会更快回答这些关键问题：

• 为什么它会错；
• 错在数据、结构还是推理策略；
• 下一次迭代该改哪里。

这也是从“能训练模型”走向“能调优模型”的关键一步。