APP: Accelerated Path Patching with Task-Specific Pruning

TL;DR

This work tackles the challenge of identifying minimal, interpretable circuits in large language models, a task hindered by the high cost of Path Patching. It introduces Accelerated Path Patching (APP), a hybrid pipeline that uses Contrastive-FLAP pruning to reduce the search space before applying circuit discovery, thereby achieving substantial runtime savings while maintaining circuit quality. Contrastive FLAP better preserves task-specific, context-sensitive heads than vanilla pruning, but pruning alone cannot replace Path Patching due to minimality constraints; APP resolves this by merging pruning-informed candidates and applying PP on a sparser model. Across multiple models and tasks, APP delivers large reductions in search space (average at least $56 ext{ extbf{%}}$) and significant GFLOPs/time speedups, with circuits that retain comparable performance to PP, enabling scalable mechanistic interpretability for larger models.

Abstract

Circuit discovery is a key step in many mechanistic interpretability pipelines. Current methods, such as Path Patching, are computationally expensive and have limited in-depth circuit analysis for smaller models. In this study, we propose Accelerated Path Patching (APP), a hybrid approach leveraging our novel contrastive attention head pruning method to drastically reduce the search space of circuit discovery methods. Our Contrastive-FLAP pruning algorithm uses techniques from causal mediation analysis to assign higher pruning scores to task-specific attention heads, leading to higher performing sparse models compared to traditional pruning techniques. Although Contrastive-FLAP is successful at preserving task-specific heads that existing pruning algorithms remove at low sparsity ratios, the circuits found by Contrastive-FLAP alone are too large to satisfy the minimality constraint required in circuit analysis. APP first applies Contrastive-FLAP to reduce the search space on required for circuit discovery algorithms by, on average, 56\%. Next, APP, applies traditional Path Patching on the remaining attention heads, leading to a speed up of 59.63\%-93.27\% compared to Path Patching applied to the dense model. Despite the substantial computational saving that APP provides, circuits obtained from APP exhibit substantial overlap and similar performance to previously established Path Patching circuits

Figures (16)

• Figure 1: Depiction of the Accelerated Path Patching (APP) Algorithm. APP reduces the search space of circuit discovery methods by successfully pruning task-irrelevant heads while preserving task-critical attention heads. APP then runs Path Patching on the remaining sparse model.
• Figure 2: Activation Patterns of context-insensitive (top) and context-senstive (bottom) heads. Left: clean activations, middle: corrupted activations, right: contrastive activations.
• Figure 3: Performance and TP over a sparsity of 0 to 1 for vanilla and contrastive FLAP. Points are the fixed sparsity used in Table \ref{['tab:exp1']}, dotted lines are the sparsity ratios found via cliff points.
• Figure 4: Comparison of attention head types identified by Contrastive FLAP and Automatic Path Patching for the IOI task in GPT-2 small.
• Figure 5: Difference of required GFLOPs across all models and tasks