注意力可视化

源文件：chapter2/attention_visualization/README.md

注意力可视化

一个交互式可视化工具，用于探索语言模型中的注意力机制。每次 Agent 运行都会生成一条独特的轨迹，可在前端查看与对比。

概览

本项目提供了一种交互方式，帮助理解语言模型在处理不同类型查询时如何分配注意力。Agent 每次运行都会生成一个新的轨迹文件，记录：

• 输入查询与模型响应
• 逐 token 的注意力权重
• 跨层、跨注意力头的注意力模式
• 注意力分布的统计分析

架构

系统采用简单的架构：

1. Agent 生成轨迹：运行 agent.py 或 main.py 生成新轨迹
2. JSON 存储：每条轨迹保存为 frontend/public/trajectories/ 下的独立 JSON 文件
3. 前端可视化：React 应用加载并通过标签页导航展示所有轨迹

快速开始（独立 CLI）

复现第 2 章中描述的注意力模式（实验 2-2）最快的办法是使用独立命令行工具 attention_cli.py。它会运行真实模型、捕获其自注意力，并直接输出一张热力图 PNG——无需前端。

# Single heatmap for the default prompt (last layer, heads averaged)
python attention_cli.py

# Custom prompt, inspect a specific layer/head, choose the output path
python attention_cli.py --prompt "北京 的 天气 怎么样" \
    --layer 0 --head 3 --output layer0_head3.png

# Generate a short continuation first, then visualize the whole sequence
python attention_cli.py --prompt "Explain attention in one sentence." \
    --max-new-tokens 40

# Compare how the attention sink emerges across layers, side by side
python attention_cli.py --compare-layers 0 13 -1 --output layer_compare.png
​

运行 python attention_cli.py --help 查看完整参数列表。关键参数：

热力图怎么看。 行是 Query 位置（发起注意力的 token），列是 Key 位置（被注意的 token）。该工具会直接基于模型自身的权重测量并打印注意力池占比（attention-sink share）——即每行注意力落到第一个 token 上的比例。在
Qwen3-0.6B 上，最后一层的注意力池通常吸收每行约 75–85% 的注意力（即第 2 章"注意力储存池 / Attention Sink"现象），而第 0 层则接近局部的对角线模式。被掩码的上三角让因果"三角形"结构一目了然：每个 token 只能注意到自己及其之前的 token。

首次运行会下载模型权重（约 1–2 GB）。推荐使用 GPU/MPS，但这些短 prompt 在 CPU 上也能跑。

交互式前端工作流

可视化过程也可以拆分为两个手动步骤：生成轨迹，再在前端查看。

第 1 步：生成轨迹

选择以下任一方式生成轨迹数据：

# Option A: Run basic attention tracking demo
python agent.py

# Option B: Run ReAct agent with tool calling (demonstrates multi-step reasoning)
python main.py
​

每次运行都会在 frontend/public/trajectories/ 下生成一个带唯一时间戳的轨迹文件。

第 2 步：启动前端

在另一个终端启动前端服务：

cd frontend
npm install  # First time only
npm run dev
​

第 3 步：查看可视化

打开浏览器访问 http://localhost:3000

前端可以一直保持运行，同时在第一个终端生成新轨迹——它们会自动出现在界面中。

项目结构

attention_visualization/
├── attention_cli.py      # Standalone CLI: prompt -> attention heatmap PNG
├── agent.py               # Core attention tracking agent
├── main.py               # ReAct agent with tool calling
├── tools.py              # Tool implementations
├── visualization.py      # Visualization utilities (heatmap / comparison)
├── config.py            # Configuration settings
├── requirements.txt     # Python dependencies
├── env.example          # Environment variable template
├── frontend/            # Next.js frontend
│   ├── pages/          # React pages
│   ├── components/     # Visualization components
│   └── public/
│       └── trajectories/  # Stored trajectory JSONs
│           ├── trajectory_YYYYMMDD_HHMMSS.json
│           └── manifest.json  # Index of all trajectories
└── attention_data/      # Additional trajectory storage
​

工作原理

1. 轨迹生成

使用 agent.py

• 运行一个基础注意力跟踪 demo，涵盖多种查询类型
• 捕获单步响应的注意力权重
• 适合理解基础注意力模式

使用 main.py

• 实现一个带工具调用能力的 ReAct Agent
• 展示带结构化思考过程的多步推理
• 展示 Agent 使用工具时注意力如何转移
• 更适合理解复杂推理模式

两个脚本都会：

• 生成带时间戳的唯一轨迹文件
• 把结果保存到 frontend/public/trajectories/
• 更新清单文件以便前端发现

2. 数据格式

每条轨迹 JSON 包含：

{
  "id": "20250914_123456",
  "timestamp": "2025-09-14 12:34:56",
  "test_case": {
    "category": "Math",
    "query": "What is 25 * 37?",
    "description": "Agent trajectory from..."
  },
  "response": "The answer is...",
  "tokens": ["What", "is", "25", ...],
  "attention_data": {
    "tokens": [...],
    "attention_matrix": [[...]],
    "num_layers": 1,
    "num_heads": 16
  },
  "metadata": {...}
}
​

3. 前端可视化

React 前端：

• 从清单中加载所有轨迹
• 提供标签页在不同运行之间切换
• 展示注意力热力图、token 分析和统计信息
• 生成新轨迹时自动更新

特性

• 多条轨迹：每次 Agent 运行都会生成新的轨迹文件
• 标签页导航：在不同 Agent 运行之间轻松切换
• 注意力热力图：token 间注意力的交互式可视化
• token 分析：查看单个 token 及其注意力模式
• 统计指标：平均注意力、最大注意力、熵
• 类别支持：查询被归类（数学、知识、推理、代码、创意）
• 持久化存储：所有轨迹都会保存，可反复查看

生成自定义轨迹

使用 agent.py

编辑 demonstrate_attention_tracking() 函数以添加自定义查询：

test_prompts = [
    ("Your custom query here", "Category"),
    # Add more queries...
]
​

使用 main.py

ReAct Agent 演示了工具使用与多步推理。请在 demonstrate_react_agent() 中编辑测试查询。

手动生成

你也可以通过编程方式使用 Agent：

from agent import AttentionVisualizationAgent

agent = AttentionVisualizationAgent()
result = agent.generate_with_attention(
    "Your query here",
    max_new_tokens=100,
    temperature=0.3,
    save_trajectory=True,
    category="Custom"
)
​

环境要求

Python

• Python 3.10+
• PyTorch
• Transformers
• 完整列表见 requirements.txt

前端

• Node.js 14+
• npm 或 yarn
• 依赖见 frontend/package.json

安装

1. 克隆仓库

2. 设置环境变量（可选）：

cp env.example .env
# Edit .env to customize model, device, and visualization settings
​

3. 安装 Python 依赖：

pip install -r requirements.txt
​

4. 安装前端依赖：

cd frontend
npm install
​

使用技巧

• 首次安装：初次运行会下载模型（约 1~2 GB）。为获得更好性能，推荐使用 GPU/MPS
• 多次运行：同时运行 agent.py 和 main.py 以查看不同的注意力模式
• 对比轨迹：用标签页界面对比模型对相似查询的处理方式
• 有工具 vs. 无工具：对比 main.py（带工具）和 agent.py（不带工具），看工具使用如何影响注意力
• 分析模式：观察以下任务中注意力焦点的差异：
  • 数学计算
  • 知识查询
  • 推理任务
  • 代码生成
  • 创意写作
• 前端自动发现：前端通过清单文件自动检测新轨迹

故障排查

前端没有显示轨迹

1. 确保至少运行过一次 agent.py 或 main.py
2. 检查 frontend/public/trajectories/ 下是否存在轨迹文件
3. 确认 manifest.json 存在且包含轨迹条目

前端无法启动

1. 确认已安装 Node.js（14+ 版本）
2. 在 frontend 目录运行 npm install
3. 检查端口冲突（默认端口 3000）

生成缓慢

• 首次运行会下载模型（约 1~2 GB）
• 如有 GPU/MPS，启用后生成更快
• 在测试查询中设置较小的 max_new_tokens 以加快演示

注意事项

• 每条轨迹都带时间戳以确保唯一性
• 清单文件最多跟踪最近 50 条轨迹
• 轨迹在不同会话间持久存在
• 前端通过清单自动发现新轨迹
• agent.py 和 main.py 都可多次运行以生成不同的轨迹