一阶模型编辑（ROME）

本仓库提供了自回归Transformer模型上的一阶模型编辑（ROME）的实现（仅支持GPU）。 我们目前支持OpenAI的GPT-2 XL（15亿参数）和EleutherAI的GPT-J（60亿参数）。EleutherAI即将发布一款200亿参数的类GPT模型，我们希望尽快支持它。

如果您遇到任何问题，请随时提交issue；我们正在积极开发此仓库，并会密切关注所有工单。

目录

1. 安装
2. 因果追踪
3. 一阶模型编辑（ROME）
4. CounterFact
5. 评估
   • 运行完整评估套件
   • 集成新的编辑方法
6. 如何引用

安装

我们建议使用conda来管理Python、CUDA和PyTorch相关的依赖，而其他内容则使用pip。要开始使用，只需安装conda并运行：

./scripts/setup_conda.sh

因果追踪

notebooks/causal_trace.ipynb展示了因果追踪技术，该技术可以被修改以应用于任何语句的处理过程。

一阶模型编辑（ROME）

notebooks/rome.ipynb演示了ROME。API非常简单；用户只需指定如下形式的请求重写即可：

request = {
    "prompt": "{} plays the sport of",
    "subject": "LeBron James",
    "target_new": {
        "str": "football"
    }
}

笔记本中还包含几个类似的示例。

CounterFact

详情即将发布！

评估

有关可用基准的说明，请参阅baselines/。

运行完整评估套件

experiments/evaluate.py可用于评估baselines/中的任何方法。 要开始（例如在GPT-2 XL上使用ROME），请运行：

python3 -m experiments.evaluate \
    --alg_name=ROME \
    --model_name=gpt2-xl \
    --hparams_fname=gpt2-xl.json

每次运行的结果都会以特定格式存储在results/<method_name>/run_<run_id>目录下：

results/
|__ ROME/
    |__ run_<run_id>/
        |__ params.json
        |__ case_0.json
        |__ case_1.json
        |__ ...
        |__ case_10000.json

要汇总结果，可以使用experiments/summarize.py：

python3 -m experiments.summarize --dir_name=ROME --runs=run_<run_id>

运行python3 -m experiments.evaluate -h或python3 -m experiments.summarize -h可获取关于命令行参数的详细信息。

集成新的编辑方法

假设您有一个新方法X，并想在CounterFact数据集上对其进行基准测试。要将X集成到我们的评估流程中：

• 将HyperParams子类化为XHyperParams，并指定所有超参数字段。请参阅ROMEHyperParameters以获取示例实现。
• 在hparams/X/gpt2-xl.json中创建一个超参数文件，并指定一些默认值。请参阅hparams/ROME/gpt2-xl.json以获取示例。
• 定义一个函数apply_X_to_model，该函数接受多个参数并返回：(i) 编辑后的模型，以及 (ii) 被编辑参数的原始权重值（以字典格式 {weight_name: original_weight_value} 表示）。请参阅rome/rome_main.py以获取示例。
• 将X添加到experiments/evaluate.py中的ALG_DICT，插入以下行：“X”: (XHyperParams, apply_X_to_model)。

最后，运行主脚本：

python3 -m experiments.evaluate \
    --alg_name=X \
    --model_name=gpt2-xl \
    --hparams_fname=gpt2-xl.json

python3 -m experiments.summarize --dir_name=X --runs=run_<run_id>

关于跨平台兼容性的说明

我们目前仅支持使用PyTorch后端编辑自回归HuggingFace模型的方法。我们正在开发一套通用方法（可在TensorFlow等框架上使用，且无需HuggingFace），这些方法将很快发布。

如何引用

@article{meng2022locating,
  title={Locating and Editing Factual Associations in {GPT}},
  author={Kevin Meng and David Bau and Alex Andonian and Yonatan Belinkov},
  journal={Advances in Neural Information Processing Systems},
  volume={35},
  year={2022}
}

ROME (Rank-One Model Editing) 快速上手指南

ROME 是一种针对自回归 Transformer 模型（如 GPT-2 XL、GPT-J）的秩一模型编辑技术，允许开发者直接修改模型内部的知识事实，而无需重新训练整个模型。

环境准备

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

• 操作系统: Linux (推荐) 或 macOS。Windows 支持可能受限，建议在使用前确认兼容性。
• 硬件: 必须配备 NVIDIA GPU。本项目仅支持 GPU 运行。
• 软件依赖:
  • Python 3.8+
  • CUDA Toolkit (版本需与 PyTorch 匹配)
  • Conda (强烈推荐使用 Conda 管理环境)
• 支持模型: 目前官方支持 gpt2-xl (1.5B) 和 EleutherAI/gpt-j-6B。

国内加速建议： 在安装依赖时，建议使用国内镜像源以提升下载速度：

• Conda: 使用清华源 (https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/)
• Pip: 使用阿里源 (https://mirrors.aliyun.com/pypi/simple/) 或清华源

安装步骤

项目推荐使用 conda 管理 Python、CUDA 和 PyTorch 依赖，其余包通过 pip 安装。

1. 克隆仓库

   git clone https://github.com/kmeng01/rome.git
   cd rome
   

2. 执行自动化安装脚本 运行官方提供的脚本来创建 Conda 环境并安装所有必要依赖：

   ./scripts/setup_conda.sh
   

   如果脚本执行失败或您需要手动配置，请参考以下逻辑：

   • 创建 conda 环境并安装 pytorch (cuda 版本)。
   • 激活环境后，运行 pip install -r requirements.txt (如果存在) 或根据 setup_conda.sh 内的命令逐个安装。

基本使用

ROME 的核心功能是通过简单的 API 调用，将模型中的某个事实（例如“勒布朗·詹姆斯打篮球”）修改为新的事实（例如“勒布朗·詹姆斯踢足球”）。

1. 交互式体验 (推荐新手)

最快速的上手方式是运行官方提供的 Jupyter Notebook 演示：

jupyter notebook notebooks/rome.ipynb

在该 Notebook 中，您可以直接运行单元格来观察模型编辑前后的效果对比。

2. 代码示例

如果您希望在 Python 脚本中使用，核心逻辑如下：

from rome import ROMEHyperParameters, apply_rome_to_model
import torch

# 1. 定义修改请求
# 目标：将 "LeBron James plays the sport of basketball" 修改为 "football"
request = {
    "prompt": "{} plays the sport of",
    "subject": "LeBron James",
    "target_new": {
        "str": "football"
    }
}

# 2. 加载模型 (以 gpt2-xl 为例)
# 注意：首次运行会自动下载模型权重
model_name = "gpt2-xl"
# 此处省略模型加载的具体 boilerplate 代码，详细请参考 notebooks/rome.ipynb
# 通常涉及 from transformers import AutoModelForCausalLM, AutoTokenizer

# 3. 应用 ROME 编辑
# edited_model 是修改后的模型，orig_weights 是被修改参数的原始权重
# params 需要从 hparams/ROME/gpt2-xl.json 加载
edited_model, orig_weights = apply_rome_to_model(
    model=model, 
    request=request, 
    hparams=params # 需预先加载 HyperParameters
)

# 4. 验证结果
# 使用 edited_model 生成文本，观察输出是否已变为 "football"

3. 运行评估套件

若需对编辑效果进行批量评估（基于 CounterFact 数据集），可使用以下命令：

python3 -m experiments.evaluate \
    --alg_name=ROME \
    --model_name=gpt2-xl \
    --hparams_fname=gpt2-xl.json

评估结果将保存在 results/ROME/run_<run_id>/ 目录下。如需生成摘要报告：

python3 -m experiments.summarize --dir_name=ROME --runs=run_<run_id>

提示：对于首次使用，建议优先阅读并运行 notebooks/rome.ipynb，其中包含了完整的上下文处理和参数配置示例。