PaperlayerView exact revision on arXiv

Source-linked AI summary

Circuit Breaking: Removing Model Behaviors with Targeted Ablation

Maximilian Li, Xander Davies, Max Nadeau

arXiv:2309.05973v2cs.CLcs.LG

TL;DR

Language models often retain behaviors that improve pre-training performance but harm downstream tasks, and these behaviors can be difficult to remove. Targeted edge ablation learns a small mask over causal pathways from examples of undesirable behavior and removes the selected edges at inference. A 12-edge ablation substantially mitigates GPT-2 toxic generation, lowering average toxicity from 0.458 to 0.328 and, for the most toxic-inducing prompts, from 0.944 to 0.567.

  • Problem

    Language models often retain behaviors that improve pre-training performance but harm downstream tasks, and these behaviors can be difficult to remove.

    Source: §1 · p. 1

  • Method

    Targeted edge ablation learns a small mask over causal pathways from examples of undesirable behavior and removes the selected edges at inference.

    Source: §1 · p. 1, §2 · p. 2, §3.3 · p. 2

  • Results

    A 12-edge ablation substantially mitigates GPT-2 toxic generation, lowering average toxicity from 0.458 to 0.328 and, for the most toxic-inducing prompts, from 0.944 to 0.567.

    Source: §4.3 · p. 4

  • Takeaways & Limitations

    Targeted edge ablation can mitigate a bad behavior across triggering inputs while causing only small amounts of damage to other model behaviors.

    Source: §6 · p. 5

    The approach uses mean ablation despite evidence that random marginal-distribution samples may be preferable, and its interventions could be made more granular.

    Source: D. Limitations · p. 10

Abstract

from arXiv · show

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.

1. Introduction

Language models can retain undesirable behaviors that help pre-training but harm downstream use. Targeted edge ablation removes small numbers of causal pathways to disable circuits responsible for such behaviors, with preliminary evidence from toxic-generation mitigation.

1. Introduction

  • Language models may exhibit hallucination, lying, manipulation, or undesirable biases because behaviors useful during pre-training can be difficult to remove during fine-tuning.Source: §1 · p. 1
  • Targeted edge ablation removes a small number of causal pathways through a model at inference time to disable circuits responsible for undesirable behaviors.The method discovers causal cuts through circuits rather than identifying the circuits themselves.Source: §1 · p. 1, Figure 1 · p. 1
  • The paper formulates behavior removal and presents preliminary targeted-ablation results for reducing toxic language generation.Source: §1 · p. 1

2. Background

Circuit analysis represents model computation as a graph and seeks subgraphs responsible for behaviors. Edge ablation instead removes specific dependencies, helping target causal pathways while reducing polysemantic side effects.

2. Background

  • Circuit analysis models computation as a directed acyclic graph and identifies subgraphs responsible for behaviors.The graph contains input source nodes, output sink nodes, and intermediate computational units.Source: §2 · p. 2
  • Edge ablation removes dependencies between computational nodes rather than entire nodes, allowing the method to target individual causal paths into and out of components.This can mitigate polysemantic behavior associated with ablating whole components.Source: §2 · p. 2
  • The experiments use zero ablation and mean ablation, which replace a source value with zero or its training-set mean, respectively.Source: §2 · p. 2

3. Targeted Ablation for Behavior Removal

The paper defines behavior removal as increasing loss on a targeted behavior while preserving performance on disjoint behaviors, then learns a sparse edge mask and applies it at inference. The approach favors limited, structure-preserving edits over expressive fine-tuning.

3. Targeted Ablation for Behavior Removal

  • Behavior removal requires efficacy on the targeted distribution and specificity on completely disjoint behaviors.The edited model should achieve high loss on the bad behavior while preserving low loss on distinct behaviors the original model handled well.Source: §3 · p. 2
  • The training and behavior distributions may overlap, so preserving training performance does not guarantee complete separation between the targeted behavior and ordinary training examples.Source: §3 · p. 2

3.3. Targeted Edge Ablation

  • Targeted edge ablation rewrites the model as a graph, learns a sparse edge mask, and ablates the selected edges during inference.The graph granularity controls the expressivity of possible ablations.Source: §3.3 · pp. 2–3
  • Each learnable edge weight ranges from 0 to 1, with zero meaning ablation, one meaning preservation, and intermediate values interpolating between preserved and ablated source values.Source: §3.3 · pp. 2–3
  • Increasing regularization during training first encourages behavior removal and then pressures the mask toward a smaller set of ablations.The final mask weights are rounded using a threshold before inference-time intervention.Source: §3.3 · pp. 2–3

3.4. Conceptual Advantages over Fine-Tuning

  • The method limits solution expressivity to reduce overfitting to a potentially mis-specified objective and relies on previously learned causal specialization.High-level edge edits also preserve most of the model’s mechanistic calculus and strictly reduce information available to computation.Source: §3.4 · p. 3

4. Removing Toxicity in GPT-2

The authors apply targeted edge ablation to GPT-2 Small, learning sparse causal cuts that reduce toxic generation while preserving other behaviors. Ablating 12 edges substantially mitigates toxicity, though it does not eradicate it.

4. Removing Toxicity in GPT-2

  • GPT-2 Small is edited to assign high loss to toxic sequences, using OpenWebText as an approximation of the training distribution and highly toxic 4chan excerpts as behavior examples.The toxic examples have Perspective API Toxicity V6 scores greater than 0.9.Source: §4 · p. 3

4.1. Learning Edge Mask Details

  • The residual-rewrite graph contains 158 nodes and 11,611 edges connecting inputs, outputs, attention heads, and MLPs, excluding communication between attention heads in the same layer.An ablated edge replaces a component contribution with zero or with that component’s mean value.Source: §4.1 · p. 3

4.2. Evaluation Metrics

  • Evaluation measures efficacy through toxic loss and generation toxicity, and specificity through ordinary and toxic-prepended OpenWebText perplexity.The study emphasizes retaining coherence after toxic prompts because this behavior may be inadvertently removed when increasing toxic-text loss.Source: §4.2 · p. 4

4.3. Results

  • 0.458 to 0.328 average toxicity, and 0.944 to 0.567 for the most toxic-inducing prompts, with 12 ablated edges.The method outperforms task arithmetic across efficacy and specificity metrics and increases incoherence less than joint fine-tuning, but does not eradicate toxicity.Source: §4.3 · p. 4, Figure 2 · p. 4

5. Related Work

The paper distinguishes circuit breaking from prior circuit discovery, model editing, pruning, and causal-ablation approaches. Its edge masks target necessary causal cuts and can modify internal component connections beyond individual weights.

Evaluation context

  • Table 1 organizes evaluation around toxic loss, toxic generation, filtered toxic cases, ordinary incoherence, and toxic-prepended incoherence.The filtered columns select samples based on the original model’s low loss or highly toxic generation.Source: Table 1 · p. 5

Causal mediation and circuit analysis

  • Circuit discovery tests whether edges are sufficient for a behavior, whereas circuit breaking seeks necessary edges forming a horizontal cut through the responsible circuit.The distinction is between preserving a vertical path and disabling a pathway that the behavior requires.Source: §5 · p. 5

Automated circuit discovery

  • The approach differs from automated circuit discovery by seeking neither individual features nor full computational circuits, but a sparse set of edges that removes a behavior.Prior automated methods include graph-crawling and gradient-based approaches.Source: §5 · p. 5

Weight-masking and model pruning

  • Unlike weight masking, edge masking can remove internal communication channels between high-level components, such as pairs of attention heads, and uses mean rather than zero ablation.The method also discourages sparsity because it seeks fewer ablations rather than model compression.Source: §5 · p. 5

Model editing to change or remove behaviors

  • Compared with model-editing methods that alter weights, activations, concepts, or representations, this work focuses on coarser ablations for removing behaviors.The supplied passage introduces this comparison without specifying the complete scope of the proposed alternative.Source: §5 · p. 5

6. Conclusion

The conclusion presents targeted edge ablation as a high-level behavioral modification method that can generalize beyond provided examples while causing only small damage elsewhere. The authors position it as a possible alternative to fine-tuning and call for further study.

6. Conclusion

  • Using a small bad-behavior dataset, the method mitigates the behavior on provided examples and generalizes to other inputs that trigger it.The conclusion states that these modifications cause only small amounts of damage to performance on other inputs.Source: §6 · p. 5
  • The authors conjecture that model editing may offer an alternative tool to fine-tuning for targeted behavioral modification and encourage future investigation.Source: §6 · p. 5

A. Additional Related Work

The paper distinguishes its setting from machine unlearning and relates it to prior backdoor-removal methods. Its target is undesirable behavior without enumerating the full set of undesirable training examples.

Unlearning

Backdoor removal

  • Prior backdoor-removal methods learn neuron masks or use Shapley importance to remove backdoors with limited data, whereas this work considers a broader edge-ablation strategy.The supplied passage says the authors believe their technique could be effective for backdoor removal, but does not provide results here.Source: A. Additional Related Work · p. 9

B. Writing models as computational graphs.

Models can be represented as computational graphs at different levels of detail, from individual operations to higher-level modules such as attention heads.

C. Ablation Types

The method considers zero and mean ablation as two ways to remove information flow along an edge, with mean ablation avoiding potentially unusual zero-valued intermediate signals.

  • Zero ablation computes a destination node as if its source value were zero, which can sometimes introduce an unusual signal.Source: C. Ablation Types · p. 9
  • Mean ablation substitutes the source node’s training-set mean and is intended to better represent withheld information than an atypical zero value.Source: C. Ablation Types · p. 9

D. Limitations

The experiments use zero and train-set-mean edge ablations, but the authors identify distributional concerns, finer-grained interventions, and stronger evaluations as limitations.

  • The ablation interventions set edge inputs to zero or their train-set mean, while random marginal-distribution samples may be preferable because mean ablation can be out of distribution.Source: D. Limitations · p. 10
  • More surgical circuit-breaking interventions could use finer-grained nodes and edges, such as separate query, key, and value nodes within attention heads.Source: D. Limitations · p. 10
  • The results could be strengthened with stronger baselines and additional methods for learning binary masks.Source: D. Limitations · p. 10

E. Additional Experiments: Breaking Digit Classification in an MLP

In an MNIST MLP, targeted edge ablation substantially reduces correct classification of a held-out digit while largely preserving classification of other digits.

Related papers

Loading 2309.05973v2…