S³martCirc: Self-supervised Smart Circuit Discovery

To unify circuit discovery and functional interpretation, we jointly model both node importance and its computational role. Specifically, for every node $n$ in the LLM, we introduce a set of learnable masking coefficients $c_{n}=\{c_{n,u},c_{n,p},c_{n,f}\}$, where $c_{n,u}$ denotes the probability that node $n$ is unimportant, $c_{n,p}$ denotes the probability that the node acts as a passthrough node, and $c_{n,f}$ denotes the probability that it acts as a functional node. These coefficients are constrained such that $c_{n,u}+c_{n,p}+c_{n,f}=1$. We optimize these coefficients through the two-stage procedure shown in Figure 2. In the first stage, Important Node Discovery, we identify nodes that contribute to task performance by distinguishing nodes assigned to either functional role introduced earlier ($c_{n,p}+c_{n,f}$) from nodes classified as unimportant ($c_{n,u}$). Nodes satisfying $c_{n,p}+c_{n,f}>c_{n,u}$ are selected as candidate circuit nodes. In the second stage, Important Node Classification, each candidate node is assigned to one of the two roles, determined using a task-dependent metric described later. During training, node masking is applied softly using the coefficients $c_{n}$, while hard assignments are used only during evaluation through the following indicator function:

A circuit contains the nodes responsible for performing the task at hand. Consequently, removing any truly important node from the circuit should substantially degrade task performance. We identify important nodes by optimizing two complementary objectives: (1) minimizing the KL divergence between the circuit’s output distribution and that of the full model to ensure that the circuit faithfully reproduces the model’s original computation; and (2) minimizing cross-entropy loss with respect to the ground-truth answer to ensure the circuit produces correct predictions.

Formally, let $P_{\text{orig}}$ denote the output probability distribution of the full LLM, and $P_{\text{mask}}$ denote the distribution obtained when unimportant nodes are masked (i.e., nodes where $c_{n,u}>c_{n,p}+c_{n,f}$). The combined objective becomes

Here, $v\in V$ represents a token in the LLM’s vocabulary, and $P_{\text{orig}}(v)$ and $P_{\text{mask}}(v)$ represent the probability assigned to token $v$ by the full model and the model when unimportant nodes are masked, respectively. $y\in V$ is the ground truth token, and $\alpha\in[0,1]$ controls the trade-off between faithfulness to the original model behavior and task performance.

As defined earlier, circuit nodes are categorized as functional or passthrough nodes. Prior work typically assigns these roles using qualitative analysis or heuristic criteria (Wang et al. 2022; Chandna et al. 2025). In contrast, we propose a quantitative metric based on the following observation: nodes that perform nontrivial computation should alter the relational structure between tokens, whereas nodes that propagate information should preserve this structure. To formalize this idea, we measure how the pairwise similarity between tokens changes after applying a node. Intuitively, passthrough nodes induce minimal similarity changes, while functional nodes produce larger deviations.

We define a token-wise similarity operator $\text{Sim}(A)$ for an activation matrix $A\in\mathbb{R}^{T\times d_{n}}$ (i.e. the node’s input or output) as $\text{Sim}(A)=\hat{A}\hat{A}^{T}\odot(\mathbf{1}\mathbf{1}^{T}-I).$ $\hat{A}$ denotes the row-wise $L_{2}$ normalization of $A$, $\mathbf{1}$ is a column vector of ones, and $I$ denotes an identity matrix. We apply element-wise multiplication with $(\mathbf{1}\mathbf{1}^{T}-I)$ to exclude self-similarity. Using this operator, we define the functional interpretation metric as the normalized change in token-wise similarity:

where $X\in\mathbb{R}^{T\times d_{m}}$ and $Y\in\mathbb{R}^{T\times d_{n}}$ denote the input and output of a node, respectively. Let $N_{i}$ denote the set of nodes identified as important (i.e., those satisfying $c_{n,p}+c_{n,f}>c_{n,u}$). We incorporate this metric into the objective for the second stage:

This objective encourages nodes that induce larger changes in similarity structure to be assigned higher $c_{n,f}$ values, while nodes that preserve similarity are assigned higher passthrough probability $c_{n,p}$. In this way, functional roles are inferred directly from measurable changes in representation, enabling a fully quantitative interpretation.

Due to the vast size of LLMs, we introduce coefficient regularization, applied to both stages, to encourage the identification of minimal circuits while maximizing task performance. The objective consists of three terms that promote (1) confident assignments by enforcing binary values, (2) valid probability distributions over node roles (i.e., coefficients sum to one), and (3) sparse circuit selection:

where $C$ is the set of all masking coefficients in the model. The first term pushes each coefficient toward binary values (0 or 1), the second term enforces that the three probabilities sum to one, and the third term directly penalizes nodes being included in the circuit, encouraging sparsity.

As illustrated in Figure 2, training alternates between the two stages to capture the interdependence between node importance and functional role, allowing bidirectional influence. However, assigning functional roles with randomly initialized coefficients is unreliable at the start of training. Therefore, we introduce a warmup phase that optimizes only the discovery objective ($L_{\text{disc.}}$) to first identify important nodes. After warmup, we alternate between both stages to train the masking coefficients.

In Table 1, we compare S³martCirc to current circuit discovery methods and observe the following:

Previous work (Wang et al. 2022) manually identified and verified a ground-truth circuit for the IOI task through extensive analysis and validation. This circuit serves as a reliable reference for evaluating whether automated methods can recover the same critical components. Accordingly, a desirable property of circuit discovery methods is the ability to recover nodes from this known IOI circuit. Figure 3 reports the percentage of nodes from the original IOI circuit that are recovered by each method within the final end-to-end discovered circuits. We evaluate this across varying circuit sizes obtained by thresholding at 50, 100, 150, and 200 nodes. Note that perfect recovery (100%) is not achievable due to the vast number of possible nodes to choose from.

Overall, S³martCirc consistently recovers a higher percent of the original circuit compared to most baselines, and its performance is comparable only to Activation Patching. In contrast, other baselines, particularly optimization-based methods, achieve substantially lower recovery rates, even when allowed to select up to 200 nodes. These results provide evidence that S³martCirc not only improves circuit discovery performance (RQ1), but also aligns more closely with human-verified mechanistic interpretations of the IOI task. This increases the credibility of its discovered circuits, particularly in settings where ground-truth labels are unavailable.

We further validate the performance of S³martCirc by comparing its functional interpretations against the ground-truth circuit for the Acronym Prediction task (García-Carrasco et al. 2024a), shown in Figure 4. In the figure, each node is labeled as layer.head (e.g., 8.11 represents layer 8, attention head 11). S³martCirc successfully recovers four of the eight nodes (50%) from the original circuit, and its functional classifications align with the mapping shown in Figure 1. Specifically, Previous Token node 1.0 is correctly identified as a passthrough node, while Letter Mover node 11.4 is correctly classified as a functional node. This alignment, again, indicates that S³martCirc not only identifies important nodes but also accurately captures their mechanistic roles, providing automated interpretations that correspond closely to those derived through human analysis.

However, an apparent discrepancy arises in the classification of nodes 8.11 and 10.10. S³martCirc classifies 8.11 as a passthrough node while classifying 10.10 as a functional node, even though prior work shows that both heads can act as either Bridge nodes or Letter Mover nodes depending on the specific intervention applied. This divergence highlights a fundamental difference in methodology: human-verified approaches evaluate node behavior under carefully isolated interventions to determine what functions a head can perform, while S³martCirc measures the dominant contribution of each head across multiple samples. In other words, S³martCirc assigns labels based on the function that most strongly influences the circuit’s overall behavior. The differing classifications therefore reflect the multi-functional nature of these heads rather than a misclassification.

To answer RQ3, we perform a parameter analysis, shown in Figure 5, on the IOI task using GPT-2 to evaluate the sensitivity of S³martCirc to two key hyperparameters: $\alpha$, which balances the two objectives in the Important Node Discovery stage, and $k$, which controls the number of steps in the Important Node Classification stage. Specifically, each iteration consists of one discovery step followed by $k$ classification steps. We evaluate both accuracy (Acc.) and the percentage of recovered nodes from the original IOI circuit across $k\in\{2,4,6,8,10\}$ and $\alpha\in\{0.0,0.25,0.5,0.75,1.0\}$.

Figure 5a demonstrates that performance generally improves as $k$ increases, and stabilizes at $k=6$. The accuracy drop decreases from -71% at $k=2$ to approximately -96% at $k=6$, and percent circuit recovery monotonically increases. The consistent improvement suggests that additional refinement steps in the classification stage improve the identification of important nodes and the circuit quality.

Figure 5b reveals that $\alpha=0.75$ achieves the best performance in terms of both accuracy and circuit recovery. This indicates that, for the IOI task, the KL divergence objective should be weighted more heavily than the task loss during circuit discovery. Nevertheless, the cross-entropy term remains important, as performance degrades sharply when it is removed ($\alpha=1.0$). This helps explain the weaker performance of baselines, which rely primarily on output distribution differences (e.g., KL divergence) to identify important nodes. Our results suggest that, for IOI, combining output faithfulness with task performance signals is more effective, as some critical components contribute to correct behavior without inducing large changes in logits.

We perform an ablation study to understand how each component of S³martCirc contributes to its overall performance. Figure 6 compares the accuracy difference achieved by the original S³martCirc method against three variants: S³martCirc-NW (No Warmup), which removes the initial warmup period and directly starts the alternating training process; S³martCirc-DF (Discovery First), which completes the entire Important Node Discovery stage before beginning Important Node Classification; and S³martCirc-IF (Interpretation First), which reverses the order by performing Important Node Classification before Important Node Discovery. These ablations isolate the contributions of the warmup period and the alternating training strategy.

S³martCirc-IF performs significantly better than S³martCirc-DF (approximately 88% on IOI and 42% on Acronym), indicating that incorporating functional information into the discovery process provides useful guidance for identifying important nodes. However, it still underperforms S³martCirc. This demonstrates that, while functional interpretation is beneficial as a prior, it is insufficient without the alternating training.

Overall, these results highlight the importance of bidirectional coupling between Important Node Discovery and Important Node Classification. The alternating optimization allows each stage to refine the other: discovery benefits from evolving functional understanding, while classification operates on progressively more focused candidate sets.