跳转至

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\) 替换后继续生成。

关键设计

  1. 事实编辑的分布建模(Distribution Modeling):

    • 功能:将单个编辑事实扩展为一个覆盖改述和逻辑蕴含的提示集合,在激活空间中形成有意义的分布
    • 核心思路:支持两种生成策略——"智能体生成"(用 GPT-4 按指令生成改述、主语别名、多跳推理问题等)和"专家生成"(人工编写)。源分布的激活通过直接将提示送入原始模型收集;目标分布的激活则通过在提示前拼接引导上下文(如 "Do not mention \(o\). Repeat: \(p_i + o^*\)")让原始模型在不修改参数的情况下输出新答案,然后收集对应激活
    • 设计动机:这是解决"单提示过拟合"问题的核心——通过显式建模编辑的影响范围,使后续的映射能覆盖改述和逻辑蕴含。消融实验表明 50 条提示即可达到 Accuracy 0.92、Generality 0.84
  2. 最优传输线性映射(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%,证明协方差匹配对泛化至关重要。闭式解无需梯度优化,计算效率高
  3. 编辑范围检测(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 映射的知识编辑范式有较强的通用性和扩展潜力,对后续研究有启发

相关论文