arXiv:2605.15183v1 [cs.LG] 14 May 2026
When Are Two Networks the Same? Tensor Similarity for Mechanistic Interpretability
ML Nissen Gonzalez∗ MARS & University Heidelberg
Melwina Albuquerque∗ MARS
Laurence Wroe∗ MARS
Jacob Meyer Cohen MARS & Stanford University
Logan Riggs Smith† Independent
Thomas Dooms† Independent
Abstract Mechanistic interpretability aims to break models into meaningful parts; verifying that two such parts implement the same computation is a prerequisite. Existing similarity measures evaluate either empirical behaviour, leaving them blind to outof-distribution mechanisms, or basis-dependent parameters, meaning they disregard weight-space symmetries. To address these issues for the class of tensor-based models, we introduce a weight-based metric, tensor similarity, that is invariant to such symmetries. This metric captures global functional equivalence and accounts for cross-layer mechanisms using an efficient recursive algorithm. Empirically, tensor similarity tracks functional training dynamics, such as grokking and backdoor insertion, with higher fidelity than existing metrics. This reduces measuring similarity and verifying faithfulness into a solved algebraic problem rather than one of empirical approximation.
1
Introduction
A central bottleneck to interpreting neural networks is validating whether the decomposed [Cunningham et al., 2023] or extracted [Kim et al., 2018, Bau et al., 2017] components preserve or capture the functional computation implemented by the original model. This is especially important for mechanistic interpretability, which aims to faithfully decompose models into constituent parts [Sharkey et al., 2025]. Such a decomposition must preserve not only the behaviour of the original network but also its internal functions. Yet current evaluations often rely entirely on input-based similarity such as comparing outputs or activations [Conmy et al., 2023, Marks et al., 2025a, Hanna et al., 2024]. This is insufficient: a decomposition can near-perfectly reconstruct activations by relying on approximate heuristics rather than the true underlying function [Olah, 2025, Ameisen et al., 2025]. The implication is that the resulting decomposition is not simply capturing an incomplete mechanism, but a different algorithm altogether. The root of this problem is the absence of a rigorous metric for mechanistic similarity. While defining such a metric is a matter of philosophical debate (when are two systems truly equal?) [Méloux et al., 2025], we argue a sensible similarity metric should satisfy three criteria: • Dataset independence: The metric should depend on the model, not on the sampled data. We can make stronger robustness claims when they apply to the full distribution, such as ruling out backdoors [Marks et al., 2025b], which by design trigger only on particular inputs. • Symmetry invariance: The metric should be invariant under transformations that leave the function unchanged. Whether a mechanism is implemented in one neuron or spread across ∗ Equal contribution. † Equal senior authorship. Correspondence to [email protected] and [email protected].
Preprint.
Figure 1: Overview of our method. The model learns a backdoor that existing similarity metrics miss when given only clean data or weights. Tensor similarity flags the change and localises it to the attacked class. ∆ measures how sharply similarity shifts from pre- to post-backdoor (subsection 3.3).
many, or distributed across layers rather than concentrated in one [Ameisen et al., 2025], should have no effect on the score. • Tractable scalability: The metric should be efficiently computable at scale, without exponential or quadratic convergence bounds (as Monte Carlo sampling incurs). Only then can it serve as an optimisation target for the largest networks. Current methods fail at least one criterion. For instance, behavioural similarity [Kornblith et al., 2019, Raghu et al., 2017, Kriegeskorte et al., 2008] measures outputs or activations empirically by comparing input-output pairs, but is not guaranteed to generalise out of distribution. It will therefore fail to detect backdoors, which are designed precisely to resist such measurements. Another common metric, matrix similarity, uses the existing parametrisation as a comparison point, making it highly susceptible to model permutations. It might succeed at studying training dynamics (where an update doesn’t arbitrarily reshuffle weights) but fail across differently trained models. In this paper, we propose a weight-based similarity that satisfies all three criteria (Figure 2), at the cost of using tensor-based neural networks [Pearce et al., 2025, Dooms et al., 2025, Riggs, 2026, Shazeer, 2020]. These networks train and function equivalently to ordinary deep networks like transformers, but use different activation functions. This trade-off is fundamental: contemporary neural networks do not appear to admit a similarity metric that tractably satisfies all three criteria. Our metric extends cosine similarity from vectors to tensor-based models, formally verifying whether two sub-networks execute identical internal logic regardless of differences in parametrisation. It decays gracefully under minor structural perturbations and, because it is rooted in linear algebra, we can attribute the contribution of each component to the total. The remainder of this paper has two parts: discussing the theory and implementation of tensorbased similarity, and validating its discriminative power empirically across a spectrum of domains. Theoretically, we cover the architectural changes required for the metric to be computable and the construction that makes it invariant under weight symmetries. Empirically, we show that tensor-based similarity exposes functional changes that other methods miss: catastrophic forgetting in vision, grokking-like phase transitions in modular arithmetic [Nanda et al., 2023], and training dynamics in 2
Figure 2: Matrix similarity computes the direct cosine similarity between (flattened) weight matrices, but is sensitive to weight-space symmetries such as permutation and scaling, severely limiting it as an accurate proxy for functional equivalence, especially for models that are trained differently. Behavioural similarity (or output similarity) mitigates this by comparing model outputs directly, making it invariant to weight-space symmetries. It depends, however, on an input distribution and is consequently blind to out-of-distribution artifacts. (Activation-based similarity, a tempting alternative, is sensitive to both the distribution and weight symmetries like cross-layer structure.) Tensor similarity is a closed-form, data-free measure of functional similarity between two models. It computes the cosine angle between their weight tensors in L2 (N (0, I)); the function space under the Gaussian metric. Tensor similarity satisfies all criteria at the cost of requiring a modified architecture. language [Hoogland et al., 2025, Olsson et al., 2022]. Together, this positions tensor similarity as a foundational tool for mechanistic interpretability3 .
2
Theory
This paper studies polynomial (or multilinear) variants of standard neural components that use multiplicative gating rather than ordinary non-linear activation functions [Dauphin et al., 2017]. In this section, we first review relevant prior work and introduce the underlying architectures, then demonstrate how the proposed similarity measure handles various architectural symmetries. The discussion assumes a fundamental familiarity with tensor algebra and the symmetric group; additional details and derivations are provided in Appendix A. 2.1
Defining multilinear transformer models
We build on literature that uses bilinear and tensor-based architectures as intrinsically interpretable models. These works show that polynomial models enable formal weight-based analysis via the learned weight tensors using bilinear MLPs [Pearce et al., 2025, Sharkey et al., 2025] (or stacks thereof [Dooms et al., 2025]) while remaining competitive with standard architectures [Pearce et al., 2025, Riggs, 2026, Shazeer, 2020]. More broadly, this line of work sits within a longer tradition of using tensor decompositions to study expressiveness of deep learning [Cohen et al., 2015]. Bilinear layers. The core building block of our architecture is a GLU [Shazeer, 2020, Dauphin et al., 2017] without gate, called a bilinear layer. It maps a lifted input, i.e. x̃ = (1, x) ∈ Rd+1 to an output in RK as A(x) = D (Lx̃) ⊙ (Rx̃) , (1) where L, R ∈ Rr×(d+1) and D ∈ RK×r are weight matrices, r is the hidden dimension, and ⊙ denotes the element-wise (Hadamard) product. The above computation can be expressed as a tensor Pd Pr contraction, A(x)i = j,k=0 h=1 Dih Lhj Rhk x̃j x̃k , and is degree-2 polynomial in x. Bilinear attention. The attention mechanism can be made multilinear by replacing the softmax attention pattern with the elementwise product of two linear attention patterns. Concretely, for a sequence of embedded tokens X ∈ Rseq×d the attention mechanism is defined as ⊤ ⊤ A(X) = (X⊤ Q⊤ L KL X) ⊙ (X QR KR X), 3 Our code is available at https://github.com/tdooms/tensor-similarity
3
(2)
where Q and K are the query and key weight matrices, respectively. Combined with a causal mask and rotary positional embeddings [Riggs, 2026], bilinear attention is a degree-5 polynomial in X. Bilinear transformers. We can combine bilinear layers and bilinear attention with a residual stream into bilinear transformers, which perform comparably to standard transformers on natural language tasks [Dooms et al., 2025, Pearce et al., 2025, Sharkey, 2023, Riggs, 2026]. Multilinear models. Equation (1) and Equation (2) can be equivalently expressed as a single contraction of a weight tensor A ∈ RK ⊗n Rd+1 with n copies of the lifted input x̃ ∈ Rd+1 as A(x) = A·in ⊗n x̃ ∈ RK , which is a polynomial of order n in x. Models of this form are multilinear models. Deep multilinear models. Multilinear models can be layered to tackle the curse of dimensionality, analogously to depth in standard neural networks. While the global weight-tensor A of a layered multilinear model A = ◦li=1 Ai can be too expensive to be instantiated for deep networks, the global tensor decomposes into local tensors according to a tree graph with shared local tensors Ai . This lowers the memory required to store the tensor from O(dn ) to O(log n) without losing access to functions of the global tensor4 . 2.2
Multilinear models enable global weight-based metrics
Having defined a multilinear transformer within the framework of deep multilinear models, we now leverage its algebraic properties for a global and tractable weight-based similarity measure. P Tensor similarity. The Frobenius inner product, ⟨A | B⟩ = Aki1 ···in Bki1 ···in , can be used to compare two multilinear models A and B with weight tensors A and B that carry the full functional information. Tensor similarity uses a slightly generalised version of this inner product by inserting a metric M (symmetric and positive-definite over RK ⊗n Rd+1 ) and dividing by the induced norms as simM (A, B) =
⟨A | M | B⟩ ∈ [−1, 1]. ∥A∥M ∥B∥M
(3)
The choice of the metric M determines which directions in weight space are considered relevant. Tractable inner product. Before we constrain the specific form of the metric, we first observe how the stratification of the global tensor of a deep multilinear model allows us to compute inner products between multilinear models efficiently. Otherwise, naïve computation of the inner product between two deep networks would require materialising the expensive global weight tensor. Gram-based recursion [Dooms et al., 2025] avoids this by processing each layer’s local tensors sequentially, without ever materialising the global tensor. Here, inner products are computed between each layer’s tensors using metric tensors that depend on the inner product of the previous layer. Tracing the last Gram matrix of this recursion recovers the inner product. Crucially, the global metric tensor M has to preserve the tree shape of the global tensor assumed by this recursion4 . 2.3
Projecting out the weight-based symmetries induces principled metrics
The aim is to determine a matrix M that leads to an induced tensor similarity which tracks the equivalence of two models, defined as: Two models are equivalent if and only if they implement functions related simply by a positive proportionality factor (i.e. A ∼ = B ⇐⇒ ∃λ > 0 : A = λB). Such an M should satisfy simM (A, B) = 1 ⇐⇒ ∃ λ > 0 : A = λB. That is, it must be invariant under weight-space transformations mapping the parameters of one model to the parameters of an equivalent model. As a first step towards a candidate metric, we identify the basic symmetry group of multilinear models. We lead with an example. 4 More details on deep multilinear models and related restrictions for the metric tensor are discussed in Appendix A.
4
Example: Let’s decompose a bilinear layer into a symmetric and antisymmetric component: Aijk = Pr sym/asym sym = 21 (Dih Lhj Rhk ±Dih Rhj Lhk ). Consider +Aasym , where Aijk h=1 Dih Lhj Rhk = A Pn asym the contribution of the antisymmetric part j,k=0 Aijk x̃j x̃k . Swapping j ↔ k changes the sign of Aasym but leaves x̃j x̃k and therefore the result invariant. The contribution to the result of the antisymmetric part is therefore zero. Thus, A(x) = Asym ·in (x̃ ⊗ x̃) for all x. Two bilinear layers with different (L, R) but equal Asym implement the same quadratic polynomial. Symmetrisation tracks functional equivalence by polarisation. The principle demonstrated in the example above generalises to polynomials of any order by the polarisation isomorphism [Ramshaw, 1989]. Denoting the group of permutations of n elements by Sn , a tensor A ∈ ⊗n Rd is said to be symmetric if it is invariant under the permutation of its indices: ∀σ ∈ Sn : Ai1 ,...,in = Aiσ(1) ,...,iσ(n) . Polarisation establishes a one-to-one correspondence between the symmetric tensors of order n, and polynomials of the same order in x4 . In other words, two multilinear models with different tensors are equivalent if and only if these tensors have the same representative in the symmetric subspace. The projection onto the symmetric subspace is called symmetrisation and denoted with Pn . With this, we restrict this analysis to metric tensors of the form M = 1 ⊗ Λ with Pn ΛPn = Λ. Such metric tensors are invariant under the action of Sn × Sn on its 2n indices, with the simplest choice here Λ = Pn . Intuitively, using Pn as a metric in input space computes the inner product of two tensors by first projecting them onto the symmetric subspace and then computing the normal Frobenius inner product there. Combining symmetrisation to invoke polarisation with the Cauchy–Schwarz inequality gives the central guarantee of tensor similarity: simPn (A, B) = 1 ⇐⇒ ∃ λ > 0 : A = λB.
(4)
Symmetry of behavioural similarity. While metric tensors of the form given by Equation 4 track functional equivalence, they can be further restricted to be even more useful. Consider the symmetries of the behavioural similarity computation by re-expressing the inner product between two outputs of multilinear models as ⟨A(x) | B(x)⟩ = ⟨A ·in ⊗n x̃ | B ·in ⊗n x̃⟩ = ⟨A | ⊗2n x̃ | B⟩. The input tensor ⊗2n x lives in the same space as Λ and is invariant under permutations of its 2n indices. By definition, it is therefore invariant under the action of the group S2n , and Sn × Sn is a subgroup of S2n . Thus the space of metrics M = 1 ⊗ Λ that are invariant under S2n is a subspace of metrics that are invariant under the action of Sn × Sn . Equation 4 therefore also applies to metrics M = 1 ⊗ Λ with P2n · Λ = Λ. However, unlike Pn which has order 2n, P2n cannot be taken as an input space metric because its order 4n is too high. This motivates the search for the Λ used in this paper by considering the Gaussian expectation value of ⊗2n x. Gaussian metric. Consider Λ = Ex∼N (0,I) ⊗2n x which has S2n symmetry. Computed using Isserlis’ theorem [Munthe-Kaas et al., 2025], it evaluates to a sum over pairwise contractions of all 2n input indices of A and B, including “self-contractions”, contracting two indices of the same tensor as ⌊n/2⌋
Λ=
X
cn,m (τ m Pn )† (τ m Pn ), with Λi1 ···i2n ∝
m=0
X
Y
δia ib ,
(5)
pairings (a,b)∈pairing
where τ m denotes the partial trace operator on m pairs of indices (which need not be specified as the tensors are symmetrised). Thus, the expected inner product of the activations of two multilinear models under a Gaussian input is equal to their weight-space inner product.5 2.4
Applying tensor similarity locally isolates components
Tensor similarity (Equation 3) with the Gaussian metric (Equation 5) is a closed form, data-free measure of functional equivalence between two models. It can be further localised to focus on specific output dimensions or differences, both of which inherit the invariance guarantees of Equation 4. 5We can compute expectation values of inner products under Gaussian distribution entirely in weight space. This result is
explained and derived in more detail in Appendix A.
5
Tensor slice similarity. Fixing an output index k yields a slice Ak ∈ ⊗n Rd+1 of the full weight tensor, capturing the weights responsible for a specific output dimension. Since M factorises as 1 ⊗ Λ, the induced metric on this slice is simply Λ defining tensor slice similarity. Tensor slice similarity focuses on functional agreement in one specific output dimension. Tensor diff similarity. The space of weight tensors is a vector space, the difference ∆ = B−C (diff ) is itself a valid tensor, and so tensor similarity of A with ∆ is well defined. Tensor diff similarity isolates similarity with respect to the functional change captured by a diff.
3
Results
Tensor similarity satisfies all desired criteria for interpretability in theory; how accurately does this translate into practice? We test four capabilities in turn: localising ground-truth changes to specific outputs (catastrophic forgetting on SVHN), tracking continuous reorganisation through training (grokking on modular addition), catching out-of-distribution changes that other metrics cannot see (backdoor injection on SVHN), and scaling all of this to language models (a two-layer bilinear attention transformer on The Pile). 3.1
Tensor similarity localises mechanism changes to specific outputs
Per-class slices pinpoint which output is responsible for a ground-truth mechanism change, not just that one occurred. We demonstrate this with continued training on SVHN [Netzer et al., 2011], a harder version of MNIST [Deng, 2012], to induce catastrophic forgetting. Training begins with a base stage on the digit subset {0, · · · , 4}, followed by stages which incrementally introduce one digit at a time, culminating in the add 9 stage on the full set {0, · · · , 9}. We further stress-test the metric’s sensitivity to dataset changes by appending three stages: control (continued training on the full set, disentangling the effect of a learning rate restart from that of any dataset change), remove 9 (training on {0, · · · , 8} to induce catastrophic forgetting), and re-add 9 (returning to the full set), as shown in Figure 3. Further setup details and additional plots are reported in Appendix B. 3.2
Tensor similarity tracks continuous reorganisation through training
Grokking is not a single phase change but a continuous reorganisation into particular Fourier frequencies. We demonstrate this by training a single bilinear layer on modular addition following Nanda et al. [2023], where the Fourier algorithm is well-established. We save K checkpoints throughout training and compute their pairwise tensor similarity to produce a K × K matrix, with further setup details reported in Appendix C. The accuracy curves in Figure 4 exhibit standard grokking: training saturates immediately, followed by validation jumping from chance to perfect (top). The similarity matrix (bottom) shows the reorganisation continuing beyond this jump: frequencies (third row) stabilise only near the end of training, when the loss also reaches its floor (second row). Tensor similarity makes this development legible without having to reverse engineer the whole model. 3.3
Tensor similarity catches out-of-distribution changes that other metrics miss
Behavioural metrics evaluate similarity on a fixed input distribution, so a mechanism active only out of distribution does not affect their score. Tensor similarity is data-free and so detects it. We inject a backdoor into SVHN training: any image stamped with a small black diamond is relabelled as digit ‘9’ (Figure 1). In this injection phase, 10% of training samples contain trigger patch in the upper-right corner, with target class 9. The remainder of the setup is equivalent to subsection 3.1, detailed in Appendix B. The model learns the backdoor at 96% attack success while train and test accuracy on clean data remain at 85% and 81%, leaving the attack invisible to standard performance metrics. We summarise each similarity matrix by ∆, the average within-block similarity minus the average across-block similarity for the pre- and post-backdoor checkpoints; higher values indicate a sharper split between the two regimes. Behavioural similarity on clean inputs remains flat (∆ = 0.09), reflecting that the attack does not register on the clean distribution. On poisoned inputs the signal returns (∆ = 0.62), but this requires knowing the trigger in advance. Matrix cosine on the weights picks up a weak shift (∆ = 0.15) yet 6
Figure 3: We report four similarity measures across the training trajectory of an SVHN model. The top row compares whole-model summaries: global tensor similarity reveals a block structure separating stages with and without digit ‘9’, while CKA on the logits and matrix similarity show much weaker, more ambiguous signal. The per-class tensor slices sharpen the picture: each exhibits exactly the block pattern predicted by when that digit was present in training, and the class ‘9’ slice groups base through add 8 with remove 9, all orthogonal to add 9, control, and re-add 9. Tensor similarity does not merely detect functional changes but can precisely localise them to the output.
collapses everything into a scalar, with no handle on which input or output is responsible. Tensor similarity flags the backdoor from clean data alone (∆ = 0.24), and the digit-‘9’ slice sharpens this to ∆ = 0.43: the same localisation pattern we saw under forgetting.
3.4
Tensor similarity scales to language models
The previous results used 1-layer bilinear models. Tensor similarity scales unchanged to deeper transformers and continues to find structure where existing metrics blur. We train a two-layer bilinear attention transformer on The Pile [Gao et al., 2020] and save 101 log-spaced checkpoints, with further setup details reported in Appendix D. Figure 5 reports five similarity measures across these checkpoints alongside the n-gram behaviour of the model, which prior work uses to track the increasingly complex structures models acquire over training, often discontinuously [Wang et al., 2024, Nguyen, 2024, Belrose et al., 2024]. Existing metrics separate only the earliest checkpoints from the rest, blurring everything beyond; tensor similarity exhibits clear block structure with sharp transitions between regimes. More broadly, behavioural metrics see only the slice of input space they are evaluated on, and a mechanism that has reorganised globally need not register on that slice if its predictions there happen to coincide [D’Amour et al., 2020]. What generalises is what holds beyond the training distribution; measuring similarity on the training distribution conflates it with memorisation. Tensor similarity integrates over the full Gaussian input space and so registers reorganisations that would be invisible to any finite probe, which is precisely the regime where mechanistic claims need to hold. 7
Figure 4: Modular addition training tracked by accuracy and loss (top), the Fourier components of the embeddings (middle), and pairwise tensor similarity between checkpoints (bottom). Three diagonal blocks in the similarity matrix correspond to initialisation, memorisation, and the converged solution.
4
Conclusion
We introduce tensor similarity, a principled metric for evaluating functional equivalence from the weights of tensor-based models such as bilinear transformers. We then apply our metric to problems in mechanistic and developmental interpretability. Summary. Measuring functional similarity from weights requires accounting for the symmetries that leave a model’s outputs invariant. Tensor similarity projects these out; combined with the polarisation isomorphism, the result is a principled measure of functional equivalence. Exploiting the link to behavioural similarity under Gaussian inputs, we further derive a Gaussian metric that yields a stricter equivalence relation. We validate the framework on four tasks: catastrophic forgetting in vision, backdoor insertion via dataset poisoning, grokking in modular arithmetic, and the formation of n-gram statistics during language modelling on The Pile. Across all four tasks, tensor similarity tracks functional changes more cleanly than existing metrics, and slicing localises that change to specific outputs or components. Limitations. The primary limitation is the reliance on multilinear models. This excludes most state-of-the-art architectures where non-polynomial non-linearities between components block the algebraic computation of functional properties from weights, particularly for cross-layer mechanisms. This limitation bites hardest if multilinear models cannot match standard architectures in performance, leaving no incentive to adopt them at scale. Even in that case, we believe multilinear models earn their place as test beds for mechanistic interpretability — they form a non-trivial class of architectures that admits effective algebraic analysis — and regardless, there is mounting evidence that multilinear models can be in the same league as state-of-the-art architectures [Pearce et al., 2025, Shazeer, 2020]. A secondary limitation is the limited scope of the analysis conducted here. While theory and the diff construction gesture towards attribution, our experiments demonstrate detection of functional changes, such as backdoors, more thoroughly than they demonstrate attribution of the specific model components that cause them. Outlook. A natural next step toward the goal of attribution, a key piece of the ambitious mechanistic interpretability agenda, is local application: computing tensor similarity between paths through a deep multilinear transformer and extending the framework of [Elhage et al., 2021]. This could include bilinear MLPs and reveal mechanisms beyond the two-layer attention-only case that originally produced induction heads. The flipside of the restriction to multilinear models is that this work adds to a growing literature showing they are much more transparent than ordinary deep networks. Their symmetries are well understood, polarisation identifies a function from its weights alone, and norms 8
Figure 5: Five similarity measures across 101 log-spaced checkpoints of a two-layer bilinear attention transformer on The Pile. We track the n-gram score as a heuristic measure of change, which shows continued changes throughout most of training. Both the Pile · residual and the Pile · logits are behavioural similarities on pre-unembedding activations and on logits over Pile tokens; Gaussian · logits is the same on logits driven by matched-σ Gaussian inputs. Tensor is the closed-form weightbased similarity; Matrix cosine is Frobenius cosine of flattened weights, as discussed in Figure 2. and correlation measures can be traced globally through the weights [Levine et al., 2019]. Combined with evidence that multilinear models match non-linear ones on standard benchmarks [Pearce et al., 2025, Cohen and Shashua, 2016, Riggs, 2026], the case builds for multilinear architectures as a foundation for trustworthy machine learning.
Acknowledgements This work was supported by the MARS (Mentorship for Alignment Research Students) program run by the Cambridge AI Safety Hub (caish.org/mars). We are grateful to the people of CAISH for making our collaboration possible, and to our research manager, Mikhail Mironov, for ensuring we didn’t veer too far off track. Additionally, we thank Ward Gauderis, Richard Mohn, and Viktor Rehnberg for technical discussions and encouragement.
References Emmanuel Ameisen, Jack Lindsey, Adam Pearce, Wes Gurnee, Nicholas L. Turner, Brian Chen, Craig Citro, David Abrahams, Shan Carter, Basil Hosmer, Jonathan Marcus, Michael Sklar, Adly Templeton, Trenton Bricken, Callum McDougall, Hoagy Cunningham, Thomas Henighan, Adam Jermyn, Andy Jones, et al. Circuit tracing: Revealing computational graphs in language models. Transformer Circuits Thread, 2025. URL https://transformer-circuits.pub/ 2025/attribution-graphs/methods.html. David Bau, Bolei Zhou, Aditya Khosla, Aude Oliva, and Antonio Torralba. Network dissection: Quantifying interpretability of deep visual representations. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 6541–6549, 2017. URL https://arxiv. org/abs/1704.05796. Nora Belrose, Quintin Pope, Lucia Quirke, Alex Mallen, and Xiaoli Fern. Neural networks learn statistics of increasing complexity, 2024. URL https://arxiv.org/abs/2402.04362. 9
Nadav Cohen and Amnon Shashua. Convolutional Rectifier Networks as Generalized Tensor Decompositions, 2016. URL https://arxiv.org/abs/1603.00162v2. Nadav Cohen, Or Sharir, and Amnon Shashua. On the Expressive Power of Deep Learning: A Tensor Analysis, 2015. URL https://arxiv.org/abs/1509.05009v3. Arthur Conmy, Augustine N. Mavor-Parker, Aengus Lynch, Stefan Heimersheim, and Adrià GarrigaAlonso. Towards automated circuit discovery for mechanistic interpretability. In Advances in Neural Information Processing Systems, 2023. URL https://arxiv.org/abs/2304.14997. Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoencoders find highly interpretable features in language models, 2023. URL https://arxiv.org/ abs/2309.08600. Alexander D’Amour, Katherine Heller, Dan Moldovan, Ben Adlam, Babak Alipanahi, Alex Beutel, Christina Chen, Jonathan Deaton, Jacob Eisenstein, Matthew D. Hoffman, Farhad Hormozdiari, Neil Houlsby, Shaobo Hou, Ghassen Jerfel, Alan Karthikesalingam, Mario Lucic, Yian Ma, Cory McLean, Diana Mincu, Akinori Mitani, Andrea Montanari, Zachary Nado, Vivek Natarajan, Christopher Nielson, Thomas F. Osborne, Rajiv Raman, Kim Ramasamy, Rory Sayres, Jessica Schrouff, Martin Seneviratne, Shannon Sequeira, Harini Suresh, Victor Veitch, Max Vladymyrov, Xuezhi Wang, Kellie Webster, Steve Yadlowsky, Taedong Yun, Xiaohua Zhai, and D. Sculley. Underspecification presents challenges for credibility in modern machine learning, 2020. URL https://arxiv.org/abs/2011.03395. Yann N. Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks, 2017. URL https://arxiv.org/abs/1612.08083. Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 29(6):141–142, 2012. Thomas Dooms, Ward Gauderis, Geraint A. Wiggins, and José Oramas. Compositionality unlocks deep interpretable models, 2025. URL https://arxiv.org/abs/2504.02667. Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https://transformer-circuits.pub/2021/framework/index.html. Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The pile: An 800gb dataset of diverse text for language modeling, 2020. URL https://arxiv.org/abs/ 2101.00027. Johnnie Gray. quimb: a python library for quantum information and many-body calculations. Journal of Open Source Software, 3(29):819, 2018. doi: 10.21105/joss.00819. Michael Hanna, Sandro Pezzelle, and Yonatan Belinkov. Have faith in faithfulness: Going beyond circuit overlap when finding model mechanisms, 2024. URL https://arxiv.org/abs/2403. 17806. Jesse Hoogland, George Wang, Matthew Farrugia-Roberts, Liam Carroll, Susan Wei, and Daniel Murfet. Loss landscape degeneracy and stagewise development in transformers, 2025. URL https://arxiv.org/abs/2402.02364. Been Kim, Martin Wattenberg, Justin Gilmer, Carrie Cai, James Wexler, Fernanda Viégas, and Rory Sayres. Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (TCAV). In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 2668–2677. PMLR, 2018. URL https://arxiv.org/abs/1711.11279. 10
Simon Kornblith, Mohammad Norouzi, Honglak Lee, and Geoffrey Hinton. Similarity of neural network representations revisited. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3519–3529. PMLR, 2019. URL https://arxiv.org/abs/1905.00414. Nikolaus Kriegeskorte, Marieke Mur, and Peter Bandettini. Representational similarity analysis — connecting the branches of systems neuroscience. Frontiers in Systems Neuroscience, 2:4, 2008. doi: 10.3389/neuro.06.004.2008. URL https://doi.org/10.3389/neuro.06.004.2008. Yoav Levine, Or Sharir, Nadav Cohen, and Amnon Shashua. Quantum Entanglement in Deep Learning Architectures. Phys. Rev. Lett., 122(6):065301, 2019. doi: 10.1103/PhysRevLett.122. 065301. URL https://link.aps.org/doi/10.1103/PhysRevLett.122.065301. Samuel Marks, Can Rager, Eric J. Michaud, Yonatan Belinkov, David Bau, and Aaron Mueller. Sparse feature circuits: Discovering and editing interpretable causal graphs in language models. In International Conference on Learning Representations, 2025a. URL https://arxiv.org/abs/ 2403.19647. Samuel Marks, Johannes Treutlein, Trenton Bricken, Jack Lindsey, Jonathan Marcus, Siddharth Mishra-Sharma, Daniel Ziegler, Emmanuel Ameisen, Joshua Batson, Tim Belonax, Samuel R. Bowman, Shan Carter, Brian Chen, Hoagy Cunningham, Carson Denison, Florian Dietz, Satvik Golechha, Akbir Khan, Jan Kirchner, Jan Leike, Austin Meek, Kei Nishimura-Gasparian, Euan Ong, Christopher Olah, Adam Pearce, Fabien Roger, Jeanne Salle, Andy Shih, Meg Tong, Drake Thomas, Kelley Rivoire, Adam Jermyn, Monte MacDiarmid, Tom Henighan, and Evan Hubinger. Auditing language models for hidden objectives, 2025b. URL https://arxiv.org/abs/2503. 10965. Hans Z. Munthe-Kaas, Olivier Verdier, and Gilles Vilmart. A short proof of Isserlis’ theorem, 2025. URL http://arxiv.org/abs/2503.01588. Maxime Méloux, Silviu Maniu, François Portet, and Maxime Peyrard. Everything, everywhere, all at once: Is mechanistic interpretability identifiable?, 2025. URL https://arxiv.org/abs/2502. 20914. Neel Nanda, Lawrence Chan, Tom Lieberum, Jess Smith, and Jacob Steinhardt. Progress measures for grokking via mechanistic interpretability, 2023. URL https://arxiv.org/abs/2301.05217. Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y. Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning 2011, 2011. URL http://ufldl.stanford.edu/ housenumbers/nips2011_housenumbers.pdf. Timothy Nguyen. Understanding transformers via n-gram statistics, 2024. URL https://arxiv. org/abs/2407.12034. Chris Olah. A toy model of mechanistic (un)faithfulness. Transformer Circuits Thread, 2025. URL https://transformer-circuits.pub/2025/faithfulness-toy-model/index.html. Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads. Transformer Circuits Thread, 2022. URL https://transformer-circuits.pub/2022/ in-context-learning-and-induction-heads/index.html. Michael T. Pearce, Thomas Dooms, Alice Rigg, José M. Oramas, and Lee Sharkey. Bilinear MLPs enable weight-based mechanistic interpretability. In International Conference on Learning Representations, 2025. URL https://arxiv.org/abs/2410.08417. Maithra Raghu, Justin Gilmer, Jason Yosinski, and Jascha Sohl-Dickstein. SVCCA: Singular vector canonical correlation analysis for deep learning dynamics and interpretability. In Advances in Neural Information Processing Systems, 2017. URL https://arxiv.org/abs/1706.05806. 11
Lyle Ramshaw. Blossoms are polar forms. Computer Aided Geometric Design, 6(4):323–358, 1989. ISSN 0167-8396. doi: 10.1016/0167-8396(89)90032-0. URL https://www.sciencedirect. com/science/article/pii/0167839689900320. Logan Riggs. Tensor-transformer variants are surprisingly performant. LessWrong, 2026. URL https://www.lesswrong.com/posts/hp9bvkiN3RzHgP9cq/ tensor-transformer-variants-are-surprisingly-performant. Lee Sharkey. A technical note on bilinear layers for interpretability, 2023. URL http://arxiv. org/abs/2305.03452. Lee Sharkey, Bilal Chughtai, Joshua Batson, Jack Lindsey, Jeff Wu, Lucius Bushnaq, Nicholas Goldowsky-Dill, Stefan Heimersheim, Alejandro Ortega, Joseph Bloom, Stella Biderman, Adria Garriga-Alonso, Arthur Conmy, Neel Nanda, Jessica Rumbelow, Martin Wattenberg, Nandi Schoots, Joseph Miller, Eric J. Michaud, Stephen Casper, Max Tegmark, William Saunders, David Bau, Eric Todd, Atticus Geiger, Mor Geva, Jesse Hoogland, Daniel Murfet, and Tom McGrath. Open problems in mechanistic interpretability, 2025. URL https://arxiv.org/abs/2501. 16496. Noam Shazeer. Glu variants improve transformer, 2020. URL https://arxiv.org/abs/2002. 05202. George Wang, Jesse Hoogland, Stan van Wingerden, Zach Furman, and Daniel Murfet. Differentiation and specialization of attention heads via the refined local learning coefficient, 2024. URL https: //arxiv.org/abs/2410.02984. Sang Michael Xie, Shibani Santurkar, Tengyu Ma, and Percy Liang. Data selection for language models via importance resampling. arXiv preprint arXiv:2302.03169, 2023.
12
A
Theory appendix
A.1
Multilinear models and their symmetries.
Multilinear models. Multilinear models take the following form: A(x) = A · ⊗n x̃ where x ∈ Rd , x̃ = (1, x), A ∈P RK ⊗n Rd+1 , and · denotes the tensor contraction along ⊗n Rd+1 . In component form: A(x)k = i∈[d]n Aki1 ...in xi1 · · · xin , which is a polynomial of order n in the components of 0 x. A thus models K ∈ N polynomials. Deep learning with multilinear models. The order of the polynomial governs the expressiveness of multilinear models. However, the size of the tensor A grows exponentially with order as O(K(d + 1)n ): this is the curse of dimensionality. To tackle this with multilinear models, we layer them. Consider a deep multilinear model with l ∈ N layers: A(x) = A(l) (A(l−1) · · · (A(1) (x)) · · · ) = A(l) · ⊗nl (A(l−1) · · · ⊗n2 (A(1) · ⊗n1 x̃)), where ni ∈ N denotes the input order of the local tensor A(i) ∈ RHi ⊗ni RHi−1 for i ∈ [l] and {Hi } the hidden dimensions with H0 = d + 1 and Hl = K. While the global tensor has order Ql n = 1 + i=1 ni , only local tensors of order ni need to be stored and this therefore constitutes a tensor decomposition of the global tensor A. Tensor networks are decompositions of tensors according to an underlying graph. The graph that underlies the decomposition of a deep multilinear model has tree shape, with the tensors within each layer being copies of each other. Thus, while the order of the modelled tensor is exponential in the number of layers, the number of coefficients that need to be stored in the local tensors grows only linearly with the number of layers. Action of the permutation group. Consider a multilinear model given by the tensor A ∈ RK ⊗n Rd+1 and the symmetric group on n elements, Sn . Its action on RK ⊗n Rd+1 permutes the copies ⊗n Rd+1 as Aki1 ...in → Akiσ(1) ...iσ(n) . Tensors that are invariant under this action are said to be symmetric. The space of these symmetric tensors, RK ⊗ symn Rd+1 , is a proper vector subspace of the vector space of generic tensors. The projection down to the symmetric subspace, called symmetrisation and denoted by Pn : ⊗n Rd+1 → symn Rd+1 , is a linear orthogonal projection defined as the mean of all the actions of the symmetric group: 1 X Asym := (1 ⊗ Pn ) · A, Asym Akiσ(1) ···iσ(n) . (6) ki1 ···in = n! σ∈Sn
The orthogonal projector Pn satisfies P2n = Pn and ⟨Pn A | B⟩ = ⟨A | Pn B⟩. Polarisation. This symmetry leaves the function A invariant, and so permuting the input copies leaves the modelled polynomial invariant. The precise statement of this is the polarisation isomorphism, usually stated for homogeneous polynomials. It states that there is a vector space isomorphism between the symmetric tensors over Rd of order n, symn Rd , and the space of homogeneous polynomials of order n in x ∈ Rd . In our case, because we consider general (not merely homogeneous) polynomials, we lift the input to x̃ = (1, x) and work with degree-n homogeneous polynomials over Rd+1 ; by polarisation these are in bijection with symn Rd+1 . This applies for each slice Ak ∈ symn Rd+1 of the multilinear model considered. The output space RK has a privileged basis such that slicing along it is principled. For a single-layer multilinear model, therefore, symmetrising the input indices of the tensor A puts it into a one-to-one correspondence with the vector of polynomials modelled, accounting for all weight-space symmetries. Example: bilinear layer. A bilinear layer is a single-layer multilinear model of order n = 2, so its weight tensor A ∈ RK ⊗2 Rd+1 encodes K quadratic polynomials. Given matrices D ∈ RK×r and L, R ∈ Rr×(d+1) with hidden dimension r ∈ N, and the order-3 delta tensor δ ∈ Rr ⊗ Rr ⊗ Rr with components δh1 h2 h3 = 1 iff h1 = h2 = h3 and zero otherwise, the weight tensor is r r X X Aki1 i2 = Dkh1 δh1 h2 h3 Lh2 i1 Rh3 i2 = Dkh Lhi1 Rhi2 . (7) h1 ,h2 ,h3 =1
h=1
13
Pr This is a rank-r CP decomposition, equivalently written A = h=1 dh ⊗ lh ⊗ rh , where dh , lh , rh are the h-th columns of D and rows of L, R respectively. The forward pass reduces to A(x̃) = D (Lx̃) ⊙ (Rx̃) , (8) where ⊙ denotes the Hadamard (elementwise) product, computable with matrix operations alone. The output depends only on the symmetric part of A. Decompose A = Asym + Aasym , where 1 A⊥ ki1 i2 = 2 (Aki1 i2 − Aki2 i1 ) is the antisymmetric complement. Since x̃ ⊗ x̃ is symmetric under i1 ↔ i2 , the contraction with Aasym vanishes for all x̃: X asym X Aki1 i2 x̃i1 x̃i2 = 12 (Aki1 i2 − Aki2 i1 )x̃i1 x̃i2 = 0. i1 ,i2 sym
Therefore A(x̃) = A
i1 ,i2
·in (x̃ ⊗ x̃), and in terms of the CP factors, 1 Asym ki1 i2 = 2
r X
Dkh Lhi1 Rhi2 + Rhi1 Lhi2 .
(9)
h=1
Two bilinear layers with different (L, R) but equal Asym implement the same polynomial; weightbased comparison is therefore only principled after symmetrisation, as established by the polarisation isomorphism above. The symmetry of deep multilinear models Layering multilinear models introduces an additional subtlety. It is too expensive to instantiate the full tensor A that decomposes according to a tree into the local tensors A(i) for i ∈ [l]. However, full permutations of all input indices of the global tensor would require instantiating the full tensor. Moreover, these full permutations would destroy the tree decomposition structure. Because leaving the decomposition space is computationally intractable, we do not consider the action of the full symmetric group. Restricting to these tree decompositions of the global tensor, we identify the symmetry group that is compatible with the tree decomposition of A to be the wreath product of local permutations: Gtree := Snl ≀ Snl−1 ≀ · · · ≀ Sn1 ⊊ Sn . Symmetrising with respect to this group symmetrises each layer in turn, resulting in the locally symmetrised A(l),sym · ⊗nl (A(l−1),sym · · · ⊗n2 (A(1),sym )). Projecting A to the symmetric subspace of Gtree by symmetrising each layer locally removes all degrees of freedom related to input permutation symmetry, while remaining within the class of multilinear models sharing the same tree decomposition. This enables principled weight-based comparison between models of the same decomposition class. Another symmetry of layered multilinear models concerns basis transformations of the hidden dimensions along which the different layers are connected. This is discussed next but concerns only functions of local tensors (not of the global tensor). Gauge symmetry. For tensor decompositions, just like matrix decompositions, the exact values of the local tensors are not unique. Inserting UU−1 in any contracted dimension leaves the global tensor invariant. Functions of the global tensor, including the modelled polynomial for multilinear models or the inner product between two decomposed tensors, are also invariant under gauge transformations. Gauge fixing is the procedure to remove these degrees of freedom. For our models, privileged output and input bases, as well as basis-dependent computations such as the elementwise product of tensors, fix a gauge. Combining this with a residual stream that propagates the privileged input basis through the layers of a deep multilinear model fixes a gauge. This is only relevant if the tensors of individual layers are analysed instead of computing functions from the global tensor such as tensor similarity. Implications for tensor similarity. The central choice in defining tensor similarity is the metric tensor M ∈ ⊗2 (⊗n Rd+1 ). Three requirements constrain this choice: it should project out weightspace symmetries, relate deviation from unity with functional difference like behavioural similarity, and preserve the class of deep multilinear models, i.e. those whose global tensors admit a special tree decomposition. When the metric projects out weight-space symmetries (for instance, when combined with layerwise symmetrisation), tensor similarity reduces to the inner product between two tensors that are in one-to-one correspondence with polynomials. The Cauchy–Schwarz inequality relates this inner product to the product of the tensor norms, precisely the quantities appearing in the denominator of tensor similarity. It further guarantees that tensor similarity equals one if and only if the tensors, and hence the polynomials, are related by a positive scalar. 14
A.2
Gram recursion: derivation, symmetrisation compatibility, and algorithms
This appendix provides the additional details used in Section 2.2. Here we: • derive the Gram recursion from the tree structure of layered networks; • prove the commutativity result: Pni commutes with the Gram step, so layer-wise symmetrisation is compatible with the recursion; • give the explicit four-matrix-product algorithm for bilinear layers (ni = 2) and state its complexity. Define the Gram matrix at layer i as the partial contraction ⊗qi G(i) = A(i) ·in G(i−1) ·in B(i) ,
G(0) = I.
(10)
The inner product is recovered as ⟨A | B⟩ = tr G(L) . Derivation of the recursion. The global weight tensor A of a layered network is a tree tensor network whose leaves are the local tensors A(1) , . . . , A(L) . The inner product ⟨A | B⟩ contracts all indices of A against the corresponding indices of B. Because the underlying tree is acyclic, contraction can proceed from the leaves upward: contracting the input legs of layer i uses only the local tensors A(i) , B(i) and the result of the contraction at layer i − 1, which is captured by G(i−1) . Initialising with G(0) = 1 ∈ ⊗2 Rd+1 (the identity on the input space) and applying this observation recursively yields (10). The inner product is the trace of the final Gram matrix because the remaining P (L) open legs at layer L are the output legs, and tracing them gives ⟨A | B⟩ = k Gkk = tr G(L) . The next paragraph shows a commutativity result that reduces the computational cost of layerwise symmetrisation during the Gram recursion. The final paragraph of this subsection gives the algorithm for the computation of the Gram recursion explicitly for bilinear layer. Nn H Commutativity. Let G ∈ ⊗2 RH and Pn the symmetrisation projector on R . Then n [Pn , ⊗ G] = 0. Proof. It suffices to show that each permutation operator Uσ , σ ∈ Sn , commutes with ⊗n G, since Pn is a linear combination of these. On product vectors, Uσ ⊗n G(x1 ⊗ · · · ⊗ xn ) = Gxσ−1 (1) ⊗ · · · ⊗ Gxσ−1 (n) = ⊗n G Uσ (x1 ⊗ · · · ⊗ xn ). Nn H Product vectors span R , so the claim follows.
□
By idempotency of Pni and the above commutativity result, Pni ·⊗ni G(i−1) ·Pni = ⊗ni G(i−1) ·Pni . It follows that the symmetrised Gram recursion requires symmetrising only one of the two layers per step: G(i),sym = A(i) ·in ⊗ni G(i−1),sym ·in Pni B(i) . (11) (i)
(i)
(i)
Bilinear algorithm. For bilinear layers (ni = 2) parameterised by LA , RA ∈ Rri ×Hi−1 , DA ∈ RHi ×ri (and likewise for B), the Gram step (10) reduces to matrix operations alone. Initialise G(0) = 1d+1 and, for each i ∈ [L], compute four matrix products, (i) (i) T (i) (i) T (LL)(i) := LA G(i−1) LB , (RR)(i) := RA G(i−1) RB , T T (i) (i) (i) (i) (LR)(i) := LA G(i−1) RB , (RL)(i) := RA G(i−1) LB , (12) two elementwise products, (i)
(i)
E∥ := (LL)(i) ⊙ (RR)(i) ,
E× := (LR)(i) ⊙ (RL)(i) ,
(13)
and one final matrix multiplication, (i)
(i)
(i)
G(i) = DA E∥ + E×
(i) T
DB
.
(14)
The two terms E∥ and E× correspond to the two terms in the symmetrisation of the bilinear layer: E∥ pairs same-role legs (L–L and R–R) while E× pairs swapped-role legs (L–R and R–L). Our experiments are implemented in Python with the tensor contraction algorithms like the above implemented using quimb Gray [2018]. 15
A.3
Derivation of the Gaussian metric
This appendix derives the Gaussian metric tensor presented in Section 2.3. Here we: • show that the expected output similarity equals ⟨A | M(p) | B⟩ where M(p) has S2n symmetry; • compute Λ(p) for a Gaussian input distribution via the Isserlis–Wick theorem; • conclude that this metric projects onto a strictly smaller subspace than the plain Sn symmetrisation. Taking the expectation over inputs of the activation inner product gives Ex ⟨A(x) | B(x)⟩ = ⟨A | M(p) | B⟩,
(15)
where M(p) = 1 ⊗ Λ(p) and (p) Λi1 ···in j1 ···jn = Ex xi1 · · · xin xj1 · · · xjn .
(16)
The tensor Λ(p) has order 2n, pairing all 2n input legs of A and B, including self-contractions of two legs belonging to the same tensor. This makes Λ(p) manifestly invariant under the action of S2n . Gaussian case via Isserlis’ theorem. For x ∼ N (0, Id ), Isserlis’ theorem [Munthe-Kaas et al., 2025] expresses the 2n-th moment as a sum over all perfect matchings of the 2n indices i1 , . . . , in , j1 , . . . , jn . For unit covariance, each matched pair (a, b) contributes δia ib . Classify pairings by m ∈ N, the number of internal A-pairs (pairs drawn entirely from {i1 , . . . , in }). A pairing with m internal A-pairs must also have m internal B-pairs (because the total number of A-indices must equal the number of B-indices), leaving n − 2m cross-pairs. Define the partial trace operator τ such that τ m T contracts m pairs of indices as (τ T)i3 ···in =
d X
Ta a i3 ···in .
(17)
a=1
All pairings of type m contribute the same traced inner product ⟨τ m A| τ m B⟩. n (2m−1)!! Counting the pairings of type m: on the A-side, choose and pair 2m of the n indices in 2m ways; the same count applies to B. The remaining n − 2m legs on each side are then cross-matched in (n − 2m)! ways. Setting 2 n cn,m = (2m − 1)!! 2 (n − 2m)!, (18) 2m and summing over m gives the main result: ⌊n/2⌋ X Ex∼N (0,Id ) ⟨A(x) | B(x)⟩ = cn,m ⟨τ m A, τ m B⟩.
(19)
m=0
Equation (19) identifies Λ(p) as a weighted sum of partial-trace operators, confirming N that M(p) is n d+1 positive definite and carries full S2n symmetry. Because the S2n -invariant subspace of R is strictly smaller than the Sn -invariant one, this Gaussian metric is a strictly finer equivalence relation than plain Sn -symmetrisation, while still satisfying the discriminative-power guarantee (4).
B
Additional catastrophic forgetting details
Dataset: The SVHN dataset comprises over 600, 000 labelled digits cropped from Google Street View images [Netzer et al., 2011]. It occupies a useful middle ground in difficulty: unlike MNIST [Deng, 2012], it is not solvable using linear models, yet it does not require the large architectures for ImageNet-scale tasks. We use a subset of approximately 73, 000 training and 26, 000 validation images, converted to greyscale. 16
Model Architecture: The model consists of a linear embedding layer, a single bilinear layer, and a linear unembedding layer. For ease of analysis, biases and normalisation layers are excluded from the model. Model Dimensions and Training: Table 1 summarises the model dimensions and training hyperparameters.
Architecture d_model d_hidden n_layer
128 256 1
Training Parameters dropout 0.0 weight decay 0.5 batch size 248 learning rate 10−3 optimizer AdamW schedule Cosine Annealing epochs 20 per stage Table 1: Model architecture and training setup for the catastrophic forgetting setup.
C
Additional modular addition details
Dataset: The task is modular addition: given inputs (a, b), predict (a + b) mod p where p = 113. Inputs are represented as two one-hot vectors of dimension 2p = 226, concatenating the encodings of a and b. All p2 = 12,769 input pairs are used, with a 60 : 40 split between train and validation sets. Model Architecture: The model consists of a single bilinear layer followed by a linear unembedding layer, with biases and normalisation layers. Model Dimensions and Training: Table 2 summarises the model dimensions and training hyperparameters.
Architecture d_input d_hidden n_layer
226 64 1
Training Parameters dropout 0.0 weight decay 0.06 batch size 512 learning rate 10−3 optimizer AdamW schedule constant steps 100,000 Table 2: Model architecture and training setup for the modular addition experiment.
17
Figure 6: Model evolution across the progressive training setup on SVHN described in subsection 3.1. Accuracy and loss follow a consistent pattern across all stages, with validation accuracy reaching 83–85% at each stage. Similarity with respect to the model at the end of the add 9 stage is shown for tensor similarity, CKA on logits, and weight cosine similarity. All three measures increase as digits are incrementally added, but tensor similarity exhibits the most pronounced drop at the remove 9 stage and the clearest recovery at re-add 9, while CKA and weight cosine show considerably weaker signals. Slice similarity highlights digit ‘9’ in a darker line and all other classes shown faint. The drop at remove 9 is localised almost entirely to digit ‘9’, confirming that catastrophic forgetting affects this class specifically. The effect is very clearly captured by tensor slice similarity, and is significantly weaker in weight cosine slice similarity.
D
Language modelling with an attention only transformer
D.1
Training details
D.1.1
Data and tokenisation
We train on the DSIR-filtered-Pile-50M subset from Xie et al. [2023]. Validation is created using a deterministic 1% held-out document split from the training set. We train a custom BPE tokeniser with vocabulary size V = 4096 on the same corpus. Documents are concatenated with EOS separators and chunked into fixed-length windows of length nctx = 512. D.1.2
Architecture
We use a decoder-only, attention-only language model with bilinear attention and no MLP blocks. The model has L = 2 layers, H = 8 attention heads, and model dimension dmodel = 256. The 18
Figure 7: Tensor diff similarities across the progressive SVHN training setup (subsection 3.1). Similarity to diff is computed with respect to a diff tensor, computed as the difference between two checkpoints of an identically-architectured bilinear model trained on a different subset of SVHN data with a different random seed. The first checkpoint is taken after training on {0, · · · , 8} and the second after subsequent fine-tuning on the full set {0, · · · , 9}, thereby isolating weight updates associated with learning digit ‘9’. Tensor similarity is shown in red and slice similarities are shown in blue, with digit ‘9’ highlighted in a darker line. The resulting pattern closely mirrors the self-similarity result of Figure 6 with catastrophic forgetting at remove 9 and subsequent recovery at re-add 9 clearly visible, demonstrating that an externally constructed diff tensor can isolate the same functional component as self-comparison.
embed and unembed are bias-free and not tied. Positional information is added using RoPE inside the attention layers. To keep the model structure bilinear, we use only a scalar normalisation method applied before the unembedding. Specifically, we normalise the sequence by a scalar s computed from the first token’s hidden state. v u model u 1 dX s=t h2 + ϵ. dmodel j=1 0,j
Weight Shape Parameters embed.weight (V, dmodel ) = (4096, 256) 1,048,576 attn.q1.weight (256, 256) 65,536 attn.q1.bias (256, ) 256 attn.k1.weight (256, 256) 65,536 attn.k1.bias (256, ) 256 attn.q2.weight (256, 256) 65,536 attn.q2.bias (256, ) 256 attn.k2.weight (256, 256) 65,536 attn.k2.bias (256, ) 256 attn.v.weight (256, 256) 65,536 attn.o.weight (256, 256) 65,536 Layer subtotal – 394,240 All layers (×L = 2) – 788,480 final_norm (tok0_batch) – 0 unembed.weight (V, dmodel ) = (4096, 256) 1,048,576 Total – 2,885,632 Table 3: Attention-only language model parameter count
19
D.1.3
Loss
For a token sequence Si = (ti1 , . . . , tiT ), the empirical next-token loss at position k is n 1X i − log softmax fw (S≤k ) tik+1 , k = 1, . . . , T − 1. ℓn,k (w) = n i=1 The training loss is T −1
ℓn (w) =
1 X ℓn,k (w), T −1
T = nctx = 512,
k=1
with held-out loss defined analogously on the validation split. D.1.4
Training
Training is performed as a single streaming pass over the shuffled, EOS-separated Pile stream. The main run uses Ttrain = 20,000 Optimiser steps, batch size B = 384, and context length nctx = 512, corresponding to approximately 4 ∗ 109 training tokens. Evaluation is performed using a cached validation split from the Pile. Optimisation uses Muon for the internal attention projection weights, while the attention output projection, embedding, and unembedding parameters are optimised with AdamW. The Muon learning rate is 0.02, and the AdamW learning rate is 3 × 10−4 , with Adam betas (0.9, 0.95), weight decay 0.1, and gradient clipping at 1.0. Both optimiser groups use a cosine schedule with a 500-step warmup and decay to 0.2 times the peak learning rate by step 20,000. The model took 5 hours to train on A10. Hyperparameter Category Value Dataset Data DSIR-filtered Pile-50M V Data 4096 nctx Data 512 B Data 384 Nsteps Data 20,000 Tokens seen Data ≈ 3.93 × 109 L Model 2 H Model 8 dmodel Model 256 Attention type Model bilinear Attention scale Model 0.35 RoPE base Model 10,000 Norm Model tok0_batch before unembedding Optimiser Optimisation Muon + AdamW Muon LR Optimisation 0.02 AdamW LR Optimisation 3 × 10−4 Weight decay Optimisation 0.1 Scheduler Optimisation cosine decay with 500-step warmup Table 4: Language-modeling transformer training, data, model, and optimisation hyperparameters.
D.2 D.2.1
Metrics tracked over training N-gram metrics
The n-gram metric measures whether the model has learned common fixed-length token patterns from the training data. First, the most frequent n-grams (t1 , . . . , tn ) are extracted from the training corpus. For each n-gram, the model is given the prefix (t1 , . . . , tn−1 ) and evaluated on its ability to predict the final token tn . Averaging this cross-entropy over frequent training n-grams gives the n-gram memorisation loss, denoted lngram (n). The same prediction task is then evaluated on validation sequences, giving ltest (n). The final score is ltest (n) ngram score(n) = . lngram (n) 20
A lower ratio indicates that the model performs similarly on validation n-gram prediction and frequent training n-gram prediction, while a higher ratio suggests stronger specialization to the frequent training patterns.
E
Robustness to non-Gaussian input distributions
In section 2.3, we showed that tensor similarity coincides with a canonical behavioural metric, cosine similarity of empirical outputs, under the assumption of a perfectly Gaussian input distribution. But how robust is this to non-Gaussian input distributions? In this appendix, we investigate the question statistically. Dataset We consider a toy task where models are trained to detect the index of the second-highest value (or 2nd argmax) in a list of four numbers sampled from one of the following 9 input distributions: • Gaussian: Standard Gaussian distribution (mean 0 and variance 1). • Half Gaussian: Absolute value after sampling from a standard Gaussian distribution. • Bimodal: Combination of two Gaussians designed to be bimodal. • Uniform: Uniform distribution on [−1, 1]. • Laplace: Laplace distribution with location parameter µ = 0 and scale parameter b = √12 . • ‘Sparse Spikes’: A distribution where there is a 75% chance of drawing 0 and a 25% chance of drawing from a Gaussian distribution with mean 1 and variance 4. • Permutations: The set of permutations of the list (1, 2, 3, 4). • Correlated Gaussian: Standard Gaussian distribution with a correlation introduced between the four inputs, defined by a fixed positive-definite correlation matrix. • ‘Gaussian and −10’: Standard Gaussian distribution for the first three inputs; a constant −10 for the last input. Example: One possible sample from the ‘Gaussian and −10’ distribution is the following: [−0.54, −0.19, 0.17, −10.00]. For this data point, the correct output for the model is 1, since the second-highest value is −0.19, which lies at index 1. Model Architecture The model consists of a toy one-layer bilinear network. No residuals, normalisation layers, or biases are included in the model. Model Dimensions and Training Table 5 summarises the model dimensions and training hyperparameters. Across all five seeds, models trained on the three distributions which are not symmetric about zero (namely Half Gaussian, Permutations, and ‘Gaussian and −10’) achieved accuracy above 90%, while models trained on the six symmetric distributions plateaued at approximately 50% (with the exception of the ‘Sparse Spikes’ models, which plateaued around 39%). This is due to a fundamental limitation in the pure, single-layer bilinear architecture used: it cannot distinguish x from −x, preventing the model from meaningfully exceeding 50% accuracy on any distribution symmetric about zero. Figure 8 shows Pearson correlations between tensor similarity and the cosine similarity of empirical outputs across different training checkpoints and seeds. While the correlation between tensor similarity and output similarity is strongest for Gaussian inputs, it exceeds 0.9 for eight out of the nine distributions tested, including distributions of varying shape and one with statistical dependence between input dimensions. The exception is the correlation across the ‘Gaussian and −10’ dataset, with a measured correlation of r = 0.61. Intuitively, the reason for this discrepancy is that here, the last index is so low that it will virtually never point to the second highest value, so the model learns to never output that index. This means that a substantial portion of network weights do not influence the ultimate result, and the weight-based similarity metric we use here is less correlated with the empirical similarity. These results suggest that the tensor similarity metric is a robust proxy for empirical similarity across a broad range of distributions, not just Gaussians, so long as the weights all meaningfully contribute toward the output. 21
Architecture input dim n rank n_layer
4 32 1
Training Parameters dropout weight decay batch size learning rate optimizer schedule train steps seeds
0.0 0.0 512 10−2 Adam constant 4096 5
Table 5: Model architecture and training setup for the distribution robustness experiment.
Figure 8: Five one-layer bilinear models were trained for each of nine input data distributions, and they were each sampled at 14 checkpoints. Pairs of these models were compared with a tensor-based similarity metric (shown on the x-axis) and a behavioural, output-based cosine simila rity metric (shown on the y-axis). Pearson correlations are displayed above each scatterplot.
22