IBCircuit: Towards Holistic Circuit Discovery with Information Bottleneck

We sincerely appreciate your time and effort in reviewing our paper and providing constructive feedback. We would like to address your questions and concerns below.

Claims and Evidence: Q1: Not sure what claims of avoiding manually constructed counterfactual activations mean. Q2: In the evaluation, don’t we need counterfactual activations from other inputs to evaluate the circuits?

Q1: We emphasize that our circuit discovery process does not require manually designing counterfactual inputs. While mean ablation also avoids explicit counterfactual input design, it underperforms the resampling ablation [1]. Additionally, we agree with the comment that "the performance will probably correlate with the similarity of the discovery and evaluation patching methods"[2]. The correlation arises from the biases introduced by different counterfactual inputs as corrupted data. In contrast, IBCircuit leverages Gaussian as corrupted activations rather than corrupted data, which eliminates such biases.

Q2: To ensure fair comparison, we follow existing baselines and utilize counterfactual activations derived from other inputs during evaluation. This demonstrates that our method achieves robust results even when assessed using ablation strategies distinct from the ablation employed in circuit discovery.

[1] Causal Scrubbing: a method for rigorously testing interpretability hypotheses.

[2] Transformer Circuit Faithfulness Metrics Are Not Robust. COLM2024.

Methods And Evaluation Criteria: Q1:Treating prior circuits as "ground-truth" risks missing valid components. Q2: KL divergence recovers far more than just task-specific behavior.

Q1: We agree with your opinion that treating manually validated circuits as ground truth may disadvantage newer discoveries. However, since baselines [1][2] also evaluate against these manually validated results as ground truth, we followed their protocol to ensure fair comparison. Beyond this Grounded in Previous Work experiment, we further demonstrate the advantages of our method through comparative analyses on metrics such as Faithfulness and Minimality.

Q2: In [1], Appendix C ("Discussion of Metrics Optimized") argues that:(1) Optimizing for low KL divergence is the simplest and most robust metric applicable across diverse tasks;(2) KL divergence is universally applicable to next-token prediction tasks because it does not require predefined output labels (unlike metrics such as logit difference, which necessitate specifying the exact tokens between which differences are computed). Based on these insights, we adopted KL divergence as our optimization metric.

[1] Towards Automated Circuit Discovery for Mechanistic Interpretability. NeurIPS 2023.

[2] Attribution Patching Outperforms Automated Circuit Discovery. NeurIPS 2023.

Experimental Designs Or Analyses: Evaluation uses counterfactual activations dissimilar to baseline discovery methods.

We would like to clarify that our evaluation is the same as the existing methods. We agree with the statement that "the performance will probably correlate with the similarity of the discovery and evaluation patching methods". For this reason, in our experiments, we adopt unified counterfactual activations derived from alternate inputs for all baseline methods, as well as our proposed approach to ensure fair comparison. The proposed method is robust and does not rely on this ablation strategy for circuit discovery. We appreciate your suggestion and will explicitly discuss this point in the revision.

Citation and writing issues mentioned in the weakness and suggestions.

Thank you for your suggestion. We will revise the corresponding part in the new version. We will discuss and cite the recommended papers.

Question: Runtime and GPU.

The IBCircuit runs the IOI task on one H20 GPU for about 1.2 hours. The official implementation of ACDC takes about 16 hours on the same device. We will add the comparison in the revised version.

We sincerely appreciate your time and effort in reviewing our paper and providing constructive feedback. We would like to address your questions and concerns below.

Claims and Evidence: Q1. Discuss the causality of circuit identification. Q2. Sampling methods in previous circuit studies IOI, Greater-than.

Q1: We appreciate your insight on causality. We would like to clarify that IBCircuit is causal. As in causal inference literature [1], our method employs controlled noise injection with learned IB weights as implicit causal interventions. And we do not need explicit data corruption. Unlike conventional ablation-based approaches, IBCircuit avoids distributional biases from discovery-phase ablations. Therefore, we can overcome the robustness limitations mentioned in [2] and find more reliable circuits.

Q2: Thank you for highlighting this detail. We acknowledge that IOI introduces mean ablation in Section 2.2, while Greater Than uses permutation sampling. We will clarify these distinctions in the revised version.

[1] A Causality-Aware Perspective on Domain Generalization via Domain Intervention.

[2] Transformer Circuit Faithfulness Metrics Are Not Robust. COLM2024.

Methods and Evaluation: Q1. Conduct more rigorous statistical tests. Q2. Experiment with larger model architectures.

Q1: We provide additional results under various ablation methods in this anonymous link, supplementing the random ablation presented in our initial draft. We also include the suggested evaluations.

Q2: We add comparison on CodeLLaMA-13B in the same anonymous link above.

Broader Scientific Literature and Essential References: Reference and comparisons with: (1)Edge pruning, (2)Sparse feature circuit.

Thank you for the valuable feedback. We include the comparison with the two baselines in this anonymous link. Due to our limited GPU resources, we have not yet completed all experiments. We will keep the remaining results updated in the next few days. Additionally, we will cite the two papers in the revision.

We sincerely appreciate your effort in reviewing our paper and providing constructive feedback. We would like to address your concerns below and revise the corresponding parts in the revised version.

Theoretical Claims: We would like to further summarize the IBCircuit. IBCircuit can be regarded as the application of the Information Bottleneck principle to circuit discovery with our proposed appropriate adaptations.

• Why use variational approximation? The mutual information is intractable and IB can not be directly applied to circuit discovery. The variational method is introduced to estimate the bound of mutual information. The problem is also analyzed in VAE[1] and VIB[2] in deep representation learning.
• Random Variable. In deep neural network[1][2], the input $X$ are random variables. Each concrete input string $x$ can be viewed as an instance sampled from the task distribution $X$, e.g. IOI task. Similarly, $Y$, $\mathcal{C}$ and $\mathcal{G}$ are also random variables, and $[x,y,c,g]$ is an concrete instance. We provide an analysis on random variables, but this formulation can be trained with instances[1][2].
• Gaussian in Appendix A.2. Gaussian is the variational prior[1][2]. Gaussian is widely used as variational priors due to the reparameterization trick and convenient KL calculation[1][2]. Following existing literatures, in IBCircuit, the prior of circuit $\mathcal{C}$ is set as Gaussian. IBCircuit will learn the posterior distribution of $\mathcal{C}$ over task $X$, and we report the mean of the learned distribution $\mathcal{C}$ as the final result.
• $I[Y; \mathcal{C}]$ Computation. For an instance $x$ in a batch, we compute the KL divergence between the output of the perturbed model $y_{\mathcal{c}}$ and the output of the clean model $y$. Then we take the average over the batch. For each instance $x$ in the batch, the corresponding circuit $c$ is generated from Gaussian with reparameterization trick[1]. Therefore, each instance in the training stage is deterministic.
• $I[\mathcal{G}; \mathcal{C}]$ Computation in Appendix A.2. We compute the KL between the posterior distribution and Guassian prior of the circuit over the batch. As the prior is Gaussian[2], we can calculate the KL as in Appendix A.2.

We will make them clear in the revision.

[1] Deep Variational Information Bottleneck. ICLR 2017.

[2] Auto-Encoding Variational Bayes. ICLR 2014.

W1: The writing around Proposition 4.2.

Thanks for the comment. We will provide the definitions of A and B in the main text. $L_{MI}(G; C)$ is defined as the upper bound of $I(G; C)$, i.e. the KL between the posterior distribution and Gaussian prior of the circuit over the batch.

W2: Comparison to Subnetwork Probing.

• Whether the IB or the variational approximation is key. As mentioned above, the key is IB. The variational approximation makes IB trainable in deep neural networks.
• What is responsible for the performance difference? We think the objective of IBCircuit is responsible for the performance difference. Gaussian perturbations and hard masking are just two different ways of modeling circuits with learnable parameters, not the key. (1) In $I[Y;\mathcal{C}]$, IBCircuit calculate KL while SP calculate cross-entropy. KL measures the distribution of the output logit vector. While cross-entropy only focuses on the ground truth, i.e. one element on the logit vector. SP ignores the logits of other tokens. (2) In $I[\mathcal{G}; \mathcal{C}]$, IBCircuit ensures the circuit does not receive irrelevant information from the whole computation graph. In contrast, SP only encourages sparsity. Even if the circuit obtained by SP is sparse, it may contain some irrelevant information.

W3&Theoretical Claims: Additional Experimental Results.

Thanks for the suggestions. We agree that additional comparisons would enhance our contribution. We include the results about (1) the comparison to IG-EAG, (2) more tasks, and (3) larger models in this anonymous link. Due to our limited GPU resources, we have not yet completed all experiments. We will keep the remaining results updated in the next few days.

Questions For Authors: What is the computational cost?

Taking the IOI task as an example, running the official ACDC implementation on a single H20 GPU takes approximately 16 hours, while EAP takes about 1 minute, and the proposed method requires roughly 1.2 hours. Although EAP has a faster runtime, it requires designing corresponding clean and corrupted input pairs to compute the gradients. Moreover, the proposed method outperforms EAP on both IOI and Great-than tasks. Furthermore, our approach is significantly more scalable compared to ACDC.

Thank you for your valuable feedback on our paper. We appreciate your time and effort in reviewing our work. We would like to address your questions and concerns below.

Weakness 1: Unclear explanation of "distorted information flow" in Eq.1: Why use input corruption instead of direct KL minimization?

Thanks for the comment. The distorted information flow is used for modeling circuits with learnable parameters. There is no explicit circuit within the LLM, thus we can not directly minimize KL. In the distorted information flow, we model the learnable circuit with Gaussian, i.e. input corruption. Therefore, IBCircuit can optimize the circuit with gradient descent. It is the common practice in variational methods [1][2]. Using input corruption to model circuit is common in optimization framework for circuit discovery. For example, the baseline Subnetwork Probing [3] uses a hard-concrete random masking as input corruption to model the circuit.

[1] Deep Variational Information Bottleneck. ICLR 2017.

[2] Auto-Encoding Variational Bayes. ICLR 2014.

[3] Low-complexity probing via finding subnetworks. NAACL 2021.

Weakness 2: The proofs presented for the theoretical bounds lack novelty. The derivation of the lower bounds appears to be a straightforward application of previous works.

Thanks for the comment. We derived a novel upper bound, and since previous subgraph discovery work has the same optimization objective as circuit discovery, we adopted the same form of lower bound.

• Upper Bound: The traditional upper bound can not be directly applied to the circuit discovery, as the random variables $\mathcal{G}$ and $\mathcal{C}$ are computational graphs. Therefore, we derive a new upper bound.

• Lower Bound: In previous studies, the lower bound derived for subgraph discovery aims to identify critical subgraphs that preserve information similar to the complete graph. This aligns with the objective of circuit discovery, where the goal is to find circuits that maintain the performance of the full pre-trained model. Therefore, we incorporate the same bound using the notation of Information Bottleneck Circuit to ensure self-consistency in the proposed method.

Weakness 3: No evidence for applicability to larger models.

Thanks for the suggestions. We agree that additional results on larger models would enhance our contribution. We include the results in this anonymous link. Due to our limited GPU resources, we have not yet completed all experiments. We will keep the remaining results updated in the next few days.