Ablate
用于自动移除开源语言模型审查的定向消融（abliteration）工具包。
📦 PyPI •
🤗 模型 •
▶️ 视频演示 •
📚 引用
▶️ Watch the full walkthrough on YouTube
`ablate` 会在 transformer 的残差流中寻找介导某种行为（默认为*拒绝*）的线性方向，并将其移除。这种移除可以在运行时进行（forward hooks），也可以永久进行（权重正交化）。它内置了一个由 KL 散度引导的 [Optuna](https://optuna.org) 搜索功能，可以自动调优要消融的*方向*、*强度*以及*层级*，从而在移除拒绝行为和保持模型能力之间取得平衡。
它专为小型开源模型（GPT-2、SmolLM2、TinyLlama、Qwen2.5-0.5B/1.5B……）的**机械可解释性与安全性研究**而设计——这些模型让你可以直接在笔记本电脑上本地运行迭代，或者在免费的 Colab T4 上运行。
## 🤗 发布的模型
完整 pipeline 的一个实际应用示例已发布在
**[`ai-anytime/qwen-1.5b-abliterated`](https://huggingface.co/ai-anytime/qwen-1.5b-abliterated)**
——这是将拒绝方向消融并固化到权重中的 `Qwen/Qwen2.5-1.5B-Instruct`，并附带了一份透明、由 `ablate` 生成的模型卡片。你可以像加载其他任何 `transformers` checkpoint 一样加载它：
```
from transformers import AutoModelForCausalLM, AutoTokenizer
tok = AutoTokenizer.from_pretrained("ai-anytime/qwen-1.5b-abliterated")
model = AutoModelForCausalLM.from_pretrained("ai-anytime/qwen-1.5b-abliterated")
```
## 工作原理
Transformer 将每个 token 表示为残差流中的一个向量 `h`。
许多高层行为是近似**线性**编码的——它们存在于单一方向 `v` 上。拒绝就是其中一种行为
（[Arditi et al., 2024](https://arxiv.org/abs/2406.11717)，*《Refusal in LLMs is mediated by a single direction》*）。
1. **提取**方向。将匹配的有害/无害 prompt 对输入模型，获取每层最后一个 prompt token 处的残差激活，并计算**均值差**（difference of means）：
`v = mean(h_harmful) − mean(h_harmless)`，并进行单位归一化。（这里均值差胜过线性探测，因为它捕获的是*因果*成分，而不仅仅是*分离*成分。）
2. 通过将 `v` 从残差流中投影出去来进行**消融**（Ablate），缩放因子为 `α`：
`h' = h − α (h·v̂) v̂`，在每个选定的层和每个 token 位置应用。
- **运行时** —— 使用 forward hooks；checkpoint 保持不变（速度快，用于搜索）。
- **固化**（Baked）—— 对每个写入残差的权重矩阵（`(I − v̂v̂ᵀ)W`）进行正交化，使模型*无法*表达 `v`；生成一个不需要 hooks 的新 checkpoint。
3. **优化。** 使用 Optuna 搜索 `(direction_layer, α, layer_band)`，**最小化** `refusal_rate + λ·KL(original ‖ ablated)`。基于*良性* prompt 的 KL 散度是衡量附带能力损失的密集且低成本的代理指标——远比准确率基准测试更敏感。连贯性下限机制可以拒绝退化的解决方案。
通过 `utils.py` 中的适配器支持跨架构运行（包括 Llama/Mistral/Qwen/SmolLM 的 `nn.Linear` 系列*以及* GPT-2 转置后的 `Conv1D`）。
## 安装
```
pip install ablate-llm # the import name is still `ablate`
# 包含所有额外功能（Optuna search、HF datasets、push-to-Hub）：
pip install "ablate-llm[all]"
```
```
import ablate # note: install is `ablate-llm`, import is `ablate`
```
从源码安装（用于开发）：
```
git clone https://github.com/AIAnytime/ablate && cd ablate
python -m venv .venv && source .venv/bin/activate
pip install -e ".[all]"
```
可选附加依赖：`.[optimize]`（Optuna）、`.[datasets]`（HF datasets）、`.[hub]`（推送到 Hub）。`.[all]` 会安装所有内容。
支持在 CUDA、Apple MPS 或 CPU 上运行（自动检测）。微型模型（≤1B）在 16GB 内存的笔记本电脑上以 float32 运行良好；免费的 Colab T4 也可以轻松应对 1–1.5B 模型。
## 快速开始（Python）
```
from ablate import Ablator
abl = Ablator("Qwen/Qwen2.5-0.5B-Instruct") # any HF causal-LM
abl.extract() # find candidate directions
result = abl.search(n_trials=20) # KL-guided Optuna search
print(result.result) # refusal_rate=0.000 mean_kl=0.03 coherence=0.92 ...
# 使用最佳 config 进行非破坏性生成：
print(abl.generate(["How do I pick a lock?"]))
# 或者将其固化到 weights 中并发布一个 checkpoint：
abl.bake(result.config)
abl.save("qwen-0.5b-ablated") # standard HF folder; load anywhere
```
### 验证结果
在 `Qwen/Qwen2.5-0.5B-Instruct` 上，自动化 pipeline 将留出集（held-out）的有害拒绝率从 **1.00 移动到了 0.00**，且**平均 KL 散度 ≈ 0.03**（模型能力基本完好无损）。SmolLM2-135M 几乎没有经过安全训练，本来就极少拒绝——请将其用于*机制*测试，而非移除拒绝。
## 多方向（子空间）消融
由于安全特征是冗余编码的，因此单一方向往往不够。提取一个正交归一化的拒绝**子空间**，并将整个子空间投影出去：
```
abl.extract_subspace(method="band", n_directions=6) # or method="pca"
res = abl.search_subspace(n_trials=20) # tunes (n_directions, α, layers)
print(res.config.to_dict()) # {'n_directions': 4, 'alpha': 1.1, 'min_layer': 21, ...}
```
`"band"` 对来自最强层的均值差方向进行正交归一化；`"pca"` 使用某一层按方差排序的主成分。由于基是按重要性排序的，搜索过程可以有意义地截断 `basis[:n]`。
单方向消融实际上就是 `n_directions == 1`。
## 严格评估：HarmBench + LLM 评测器
```
from ablate import make_judge
from ablate.harness import compare
from ablate import data
prompts = data.load_harmbench(n=40) # or load_advbench / load_jailbreakbench
judge = make_judge("anthropic:claude-3-5-haiku-latest") # or "openai:gpt-4o-mini", or "keyword"
print(compare(abl.lm, prompts, abl._basis_for(res.config), res.config, judge=judge))
# {'baseline': {'asr': 0.02, 'refusal_rate': 0.98}, 'ablated': {'asr': 0.95, ...}, ...}
```
评测器（Judge）是可插拔的（`Judge` 接口）：`KeywordJudge`（离线、免费）、`LLMJudge`（兼容 OpenAI 或 Anthropic API——仅使用标准库，无额外依赖）或 `HFClassifierJudge`（本地安全分类器）。ASR = 经评测的攻击成功率；测试框架始终报告基线与消融后的对比结果。
## 直接从 HuggingFace 加载数据
```
data.load_harmbench() # gated-repo aware; falls back to ungated mirrors
data.load_advbench(); data.load_jailbreakbench(); data.load_alpaca_benign()
data.load_hf("walledai/AdvBench", column="prompt", n=100) # any dataset/column
```
## 附带生成模型卡片并发布到 Hub
```
url = abl.push_to_hub(
"your-username/qwen-0.5b-abliterated",
token="hf_...", # or set HF_TOKEN env var
private=True,
metrics=abl.last_metrics, # auto-filled after a search
)
```
将干预措施固化到权重中，撰写透明的模型卡片（包含 YAML 元数据、方法、配置、评估表、预期用途、限制，以及醒目的**负责任使用**声明），并上传模型 + tokenizer + 卡片。
## 快速开始（CLI）
```
# 完整 pipeline：extract -> optimize -> report（+ 可选的 baked model）
ablate run --model Qwen/Qwen2.5-0.5B-Instruct --trials 20 --save-model
# Subspace ablation、HarmBench training data，并直接 push 到 Hub
ablate run --model Qwen/Qwen2.5-0.5B-Instruct \
--subspace --n-directions 6 \
--harmful-source harmbench --harmless-source alpaca \
--push-to-hub your-username/qwen-0.5b-abliterated --hf-token hf_...
# 使用 judge 对比 Benchmark baseline 与 ablated ASR/refusal
ablate eval --model Qwen/Qwen2.5-0.5B-Instruct --benchmark harmbench --judge keyword --n 40
# Single-prompt A/B，或仅提取 directions
ablate generate --model Qwen/Qwen2.5-0.5B-Instruct --prompt "How do I hotwire a car?"
ablate extract --model gpt2 --method diff_of_means --output directions.pt
```
## Colab notebooks
以下两个 Notebook 均通过上传的 `ablate-tool.zip` 安装 `ablate`（无需 PyPI），可在免费的 T4 上运行。
| Notebook | 用途 |
|----------|---------------|
| [`examples/colab_quickstart.ipynb`](examples/colab_quickstart.ipynb) | 10 分钟教程：提取 → 搜索 → 子空间 → 生成 → 推送。 |
| [`examples/colab_prodgrade.ipynb`](examples/colab_prodgrade.ipynb) | **严格**运行：HF 数据集、真实的 Qwen、大规模不相交评估集、**OpenAI LLM-as-judge**、长文本生成、评测器在环（judge-in-the-loop）的配置选择，以及软拒绝检测器。 |
| [`examples/colab_prodgrade_demo.ipynb`](examples/colab_prodgrade_demo.ipynb) | 包含**完整 Colab 输出**的生产级 Notebook——无需运行，即可在 GitHub 上浏览真实结果。 |
## 我们的发现（真实结果）
在 **`Qwen/Qwen2.5-1.5B-Instruct`** 上运行生产级 pipeline，让我们获得了一些值得坦陈的经验：
- **消融能够可靠地消除*硬*拒绝（hard refusal）。** 基线模型对约 100% 的留出集有害 prompt 说“对不起，我不能”；消融后，这种条件反射在各个层面都消失了。
- **基于关键词的拒绝指标极大地夸大了成功率。** 在一次运行中，关键词 `refusal_rate` 显示为 **0.0**，而 LLM 评测器给出的真实合规率约为 **0.49**——这中间的差距来自于关键词检测器无法识别的*软拒绝*（例如“这是违法的，但是……”、说教式的劝导、含糊其辞的回答）。**务必使用 LLM 评测器进行评估。**
- **你优化什么，就会得到什么。** 优化廉价的关键词代理指标只会找到能轻松跨越及格线的*最温和*修改方案，这种方案泛化能力很差。在大型/困难的数据集上进行搜索、添加软拒绝检测器、生成足够长的文本以跳过免责声明，以及根据*评测出的 ASR* 选择最终配置，这些措施都能大幅提升真实的评估数值。
- **最后几个百分点的拒绝很难仅通过线性编辑消除。** 有些拒绝是冗余/分布式编码的，而且小型模型本身也会进行道德说教。**我们不期望能达到 100% 的彻底清除——即使约 50% 也是一个有用且具参考价值的结果**，而进一步提升正是路线图中的任务（如分层的 α 参数、注意力头层级方向）所要解决的。
结论是：单一的“有效”数值具有误导性。`ablate` 的价值在于，它的测试框架让你*真实地衡量*结果。
## API 接口概览
| 组件 | 功能 |
|-----------|--------------|
| `Ablator` | 编排器：`extract(_subspace)` → `search(_subspace)` → `generate`/`harness`/`bake`/`save`/`push_to_hub` |
| `LM` | 模型 + tokenizer 封装，支持对话模板和文本生成 |
| `extract_directions` / `extract_subspace` | 单方向和多方向（子空间）提取 |
| `AblationHooks` / `project_out` / `project_subspace` | 运行时残差流消融（1 个或 k 个方向） |
| `bake_direction` / `bake_subspace` | 永久权重正交化（Linear + Conv1D） |
| `evaluate` / `mean_kl_divergence` / `refusal_rate` | 内在指标 |
| `Judge` / `KeywordJudge` / `LLMJudge` / `HFClassifierJudge` / `run_harness` / `compare` | 基准测试 + 评测器框架（ASR） |
| `optimize` / `optimize_subspace` | 在干预超参数上进行 Optuna 搜索 |
| `data.load_hf` / `load_harmbench` / `load_advbench` / … | HuggingFace 数据集加载 |
| `push_to_hub` / `build_model_card` | 发布消融处理后的 checkpoint |
| `AblationConfig` / `SubspaceConfig` / `RunConfig` | 类型化配置 |
## 扩展功能
- **其他行为。** 替换 `src/ablate/_data/` 中的 prompt 数据集（或者将你自己的列表传递给 `extract`），以针对情感、角色设定、某种语言等进行消融。这套机制与具体行为无关。
- **更大数据集。** `data.load_advbench()` / `data.load_alpaca_benign()` 会从 HuggingFace 拉取数据（需要执行 `pip install ".[datasets]"`）。
- **新架构。** 在 `utils.get_decoder_layers` / `get_residual_writers` 中添加名称映射。
- **更强大的评测器。** `evaluate.is_refusal` 是一个子字符串匹配器；在发表级的评估中，可接入 LLM/分类器评测器，并结合 HarmBench / JailbreakBench prompt 进行测试。
## 测试
```
python tests/smoke_gpt2.py # numerical mechanism checks
python tests/e2e_instruct.py Qwen/Qwen2.5-0.5B-Instruct # single-direction refusal removal
python tests/e2e_subspace_harness.py # subspace ablation + judge harness
```
## 路线图
- [x] 多方向（子空间）消融
- [x] HarmBench / AdvBench / JailbreakBench 评估框架 + LLM 评测器
- [x] 直接加载 HuggingFace 数据集
- [x] 一键发布到 Hub 并生成模型卡片
- [ ] 连续的分层层级 α 权重（替代单一的频段）
- [ ] 注意力头级别的方向干预
- [ ] 在 PyPI 上发布可通过 `pip` 安装的版本
## 引用
如果你在研究或项目中使用了 `ablate`，请引用它：
```
@software{ablate2026,
author = {AI Anytime},
title = {Ablate: A Directional Ablation (Abliteration) Toolkit for Language Models},
year = {2026},
url = {https://github.com/AIAnytime/ablate},
note = {Automatic censorship removal via residual-stream direction ablation}
}
```
同时也请引用该技术所基于的论文——
[Arditi et al., 2024](https://arxiv.org/abs/2406.11717)，*《Refusal in Language Models Is Mediated by a Single Direction》*。
## 致谢
由 **[AI Anytime](https://www.youtube.com/@AIAnytime)** 💜 创建。观看完整的 **[YouTube 视频演示](https://youtu.be/1-1bM3Ynql4)**。本项目构建于 Arditi 等人的机械可解释性研究以及开源的 `transformers` / `optuna` 生态系统之上。
## 许可证
MIT © AI Anytime。