PAHQ: Accelerating Automated Circuit Discovery through Mixed-Precision Inference Optimization

TL;DR

PAHQ presents a training-free, per-edge mixed-precision approach to accelerate Automated Circuit Discovery (ACDC) by preserving FP32 precision only for the active edge while lowering precision elsewhere, coupled with a three-stream CUDA scheduler to hide data transfer latency. This alignment between activation patching and mixed-precision quantization enables substantial reductions in runtime (up to ~80%) and memory use (up to ~30%) with minimal faithfulness loss compared to unaccelerated ACDC, and shows strong performance relative to RTN-Q and EAP across standard circuit-discovery tasks. The method integrates seamlessly with existing circuit-discovery workflows, requires no extra training, and scales to larger GPT-2 variants, offering a practical pathway for accelerated mechanistic interpretability research. Code and implementation details are made available, highlighting its potential as a plug-and-play optimization for interpretable analysis of large language models.

Abstract

Circuit discovery, which involves identifying sparse and task-relevant subnetworks in pre-trained language models, is a cornerstone of mechanistic interpretability. Automated Circuit Discovery (ACDC) has emerged as a pivotal methodology in circuit discovery, but its application to large language models is severely limited by computational inefficiency and prohibitively high memory requirements. Although several accelerated approaches have been proposed, they primarily rely on linear approximations to ACDC, which significantly compromises analytical faithfulness. Our proposed method for accelerating automated circuit discovery, Per Attention Head Quantization (PAHQ), takes a fundamentally different approach by optimizing the efficiency of each individual patching operation. PAHQ leverages a fundamental alignment between activation patching and mixed-precision quantization (MPQ): interpretability analysis through patching essentially performs targeted ablation studies. Therefore, we can maintain high precision exclusively for investigated components while safely reducing precision elsewhere in the network. PAHQ-accelerated ACDC reduces runtime by up to 80\% and memory consumption by up to 30\% compared to unaccelerated ACDC while maintaining faithfulness. Importantly, our method readily integrates with existing edge-based circuit discovery techniques by modifying the attention computation mechanism. This training-free approach provides a practical and novel pathway for accelerating mechanistic interpretability methods. Our code is available at https://github.com/626619403/PAHQ.

Figures (5)

• Figure 1: ROC curves. RTN-quantization degrades performance.
• Figure 2: Three-Stream Parallel Scheduler. We overlap high-precision weight transfers with low-precision computations to mask transfer latency.
• Figure 3: A line graph showing the change in the number of edges in the ACDC circuit before and after quantization using our method as a function of step size. The graph demonstrates that after applying our quantization method, the ACDC circuit does not experience significant changes in terms of sparsity or pruning speed. Although the number of algorithm steps slightly increases, our method still far exceeds the original algorithm in terms of speed.
• Figure 4: The results of quantization strategy comparison.
• Figure 5: The full algorithm and implementation of PAHQ.