[论文解读] Circuit Breaking: Removing Model Behaviors with Targeted Ablation

Maximilian Xiling Li, Xander Davies+1 more|arXiv (Cornell University)|Sep 12, 2023

Adversarial Robustness in Machine Learning被引用 5

一句话总结

本文提出了一种有针对性的边缘消融方法——通过识别并消融模型计算图中少量关键因果路径，来消除语言模型中的不良行为。在GPT-2中，仅消融12条边即可平均将毒性生成降低28%，在最毒性的提示下降低40%，且对其他输入的性能下降极小。

ABSTRACT

Language models often exhibit behaviors that improve performance on a pre-training objective but harm performance on downstream tasks. We propose a novel approach to removing undesirable behaviors by ablating a small number of causal pathways between model components, with the intention of disabling the computational circuit responsible for the bad behavior. Given a small dataset of inputs where the model behaves poorly, we learn to ablate a small number of important causal pathways. In the setting of reducing GPT-2 toxic language generation, we find ablating just 12 of the 11.6K causal edges mitigates toxic generation with minimal degradation of performance on other inputs.

研究动机与目标

• 解决尽管经过微调，语言模型中仍持续存在的有害行为问题。
• 开发一种方法，专门禁用不良行为，同时不降低在无关任务上的性能。
• 通过计算图消融，识别出导致不良行为的最小、因果性路径。
• 通过保留不相交、无关行为的性能，实现行为移除的高特异性。
• 探索一种无需完整微调或权重修改的推理后编辑方法。

提出的方法

• 在选定的粒度下，将模型重写为有向无环图（DAG），将计算表示为节点，边表示依赖关系。
• 使用可微目标学习边上的二值掩码，该目标惩罚消融操作，同时在不良行为输入上增加损失。
• 在推理时应用零消融或均值消融，以禁用所选边上的因果影响。
• 使用损失函数优化消融掩码，以平衡不良行为输入上损失的增加与训练数据上性能的保持。
• 使用早停和正则化，防止对小规模不良行为数据集过拟合。
• 在推理时应用消融，以禁用特定路径，而无需修改模型权重。

Figure 1: In targeted ablation, we (1) rewrite our model as a computation graph of a desired granularity, (2) learn a binary mask over edges while regularizing to penalize ablations, and (3) ablate edges at inference time to avoid the target bad behavior.

实验结果

研究问题

• RQ1 少量有针对性的边缘消融是否能有效缓解语言模型中的特定不良行为？
• RQ2 与微调和任务算术相比，有针对性的边缘消融在有效性与特异性方面表现如何？
• RQ3 消融因果路径是否能在不降低无关、不相交任务性能的前提下打破不良行为？
• RQ4 消融稀疏性与非目标行为性能下降之间的权衡如何？
• RQ5 边缘消融是否在行为移除方面优于基于权重的编辑或激活投影方法？

主要发现

• 在GPT-2中仅消融12条边，即可将有毒提示上的平均毒性得分从0.458降至0.328，相对降低28%。
• 对于最易诱发毒性的提示，毒性得分从0.944降至0.567，相对降低39.9%。
• 该方法在所有有效性和特异性指标上均优于任务算术，包括在有毒提示后更低的不连贯性。
• 使用84条边的均值消融，实现了与毒性生成缓解效果相当的水平，同时保持了其他输入的性能。
• 与联合微调相比，该方法在无关行为上的性能退化显著更小，证明了其高特异性。
• 该方法成功识别出导致毒性生成的最小因果切割路径，证实了其基于机制的方法论。

Figure 2: Ablating GPT-2 Small to remove toxicity. Left: Grey nodes are attention heads, and purple nodes are MLPs. Computation proceeds upwards, with horizontal alignment corresponding to layers. The computational graph has 11,611 edges; red edges are the 12 ablations learned to remove toxicity. Ri