一个通过干预来 理解 和 改进 PyTorch 模型的库

pyvene 是一个开源的 Python 库，用于对 PyTorch 模型的内部状态进行干预。干预在许多 AI 领域中都是一项重要操作，包括模型编辑、可控生成、鲁棒性以及可解释性等。

pyvene 具有许多使干预变得简单的特性：

• 干预是基本的原语，以字典形式表示，因此可以本地保存，并通过 HuggingFace 作为可序列化的对象进行共享。
• 干预可以组合和自定义：你可以在多个位置运行它们，针对任意数量的神经元（或其他粒度级别），并行或按顺序执行，甚至应用于生成式语言模型的解码步骤等。
• 干预可以直接应用于任何 PyTorch 模型！无需从头定义新的模型类，各种架构（RNN、ResNet、CNN、Mamba 等）都可以轻松实现干预。

pyvene 目前仍在积极开发中，并不断改进 🫡

[!IMPORTANT] 请访问 https://stanfordnlp.github.io/pyvene/ 阅读 pyvene 的文档！

安装

要安装最新稳定版的 pyvene：

pip install pyvene

或者，如果你想安装前沿版本，可以克隆仓库并安装：

git clone git@github.com:stanfordnlp/pyvene.git
cd pyvene
pip install -e .

需要更新时，只需在克隆的目录中运行 git pull 即可。

我们建议这样导入该库：

import pyvene as pv

引用

如果你使用了这个仓库，请考虑引用我们的库论文：

@inproceedings{wu-etal-2024-pyvene,
    title = "pyvene: A Library for Understanding and Improving {P}y{T}orch Models via Interventions",
    author = "Wu, Zhengxuan and Geiger, Atticus and Arora, Aryaman and Huang, Jing and Wang, Zheng and Goodman, Noah and Manning, Christopher and Potts, Christopher",
    editor = "Chang, Kai-Wei and Lee, Annie and Rajani, Nazneen",
    booktitle = "Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 3: System Demonstrations)",
    month = jun,
    year = "2024",
    address = "Mexico City, Mexico",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.naacl-demo.16",
    pages = "158--165",
}

星标历史

pyvene 快速上手指南

pyvene 是一个开源 Python 库，旨在通过对 PyTorch 模型内部状态进行干预（Interventions），来帮助开发者理解和改进模型。它广泛应用于模型编辑、引导（Steering）、鲁棒性分析及可解释性研究。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：建议 Python 3.8 及以上
• 前置依赖：
  • PyTorch
  • Hugging Face Transformers (可选，用于加载预训练模型)
  • 其他标准科学计算库 (numpy, etc.)

提示：国内用户建议在安装前配置 pip 镜像源（如清华源或阿里源）以加速下载。

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

方式一：安装稳定版（推荐）

通过 pip 直接安装最新稳定版本：

pip install pyvene

方式二：安装开发版

如需体验最新功能或贡献代码，可克隆仓库进行安装：

git clone git@github.com:stanfordnlp/pyvene.git
cd pyvene
pip install -e .

更新时只需在克隆目录中运行 git pull 即可。

导入库

建议在代码中使用以下别名导入：

import pyvene as pv

基本使用

pyvene 的核心是将“干预”作为基本原语。干预配置以字典形式定义，支持序列化并通过 HuggingFace 共享。以下是一个最简化的概念示例，展示如何定义并应用干预：

import pyvene as pv
import torch
import transformers

# 1. 加载一个预训练模型 (以 GPT-2 为例)
model_name = "gpt2"
model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

# 2. 定义干预配置
# 指定干预位置 (layer, component) 和干预来源 (source)
# 这里演示一个简单的占位符配置，实际使用中需根据具体任务设定
interventions = [
    pv.Intervention(
        layer=0, 
        component="h[0].mlp", 
        source="constant", 
        constant_value=torch.zeros_like(model.h[0].mlp.c_proj.weight)
    )
]

# 3. 创建 IntervenableModel 包装器
# 将原始模型转换为支持干预的模型
pv_model = pv.IntervenableModel(interventions, model)

# 4. 执行干预推理
# 准备输入数据
inputs = tokenizer("Hello, world!", return_tensors="pt")

# 运行带干预的前向传播
# unit_locations 指定干预作用的具体位置（如特定 token 或神经元）
_, output = pv_model(
    base=inputs, 
    unit_locations={"sources->base": (None, 0)} 
)

print(tokenizer.decode(output.logits[0, -1, :].argmax(), skip_special_tokens=True))

核心特性说明：

• 灵活组合：支持在多个位置、任意神经元集合上并行或串行运行干预。
• 通用兼容：无需重新定义模型类，即可作用于 RNN、ResNet、CNN、Mamba 等各种 PyTorch 架构。
• 生成式支持：专门优化了对生成式语言模型解码步骤的干预支持。

更多高级用法和详细 API 文档，请访问 官方文档。