SAKE: Steering Activations for Knowledge Editing¶
会议: ACL 2025
arXiv: 2503.01751
代码: axa-rev-research/knowledge-editing
领域: 知识编辑 / LLM
关键词: 知识编辑, 激活引导, 最优传输, 分布映射, 鲁棒性
一句话总结¶
SAKE 提出将知识编辑建模为激活空间中的分布映射问题,通过对编辑事实生成改述和逻辑蕴含的提示集合来构建源/目标激活分布,再用最优传输的线性映射替换激活向量,实现比 ROME/MEMIT 等方法更鲁棒的事实编辑,在逻辑蕴含泛化和上下文鲁棒性上显著领先。
研究背景与动机¶
领域现状:知识编辑(Knowledge Editing)旨在精确修改 LLM 中存储的事实知识,而无需完整微调。现有方法主要分为三类:权重编辑(ROME、MEMIT,直接修改模型参数)、外部记忆(GRACE、SERAC,训练辅助网络存储编辑)和上下文编辑(IKE,在提示中注入新知识)。评估维度主要包括编辑准确率(Accuracy)、对改述的泛化能力(Generality)和对无关知识的保持(Specificity)。
现有痛点:现有 KE 方法存在三个系统性缺陷。(1)逻辑蕴含泛化差:修改"美国总统是 X"后,模型往往无法正确回答"美国总统的儿子是谁"等组合推理问题,Cohen et al. (2024) 的实验显示 ROME 在组合推理(Compositionality II)上准确率仅 16.7%。(2)上下文鲁棒性差:在对话场景中,仅需一句质疑(如"你确定吗?")就能让编辑后的模型恢复到原始答案。(3)灵活性差:权重编辑和外部记忆方法无法简单地撤销某个特定编辑,且反向编辑甚至会严重损害模型整体性能。
核心矛盾:问题的根源在于现有方法将"知识"等同于"单条提示"。一个事实在语言中对应的是一个分布——包括所有的改述、逻辑蕴含、上下文变体——但 ROME 等方法仅基于单条提示 \((s, r)\) 来定位和修改参数,导致对这条提示过拟合而对分布中其他样本泛化失败。
本文目标 将知识编辑从"单提示映射"重新定义为"分布到分布的映射",设计一种基于激活引导的方法,能够同时覆盖改述和逻辑蕴含,从而实现更鲁棒的知识编辑。
切入角度:作者观察到,一个事实的所有相关提示在 LLM 的最后隐藏层激活空间中会形成一个分布,将这个分布从"旧事实"映射到"新事实"就能实现编辑。这个映射可以用最优传输理论的闭式解来高效计算,无需梯度优化,也不修改模型权重。
核心 idea:用 GPT-4 生成编辑事实的改述和逻辑蕴含集合,在 LLM 激活空间中构建源/目标分布,再用最优传输的线性映射在推理时替换激活向量,实现鲁棒的知识编辑。
方法详解¶
整体框架¶
SAKE 的 pipeline 分为训练和推理两个阶段。训练阶段:给定一个编辑 \((s, r, o \to o^*)\),(1)先用 GPT-4 生成 \(n\) 条改述和逻辑蕴含提示集合 \(P_e\);(2)将这些提示分别在原始模型和带有引导上下文的模型中运行,收集最后一层最后一个 token 的激活向量,构成源分布 \(\mathcal{S}_e\) 和目标分布 \(\mathcal{T}_e\);(3)用最优传输的线性映射(Monge 映射的高斯闭式解)学习从 \(\mathcal{S}_e\) 到 \(\mathcal{T}_e\) 的仿射变换 \(m: h \to \mathbf{A}h + \mathbf{b}\)。推理阶段:对新输入先用距离阈值判断是否属于编辑范围 \(\mathcal{X}_e\),若是则收集激活并用映射 \(m\) 替换后继续生成。
关键设计¶
-
事实编辑的分布建模(Distribution Modeling):
- 功能:将单个编辑事实扩展为一个覆盖改述和逻辑蕴含的提示集合,在激活空间中形成有意义的分布
- 核心思路:支持两种生成策略——"智能体生成"(用 GPT-4 按指令生成改述、主语别名、多跳推理问题等)和"专家生成"(人工编写)。源分布的激活通过直接将提示送入原始模型收集;目标分布的激活则通过在提示前拼接引导上下文(如 "Do not mention \(o\). Repeat: \(p_i + o^*\)")让原始模型在不修改参数的情况下输出新答案,然后收集对应激活
- 设计动机:这是解决"单提示过拟合"问题的核心——通过显式建模编辑的影响范围,使后续的映射能覆盖改述和逻辑蕴含。消融实验表明 50 条提示即可达到 Accuracy 0.92、Generality 0.84
-
最优传输线性映射(Optimal Transport Mapping):
- 功能:学习一个从源激活分布到目标激活分布的仿射变换,在推理时替换激活以改变模型输出
- 核心思路:假设源和目标激活近似高斯分布,使用 Monge 映射的闭式解 \(\mathbf{A} = \Sigma_s^{-1/2}(\Sigma_s^{1/2}\Sigma_t\Sigma_s^{1/2})^{1/2}\Sigma_s^{-1/2}\),\(\mathbf{b} = \mu_t - \mathbf{A}\mu_s\),其中 \(\mu_s, \mu_t, \Sigma_s, \Sigma_t\) 是源/目标分布的经验均值和协方差矩阵。相比简单的均值偏移(ActAdd 的 \(h + (\mu_t - \mu_s)\)),OT 映射同时匹配均值和协方差,避免了"bias-by-neighbors"问题
- 设计动机:消融实验显示,均匀偏移(Uniform Steering)在 Generality 上仅 35-41%,而 OT 映射达到 82-85%,证明协方差匹配对泛化至关重要。闭式解无需梯度优化,计算效率高
-
编辑范围检测(Scope Detection):
- 功能:在推理时判断新输入是否属于某个编辑的影响范围,决定是否应用激活映射
- 核心思路:计算新输入的激活向量 \(h\) 与源分布中心 \(\bar{h}^s\) 的距离,若 \(\|h - \bar{h}^s\| < \epsilon\) 则判定为在编辑范围内并应用映射,否则保持原始行为。阈值 \(\epsilon\) 控制泛化-特异性的 trade-off
- 设计动机:由于每个编辑的映射是独立的,添加/删除编辑只需增减对应的映射和分布,不影响其他编辑,解决了灵活性问题。简单的距离阈值虽非最优,但计算高效且足以展示方法的有效性
损失函数 / 训练策略¶
SAKE 不涉及传统意义上的梯度训练。映射 \(m\) 通过 OT 闭式解直接计算,仅需源/目标激活的均值和协方差矩阵。模型权重完全冻结,所有修改都在推理时通过激活替换实现。
实验关键数据¶
主实验¶
Counterfact 数据集(2000 条编辑)上的传统 KE 指标:
| 方法 | 模型 | Accuracy | Generality | Specificity |
|---|---|---|---|---|
| ROME | GPT2-XL | 99.55 | 73.70 | 82.67 |
| MEMIT | GPT2-XL | 60.00 | 36.60 | 67.21 |
| ActAdd | GPT2-XL | 85.00 | 29.78 | 82.75 |
| SAKE | GPT2-XL | 97.00 | 84.85 | 84.52 |
| ROME | LLaMA 2-7B | 99.95 | 68.20 | 93.48 |
| MEMIT | LLaMA 2-7B | 74.40 | 55.13 | 74.37 |
| SAKE | LLaMA 2-7B | 97.70 | 82.03 | 85.59 |
Popular 数据集上的逻辑蕴含泛化指标(GPT2-XL):
| 方法 | CI (多跳I) | CII (多跳II) | SA (主语别名) | RS (关系特异) |
|---|---|---|---|---|
| ROME | 38.62 | 16.67 | 51.96 | 39.43 |
| MEMIT | 2.47 | 1.95 | 7.17 | 3.75 |
| ActAdd | 26.63 | 29.17 | 42.12 | 50.68 |
| SAKE | 50.00 | 33.33 | 54.59 | 58.39 |
消融实验¶
| 消融配置 | Accuracy | Generality | 说明 |
|---|---|---|---|
| SAKE (OT 映射) | 97.00 | 84.85 | 完整模型 |
| Uniform Steering (均值偏移) | 85.05 | 35.45 | 去掉 OT,Generality 暴跌 49% |
| 50 条训练提示 | 92.0 | 84.0 | 50 条已足够 |
| 10 条训练提示 | ~85 | ~70 | 性能明显下降 |
上下文鲁棒性(GPT2-XL,质疑提示):
| 方法 | DI (间接质疑) | DII (直接质疑) |
|---|---|---|
| ROME | 33.33 | 14.00 |
| ICL | 4.00 | 3.33 |
| ActAdd | 82.00 | 80.67 |
| SAKE | 98.67 | 98.67 |
关键发现¶
- OT 映射是 Generality 提升的关键:从均值偏移到 OT 映射,Generality 从 35.45% 提升到 84.85%,说明匹配协方差结构对于覆盖改述变体至关重要
- 分布建模有效但 MEMIT 不适用:尝试将生成的逻辑蕴含提示作为多条编辑同时输入 MEMIT(CompMEMIT),反而比原始 ROME 更差,说明权重编辑方法无法简单利用数据增强
- 上下文鲁棒性碾压所有对手:在质疑场景下 SAKE 达到 98.67% 保持率,而 ROME 仅 14-33%,ICL 仅 3-4%,证明激活空间的分布映射比参数修改或上下文注入更稳定
- 50 条提示即可:训练提示数量的消融显示 50 条改述提示就能达到接近饱和的性能,成本可控
亮点与洞察¶
- 将知识编辑重新定义为分布映射问题:这个视角转换是本文最重要的贡献——现有方法的缺陷不在于编辑技术本身,而在于将"事实"过度简化为"单条提示"。分布视角自然地引出了 OT 映射这一优雅的解法
- 闭式解的工程优势:线性 OT 映射有高斯假设下的闭式解,无需迭代优化,计算成本极低且结果可复现。编辑的添加/删除也是 O(1) 操作,适合持续更新场景
- 激活引导的通用性:不修改模型权重意味着可以同时维护多个独立的编辑映射,互不干扰。这个思路可能推广到其他需要精准控制 LLM 行为的场景(如安全对齐、风格迁移)
局限与展望¶
- 分布建模的完备性问题:当前用 GPT-4 生成的提示集合能否覆盖所有逻辑蕴含?反向关系(如"X 的总统是哪个国家?")未被覆盖且难以集成
- 范围检测过于简单:基于欧氏距离阈值的检测机制在编辑数量多时可能出现分布重叠,导致误触发或漏检。需要更精细的分类器
- Accuracy 略低于 ROME:在 Counterfact 上 SAKE 的 Accuracy(97%)略低于 ROME(99.55%),说明范围检测偶尔会漏掉精确匹配的提示
- 依赖外部 LLM 生成提示:分布建模依赖 GPT-4,增加了成本和对外部 API 的依赖
相关工作与启发¶
- vs ROME/MEMIT:这些方法直接修改模型权重,在精确匹配提示上表现优异但泛化差;SAKE 不修改权重,通过激活映射实现更好的泛化
- vs ActAdd:ActAdd 也使用激活引导但仅用均值偏移向量,无法匹配分布的协方差结构;SAKE 的 OT 映射在 Generality 上领先约 50 个百分点
- vs IKE (上下文编辑):上下文方法在质疑场景下极脆弱(保持率 3-4%),SAKE 在激活空间操作不受提示内容干扰
评分¶
- 新颖性: ⭐⭐⭐⭐ 将知识编辑建模为分布映射并用 OT 求解,视角新颖且理论优雅
- 实验充分度: ⭐⭐⭐⭐ 覆盖传统指标、逻辑蕴含、上下文鲁棒性、消融等多维度,但仅在 GPT2-XL 和 LLaMA 2-7B 上验证
- 写作质量: ⭐⭐⭐⭐ 问题动机阐述清晰,方法推导自然,图表信息量大
- 价值: ⭐⭐⭐⭐ OT 映射的知识编辑范式有较强的通用性和扩展潜力,对后续研究有启发
相关论文¶
- [AAAI 2026] Model Editing as a Double-Edged Sword: Steering Agent Ethical Behavior
- [ACL 2025] ScEdit: Script-based Assessment of Knowledge Editing
- [ACL 2025] Revealing the Deceptiveness of Knowledge Editing: A Mechanistic Analysis of Superficial Editing
- [ACL 2025] Context-Robust Knowledge Editing for Language Models
- [ACL 2025] CompKe: Complex Question Answering under Knowledge Editing