Tensor Similarity: A Weight-Based Metric for Comparing Neural Network Mechanisms

Paper: 2605.15183 Authors: ML Nissen Gonzalez, Melwina Albuquerque, Laurence Wroe, Jacob Meyer Cohen, Logan Riggs Smith, Thomas Dooms Categories: cs.LG

The Gap

Mechanistic interpretability researchers want to verify that two network components implement the same computation. Existing approaches fall into two camps: behavioral metrics (comparing outputs on test data) miss out-of-distribution mechanisms, while parameter-based metrics (comparing raw weights) break when you permute neurons or apply other weight-space symmetries that don’t change the function.

The field has been stuck choosing between “blind to mechanisms” and “blind to symmetries.” For tensor-based models (MLPs, attention layers, convolutions), there’s been no way to ask: “Do these weights implement the same mathematical transformation?” while respecting that W and PW (for permutation P) can be functionally identical.

Problem: Compare two network components
   |
   v
Existing approaches split:
   |
   +---> Behavioral metrics -----> Blind to OOD mechanisms
   |                                (only see test distribution)
   |
   +---> Parameter metrics -------> Blind to symmetries
                                    (permute neurons = different)
   |
   v
Gap: Need weight-based + symmetry-aware
   |
   v
Method: Tensor similarity (recursive invariant)
   |
   v
Evidence: Tracks grokking, backdoors better than baselines
   |
   v
Conclusion: Functional equivalence becomes algebraic problem

The Increment

One sentence: Before this paper, comparing network mechanisms meant choosing between empirical approximation (behavior) or symmetry-blind comparison (raw weights); after, there’s an algebraic solution that respects weight-space structure.

Core Mechanism

Tensor similarity treats each layer as a multilinear map (a tensor) and compares them by finding the optimal alignment under all valid symmetries. For a single layer, this means: given weight matrices W₁ and W₂, find permutation matrices P and Q such that ||W₁ - P W₂ Q^T|| is minimized. The metric is the residual after optimal alignment.

For multi-layer networks, the algorithm works recursively. Start at the output layer and work backward. At each layer, the alignment from the previous layer constrains the current layer’s symmetries—if you permuted neurons in layer k, layer k-1 must respect that permutation. This captures cross-layer mechanisms: a circuit spanning multiple layers will only match if the entire computational path aligns.

The computation uses the Hungarian algorithm for optimal bipartite matching at each layer, making it polynomial-time. The result is a single scalar: zero means functionally identical (up to symmetries), larger values mean structural divergence.

Input: Two networks N1, N2
   |
   v
Layer L (output):
   W1[L], W2[L]
   |
   v
Find optimal P, Q:
   min ||W1[L] - P*W2[L]*Q^T||
   |
   v
Layer L-1:
   W1[L-1], W2[L-1]
   Constrained by Q from layer L
   |
   v
Recurse backward to input
   |
   v
Output: Scalar similarity (0 = identical)

Think of it like comparing two Rubik’s cubes. Behavioral metrics check if they produce the same colors when you apply test sequences—but miss internal mechanisms. Raw parameter comparison checks if stickers are in identical positions—but fails when one cube is rotated (a symmetry). Tensor similarity solves the cube alignment problem: it finds the rotation that makes the cubes maximally similar, then measures the remaining difference. The recursive part handles that each face’s rotation constrains adjacent faces—just as each layer’s permutation constrains neighboring layers.

Key Concepts

• Weight-space symmetries: Neural networks have redundancy in how they represent the same function. Permuting all neurons in a hidden layer (swapping rows in one weight matrix and columns in the next) changes the parameters but not the computation. Other symmetries include scaling (multiply one layer by α, divide the next by α) and sign flips for certain activations. These symmetries form a group—compositions of symmetries are symmetries. A good similarity metric must be invariant to this group, meaning symmetric networks score as identical. Existing parameter metrics fail this: they see permuted networks as different even when functionally identical.

• Tensor decomposition perspective: A neural network layer with weight matrix W ∈ R^(m×n) defines a bilinear map: (x, y) → x^T W y. This is a rank-2 tensor. Multi-layer networks compose these tensors. Tensor similarity leverages this structure: instead of treating weights as arbitrary arrays, it respects their multilinear nature. The alignment problem becomes: find the change of basis (permutations) that makes two tensors closest. This is why the method works recursively—tensor composition has a natural layer-by-layer structure.

• Grokking as a case study: Grokking is when a network suddenly generalizes after long overfitting. Behavioral metrics can’t distinguish pre-grok from post-grok on training data (both fit perfectly), and test accuracy jumps discontinuously. Tensor similarity reveals the transition smoothly: the weight structure gradually aligns with the true algorithmic solution, even before test accuracy improves. This shows the metric captures mechanistic changes invisible to behavior-based approaches.

Framework Shift

Before (mainstream approach):        After (this paper):

Network 1 ---> [Test Data] <--- Network 2    Network 1 ---> [Tensor Align] <--- Network 2
                   |                                              |
                   v                                              v
            Compare outputs                              Compare aligned weights
                   |                                              |
                   v                                              v
         Empirical similarity                           Algebraic similarity
         (distribution-dependent)                       (distribution-free)
         
         Blind to:                                      Captures:
         - OOD mechanisms                               - Global function structure
         - Internal structure                           - Cross-layer circuits
                                                        - Symmetry-invariant identity

From empirical approximation to algebraic verification, the core shift is treating functional equivalence as a solved problem in multilinear algebra rather than a statistical estimation task.

Expert Assessment

Problem choice: Real gap. Mechanistic interpretability is bottlenecked by the inability to verify component equivalence—you can’t build a library of circuits if you can’t tell when two circuits are the same. The symmetry-blindness of parameter metrics is a known pain point, not a manufactured problem.

Method maturity: Elegant use of existing tools (Hungarian algorithm, tensor algebra) rather than inventing new machinery. The recursive structure is the key insight—it’s not obvious that layer-wise alignment should compose into a global metric, but the tensor perspective makes it natural. One concern: the method assumes tensor structure (linear layers), so it doesn’t directly handle normalization layers, residual connections, or other architectural quirks without extension.

Experimental integrity: Baselines are fair (CKA, linear CKA, parameter distance). The grokking and backdoor experiments are well-chosen stress tests—they probe whether the metric sees mechanistic changes invisible to behavior. Numbers look solid. Minor quibble: would like to see failure cases or adversarial examples where tensor similarity gives misleading results.

Writing quality: Abstract and intro are crisp. The recursive algorithm description (Section 3) could use a worked example—readers will struggle to visualize the constraint propagation without one. The related work section undersells how this connects to circuit discovery (Olah et al.) and modular arithmetic grokking (Power et al.)—those connections would strengthen the motivation.

Verdict: strong accept — Solves a real problem with a principled method, strong empirical validation, and immediate utility for interpretability research.

Takeaways

If you’re comparing learned representations (not just in neural nets—think matrix factorizations, tensor decompositions, any parameterized multilinear map), steal the alignment-before-comparison principle. Don’t compare raw parameters; find the symmetry group, align optimally, then measure residuals.

For interpretability practitioners: this gives you a way to build circuit libraries. Train multiple networks on the same task, extract suspected circuits, use tensor similarity to cluster them. Circuits that score near-zero are mechanistically identical—you’ve found a canonical implementation.

For training dynamics researchers: tensor similarity tracks mechanistic phase transitions (like grokking) that behavioral metrics miss. Use it as a diagnostic: if weights are changing but tensor similarity to the final solution isn’t decreasing, you’re in a plateau.

The recursive constraint propagation trick generalizes beyond neural nets. Any hierarchical system where local symmetries compose (think: modular code with interface constraints, hierarchical Bayesian models) can use this pattern—align top-down, propagating constraints from parent to child.