APP: Accelerated Path Patching with Task-Specific Pruning

TL;DR

This paper introduces APP, a hybrid method combining contrastive attention head pruning and Path Patching to significantly accelerate circuit discovery in models, maintaining accuracy while reducing computational costs.

Contribution

The paper presents Contrastive-FLAP, a novel pruning algorithm for task-specific heads, and integrates it with Path Patching to speed up circuit discovery by over 56% with minimal loss in circuit quality.

Findings

Contrastive-FLAP preserves task-specific heads better than traditional pruning.

APP reduces search space by 56% on average.

Speed-up of 59.63%-93.27% in circuit discovery compared to dense Path Patching.

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