Organizations: Department of Systems Engineering and Engineering Management, The Chinese University of Hong Kong, Hong Kong, China · DAMO Academy, Alibaba Group, Hangzhou, China · Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Department of Computer Science, Hong Kong Baptist University, Hong Kong, China · School of Artificial Intelligence, Shenzhen University, Shenzhen, China · Tencent, Shenzhen, China · Hupan Lab, Hangzhou, China
Circuit discovery has recently attracted attention as a potential research direction to explain the non-trivial behaviors of language models. It aims to find the computational subgraphs, also known as circuits, within the model that are responsible for solving specific tasks. However, most existing studies overlook the holistic nature of these circuits and require designing specific corrupted activations for different tasks, which is inaccurate and inefficient. In this work, we propose an end-to-end approach based on the principle of Information Bottleneck, called IBCircuit, to identify informative circuits holistically. IBCircuit is an optimization framework for holistic circuit discovery and can be applied to any given task without tediously corrupted activation design. In both the Indirect Object Identification (IOI) and Greater-Than tasks, IBCircuit identifies more faithful and minimal circuits in terms of critical node components and edge components compared to recent related work.
Figures & tables
Figure 1: (a) Transformer blocks from the perspective of the residual stream. (b) Adding Gaussian noise to the activations of attention heads using node-wise IB weights and optimizing through the Information Bottleneck. (c) Selecting attention heads with less noise as the node-level circuit. (d) Adding Gaussian noise to the activations of source nodes using edge-wise IB weights and optimizing through the Information Bottleneck. (e) Selecting edges with less noise as the edge-level circuit.
Figure 2: ROC curves of SP, ACDC, AP and IBCircuit identifying model components from previous work, across IOI circuit and Greater-Than circuit.
Figure 3: Comparison of the impact of different trade-off coefficients α on IBCircuit in implementing the IOI task.
Figure 4: Comparison of IBCircuit and ablation variants in terms of Greater Probability and KL Divergence metrics.
Figure 5: Comparison of IBCircuit and related methods in terms of Logit Difference and Greater Probability metrics under different node number thresholds. Higher metric scores and fewer nodes correspond to better circuits.
Figure 6: Comparison of IBCircuit and related methods in terms of KL Divergence metric under different node number thresholds. Lower metric scores and fewer nodes correspond to better circuits.
Figure 7: Comparison of IBCircuit and related methods in terms of Logit Difference and Greater Probability metrics under different edge number thresholds. Higher metric scores and fewer edges correspond to better circuits.
Figure 8: Comparison of IBCircuit and related methods in terms of KL Divergence metric under different edge number thresholds. Lower metric scores and fewer edges correspond to better circuits.
Figure 9: Comparison of IBCircuit and related methods in terms of Logit Difference metric under different edge number thresholds. Higher metric scores and fewer edges correspond to better circuits.
Max Planck Institute for Informatics, Saarland Informatics Campus, Saarbrücken, Germany · CISPA Helmholtz Center for Information Security, Saarbrücken, Germany