电路稀疏性可视化工具与模型

用于检查 Gao et al. 2025 提出的稀疏电路模型的工具。提供运行推理的代码以及一个 Streamlit 仪表板，允许您与通过剪枝发现的任务特定电路进行交互。注意：本 README 是由 AI 生成并经过轻微编辑的。

安装

pip install -e .

启动可视化工具

从项目根目录启动 Streamlit 应用：

streamlit run circuit_sparsity/viz.py

该应用会从 openaipublic 网页加载数据，并在本地缓存。当可视化工具加载时，您可以使用左侧栏中的控件选择模型、数据集、剪枝扫描和节点预算 k。图表使用 Plotly 渲染；大多数元素都是可交互的，支持悬停和点击探索。

Streamlit 电路可视化工具的示例视图（wte/wpe 选项卡），显示节点消融差异和激活预览：

运行模型前向传播

Transformer 的定义位于 circuit_sparsity.inference.gpt 中。该模块导出了：

• GPTConfig / GPT：适用于 CPU/GPU 推理的轻量级 GPT 实现。
• load_model(model_dir, cuda=False)：便捷的加载函数，需要 models/... 目录下的 beeg_config.json 和 final_model.pt 文件对。

示例用法（改编自 tests/test_gpt.py）：

from circuit_sparsity.inference.gpt import GPT, GPTConfig, load_model
from circuit_sparsity.inference.hook_utils import hook_recorder
from circuit_sparsity.registries import MODEL_BASE_DIR

config = GPTConfig(block_size=8, vocab_size=16, n_layer=1, n_head=1, d_model=8)
model = GPT(config)
logits, loss, _ = model(idx, targets=targets)

# 获取激活值
with hook_recorder() as rec:
    model(idx)

# rec 是一个字典，形如 {"0.attn.act_in": tensor(...), ...}

pretrained = load_model(f"{MODEL_BASE_DIR}/models/<model_name>", cuda=False)

运行测试：

pytest tests/test_gpt.py

数据布局

项目资产位于 https://openaipublic.blob.core.windows.net/circuit-sparsity 下，结构如下：

• models/<model_name>/
  • beeg_config.json：序列化的 GPTConfig，用于重建模型。
  • final_model.pt：circuit_sparsity.inference.gpt.load_model 使用的检查点。
• viz/<experiment>/<model_name>/<task_name>/<sweep>/<k>/
  • viz_data.pkl：viz.py 加载的主要数据包（包含电路掩码、激活值、样本、重要性等）。
  • 其他每轮输出（掩码、直方图、样本桶）会在预处理脚本生成后存储在同一目录树下。
• train_curves/<model_name>/progress.json：仪表板摘要表使用的训练指标。
• 其他实验特定目录（例如 csp_yolo1/、csp_yolo2/）存放准备剪枝运行时产生的原始数据。

viz.py 和 registries.py 中引用的文件路径假定采用此布局。如果重新定位数据，请更新 registries.py。

模型

我们发布了论文中所有用于获得结果的模型。完整模型列表请参见 registries.py。确切的训练超参数可在 [todo] 中找到。

• csp_yolo1：用于 single_double_quote 定性结果的模型。这是一个总参数量为 1.18 亿的模型。该模型相对较旧，采用的训练方法与论文中不完全一致；特别是同时使用多个 L0 值进行训练的方法。
• csp_yolo2：用于 bracket_counting 和 set_or_string_fixedvarname 定性结果的模型。这是一个总参数量为 4.75 亿的模型。
• csp_sweep1_*：这些模型用于获得图 3 的结果。名称表示模型大小（相对于任意基准大小的“扩展因子”）、权重 L0 和激活稀疏度水平（afrac）。
• csp_bridge1：用于获得论文结果的桥梁模型。
• csp_bridge2：另一款桥梁模型。
• dense1_1x：在我们的数据集上训练的密集模型。
• dense1_2x：在我们的数据集上训练的密集模型，宽度是前者的两倍。
• dense1_4x：在我们的数据集上训练的密集模型，宽度是前者的四倍。

剪枝扫描 ID

• prune_v2：256 次 CARBS 迭代，批量大小为 16，采用非常古老的（未发表的）剪枝算法，目标是固定的 k 值而非固定的目标损失。
• prune_v3：256 次 CARBS 运行，批量大小为 64，训练 32 个周期，采用较早的（未发表的）算法，目标是固定的目标损失。
• prune_v4：768 次 CARBS 运行，批量大小为 64，训练 48 个周期，采用已发表的算法，目标是固定的目标损失。
• prune_v5_logitscaling：256 次 CARBS 运行，批量大小为 32，训练 32 个周期，采用带有 logits 缩放的已发表算法，目标是固定的目标损失。

其他实用工具

• per_token_viz_demo.py：用于标记级别可视化的小型示例。
• clear_cache.py：删除本地缓存的 blobstore 文件副本（包括 Streamlit/viz 缓存和 tiktoken 缓存）；如果您需要重新获取最新数据，请运行此脚本。

该项目依赖于 Streamlit、Plotly、matplotlib、seaborn 和 torch（完整依赖列表请参阅 pyproject.toml）。

Circuit Sparsity 快速上手指南

本指南帮助开发者快速部署并使用 Circuit Sparsity 工具，用于可视化和分析稀疏电路模型（基于 Gao et al. 2025 的研究）。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python 版本: 推荐 Python 3.8 或更高版本
• 核心依赖:
  • torch (PyTorch)
  • streamlit (用于可视化仪表盘)
  • plotly, matplotlib, seaborn (用于绘图)
  • pytest (可选，用于运行测试)

提示：国内用户建议使用清华或阿里镜像源加速依赖安装。

安装步骤

1. 克隆项目代码库到本地（假设您已获取源码）。
2. 进入项目根目录，使用 pip 进行可编辑安装：

pip install -e .

若需加速安装，可使用国内镜像源：

pip install -e . -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

1. 启动可视化仪表盘 (Visualizer)

这是最常用的功能，允许您交互式地探索不同任务下的稀疏电路结构。

在项目根目录下运行以下命令启动 Streamlit 应用：

streamlit run circuit_sparsity/viz.py

使用说明：

• 应用启动后会自动从 OpenAI 公共存储加载数据并缓存在本地。
• 在左侧边栏中，您可以选择：
  • Model: 模型名称（如 csp_yolo1, csp_bridge1 等）。
  • Dataset: 数据集类型。
  • Pruning Sweep: 剪枝策略版本（如 prune_v4）。
  • Node Budget (k): 节点预算数量。
• 主界面将渲染 Plotly 交互图表，支持悬停查看详细信息和点击探索节点消融差异及激活预览。

2. 运行模型推理 (代码示例)

如果您希望通过代码直接加载模型并进行前向传播或提取激活值，可以参考以下示例。

加载预训练模型并推理：

from circuit_sparsity.inference.gpt import GPT, GPTConfig, load_model
from circuit_sparsity.inference.hook_utils import hook_recorder
from circuit_sparsity.registries import MODEL_BASE_DIR

# 方式 A: 从零初始化一个小型 GPT 模型
config = GPTConfig(block_size=8, vocab_size=16, n_layer=1, n_head=1, d_model=8)
model = GPT(config)
# logits, loss, _ = model(idx, targets=targets)

# 方式 B: 加载预训练模型 (推荐)
# 请替换 <model_name> 为具体的模型文件夹名称，例如 csp_yolo1
pretrained = load_model(f"{MODEL_BASE_DIR}/models/<model_name>", cuda=False)

# 提取激活值 (使用 Hook 机制)
with hook_recorder() as rec:
    pretrained(idx)

# rec 是一个字典，包含类似 {"0.attn.act_in": tensor(...)} 的激活数据
print(rec.keys())

运行测试验证环境：

pytest tests/test_gpt.py

3. 其他实用工具

• 清除缓存：如果需要重新拉取最新的数据工件或解决缓存冲突，运行：

  python clear_cache.py
  

• Token 级可视化演示：参考 per_token_viz_demo.py 获取最小化示例。