Neuron Alignment via Permutations

• The paper details how neuron alignment recovers permutation-invariant structures by matching neurons with cost minimization techniques such as the Hungarian algorithm.
• It employs methods like Euclidean and cosine distance cost matrices and soft assignment relaxations (Sinkhorn) to robustly align neurons across diverse architectures.
• Applications include model merging, watermark extraction, federated learning, and even mapping biological neuron patterns to reveal functional groupings.

Neuron alignment via permutations refers to the process of explicitly finding, recovering, or exploiting the permutation-invariant structure of neural representations by matching the indices of individual neurons across networks, weight matrices, or data modalities. This alignment is fundamental for resolving the permutation symmetry inherent in multilayer neural networks—where permuting neurons within a hidden layer (with appropriate compensatory permutations on adjacent weights) produces functionally equivalent models. Methods for neuron alignment via permutations are now central to diverse areas: model merging, robust watermark extraction, multi-model fusion, mode connectivity, and representational similarity analysis.

1. Mathematical Foundations of Permutation Symmetry

Permutation symmetry in neural networks means that any network with layers of size $N_1, N_2, \ldots, N_L$ admits the action of the product group $\Pi = S_{N_1} \times \ldots \times S_{N_{L-1}}$, where $S_{N}$ is the symmetric group on $N$ elements. For weights $\{W_\ell, b_\ell\}$ and permutation matrices $P_\ell \in \{0,1\}^{N_\ell \times N_\ell}$, the function computed by the network remains invariant if each hidden layer's output is permuted by $P_\ell$, with all associated weight matrices and biases appropriately re-indexed or transformed: $$W_\ell \mapsto P_\ell W_\ell P_{\ell-1}^\top, \quad b_\ell \mapsto P_\ell b_\ell.$$ This invariance leads to equivalence classes (“orbits”) of network parameterizations, which are critical sources of redundancy and non-convexity in the parameter space (Sharma et al., 9 Apr 2024).

Permutation-based neuron alignment seeks to find a set of permutations $\{P_\ell\}$ such that, given two functionally similar networks $f(\cdot;\theta)$ and $f'(\cdot;\theta')$, their respective neurons and/or weights (and occasionally activations) are brought into correspondence, minimizing a cost such as

$$\sum_{\ell=1}^{L} \Vert W_\ell - P_\ell W'_\ell Q_{\ell-1} \Vert_F^2$$

subject to $P_\ell, Q_\ell$ being proper permutation matrices (Trias et al., 2023).

2. Algorithms for Neuron Alignment

The core alignment problem is to solve an assignment of neurons (indices) across two layers, or across whole networks. In the most direct scenario (layers of equal width, $N$), this is formalized as seeking a permutation $P^*$ minimizing the Frobenius norm: $$P^* = \arg\min_{P \in \mathcal{P}_N} \| W - W'^P \|_F^2$$ where $\mathcal{P}_N$ is the set of $N \times N$ permutation matrices, and $W$, $W'$ are corresponding weight matrices (Li et al., 2021).

Common steps:

1. Cost Matrix Computation: Compute pairwise costs between neuron “signatures” (weight vectors, or activation vectors across a data batch). Popular choices:
   • Squared Euclidean distance: $C_{ij} = \|w_{i} - w'_j\|^2$
   • Cosine similarity: $C_{ij} = 1 - \frac{w_i^\top w'_j}{\|w_i\|\|w'_j\|}$
   • Cross-correlation or canonical correlation between activations (Uriot et al., 2020, Tatro et al., 2020).
2. Layerwise Assignment: Solve the linear assignment problem (LAP) using the Hungarian (Kuhn–Munkres) algorithm (complexity $O(N^3)$ per layer). This can be relaxed to doubly-stochastic matrices and solved via Sinkhorn-Knopp for “soft” assignments, especially when sizes differ (Khosla et al., 2023).
3. Permutation Application: Apply found permutations to permute $W'$, $b'$, and all related parameters/outputs, restoring a common index basis for subsequent operations.

Extensions:

• Soft/permutation-relaxed matching via optimal transport: for non-equal sizes or partial matches (Khosla et al., 2023).
• Clustering matrices by optimal permutations for whole-graph or cross-modal alignment (Morone, 2021).
• Trainable Feedforward Alignment: Deep-Align merges assignment logic into learned, group-equivariant neural networks for fast inference and large-scale applications (Navon et al., 2023).

3. Applications

3.1 Watermark Robustness and Integrity Verification

White-box DNN watermarking schemes embed signatures into specific layers’ weights or activations (Li et al., 2021). An adversary can permute neuron order in a protected layer, rendering the watermark irrecoverable via standard extraction. Neuron alignment recovers the original ordering by solving (via the Hungarian algorithm) for permutations that best match the suspect model’s layer weights to the owner’s reference, fully restoring watermark verification accuracy against permutation-only attacks and substantially improving robustness under additional perturbations (fine-tuning, pruning).

Attack Type	Success w/o Alignment	Success w/ Alignment (e.g., Residual backend)
Neuron Permutation (NP)	0.0%	99.1%
Fine-Tuning + Permutation	0.0%	84.6%
Neuron-Pruning + Permutation	0.0%	79.8%

(Li et al., 2021)

3.2 Model Fusion, Mode Connectivity, and Federated Aggregation

Permutation symmetry is a barrier to simple model fusion since weights from independently trained networks are non-aligned in index space. Recent algorithms first perform neuron alignment across all intermediate layers, then apply arithmetic averaging or mode-connecting curves (e.g., Bézier) (Tatro et al., 2020, Luenam et al., 18 Jun 2025, Li et al., 2 Feb 2024). In federated learning or model soups, explicit alignment enables “zero-shot” fusion, outperforming prior rebasin techniques, especially under data heterogeneity or non-IID settings (Luenam et al., 18 Jun 2025, Li et al., 2 Feb 2024). Training-time approaches such as TNA-PFN further break permutation symmetry by constraining all models to a common permutation subspace using a universal mask, achieving barrier-free interpolation and robust fusion without post-hoc alignment (Li et al., 2 Feb 2024).

3.3 Neuroevolution, Safe Crossover, and Representation Metric Analysis

In neuroevolution, naive parameter crossover fails due to the “competing conventions problem”—permutation misalignment of parent networks. Safe crossover strategies align neurons via cross-correlation or canonical correlation analysis before recombination, ensuring that arithmetic blending between parents produces offspring with meaningful inherited features (Uriot et al., 2020).

Permutation-based alignment metrics (strict one-to-one, semi-matching, soft-matching) enable fine-grained comparison of network representations—crucially, they capture single-unit tuning structure missed by rotation-invariant metrics (CKA, CCA, Procrustes). This has implications for model-brain comparisons, representation analysis, and suggest that soft-matching alignment (i.e., optimal transport between networks' neuron activations) is strictly more informative when neuron identity is meaningful (Khosla et al., 2023, Longon et al., 3 Oct 2025).

Metric Type	Permutation Sensitivity	Captures Neuron Identity?	Symmetric?	Triangle Inequality?
Permutation Matching	Yes	Yes	Yes	Yes
Soft Matching (OT)	Yes	Yes	Yes	Yes
CCA/CKA	No	No	Yes	Yes

(Khosla et al., 2023)

4. Theoretical Guarantees and Limitations

4.1 Uniqueness and Robustness

Uniqueness of assignment is guaranteed when weight (or activation) columns are pairwise well-separated: if the minimum inter-column distance is $\Delta > 2\varepsilon$ (with adversary perturbation at most $\varepsilon$), exact recovery holds (Li et al., 2021, Trias et al., 2023). Cosine similarity-based alignment is proven robust to extensive integrity attacks (fine-tuning, quantization, pruning up to high levels), maintaining >99% recovery unless the network's own accuracy is catastrophically degraded (Trias et al., 2023).

4.2 Mode Connectivity and Loss Barriers

Theoretical analysis (e.g., for linear mode connectivity) shows that, after alignment, interpolation barriers in loss or accuracy are sharply reduced or entirely eliminated; critical in both pairwise and simultaneous multi-model fusion (Sharma et al., 9 Apr 2024, Li et al., 2 Feb 2024). In the large-width regime, empirical evidence suggests the possibility of “strong linear connectivity” where a single permutation aligns all models in a set, not just pairwise (Sharma et al., 9 Apr 2024).

4.3 Superposition Interference

Permutation-based alignment critically depends on the assumption that features are not in superposition (i.e., not distributed as different linear combinations in different models). When superposition arrangements differ across models, alignment metrics are systematically deflated. Disentangling superposition using sparse overcomplete codes (e.g., SAEs) restores near-perfect alignment in toy models and boosts DNN→DNN and DNN→brain alignment scores (Longon et al., 3 Oct 2025).

5. Extensions, Relaxations, and Open Problems

• Soft-Permutation and Optimal Transport: Allow assignment matrices to be doubly-stochastic (“soft” matching), enabling generalization to networks with different widths and yielding Wasserstein (OT) metrics with desirable theoretical properties (symmetry, triangle inequality) (Khosla et al., 2023).
• Activation- vs Weight-based Matching: Cross-correlation on post-activations generally outperforms purely weight-based alignment, especially when architectures are widened or the activation code is highly distributed (Tatro et al., 2020, Sharma et al., 9 Apr 2024).
• Trainable Alignment Architecture: Deep-Align leverages group-equivariant architectures to accelerate and generalize alignment, simulating classical assignment but requiring only a forward pass at inference (Navon et al., 2023).
• Clustering via Optimal Permutations: When aligning networks to specific structural motifs or biological data (e.g., C. elegans connectome), permutation alignment can be driven by geometric patterns, not just minimum $\ell_2$ distance or maximum correlation (Morone, 2021).

Open challenges include extensions to networks of varying widths, dynamic architectures, or with fundamentally different structures (e.g., transformer vs convolutional, beyond permutation-invariance), as well as more efficient approximation algorithms for large-scale weight or activation matching, and understanding the global landscape of near-optimal permutations (Navon et al., 2023, Sharma et al., 9 Apr 2024).

6. Empirical Results and Best Practices

Experimental work demonstrates that permutation-based neuron alignment:

• Accurately and robustly recovers neuron correspondence across architectures, datasets, and even under moderate integrity attacks (Li et al., 2021, Trias et al., 2023).
• Enables barrier-free merging and ensembling in both centralized and federated learning (Li et al., 2 Feb 2024, Luenam et al., 18 Jun 2025).
• Clarifies representational convergence at the neuron level, distinct from merely subspace alignment (Khosla et al., 2023).
• Strongly benefits from activation-over-weight cost matrices and subset selection for large datasets (Tatro et al., 2020, Khosla et al., 2023).

Best practices include matching by post-activation cross-correlation with the Hungarian algorithm, relaxing to Sinkhorn or soft-assignment for large or mismatched layers, and prescreening for superposition via sparse code disentanglement (Khosla et al., 2023, Longon et al., 3 Oct 2025). For maximal efficiency in deployment or integration into pipelines such as federated learning, training-time approaches (e.g., TNA-PFN) can circumvent the costs of post-hoc alignment altogether (Li et al., 2 Feb 2024).

7. Biological and Neuroscientific Contexts

Beyond artificial neural networks, permutation alignment has been systematically applied to biological data, such as clustering neuron activity or connectome graphs in C. elegans. By encoding block-diagonal, nested, or banded patterns in template matrices and seeking permutations that optimally match experimental adjacency matrices to these templates, hidden modular or functional groupings are revealed—e.g., distinguishing interneurons, motor, sensory, and polymodal neuron types (Morone, 2021).

Permutation alignment in this context utilizes exactly the same mathematical tools (Frobenius/min-norm assignment, Sinkhorn projection, etc.), providing a theoretical and computational bridge between deep learning and neuroscience-domain clustering.

---

References are to arXiv preprints using the following identifiers: (Li et al., 2021, Khosla et al., 2023, Trias et al., 2023, Luenam et al., 18 Jun 2025, Li et al., 2 Feb 2024, Uriot et al., 2020, Navon et al., 2023, Sharma et al., 9 Apr 2024, Longon et al., 3 Oct 2025, Tatro et al., 2020, Morone, 2021).