Technical Summary: Tensor Similarity for Mechanistic Interpretability

Problem Statement

Mechanistic interpretability seeks to decompose neural networks into meaningful, constituent parts to understand their internal computations. A critical prerequisite for this field is the ability to rigorously verify whether two decomposed components (or two different models) implement the same underlying function. Current evaluation methods suffer from two primary limitations:

1. Empirical/Behavioral Similarity: Methods that compare model outputs or activations on specific datasets are blind to out-of-distribution (OOD) mechanisms. They cannot detect issues like backdoors that trigger only on specific inputs or fail to generalize if the decomposition relies on approximate heuristics rather than the true function.
2. Parameter-Based Similarity: Methods that compare weight matrices (e.g., matrix cosine similarity) are highly sensitive to weight-space symmetries (such as permutation and scaling). They fail to recognize functional equivalence when models are trained differently or when mechanisms are distributed across layers rather than concentrated in single neurons.

There is a lack of a rigorous, dataset-independent metric that is invariant to these symmetries and capable of capturing global functional equivalence.

Methodology

The authors propose Tensor Similarity, a weight-based metric designed for tensor-based neural networks (specifically multilinear models like bilinear transformers). This approach trades the use of standard non-linear activation functions for polynomial (multilinear) architectures to enable algebraic analysis of weights.

Core Architectural Shift

The method relies on multilinear models where activation functions are replaced by multiplicative gating (e.g., bilinear layers and bilinear attention).

• Bilinear Layers: Map a lifted input $\tilde{x} = (1, x)$ to an output via a degree-2 polynomial: $A(x) = D((L\tilde{x}) \odot (R\tilde{x}))$.
• Multilinear Representation: These layers can be expressed as a contraction of a weight tensor $A$ with $n$ copies of the input: $A(x) = A \cdot \text{in}^{\otimes n} \tilde{x}$.
• Deep Models: Deep multilinear networks are constructed as tree-structured decompositions of global tensors, allowing for efficient computation without materializing the exponentially large global tensor.

The Tensor Similarity Metric

The metric is defined as a generalized inner product between two weight tensors $A$ and $B$, normalized by their induced norms:
$$\text{sim}_M(A, B) = \frac{\langle A | M | B \rangle}{\|A\|_M \|B\|_M}$$
where $M$ is a metric tensor.

To ensure the metric satisfies the criteria of symmetry invariance and dataset independence, the authors derive a specific form for $M$:

1. Symmetrisation: Based on the polarisation isomorphism, the authors establish that two multilinear models are functionally equivalent if and only if their weight tensors have the same representative in the symmetric subspace. The metric projects tensors onto this symmetric subspace to remove degrees of freedom related to input permutation symmetries.
2. Gaussian Metric: To further refine the metric and relate it to behavioral similarity, the authors define $M$ using the expectation of the input tensor under a standard Gaussian distribution ($x \sim \mathcal{N}(0, I)$). Using Isserlis' theorem, this results in a closed-form metric that computes the inner product of activations over the entire Gaussian input space, rather than a sampled dataset.
3. Efficient Computation: A Gram-based recursive algorithm allows the computation of this similarity for deep networks by processing local tensors layer-by-layer, avoiding the need to instantiate the full global tensor.

Key Contributions

1. A Principled Weight-Based Metric: The paper introduces a metric that is invariant to weight-space symmetries (permutations, scaling) and independent of the input dataset. It formally verifies functional equivalence by checking if two models implement the same polynomial logic.
2. Closed-Form, Data-Free Analysis: Unlike behavioral metrics, tensor similarity does not require sampling data. It integrates over the full Gaussian input space, making it capable of detecting OOD mechanisms (like backdoors) that are invisible to standard performance metrics.
3. Localisation Capabilities: The metric can be "sliced" to focus on specific output dimensions or "diff" tensors to isolate functional changes between checkpoints, allowing for the attribution of specific mechanisms to specific outputs.
4. Scalability: The recursive algorithm ensures the metric is computationally tractable for deep networks, avoiding exponential or quadratic convergence bounds associated with Monte Carlo sampling.

Results

The authors validate tensor similarity across four distinct domains, demonstrating its superiority over existing metrics (CKA, matrix cosine, behavioral similarity):

1. Catastrophic Forgetting (Vision): In a progressive training setup on SVHN, tensor similarity precisely localizes the loss of knowledge for specific digits (e.g., digit '9') when the dataset is reduced. It shows a sharp drop in similarity for the affected class while remaining stable for others, whereas behavioral metrics show ambiguous signals.
2. Grokking (Modular Arithmetic): In modular addition tasks, tensor similarity tracks the continuous reorganization of the model into specific Fourier frequencies. It reveals that structural reorganization continues even after accuracy jumps to perfection, providing a more legible view of training dynamics than accuracy curves alone.
3. Backdoor Detection (OOD): When a backdoor is injected into SVHN (triggering on a black diamond patch), standard behavioral metrics on clean data fail to detect the change. Tensor similarity, however, flags the backdoor from clean data alone and localizes the change to the specific attacked class ('9'), demonstrating its ability to detect OOD mechanisms.
4. Language Modeling: Applied to a two-layer bilinear attention transformer trained on The Pile, tensor similarity reveals clear block structures and sharp transitions in training regimes that existing metrics blur. It tracks the formation of complex n-gram statistics more effectively than behavioral proxies.

Significance and Claims

The paper positions tensor similarity as a foundational tool for mechanistic interpretability. Its primary significance lies in reducing the problem of measuring similarity and verifying faithfulness from an empirical approximation problem to a solved algebraic problem.

• Faithfulness: By relying on the algebraic properties of multilinear models, the metric ensures that decompositions capture the true underlying function rather than approximate heuristics.
• Robustness: It provides stronger robustness claims by applying to the full distribution, effectively ruling out mechanisms that only trigger on specific inputs (like backdoors).
• Limitations: The authors modestly acknowledge that the method currently relies on multilinear models, which excludes most state-of-the-art architectures using non-polynomial non-linearities. However, they argue that multilinear models serve as a vital testbed for interpretability and that evidence suggests they can match the performance of standard architectures.
• Future Outlook: The work suggests that the restriction to multilinear models is a necessary trade-off for tractable algebraic analysis, but it paves the way for future attribution methods (e.g., computing similarity between paths in deep transformers) and supports the view that multilinear architectures are more transparent than ordinary deep networks.