Source-linked AI summary

Can Circuit Alignment Predict OOD Generalization?

Ayan Banerjee, Abhra Chaudhuri, Josep Llados, Umapada Pal, Anjan Dutta

arXiv:2609.31996v1cs.LG

TL;DR

Predicting OOD generalization from weights alone remains an open problem, and activation-based similarity metrics miss structural rerouting. The paper introduces CAS to compare class-specific circuits across domains and proves a Monte Carlo ranking guarantee. Across 48 learners on three benchmarks, CAS outperforms the reported baselines, with a PACS Spearman correlation of 0.88.

  • Problem

    Weight-only OOD prediction lacks a formally grounded answer, while activation-based metrics are provably blind to structural rerouting that distribution shift induces.

    Source: §1 · p. 1
  • Method

    CAS compares class-specific circuits via graph kernels, separates same-class coherence from cross-class confusion, and estimates population rankings with Monte Carlo sampling.

    Source: §1 · p. 2
  • Results

    Across 48 learners on three benchmarks, CAS attains Spearman ρS ≥0.77 in every setting; on PACS it reaches ρS = 0.88 versus 0.58 (CKA), 0.23 (SVCCA), and 0.14 (RSA).

    Source: §1 · p. 2, §5 · p. 9
  • Takeaways & Limitations

    The results support circuit alignment as a structural basis for ranking OOD robustness, with a Monte Carlo pairwise inversion error of O(1/M).

    Source: §1 · p. 2

    Finite-accuracy monotonicity is empirically verified rather than proved; the framework targets classification with bounded class counts, and ACE requires source-domain adapter training.

    Source: §5 · p. 10

Abstract

from arXiv · show

Can out-of-distribution (OOD) generalization be predicted from a trained model's weights alone, without any target-domain data? Existing representational similarity metrics (CKA, SVCCA, RSA) compare activations rather than forecast generalization. We show they are provably insensitive to structural rerouting in the computational graph, the very change distribution shift induces. We close this gap with the Circuit Alignment Score (CAS), which compares class-specific circuits across domains via graph kernels, decomposed into same-class coherence and cross-class confusion. Casting CAS as a Lebesgue integral over the domain distribution, we prove its Monte Carlo estimate recovers the ground-truth ranking of learners by OOD accuracy, with pairwise inversion error vanishing at rate $O(1/M)$, where $M$ is the number of sampled domains. Across $48$ learners on PACS, CAS attains $0.88$ rank correlation with OOD accuracy, versus $0.58$ (CKA), $0.23$ (SVCCA), and $0.14$ (RSA), with similar trends on other benchmarks and even against data-dependent methods, making it the first provably consistent predictor of distributional robustness requiring neither target-domain data nor labels. The code is available at: https://github.com/ayanban011/ACE

1 Introduction

The paper asks whether OOD generalization can be predicted from weights without target-domain data, while activation-similarity metrics miss structural rerouting. It proposes CAS to compare class-specific circuits and gives theoretical and empirical support for using it to rank learners.

1 Introduction

  • In the limit of perfect OOD accuracy, learners are necessarily circuit-robust: their class-specific circuits are preserved across domains and remain distinct between classes.This grounds circuit invariance as a structural proxy for OOD robustness rather than a heuristic.Source: §1 · p. 2
  • Any weight-only OOD metric must detect structural rerouting, distinguish same-class from cross-class structure, and respond monotonically to graded perturbation.CKA, SVCCA, and RSA each fail at least two of these conditions by construction.Source: §1 · p. 2
  • CAS compares circuits via graph kernels, forms a class-conditional similarity matrix, and aggregates the gap between mean-diagonal and mean-off-diagonal similarities.The two components capture same-class circuit drift and cross-class circuit entanglement.Source: §1 · p. 2, Figure 1 · p. 2
  • The M-sample Monte Carlo estimate of population CAS has pairwise ranking inversion error vanishing at O(1/M), yielding convergence to the ground-truth OOD ranking when monotonicity holds.The paper reports that monotonicity condition as empirically verifiable.Source: §1 · p. 2
  • On PACS, CAS reaches Spearman ρS = 0.88 versus 0.58 (CKA), 0.23 (SVCCA), and 0.14 (RSA); trends extend to Office-Home and DomainNet.The comparison uses 48 learners, and the reported gap is consistent across target domains.Source: §1 · p. 2

2 Related Work

Prior OOD predictors span training dynamics, influence, margins, bounds, meta-learning, and data-centric estimates, but require target data or assumptions about cross-domain variance. Circuit-tracing work leaves cross-domain circuit stability and comparability open, while activation-based similarity metrics can conflate robust and non-robust learners.

Predictors of OOD Generalization

  • Existing OOD predictors include training-loss trends, influence functions [67], margin-based predictors [24] [33], PAC-Bayes bounds [32] [41], meta-learning [61], and data-centric estimators [6] [18] [68].These approaches require either prediction-variance preservation across domains or unlabeled target data; CAS predicts rankings from source-domain circuits alone.Source: §2 · p. 2

Representation Similarity Metrics

  • CKA [29], SVCCA [44], and RSA compare activation geometry, and activation-factoring metrics provably conflate circuit-robust and non-robust learners.Prior work [2] reports inconsistent correlations between representational metrics and OOD accuracy.Source: §2 · p. 3

Circuit Tracing

  • Circuit tracing uses methods including attribution patching [53], sparse autoencoders [54], probing, causal tracing [37], and neuron-level analysis [48].These methods assume a fixed data distribution; CAS instead defines a metric on extracted circuits that is agnostic to the extractor.Source: §2 · p. 3

3 Circuit Alignment Score (CAS)

The section derives necessary criteria for predicting OOD generalization from weights, then constructs CAS to meet them. It proves ranking recovery under an empirically verified monotonicity assumption, with pairwise inversion error O(1/M).

3.1 What Does Predicting OOD Generalization from Weights Require?

  • Perfect OOD accuracy requires class-specific circuits to be preserved across domains and distinct between classes, supporting circuit robustness as a structural proxy.The section identifies circuit drift and circuit entanglement as distinct ways this structure can break down.Source: §3.1 · p. 3
  • CKA, SVCCA, and RSA fail P1 and P2, while CAS satisfies all three necessary conditions by construction.P1 requires sensitivity to structural rerouting, P2 separate same-class from cross-class structure, and P3 monotonic response to perturbation.Source: §3.1 · pp. 3–4

3.2 Construction of Circuit Alignment Score (CAS)

  • CAS compares circuit graphs with a graph kernel, organizes similarities by class, and subtracts mean off-diagonal similarity from mean diagonal similarity.Diagonal entries capture same-class preservation; off-diagonal entries capture cross-class entanglement.Source: §3.2 · p. 4
  • CAS is high when same-class circuits are preserved and different-class circuits remain distinct, and its components expose class vulnerabilities and confusion topology.The matrix supports per-class failure analysis beyond the aggregate scalar score.Source: §3.2 · pp. 4–5

3.3 Consistency: From CAS to OOD Ranking

  • The population CAS averages source-to-domain CAS over the domain distribution, and its Monte Carlo estimate converges almost surely as sampled domains increase.The convergence rate for preserving pairwise rankings depends on the gap between learners' population CAS values relative to domain-wise variability.Source: §3.3 · p. 5
  • Under the monotonicity assumption, empirical CAS rankings recover OOD-accuracy rankings with pairwise inversion error O(1/M).The assumption is empirically verified across 48 learners on PACS, Office-Home, and DomainNet, rather than proved from first principles.Source: §3.3 · p. 6

4 Experimental Results

Across three benchmarks, the experiments assess CAS as an OOD-accuracy predictor, test its ranking consistency and response to controlled domain shifts, and compare it with other prediction methods.

4 Experimental Results

  • Experiments span PACS, Office-Home, and DomainNet, use Spearman rank correlation as the primary evaluation metric, and extract class-specific circuits with ACE.The benchmarks range from PACS’s four domains and seven classes to DomainNet’s six domains and 345 classes.Source: §4 · p. 6

4.1 CAS as OOD Predictors

  • ρS = 0.93 for CAS, compared with 0.58 for CKA, 0.23 for SVCCA, and 0.14 for RSA, in the PACS learner comparison.CAS shows a near-monotone relationship with OOD accuracy, while the activation-based metrics show weak trends with high variance.Source: Figure 2 · p. 7, §4.1 · p. 7

4.2 Consistency of CAS Rankings

  • At practical M values, CAS rankings are near convergence; TK correlations are 0.89, 0.81, and 0.79 on PACS, Office-Home, and DomainNet.Pairwise inversion probability follows the O(1/M) trend and lies at or below the theoretical bound, with M=3 for PACS and Office-Home and M=5 for DomainNet.Source: §4.2 · pp. 7–8, Figure 3 · p. 8

4.3 Evolution of CAS with Domain Perturbations

  • Across four architectures and three graph kernels, CAS decreases monotonically as LoRA photo weight increases, recovering the known perturbation order without reversals.The interpolation preserves semantic content while changing domain-specific appearance; the relative ordering holds across configurations despite different architecture-specific slopes.Source: §4.3 · pp. 8–9, Figure 5 · p. 9

4.4 Comparison with OOD Accuracy Prediction Methods

  • CAS achieves ρS = 0.88 and MAE = 2.14%; ALine-D is the closest competitor at ρS = 0.72 and MAE = 4.85%.Table 3 compares prediction methods on PACS; CAS has the highest ρS and lowest MAE across all four target-domain splits.Source: Table 3 · p. 9, §4.4 · p. 9

5 Conclusion and Discussions

The paper presents CAS as a theoretically motivated, weight-derived OOD predictor and reports strong results across benchmarks. It also identifies limits on the guarantee and current framework.

5 Conclusion and Discussions

  • Across 48 learners on three benchmarks, CAS attains Spearman ρS ≥ 0.77 in every setting, exceeding weight-only and target-data-using baselines.The framework has a Monte Carlo consistency guarantee with O(1/M) pairwise inversion error.Source: §5 · p. 9
  • The class-conditional similarity matrix exposes per-class vulnerability and confusion topology beyond the scalar CAS score.Alternative circuit extractors can be used with the same metric.Source: §5 · p. 10
  • Finite-accuracy monotonicity is empirically verified rather than proved; the framework is currently specialized to classification with bounded class counts.ACE also requires source-domain adapter training, though this is lighter than retraining-based baselines.Source: §5 · p. 10
  • CAS is proposed as a model-selection aid alongside conventional validation, not as a substitute, because high CAS does not establish correctness on individual predictions.The paper identifies settings with unavailable target data, including clinical pre-deployment, low-resource, and federated systems.Source: §5 · p. 10

A.1 Circuit-level robustness as a structural proxy

The section defines circuit robustness through preservation of same-class circuits and separation of different-class circuits across domains. It motivates CAS by showing why predictive metrics must detect structural rerouting, preserve class-conditional structure, and respond monotonically to perturbation.

A.1 Circuit-level robustness as a structural proxy

  • Perfect OOD accuracy implies circuit robustness under shared-support overlap: same-class circuits are causally equivalent across domains, while different-class circuits remain functionally distinguishable.The equivalence allows functionally redundant rerouting; the kernel-based bounds are stated for domains modulo a null set.Source: §A.1 · pp. 15–16

A.2 Structural Sensitivity

  • Activation-factorized metrics can assign identical values to robust and rerouted non-robust learners, so predictive metrics must be structurally sensitive to circuit topology.CKA, SVCCA, and RSA factor through activation embeddings, whereas CAS compares circuit graphs directly.Source: §A.2 · p. 16

A.3 Class-Conditional Resolution

  • A metric that ignores similarity-matrix positions can conflate robust circuits with configurations exhibiting drift and entanglement; CAS separates diagonal from off-diagonal similarities.CKA, SVCCA, and RSA do not form a per-class similarity matrix, while CAS uses the gap between mean-diagonal and mean-off-diagonal values.Source: §A.3 · p. 17

A.4 Semantic Consistency

  • A semantically consistent metric must not increase when same-class similarity falls and cross-class similarity rises; CAS satisfies this monotonicity by construction.The section defines stronger perturbation as simultaneously reducing same-class similarity and increasing mean cross-class similarity for every class.Source: §A.4 · pp. 17–18

B CAS Construction: Full Derivation and Proofs

  • The construction uses a graph kernel, a class-conditional similarity matrix, and a scalar aggregation, with the kernel operating on circuit topology rather than activations.The treelet, random-walk, and optimal-transport-based kernels are used; the treelet kernel is the default because it gives the cleanest diagonal/off-diagonal separation in the experiments.Source: §B.1 · p. 18

B.1 Design rationale

  • CAS aggregates a class-conditional similarity matrix using uniform weights, yielding a bounded scalar that is class-permutation-invariant and monotone in the matrix entries.The section characterizes CAS as the unique class-permutation-invariant, monotone-in-S, [−1, 1]-bounded scalar aggregation.Source: §B.1 · p. 18

B.2 Proof of CAS Soundness

The soundness argument shows that CAS satisfies structural sensitivity, class-conditional resolution, and semantic consistency. The section also records boundedness, symmetry, and self-similarity properties.

B.2 Proof of CAS Soundness

  • CAS satisfies structural sensitivity because graph kernels compare circuit topology directly, distinguishing circuits with different topology even when their activation footprints match.The cited argument relies on the kernel’s topology-respecting property.Source: §B.2 · p. 19
  • CAS satisfies class-conditional resolution by separately aggregating diagonal and off-diagonal similarities, so it is not invariant to arbitrary permutations of the similarity matrix.Its same-class and cross-class components can be reported independently.Source: §B.2 · p. 19
  • CAS is non-decreasing in diagonal similarities and non-increasing in off-diagonal similarities, so stronger perturbation cannot yield a higher CAS value.The proof establishes the inequality term by term from the signs of the aggregation coefficients.Source: §B.2 · p. 19

B.3 Algebraic properties of CAS

  • CAS lies in [−1, 1], attaining +1 under perfect class-conditional preservation and −1 under complete entanglement.The extrema occur when diagonal similarities are respectively all 1 or all 0, with off-diagonal similarities respectively all 0 or all 1.Source: §B.3 · p. 19
  • If the base kernel is symmetric, swapping the two circuit families leaves CAS unchanged.The proof notes that swapping transposes the similarity matrix, preserving its diagonal and the set of off-diagonal entries.Source: §B.3 · p. 19
  • When each class circuit is more similar to itself than to other classes, self-comparison has positive CAS; it equals 1 only with perfect within-class similarity and zero cross-class similarity.The stated equality condition requires κ(C(i), C(i)) = 1 and κ(C(i), C(j)) = 0 for all i ≠ j.Source: §B.3 · p. 20

C Consistency: Full Proofs

The section analyzes how empirical CAS rankings relate to population rankings and how population CAS connects to OOD accuracy. It reports faster inversion-bound convergence for comparisons across training algorithms, while noting that the accuracy-alignment assumption is not proved for finite-accuracy regimes.

C Consistency: Full Proofs

  • The appendix extends the consistency analysis by comparing inter-algorithmic and intra-algorithmic separation and relating the structural premise to Assumption 1.It also gives full proofs for the consistency results.Source: C Consistency: Full Proofs · p. 20

C.2 Inter- vs intra-algorithmic separation

  • Both inter- and intra-algorithmic pairwise inversion bounds vanish at rate 1/M, but the inter-algorithmic bound is tighter by a factor of (∆intra/∆inter)^2.The paper reports ∆inter ≫ ∆intra empirically, and says CAS distinguishes differently trained learners with fewer sampled domains than learners differing only by initialization or regularization.Source: §C.2 · pp. 20–21

C.3 Proof of ranking recovery

  • For a fixed learner pool, the probability that any pair is inverted is O(1/M), so the empirical Kendall ranking converges in probability to the ground-truth ranking as sampled domains increase.The proof applies a union bound over pairs after bounding each pair’s inversion probability.Source: §C.3 · p. 21

C.4 Relationship between Assumption 1 and the structural premise

  • Assumption 1 bridges population CAS and OOD accuracy, but the paper does not derive it from first principles in finite-accuracy regimes.The authors cite the perfect-accuracy limiting case, causal-invariance theory, and empirical verification as support.Source: §C.4 · p. 21
  • Across 48 learners on PACS, Office-Home, and DomainNet, CAS and OOD accuracy have mean Spearman rank correlation ρS = 0.85 across benchmarks and target domains.Reported correlations range from ρS = 0.77 on OfficeHome Real World to ρS = 0.93 on PACS Cartoon; inversions occur primarily among learners with very similar OOD accuracies.Source: §C.4 · p. 21
  • The paper leaves a formal bound linking inversion-set size to the OOD-accuracy gap for future work.The proposed strengthening would relate |g(ℓi)−g(ℓj)| to the inversion set’s size.Source: §C.4 · p. 21

D Implementation Details

The implementation uses frozen pretrained backbones with trainable adapters and a task head, then extracts class-specific circuits through causal activation patching. Graph kernels compare those circuits, while baseline metrics and evaluation follow defined protocols.

D Implementation Details

  • Adapters are placed after late representation stages, and ViT circuits include attention heads as well as adapter neurons to capture high-level task computation.Backbones remain frozen; only adapters and the classification head are optimized.Source: D Implementation Details · pp. 21–22, Table 4 · p. 22

D.1 Baseline Protocols for OOD Accuracy Prediction

  • Baselines compare penultimate-layer features, while leave-one-domain-out evaluation correlates each metric with held-out-target accuracy across the learner pool.ATC, ProjNorm, and ALine-D use their specified target-data protocols.Source: D Implementation Details · p. 23, §D.1 · p. 23

E Adaptive Circuit Tracing (ACE)

  • ACE extracts class-specific circuits by scoring units with activation patching, retaining top-K units, then estimating directed inter-layer edges from causal activation changes.Positive importance promotes the correct class; negative importance suppresses it, and edges are retained when their absolute weight exceeds ϵ.Source: E Adaptive Circuit Tracing (ACE) · p. 24
  • Circuit extraction scales with adapter and attention units rather than total backbone parameters, with ViT attention patching adding 36 forward passes.Cached outputs keep edge-estimation cost linear in K rather than quadratic.Source: E Adaptive Circuit Tracing (ACE) · p. 25

F Choice of Graph Kernel

  • The experiments use the Treelet Kernel, which captures local hierarchical sparse substructures alongside the complementary global alignment and path-based notions of OT and RWK.The compared kernels produce symmetric circuit similarities; the Treelet Kernel is used for CAS throughout the experiments.Source: F Choice of Graph Kernel · pp. 25, 27, Figure 8 · p. 27

G Ranking pairwise domain similarity with CAS

CAS ranks PACS domain pairs in agreement with independent measures of domain dissimilarity, whereas representational baselines produce incoherent orderings. The section also estimates extraction noise to distinguish it from cross-domain circuit divergence.

G Ranking pairwise domain similarity with CAS

H Decomposing Circuit Differences: Noise vs. Signal

  • Within-domain resampling estimates extraction noise, allowing observed cross-domain circuit dissimilarity to be decomposed into noise and signal components.The decomposition assumes independent extraction noise; the signal-to-noise ratio indicates whether observed differences exceed the noise floor.Source: H Decomposing Circuit Differences:… · p. 31

H.1 Circuit Fragility Test for PACS dataset

  • On PACS, within-domain variability is 0.096 ± 0.021 versus cross-domain dissimilarity of 0.501 ± 0.093, and 47/48 learners have SNR > 1.The section reports lower within-domain variability for IRM, CORAL, and ViTs than for ERM and MLPs.Source: §H.1 · p. 32, Table 8 · p. 32

H.2 Circuit Fragility Test for OfficeHome and DomainNet

Across Office-Home and DomainNet, cross-domain circuit signal remains larger than extraction noise despite increasing benchmark difficulty. Noise correction improves CAS ranking correlation, with the largest reported gain on DomainNet.

H.2 Circuit Fragility Test for OfficeHome and DomainNet

  • Office-Home's mean extraction noise is 0.118 ± 0.027, while cross-domain dissimilarity is 0.540 ± 0.102, yielding mean SNR 3.58 and 46/48 learners with SNR > 1.The two exceptions are MLP-Mixer models trained with ERM.Source: §H.2 · p. 32
  • On Office-Home, noise correction raises CAS's Spearman correlation from ρS = 0.790 to ρS = 0.830.The reported objective-wise noise ordering is IRM (0.089), CORAL (0.108), DANN (0.119), then ERM (0.146).Source: §H.2 · p. 33
  • On DomainNet, mean cross-domain dissimilarity (∆obs = 0.559 ± 0.108) exceeds the noise floor by nearly 4×, while noise correction raises mean Spearman ρS from 0.725 to 0.775 (+0.050).Mean SNR is 2.94, with 44/48 learners above 1; all four failures were MLP-Mixer models trained with ERM.Source: §H.2 · p. 33

I Circuit Extraction Hyperparameters

Ablations of retained neurons, adapter bottleneck ratio, and edge threshold show interior optima for CAS-based OOD ranking. Optimal top-K grows with benchmark label-space size, while the preferred adapter ratio and edge threshold remain consistent.

I Circuit Extraction Hyperparameters

  • On PACS, CAS–OOD correlation is maximized at K = 30, r = 4, and ϵ = 10−4; settings on either side degrade as circuits become too sparse, dense, constrained, or noisy.The ablation reports CAS = 0.12 and ρS = 0.71 at K = 5, and ρS = 0.83 at K = 164.Source: I Circuit Extraction Hyperparameters · p. 34, Table 12 · p. 34
  • Office-Home's optimal top-K is 60, while r = 4 and ϵ = 10−4 remain preferred; ρS peaks at 0.82.The passage attributes the higher K than PACS to the larger 65-class label space.Source: I Circuit Extraction Hyperparameters · p. 34
  • DomainNet's optimal top-K is 164 with ρS = 0.79, while r = 4 and ϵ = 10−4 remain unchanged.Correlation declines to 0.77 at K = 200, and the passage reports that optimal edge counts scale approximately linearly with label cardinality.Source: I Circuit Extraction Hyperparameters · pp. 34–35

J CAS vs OOD generalization

Across PACS, Office-Home, and DomainNet, CAS tracks OOD accuracy more consistently than activation-similarity baselines. The relationship remains near-monotone across target domains, with benchmark-specific variability.

J CAS vs OOD generalization

  • PACS CAS correlations with leave-one-domain-out OOD accuracy are ρS = 0.91 on ArtPainting, 0.93 on Photo, and 0.88 on Sketch.CKA saturates at ≈0.70–0.88, while SVCCA and RSA show diffuse scatter and sometimes invert the ranking.Source: J CAS vs OOD generalization · p. 35
  • Office-Home CAS correlations with OOD accuracy are ρS = 0.81 on Art, 0.76 on Clipart, 0.82 on Product, and 0.77 on Real World.CKA saturates at ≈0.70–0.88; SVCCA and RSA yield inconsistent orderings.Source: J CAS vs OOD generalization · p. 35
  • On DomainNet, CAS preserves a monotone relationship with OOD accuracy across six targets despite accuracy ranging from 5%–70%.Quickdraw is the weakest within-objective cluster: IRM models differ by up to 15 points at similar CAS values.Source: J CAS vs OOD generalization · p. 35

K Circuit Minimality and TopK Selection

Circuit minimality follows a two-phase pattern: a small core sub-circuit captures most performance, while additional neurons yield diminishing returns. CAS remains stable near the task-level minimal circuit size.

K Circuit Minimality and TopK Selection

  • Circuit minimality retains the smallest sufficient neuron subset that preserves near-optimal accuracy, with performance following a consistent two-phase pattern across benchmarks.The K sweep is the basis for examining this pattern.Source: K Circuit Minimality and TopK Selection · p. 35
  • Accuracy rises steeply as small K recruits the core invariant sub-circuit, with elbow ranges of K ≈10–16 on MNIST, 25–40 on PACS, 50–75 on Office-Home, and K*=164 on DomainNet.The heavy-tailed importance distribution means a small fraction of neurons carries most causal signal.Source: K Circuit Minimality and TopK Selection · pp. 35–36
  • Beyond the elbow, adding neurons yields marginal gains of < 1–2% per +10 neurons, and peripheral nodes do not alter core topology.Their contributions are down-weighted by K and largely cancel in CAS, explaining its robustness once K* is reached.Source: K Circuit Minimality and TopK Selection · p. 36
  • DomainNet has a second inflection at K ≈80–120 as domain-bridging mid-level features are recruited before performance plateaus at K*.These features span extreme stylistic gaps such as Quickdraw versus Real.Source: K Circuit Minimality and TopK Selection · p. 36
  • All domains within a dataset reach near-optimal accuracy at the same K*, indicating task-level minimality.Evaluating CAS at {0.5K*, K*, 1.5K*, 2K*} changes ρS by < 0.03; below 0.5K* circuits lose discriminative structure, while above 2K* noise edges accumulate but kernel down-weighting keeps CAS stable.Source: K Circuit Minimality and TopK Selection · p. 36
Loading 2609.31996v1…