# Equivariant Deep Weight Space Alignment

Aviv Navon
Bar-Ilan University
<EMAIL_ADDRESS>
&Aviv Shamsian∗
Bar-Ilan University
<EMAIL_ADDRESS>
&Ethan Fetaya
Bar-Ilan University
&Gal Chechik
Bar-Ilan University
NVIDIA Research
&Nadav Dym
Technion
&Haggai Maron
Technion
NVIDIA Research
Equal contributor

###### Abstract

Permutation symmetries of deep networks make simple operations like model
averaging and similarity estimation challenging. In many cases, aligning the
weights of the networks, i.e., finding optimal permutations between their
weights, is necessary. More generally, weight alignment is essential for a
wide range of applications, from model merging, through exploring the
optimization landscape of deep neural networks, to defining meaningful
distance functions between neural networks. Unfortunately, weight alignment is
an NP-hard problem. Prior research has mainly focused on solving relaxed
versions of the alignment problem, leading to either time-consuming methods or
sub-optimal solutions. To accelerate the alignment process and improve its
quality, we propose a novel framework aimed at learning to solve the weight
alignment problem, which we name Deep-Align. To that end, we first demonstrate
that weight alignment adheres to two fundamental symmetries and then, propose
a deep architecture that respects these symmetries. Notably, our framework
does not require any labeled data. We provide a theoretical analysis of our
approach and evaluate Deep-Align on several types of network architectures and
learning setups. Our experimental results indicate that a feed-forward pass
with Deep-Align produces better or equivalent alignments compared to those
produced by current optimization algorithms. Additionally, our alignments can
be used as an initialization for other methods to gain even better solutions
with a significant speedup in convergence.

##  1 Introduction

The space of deep network weights has a complex structure since networks
maintain their function under certain permutations of their weights. This fact
makes it hard to perform simple operations over deep networks, such as
averaging their weights or estimating similarity. It is therefore highly
desirable to “align” networks - find optimal permutations between the weight
matrices of two networks. Weight Alignment is critical to many tasks that
involve weight spaces. One key application is model merging and editing
(Ainsworth et al., 2022; Wortsman et al., 2022; Stoica et al., 2023; Ilharco
et al., 2022), in which the weights of two or more models are (linearly)
combined into a single model to improve their performance or enhance their
capabilities. Weight alignment algorithms are also vital to the study of the
loss landscape of deep networks (Entezari et al., 2022), a recent research
direction that has gained increasing attention. Moreover, weight alignment
induces an invariant distance function on the weight space that can be used
for clustering and visualization.

Since weight alignment is NP-hard (Ainsworth et al., 2022), current approaches
rely primarily on local optimization of the alignment objective which is time-
consuming and may lead to suboptimal solutions. Therefore, identifying methods
with faster run time and improved alignment quality is an important research
objective. A successful implementation of such methods would allow
practitioners to perform weight alignment in real-time, for example, when
merging models in federated or continual learning setups, or to perform
operations that require computing many alignments in a reasonable time, such
as weight space clustering.

Following a large body of works that suggested _learning_ to solve
combinatorial optimization problems using deep learning architectures (Khalil
et al., 2017; Bengio et al., 2021; Cappart et al., 2021), we propose the first
learning-based approach to weight alignment, called Deep-Align. Deep-Align is
a neural network with a specialized architecture to predict high-quality
weight alignments for a given distribution of data. A major benefit of our
approach is that after a model has been trained, predicting the alignment
between two networks amounts to a simple feed-forward pass through the network
followed by an efficient projection step, as opposed to solving an
optimization problem in other methods.

This paper presents a principled approach to designing a deep architecture for
the weight alignment problem. We first formulate the weight-alignment problem
and prove it adheres to a specific equivariance structure. We then propose a
neural architecture that respects this structure, based on newly suggested
equivariant architectures for deep-weight spaces (Navon et al., 2023) called
Deep Weight Space Networks (DWSNets). The architecture is based on a Siamese
application of DWSNets to a pair of input networks, mapping the outputs to a
lower dimensional space we call _activation space_ , and then using a
generalized outer product layer to generate candidates for optimal
permutations.

Theoretically, we prove that our architecture can approximate the Activation
Matching algorithm Tatro et al. (2020); Ainsworth et al. (2022), which
computes the activations of the two networks on some pre-defined input data
and aligns their weights by solving a sequence of linear assignment problems.
This theoretical analysis suggests that Deep-Align can be seen as a learnable
generalization of this algorithm. Furthermore, we show that Deep-Align has a
valuable theoretical property called _Exactness_ , which guarantees that it
always outputs the correct alignment when there is a solution with zero
objective.

Obtaining labeled training data is one of the greatest challenges when
learning to solve combinatorial optimization problems. To address this
challenge, we generate labeled examples on the fly by applying random
permutations and noise to our unlabeled data. We then train our network using
a combination of supervised and unsupervised loss functions _without_ relying
on any labeled examples.

Our experimental results indicate that Deep-Align produces better or
comparable alignments relative to those produced by slower optimization-based
algorithms, when applied to both MLPs and CNNs. Furthermore, we show that our
alignments can be used as an initialization for other methods that result in
even better alignments, as well as significant speedups in their convergence.
Lastly, we show that our trained networks produce meaningful alignments even
when applied to out-of-distribution weight space data.

Previous work. Several algorithms have been proposed for weight-alignment
(Tatro et al., 2020; Ainsworth et al., 2022; Peña et al., 2023; Akash et al.,
2022). Ainsworth et al. (2022) presented three algorithms: Activation
Matching, Weight Matching, and straight-through estimation. Peña et al. (2023)
improved upon these algorithms by incorporating a Sinkhorn-based projection
method. In part, these works were motivated by studying the loss landscapes of
deep neural networks. It was conjectured that deep networks exhibit a property
called _linear mode connectivity_ : for any two trained weight vectors (i.e.,
a concatenation of all the parameters of neural architecture), a linear
interpolation between the first vector and the optimal alignment of the
second, yields very small increases in the loss (Entezari et al., 2022;
Garipov et al., 2018; Draxler et al., 2018; Freeman & Bruna, 2016; Tatro et
al., 2020). Another relevant research direction is the growing area of
research that focuses on applying neural networks to neural network weights.
Early methods proposed using simple architectures (Unterthiner et al., 2020;
Andreis et al., 2023; Eilertsen et al., 2020). Several recent papers exploit
the symmetry structure of the weight space in their architectures (Navon et
al., 2023; Zhou et al., 2023a; b; Zhang et al., 2023). A comprehensive survey
of relevant previous work can be found in Appendix A.

##  2 Preliminaries

Equivariance Let $G$ be a group acting on $\mathcal{V}$ and $\mathcal{W}$. We
say that a function $L:\mathcal{V}\rightarrow\mathcal{W}$ is _equivariant_ if
$L(\rho_{1}(g)v)=\rho_{2}(g)L(v)$ for all $v\in\mathcal{V},g\in G$.

MultiLayer Perceptrons and weight spaces. The following definition follow the
notation in Navon et al. (2023). An $M$-layer MultiLayer Perceptron (MLP)
$f_{v}$ is a parametric function of the following form:

$f(x)=x_{M},\quad x_{m+1}=\sigma(W_{m+1}x_{m}+b_{m+1}),\quad x_{0}=x$ (1)

Here, $x_{m}\in\mathbb{R}^{d_{m}}$, $W_{m}\in\mathbb{R}^{d_{m}\times
d_{m-1}}$, $b_{m}\in\mathbb{R}^{d_{m}}$, and $\sigma$ is a pointwise
activation function. Denote by $v=[W_{m},b_{m}]_{m\in[M]}$ the concatenation
of all (vectorized) weight matrices and bias vectors. We define the _weight-
space_ of an $M$-layer MLP as:
$\mathcal{V}=\bigoplus_{m=1}^{M}\left(\mathcal{W}_{m}\oplus\mathcal{B}_{m}\right)$,
where $\mathcal{W}_{m}:=\mathbb{R}^{d_{m}\times d_{m-1}}$
$\mathcal{B}_{m}=\mathbb{R}^{d_{m}}$ and $\bigoplus$ denotes the direct sum of
vector spaces. A vector in this space represents all the learnable parameters
on an MLP. We define the _activation space_ of an MLP as
$\mathcal{A}=\bigoplus_{m=1}^{M}\mathbb{R}^{d_{m}}:=\bigoplus_{m=1}^{M}\mathcal{A}_{m}$.
The activation space, as its name implies, represents the concatenation of
network activations at all layers. In other words, $\mathcal{A}_{m}$ is the
space in which $x_{m}$ resides.

Symmetries of weight spaces. The permutation symmetries of the weight space
are a result of the equivariance of pointwise activations: for every
permutation matrix $P$ we have that $P\sigma(x)=\sigma(Px)$. Thus for example,
a shallow network defined by weight matrices $W_{1},W_{2}$ will represent the
same function as the network defined by $PW_{1},W_{2}P^{T}$, since the
permutations cancel each other. The same idea can be used to identify
permutation symmetries of general MLPs of depth $M$. In this case, the weight
space’s symmetry group is the direct product of symmetric groups for each
intermediate dimension $m\in[1,M-1]$ namely, $S_{d_{1}}\times\cdots\times
S_{d_{M-1}}$. For clarity, we formally define the symmetry group as a product
of matrix groups: $G=\Pi_{d_{1}}\times\cdots\times\Pi_{d_{M-1}},$ where
$\Pi_{d}$ is the group of $d\times d$ permutation matrices (which is
isomorphic to $S_{d}$). For $v\in\mathcal{V}$, $v=[W_{m},b_{m}]_{m\in[M]}$, a
group element $g=(P_{1},\dots,P_{M-1})$ acts on $v$ via a group action
$v^{\prime}=g_{\\#}(v)$, where
$v^{\prime}=[W^{\prime}_{m},b^{\prime}_{m}]_{m\in[M]}$ is defined by:

$\displaystyle W_{1}^{\prime}$
$\displaystyle=P_{1}W_{1},W_{M}^{\prime}=W_{M}P^{T}_{M-1},\text{ and
}W_{m}^{\prime}=P_{m}W_{m}P^{T}_{m-1},\forall m\in[2,M-1]$
$\displaystyle~{}b_{1}^{\prime}$
$\displaystyle=P_{1}b_{1},~{}b_{M^{\prime}}=b_{M},\text{ and
}b_{m}^{\prime}=P_{m}b_{m},\forall m\in[2,M-1].$

By construction, $v$ and $v^{\prime}=g_{\\#}(v)$ define the same function
$f_{v}=f_{v^{\prime}}$. The group product $g\cdot g^{\prime}$ and group
inverse $g^{-1}=g^{T}$ are naturally defined as the elementwise matrix product
and transpose operations $g\cdot
g^{\prime}=(P_{1}P_{1}^{\prime},\ldots,P_{M}P_{M}^{\prime}),\quad
g^{T}=(P_{1}^{T},\ldots,P_{m}^{T}).$ Note that the elementwise product and
transpose operations are well defined even if the $P_{m}$ and $P_{m}^{\prime}$
matrices are not permutations.

##  3 The weight alignment problem and its symmetries

The weight alignment problem. Given an MLP architecture as in equation 1 and
two weight-space vectors $v,v^{\prime}\in\mathcal{W}$, where
$v=[W_{m},b_{m}]_{m\in[M]},v^{\prime}=[W^{\prime}_{m},b^{\prime}_{m}]_{m\in[M]},$
the weight alignment problem is defined as the following optimization problem:

$\mathcal{G}(v,v^{\prime})=\text{argmin}_{k\in
G}\|v-k_{\\#}v^{\prime}\|^{2}_{2}$ (2)

In other words, the problem seeks a sequence of permutations
$k=(P_{1},\dots,P_{M-1})$ that will make $v^{\prime}$ as close as possible to
$v$. The optimization problem in equation 2 always admits a minimizer since
$G$ is finite. For some $(v,v^{\prime})$ it may have several minimizers, in
which case $\mathcal{G}(v,v^{\prime})$ is a set of elements. To simplify our
discussion we will sometimes consider the domain of $\mathcal{G}$ to be only
the set $\mathcal{V}^{2}_{\text{unique}}$ of pairs $(v,v^{\prime})$ for which
a unique minimizer exists. On this domain we can consider $\mathcal{G}$ as a
function to the unique minimizer in $G$, that is
$\mathcal{G}:\mathcal{V}^{2}_{\text{unique}}\rightarrow G$.

Our goal in this paper is to devise an architecture that can learn the
function $\mathcal{G}$. As a guiding principle for devising this architecture,
we would like this function to be equivariant to the symmetries of
$\mathcal{G}$. We describe these symmetries in the following subsection.

Figure 1:  The equivariance structure of the alignment problem. The function
$\mathcal{G}$ takes as input two weight space vectors $v,v^{\prime}$ and
outputs a sequence of permutation matrices that aligns them denoted
$\mathcal{G}(v,v^{\prime})$. In case we reorder the input using
$(g,g^{\prime})$ where
$g=(P_{1},P_{2}),g^{\prime}=(P^{\prime}_{1},P^{\prime}_{2})$, the optimal
alignment undergoes a transformation, namely
$\mathcal{G}(g_{\\#}v,g_{\\#}^{\prime}v^{\prime})=g\cdot\mathcal{G}(v,v^{\prime})\cdot
g^{\prime T}$ .

The symmetries of $\mathcal{G}$. One important property of the function
$\mathcal{G}$ is that it is equivariant to the action of the group $H=G\times
G$ which consists of two independent copies of the permutation symmetry group
for the MLP architecture we consider. Here, the action $h=(g,g^{\prime})\in H$
on the input space $\mathcal{V}\times\mathcal{V}$ is simply
$(v,v^{\prime})\mapsto(g_{\\#}v,g_{\\#}^{\prime}v^{\prime})$, and the action
of $h=(g,g^{\prime})\in H$ on an element $k\in G$ in the output space is given
by $g\cdot k\cdot{g^{\prime}}^{T}$. This equivariance property is summarized
and proved in the proposition below and visualized using the commutative
diagram in Figure 1: applying $\mathcal{G}$ and then $(g,g^{\prime})$ results
in exactly the same output as applying $(g,g^{\prime})$ and then
$\mathcal{G}$.

######  proposition 1.

The map $\mathcal{G}$ is $H$-equivariant, namely, for all
$(v,v^{\prime})\in\mathcal{V}^{2}_{\text{unique}}$ and $(g,g^{\prime})\in H$,

$\mathcal{G}(g_{\\#}v,g_{\\#}^{\prime}v^{\prime})=g\cdot\mathcal{G}(v,v^{\prime})\cdot
g^{\prime T}$

The function $\mathcal{G}$ exhibits another interesting property: swapping the
order of the inputs $v,v^{\prime}$ corresponds to inverting the optimal
alignment $\mathcal{G}(v,v^{\prime})$ :

######  proposition 2.

Let $(v,v^{\prime})\in\mathcal{V}^{2}_{\text{unique}}$ then
$\mathcal{G}(v^{\prime},v)=\mathcal{G}(v,v^{\prime})^{T}$.

We conclude by noting that, although for simplicity we focused on the case
where $(v,v^{\prime})\in\mathcal{V}^{2}_{\text{unique}}$, we can also state
analogous claims for the general case where multiple minimizers are possible.
In this case we will have that the equalities
$g\cdot\mathcal{G}(v,v^{\prime})\cdot g^{\prime
T}=\mathcal{G}(gv,g^{\prime}v^{\prime})$ and
$\mathcal{G}(v,v^{\prime})^{T}=\mathcal{G}(v^{\prime},v)$ still hold as
equalities between subsets of $G$.

Extension to other optimization objectives. In Appendix B we show that the
equivariant structure of the function $\mathcal{G}$ occurs not only for the
objective in equation 2, but also when the objective
$\|v-k_{\\#}v^{\prime}\|^{2}_{2}$ is replaced with any scalar function
$E(v,k_{\\#}v^{\prime})$ that satisfies the following properties: (1) $E$ is
invariant to the action of $G$ on both inputs; and (2) $E$ is invariant to
swapping its arguments.

##  4 Deep-Align

###  4.1 Architecture

Here, we define a neural network architecture $F=F(v,v^{\prime};\theta)$ for
learning the weight-alignment problem. The output of $F$ will be a sequence of
square matrices $(P_{1},\dots,P_{M-1})$ that represents a (sometimes
approximate) group element in $G$. In order to provide an effective inductive
bias, we will ensure that our architecture meets both properties: 1,2, namely
$F(g_{\\#}v,g_{\\#}^{\prime}v^{\prime})=g\cdot
F(v,v^{\prime})\cdot{g^{\prime}}^{T}\text{ and
}F(v,v^{\prime})=F(v^{\prime},v)^{T}.$ The architecture we propose is composed
of four functions:

$F=F_{proj}\circ F_{prod}\circ F_{\mathcal{V}\rightarrow\mathcal{A}}\circ
F_{DWS}:\mathcal{V}\times\mathcal{V}^{\prime}\rightarrow\bigoplus_{m=1}^{M-1}\mathbb{R}^{d_{m}\times
d_{m}},$

where the equivariance properties we require are guaranteed by constructing
each of the four functions composing $F$ to be equivariant with respect to an
appropriate action of $H=G\times G$ and the transposition action
$(v,v^{\prime})\mapsto(v^{\prime},v)$. In general terms, we choose $F_{DWS}$
to be a siamese weight space encoder, $F_{\mathcal{V}\rightarrow\mathcal{A}}$
is a siamese function that maps the weight space to the activation space,
$F_{prod}$ is a function that performs (generalized) outer products between
corresponding activation spaces in both networks and $F_{proj}$ performs a
projection of the resulting square matrices on the set of doubly stochastic
matrices (the convex hull of permutation matrices). The architecture is
illustrated in Figure 2. We now describe our architecture in more detail.

Figure 2: Our architecture is a composition of four blocks: The first block,
$F_{DWS}$ generates weight space embedding for both inputs. The second block
$F_{\mathcal{V}\rightarrow\mathcal{A}}$ maps these to the activation spaces.
The third block, $F_{Prod}$, generates square matrices by applying an outer
product between the activation vector of one network to the activation vectors
of the other network. Lastly, the fourth block, $F_{Proj}$ projects these
square matrices on the (convex hull of) permutation matrices.

Weight space encoder .
$F_{DWS}:\mathcal{V}\times\mathcal{V}^{\prime}\rightarrow\mathcal{V}^{d}\times\mathcal{V}^{\prime
d}$, where $d$ represents the number of feature channels, is implemented as a
Siamese DWSNet (Navon et al., 2023). This function outputs two weight-space
embeddings in $\mathcal{V}^{d}$, namely,
$F_{DWS}(v,v^{\prime})=(\mathcal{E}(v),\mathcal{E}(v^{\prime})),$ for a DWS
network $\mathcal{E}$. The Siamese structure of the network guarantees
equivariance to transposition, while the $G$-equivariance of DWSNet implies
equivariance with respect to the action of $G\times G$, that is
$(\mathcal{E}(g_{\\#}v),\mathcal{E}(g_{\\#}^{\prime}v^{\prime}))=(g_{\\#}\mathcal{E}(v),g_{\\#}^{\prime}\mathcal{E}(v^{\prime}))$.

Mapping the weight space to the activation space. The function
$F_{\mathcal{V}\rightarrow\mathcal{A}}:\mathcal{V}^{d}\times\mathcal{V}^{\prime
d}\rightarrow\mathcal{A}^{d}\times\mathcal{A}^{\prime d}$ maps the weight
spaces $\mathcal{V}^{d},\mathcal{V}^{\prime d}$ to the corresponding
Activation Spaces (see preliminaries section). There are several ways to
implement $F_{\mathcal{V}\rightarrow\mathcal{A}}$. As the bias space,
$\mathcal{B}=\bigoplus_{m=1}^{M}\mathcal{B}_{m}$, and the activation space
have a natural correspondence between them, perhaps the simplest way, which we
use in this paper, is to map a weight space vector $v=(w,b)\in\mathcal{V}^{d}$
to its bias component $b\in\mathcal{B}^{d}$. This operation is again
equivariant to transposition and the action of $G\times G$, where the action
of $G\times G$ on the input space is the more complicated action (by
$(g_{\\#},g_{\\#}^{\prime})$) on $\mathcal{V}\times\mathcal{V}$ and the action
on the output space is the simpler action of $G\times G$ on the activation
spaces.

Generalized outer product. $F_{prod}:\mathcal{A}^{d}\times\mathcal{A}^{\prime
d}\rightarrow\bigoplus_{m=1}^{M}\mathbb{R}^{d_{m}\times d_{m}}$ is a function
that takes the activation space features and performs a _generalized outer
product_ operation as defined below:

$F_{prod}(a,a^{\prime})_{m,i,j}=\phi([a_{m,i},a^{\prime}_{m,j}])$

where the subscripts $m,i,j$ represent the $(i,j)$-th entry of the $m$-th
matrix, and $a_{m,i},a^{\prime}_{m,j}\in\mathbb{R}^{d}$ are the rows of
$a,a^{\prime}$. Here, the function $\phi$ is a general (parametric or
nonparametric) symmetric function in the sense that $\phi(a,b)=\phi(b,a)$. In
this paper, we use $\phi(a,b)=s^{2}\langle a/\|a\|_{2},b/\|b\|_{2}\rangle$
where $s$ is a trainable scalar scaling factor. The equivariance with respect
to the action of $G\times G$ and transposition is guaranteed by the fact that
$\phi$ is applied elementwise, and is symmetric, respectively.

Projection layer. The output of $F_{prod}$ is a sequence of matrices
$Q_{1},\ldots,Q_{M-1}$ which in general will not be permutation matrices. To
bring the outputs closer to permutation matrices, $F_{proj}$ implements a
approximate projection onto the convex hull of the permutation matrices, i.e.,
the space of doubly stochastic matrices. In this paper, we use two different
projection operations, depending on whether the network is in training or
inference mode. At training time, to ensure differentiability, we implement
$F_{proj}$ as an approximation of a matrix-wise projection $Q_{m}$ to the
space of doubly stochastic matrices using several iterations of the well-known
Sinkhorn projection (Mena et al., 2018; Sinkhorn, 1967). Since the set of
doubly stochastic matrices is closed under the action of $G\times G$ on the
output space, and under matrix transposition, and since the Sinkhorn
iterations are composed of elementwise, row-wise, or column-wise operations,
we see that this operation is equivariant as well. At inference time, we
obtain permutation matrices from $Q_{i}$ by finding the permutation matrix
$P_{i}$ which has the highest correlation with $Q_{i}$, that is
$P_{i}=\operatorname*{arg\,max}_{P\in S_{d_{i}}}\langle Q_{i},P\rangle,$ where
the inner product is the standard Frobenious inner product. The optimization
problem, known as the _linear assignment problem_ can be solved using the
Hungarian algorithm.

As we carefully designed the components of $F$ so that they are all
equivariant to transposition and the action of $G\times G$, we obtain the
following proposition:

######  proposition 3.

The architecture $F$ satisfies the conditions specified in 1,2, namely for all
$(v,v^{\prime})\in\mathcal{V}\times\mathcal{V}$ and $(g,g^{\prime})\in H$ we
have: $F(g_{\\#}v,g_{\\#}^{\prime}v^{\prime})=g\cdot F(v,v^{\prime})\cdot
g^{\prime T}\text{ and }F(v,v^{\prime})=F(v^{\prime},v)^{T}.$

###  4.2 Data generation and Loss functions

Generating labeled data for the weight-alignment problem is hard due to the
intractability of the problem. Therefore, we propose a combination of both
unsupervised and supervised loss functions where we generate labeled examples
synthetically from unlabeled examples, as specified below.

Data generation. Our initial training data consists of a finite set of weight
space vectors $D\subset\mathcal{V}$. From that set, we generate two datasets
consisting of pairs of weights for the alignment problem. First, we generate a
labeled training set, $D_{\text{labeled}}=\\{(v^{j},v^{\prime
j},t^{j})\\}_{j=1}^{N_{\text{labeled}}}$ for
$t^{j}=(T^{j}_{1},\dots,T^{j}_{M-1})\in G$. This is done by sampling $v^{j}\in
D$ and defining $v^{\prime j}$ as a permuted and noisy version of $v^{j}$.
More formally, we sample a sequence of permutations $t\in G$ and define
$v^{\prime j}=t_{\\#}f_{\text{aug}}(v^{j})$, where $f_{\text{aug}}$ applies
several weight-space augmentations, like adding binary and Gaussian noise,
scaling augmentations for ReLU networks, etc. We then set the label of this
pair to be $t$. In addition, we define an unlabeled dataset
$D_{\text{unlabeled}}=\\{(v^{j},v^{\prime
j})\\}_{j=1}^{{N}_{\text{unlabeled}}}$ where $v^{j},v^{\prime
j}\in\mathcal{V}$.

Loss functions. The datasets above are used for training our architecture
using the following loss functions. The labeled training examples in
$D_{\text{labeled}}$ are used by applying a cross-entropy loss for each row
$i=1,\dots,d_{m}$ in each output matrix $m=1,\dots,M-1$. This loss is denoted
as $\ell_{\text{supervised}}(F(v,v^{\prime};\theta),t)$. The unlabeled
training examples are used in combination with two unsupervised loss
functions. The first loss function aims to minimize the alignment loss in
equation 2 directly by using the network output $F(v,v^{\prime};\theta)$ as
the permutation sequence. This loss is denoted as
$\ell_{\text{alignment}}(v,v^{\prime},\theta)=\|v-F(v,v^{\prime};\theta)_{\\#}v^{\prime}\|_{2}^{2}$.
The second unsupervised loss function aims to minimize the original loss
function used to train the input networks on a line segment connecting the
weights $v$ and the transformed version of $v^{\prime}$ using the network
output $F(v,v^{\prime};\theta)$ as the permutation sequence. Concretely, let
$\mathcal{L}$ denote the original loss function for the weight vectors
$v,v^{\prime}$, the loss is defined as
$\ell_{\text{LMC}}(v,v^{\prime},\theta)=\mathcal{L}(\lambda
v+(1-\lambda)F(v,v^{\prime};\theta)_{\\#}v^{\prime})$ for $\lambda$ sampled
uniformly $\lambda\sim U(0,1)$111This loss function satisfies the properties
as described in Section 3 when taking expectation over $\lambda$.. This loss
is similar to the STE method in Ainsworth et al. (2022) and the differentiable
version in Peña et al. (2023). Our final goal is to minimize the parameters of
$F$ with respect to a linear (positive) combination of
$\ell_{\text{alignment}},\ell_{\text{LMC}}$ and $\ell_{\text{supervised}}$
applied to the appropriate datasets described above.

##  5 Theoretical analysis

Relation to the activation matching algorithm. In this subsection, we prove
that our proposed architecture can simulate the activation matching algorithm,
a heuristic for solving the weight alignment problem suggested in Ainsworth et
al. (2022). In a nutshell, this algorithm works by evaluating two neural
networks on a set of inputs and finding permutations that align their
activations by solving a linear assignment problem using the outer product
matrix of the activations as a cost matrix for every layer $m=1,\dots,M-1$.

######  proposition 4.

(Deep-Align can simulate activation matching) For any compact set
$K\subset\mathcal{V}$ and $x_{1},\dots,x_{N}\in\mathbb{R}^{d_{0}}$, there
exists an instance of our architecture $F$ and weights $\theta$ such that for
any $v,v^{\prime}\in K$ for which the activation matching algorithm has a
single optimal solution $g\in G$ and another minor assumption specified in the
appendix, $F(v,v^{\prime};\theta)$ returns $g$.

This result offers an interesting interpretation of our architecture: the
architecture can simulate activation matching while optimizing the input
vectors $x_{1},\dots,x_{N}$ as a part of their weights $\theta$.

Exactness. We now discuss the _exactness_ of our algorithms. An alignment
algorithm is said to be exact on some input $(v,v^{\prime})$ if it can be
proven to successfully return the correct minimizer
$\mathcal{G}(v,v^{\prime})$. For NP-hard alignment problems such as weight
alignment, exactness can typically be obtained when restricting it to ‘tame’
inputs $(v,v^{\prime})$. Examples of exactness results in the alignment
literature can be found in Aflalo et al. (2015); Dym & Lipman (2017); Dym
(2018). The following proposition shows that (up to probability zero events)
when $v,v^{\prime}$ are exactly related by some $g\in G$, our algorithm will
retrieve $g$ _exactly_ :

######  proposition 5 (Deep-Align is exact for perfect alignments).

Let $F$ denote the Deep-Align architecture with non-constant analytic
activations and $d\geq 2$ channels. Then, for Lebesgue almost every
$v\in\mathcal{V}$ and parameter vector $\theta$, and for every $g\in G$, we
have that $F(v,g_{\\#}v,\theta)=g$.

##  6 Experiments

Table 1: MLP image classifiers: Results on aligning MNIST and CIFAR10 MLP image classifiers. | MNIST (MLP) | CIFAR10 (MLP)
---|---|---
| Barrier $\downarrow$ | AUC $\downarrow$ | Barrier $\downarrow$ | AUC $\downarrow$
Naive | $2.007\pm 0.00$ | $0.835\pm 0.00$ | $0.927\pm 0.00$ | $0.493\pm 0.00$
Weight Matching | $0.047\pm 0.00$ | $0.011\pm 0.00$ | $0.156\pm 0.00$ | $0.068\pm 0.00$
Activation Matching | $0.024\pm 0.00$ | $0.007\pm 0.00$ | $0.066\pm 0.00$ | $0.024\pm 0.00$
Sinkhorn | $0.027\pm 0.00$ | $0.002\pm 0.00$ | $0.183\pm 0.00$ | $0.072\pm 0.00$
Deep-Align | $0.005\pm 0.00$ | $\mathbf{0.000\pm 0.00}$ | $0.078\pm 0.01$ | $0.029\pm 0.00$
Deep-Align \+ Sinkhorn | $\mathbf{0.000\pm 0.00}$ | $\mathbf{0.000\pm 0.00}$ | $\mathbf{0.037\pm 0.00}$ | $\mathbf{0.004\pm 0.00}$
Table 2: CNN image classifiers: Results on aligning CIFAR10 and STL10 CNN image classifiers. | CIFAR10 (CNN) | STL10 (CNN) | Runtime (Sec) $\downarrow$
---|---|---|---
| Barrier $\downarrow$ | AUC $\downarrow$ | Barrier $\downarrow$ | AUC $\downarrow$
Naive | $1.124\pm 0.01$ | $0.524\pm 0.00$ | $1.006\pm 0.00$ | $0.650\pm 0.00$ | —
Weight Matching | $0.661\pm 0.02$ | $0.178\pm 0.01$ | $0.859\pm 0.00$ | $0.453\pm 0.00$ | $0.21$
Activation Matching | $0.238\pm 0.01$ | $\mathbf{0.000\pm 0.00}$ | $0.479\pm 0.00$ | $0.250\pm 0.00$ | $7.52$
Sinkhorn | $0.313\pm 0.01$ | $\mathbf{0.000\pm 0.00}$ | $0.366\pm 0.00$ | $0.163\pm 0.00$ | $79.81$
Deep-Align | $0.237\pm 0.01$ | $\mathbf{0.000\pm 0.00}$ | $0.382\pm 0.01$ | $0.182\pm 0.00$ | $0.44$
Deep-Align \+ Sinkhorn | $\mathbf{0.081\pm 0.00}$ | $\mathbf{0.000\pm 0.00}$ | $\mathbf{0.232\pm 0.00}$ | $\mathbf{0.097\pm 0.00}$ | $80.25=0.44+79.81$

In this section, we evaluate Deep-Align on the task of aligning and merging
neural networks. To support future research and the reproducibility of the
results, we will make our source code and datasets publicly available upon
publication.

Evaluation metrics. We use the standard evaluation metrics for measuring model
merging (Ainsworth et al., 2022; Peña et al., 2023): Barrier and Area Under
the Curve (AUC). For two inputs $v,v^{\prime}$ the Barrier is defined by
$\max_{\lambda\in[0,1]}\psi(\lambda)\equiv\mathcal{L}(\lambda
v+(1-\lambda)v^{\prime})-(\lambda\mathcal{L}(v)+(1-\lambda)\mathcal{L}(v^{\prime}))$
where $\mathcal{L}$ denote the loss function on the original task. Similarly,
the AUC is defined as the integral of $\psi$ over $[0,1]$. Lower is better for
both metrics. Following previous works (Ainsworth et al., 2022; Peña et al.,
2023), we bound both metrics by taking the maximum between their value and
zero.

Compared methods. We compare the following approaches: (1) Naive: where two
models are merged by averaging the models’ weights without alignment. The (2)
Weight matching and (3) Activation matching approaches proposed in Ainsworth
et al. (2022). (4) Sinkhorn (Peña et al., 2023): This approach directly
optimizes the permutation matrices using the task loss on the line segment
between the aligned models (denoted $\mathcal{C}_{Rnd}$ in Peña et al.
(2023)). (5) Deep-Align: Our proposed method described in Section 4. (5)
_Deep-Align \+ Sinkhorn_: Here, the output from the Deep-Align is used as an
initialization for the Sinkhorn method.

Experimental details. Our method is first trained on a dataset of weight
vectors and then applied to unseen weight vectors at test time, as is standard
in learning setups. In contrast, baseline methods are directly optimized using
the test networks. For the _Sinkhorn_ and _Deep-Align \+ Sinkhorn_ methods, we
optimize the permutations for 1000 iterations. For the _Activation Matching_
method, we calculate the activations using the entire train dataset. We repeat
all experiments using 3 random seeds and report each metric’s mean and
standard deviation. For full experimental details see Appendix E.

(a) MNIST MLPs.

(b) CIFAR10 MLPs.

(c) CIFAR10 CNNs.

(d) STL10 CNNs.

Figure 3:  Merging image classifiers: the plots illustrate the values of the
loss function used for training the input networks when evaluated on a line
segment connecting $v$ and $g_{\\#}v^{\prime}$, where $g$ is the output of
each method. Values are averaged over all test images and networks and 3
random seeds.

###  6.1 Results

Aligning classifiers. Here, we evaluate our method on the task of aligning
image classifiers. We use four network datasets. Two datasets consist of MLP
classifiers for MNIST and CIFAR10, and two datasets consist of CNN classifiers
trained using CIFAR10 and STL10. This collection forms a diverse benchmark for
aligning NN classifiers. The results are presented in Figure 6, Table 1 and
Table 2. The alignment produced through a feed-forward pass with Deep-Align
performs on par or outperforms all baseline methods. Initializing the Sinkhorn
algorithm with our alignment (Deep-Align \+ Sinkhorn) further improves the
results, and significantly outperforms all other methods.

Aligning INRs. We use two datasets consisting of implicit neural
representations (INRs). The first consists of Sine waves INRs of the form
$f(x)=\sin(ax)$ on $[-\pi,\pi]$, where $a\sim U(0.5,10)$, similarly to the
data used in Navon et al. (2023). We fit two views (independently trained
weight vectors) for each value of $a$ starting from different random
initializations and the task is to align and merge the two INRs. We train our
network to align pairs of corresponding views. The second dataset consists of
INRs fitted to CIFAR10 images. We fit five views per image. The results are
presented in Figure 4. Deep-Align, performs on par or outperforms all baseline
methods. Moreover, using the output from the Deep-Align to initialize the
Sinkhorn algorithm further improves this result, with a large improvement over
the Sinkhorn baseline with random initialization.

#### Generalization to out-of-distribution data (OOD).

Here, we evaluate the generalization capabilities of Deep-Align under
distribution shift. We use the Deep-Align model trained on CIFAR10 CNN image
classifiers and evaluate the generalization on two datasets. The first dataset
consists of CNN classifiers trained on a version of CIFAR10 in which each
image is rotated by a rotation degree sampled uniformly from $U(-45,45)$. The
second dataset consists of CNN image classifiers trained on the STL10 dataset.
Importantly, we note that DEEP-ALIGN is evaluated on a distribution of models
that is different than the one observed during training. In contrast, the
baselines directly solve an optimization problem for each model pair within
the test datasets. While Deep-Align significantly outperforms the Naive and WM
baselines, it falls short in comparison to the Sinkhorn and AM methods, both
of which are directly optimized using data from the new domain (NNs and
images). Employing Deep-Align as an initialization for the Sinkhorn method
consistently proves beneficial, with the Deep-Align \+ Sinkhorn approach
yielding the most favorable results in terms of the barrier metric.

(a) Sine Wave INRs.

(b) CIFAR10 INRs.

Figure 4: Aligning INRs: The test barrier vs. the number of Sinkhorn iterations ( relevant only for Sinkhorn or Deep-Align \+ Sinkhorn), using (a) sine wave and (b) CIFAR10 INRs. Deep-Align outperforms baseline methods or achieves on-par results. Our Deep-Align \+ Sinkhorn further improves this result and significantly improves over Sinkhorn with random initialization, both in terms of test barrier and convergence speed. Table 3: Aligning OOD image classifiers, using a Deep-Align network trained on CIFAR10. | Rotated CIFAR10 | STL10
---|---|---
| Barrier $\downarrow$ | AUC $\downarrow$ | Barrier $\downarrow$ | AUC $\downarrow$
Naive | $1.077\pm 0.01$ | $0.714\pm 0.00$ | $1.006\pm 0.00$ | $0.650\pm 0.00$
Weight Matching | $0.945\pm 0.02$ | $0.550\pm 0.01$ | $0.859\pm 0.00$ | $0.453\pm 0.00$
Activation Matching | $0.586\pm 0.00$ | $0.336\pm 0.00$ | $0.479\pm 0.00$ | $0.250\pm 0.00$
Sinkhorn | $0.596\pm 0.01$ | $0.321\pm 0.00$ | $0.366\pm 0.00$ | $\mathbf{0.163\pm 0.00}$
Deep-Align | $0.769\pm 0.01$ | $0.453\pm 0.00$ | $0.686\pm 0.01$ | $0.373\pm 0.01$
Deep-Align \+ Sinkhorn | $\mathbf{0.430\pm 0.01}$ | $\mathbf{0.245\pm 0.00}$ | $\mathbf{0.357\pm 0.00}$ | $0.165\pm 0.00$

#### Time comparison.

Prior methods for weight matching, which rely on optimization, often suffer
from exhaustive runtime which may be impractical for real-time applications.
In contrast, once trained, Deep-Align is able to produce high-quality weight
alignments through a single forward pass and an efficient projection step.
Here, we compare Deep-Align to baselines by measuring the time required to
align a pair of models in the CIFAR10 CNN classifiers dataset, and report the
averaged alignment time using 1000 random pairs on a single RTX 2080-Ti Nvidia
GPU. The results are presented in Table 2. Deep-Align is significantly faster
than Sinkhorn and Activation Matching while achieving comparable results.
Furthermore, Deep-Align is on par with Weight Matching w.r.t runtime, yet it
consistently generates better weight alignment solutions.

#### Aligning networks from disjoint datasets.

Figure 5: Merging networks trained on distinct subsets of CIFAR10.

Following Ainsworth et al. (2022), we experiment with aligning networks
trained on disjoint datasets. One major motivation for such a setup is
Federated learning (McMahan et al., 2017). In Federated Learning, the goal is
to construct a unified model from multiple networks trained on separate and
distinct datasets.

To that end, we split the CIFAR10 dataset into two splits. The first consists
of $95\%$ images from classes $0$-$4$ and $5\%$ of classes $5$-$9$, and the
second split is constructed accordingly with $95\%$ of classes $5$-$9$. We
train the Deep-Align model to align CNN networks trained using the different
datasets. For Sinkhorn and Activation Matching, we assume full access to the
training data in the optimization stage. For DEEP-ALIGN, we assume this data
is accessible in the training phase. The results are presented in Figure 5.
Our approach, Deep-Align, along with the Sinkhorn and Activation Matching
approaches, are able to align and merge the networks to obtain a network with
lower loss compared to the original models. However, our approach is
significantly more efficient at inference.

##  7 Conclusion

We investigate the challenging problem of weight alignment in deep neural
networks. The key to our approach, Deep-Align, is an equivariant architecture
that respects the natural symmetries of the problem. At inference time Deep-
Align can align unseen network pairs without the need for performing expensive
optimization. Deep-Align, performs on par or outperforms optimization-based
approaches while significantly reducing the runtime or improving the quality
of the alignments. Furthermore, we demonstrate that the alignments of our
method can be used to initialize optimization-based approaches. One limitation
of our approach is the need for training a network. Although this can be a
relatively time-consuming process, we only have to perform it once for each
weight distribution. Furthermore, this procedure does not require labeled
examples. To summarize, Deep-Align is the first architecture designed for
weight alignment. It demonstrates superior performance over existing methods.
The generalization capabilities of Deep-Align make it a promising and
practical solution for applications that require weight alignment.

##  8 Acknowledgements

This study was funded by a grant to GC from the Israel Science Foundation (ISF
737/2018), and by an equipment grant to GC and Bar-Ilan University from the
Israel Science Foundation (ISF 2332/18).

## References

  * Aflalo et al. (2015) Yonathan Aflalo, Alexander Bronstein, and Ron Kimmel.  On convex relaxation of graph isomorphism.  _Proceedings of the National Academy of Sciences_ , 112(10):2942–2947, 2015.
  * Ainsworth et al. (2022) Samuel K Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa.  Git re-basin: Merging models modulo permutation symmetries.  _arXiv preprint arXiv:2209.04836_ , 2022.
  * Akash et al. (2022) Aditya Kumar Akash, Sixu Li, and Nicolás García Trillos.  Wasserstein barycenter-based model fusion and linear mode connectivity of neural networks.  _arXiv preprint arXiv:2210.06671_ , 2022.
  * Andreis et al. (2023) Bruno Andreis, Soro Bedionita, and Sung Ju Hwang.  Set-based neural network encoding.  _arXiv preprint arXiv:2305.16625_ , 2023.
  * Bengio et al. (2021) Yoshua Bengio, Andrea Lodi, and Antoine Prouvost.  Machine learning for combinatorial optimization: a methodological tour d’horizon.  _European Journal of Operational Research_ , 290(2):405–421, 2021.
  * Cappart et al. (2021) Quentin Cappart, Didier Chételat, Elias B. Khalil, Andrea Lodi, Christopher Morris, and Petar Velickovic.  Combinatorial optimization and reasoning with graph neural networks.  _CoRR_ , abs/2102.09544, 2021.  URL https://arxiv.org/abs/2102.09544.
  * Draxler et al. (2018) Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht.  Essentially no barriers in neural network energy landscape.  In _International conference on machine learning_ , pp. 1309–1318. PMLR, 2018.
  * Dym (2018) Nadav Dym.  Exact recovery with symmetries for the doubly stochastic relaxation.  _SIAM Journal on Applied Algebra and Geometry_ , 2(3):462–488, 2018.
  * Dym & Lipman (2017) Nadav Dym and Yaron Lipman.  Exact recovery with symmetries for procrustes matching.  _SIAM Journal on Optimization_ , 27(3):1513–1530, 2017.
  * Eilertsen et al. (2020) Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman.  Classifying the classifier: dissecting the weight space of neural networks.  _arXiv preprint arXiv:2002.05688_ , 2020.
  * Entezari et al. (2022) Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur.  The role of permutation invariance in linear mode connectivity of neural networks.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=dNigytemkL.
  * Fey et al. (2020) Matthias Fey, Jan E Lenssen, Christopher Morris, Jonathan Masci, and Nils M Kriege.  Deep graph matching consensus.  _arXiv preprint arXiv:2001.09621_ , 2020.
  * Freeman & Bruna (2016) C Daniel Freeman and Joan Bruna.  Topology and geometry of half-rectified network optimization.  _arXiv preprint arXiv:1611.01540_ , 2016.
  * Garipov et al. (2018) Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson.  Loss surfaces, mode connectivity, and fast ensembling of dnns.  _Advances in neural information processing systems_ , 31, 2018.
  * Ilharco et al. (2022) Gabriel Ilharco, Marco Tulio Ribeiro, Mitchell Wortsman, Suchin Gururangan, Ludwig Schmidt, Hannaneh Hajishirzi, and Ali Farhadi.  Editing models with task arithmetic.  _arXiv preprint arXiv:2212.04089_ , 2022.
  * Jordan et al. (2022) Keller Jordan, Hanie Sedghi, Olga Saukh, Rahim Entezari, and Behnam Neyshabur.  Repair: Renormalizing permuted activations for interpolation repair.  _arXiv preprint arXiv:2211.08403_ , 2022.
  * Khalil et al. (2017) Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song.  Learning combinatorial optimization algorithms over graphs.  _Advances in neural information processing systems_ , 30, 2017.
  * Lim et al. (2022) Derek Lim, Joshua Robinson, Lingxiao Zhao, Tess Smidt, Suvrit Sra, Haggai Maron, and Stefanie Jegelka.  Sign and basis invariant networks for spectral graph representation learning.  _arXiv preprint arXiv:2202.13013_ , 2022.
  * Loshchilov & Hutter (2017) Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  _arXiv preprint arXiv:1711.05101_ , 2017.
  * McMahan et al. (2017) Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas.  Communication-efficient learning of deep networks from decentralized data.  In _Artificial intelligence and statistics_ , pp. 1273–1282. PMLR, 2017.
  * Mena et al. (2018) Gonzalo Mena, David Belanger, Scott Linderman, and Jasper Snoek.  Learning latent permutations with gumbel-sinkhorn networks.  _arXiv preprint arXiv:1802.08665_ , 2018.
  * Mityagin (2015) Boris Mityagin.  The zero set of a real analytic function.  _arXiv preprint arXiv:1512.07276_ , 2015.
  * Navon et al. (2023) Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron.  Equivariant architectures for learning in deep weight spaces.  _arXiv preprint arXiv:2301.12780_ , 2023.
  * Nowak et al. (2017) Alex Nowak, Soledad Villar, Afonso S Bandeira, and Joan Bruna.  A note on learning algorithms for quadratic assignment with graph neural networks.  _stat_ , 1050:22, 2017.
  * Nowak et al. (2018) Alex Nowak, Soledad Villar, Afonso S Bandeira, and Joan Bruna.  Revised note on learning quadratic assignment with graph neural networks.  In _2018 IEEE Data Science Workshop (DSW)_ , pp. 1–5. IEEE, 2018\.
  * Peña et al. (2023) Fidel A Guerrero Peña, Heitor Rapela Medeiros, Thomas Dubail, Masih Aminbeidokhti, Eric Granger, and Marco Pedersoli.  Re-basin via implicit sinkhorn differentiation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 20237–20246, 2023.
  * Selsam et al. (2018) Daniel Selsam, Matthew Lamm, Benedikt Bünz, Percy Liang, Leonardo de Moura, and David L Dill.  Learning a sat solver from single-bit supervision.  _arXiv preprint arXiv:1802.03685_ , 2018.
  * Sinkhorn (1967) Richard Sinkhorn.  Diagonal equivalence to matrices with prescribed row and column sums.  _The American Mathematical Monthly_ , 74(4):402–405, 1967.
  * Stoica et al. (2023) George Stoica, Daniel Bolya, Jakob Bjorner, Taylor Hearn, and Judy Hoffman.  Zipit! merging models from different tasks without training.  _arXiv preprint arXiv:2305.03053_ , 2023.
  * Tatro et al. (2020) Norman Tatro, Pin-Yu Chen, Payel Das, Igor Melnyk, Prasanna Sattigeri, and Rongjie Lai.  Optimizing mode connectivity via neuron alignment.  _Advances in Neural Information Processing Systems_ , 33:15300–15311, 2020.
  * Unterthiner et al. (2020) Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and Ilya Tolstikhin.  Predicting neural network accuracy from weights.  _arXiv preprint arXiv:2002.11448_ , 2020.
  * Vesselinova et al. (2020) Natalia Vesselinova, Rebecca Steinert, Daniel F Perez-Ramirez, and Magnus Boman.  Learning combinatorial optimization on graphs: A survey with applications to networking.  _IEEE Access_ , 8:120388–120416, 2020.
  * Wortsman et al. (2022) Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes, Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al.  Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time.  In _International Conference on Machine Learning_ , pp. 23965–23998. PMLR, 2022.
  * Yan et al. (2020) Junchi Yan, Shuang Yang, and Edwin R Hancock.  Learning for graph matching and related combinatorial optimization problems.  In _Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20_ , pp. 4988–4996. International Joint Conferences on Artificial Intelligence Organization, 2020\.
  * Yu et al. (2019) Tianshu Yu, Runzhong Wang, Junchi Yan, and Baoxin Li.  Learning deep graph matching with channel-independent embedding and hungarian attention.  In _International conference on learning representations_ , 2019.
  * Zanfir & Sminchisescu (2018) Andrei Zanfir and Cristian Sminchisescu.  Deep learning of graph matching.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pp. 2684–2693, 2018.
  * Zhang et al. (2023) David W Zhang, Miltiadis Kofinas, Yan Zhang, Yunlu Chen, Gertjan J Burghouts, and Cees GM Snoek.  Neural networks are graphs! graph neural networks for equivariant processing of neural networks.  2023\.
  * Zhou et al. (2023a) Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, J Zico Kolter, and Chelsea Finn.  Permutation equivariant neural functionals.  _arXiv preprint arXiv:2302.14040_ , 2023a.
  * Zhou et al. (2023b) Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn.  Neural functional transformers.  _arXiv preprint arXiv:2305.13546_ , 2023b.

##  Appendix A More previous work

Weight space alignment. Several algorithms have been proposed for weight-
alignment (Tatro et al., 2020; Ainsworth et al., 2022; Peña et al., 2023;
Akash et al., 2022). Ainsworth et al. (2022) presented three algorithms:
Activation Matching, Weight Matching, and straight-through estimation. Peña et
al. (2023) improved upon these algorithms by incorporating a Sinkhorn-based
projection method. In part, these works were motivated by studying the loss
landscapes of deep neural networks. It was conjectured that deep networks
exhibit a property called _linear mode connectivity_ : for any two trained
weight vectors, a linear interpolation between the first vector and the
optimal alignment of the second, yields very small increases in the loss
(Entezari et al., 2022; Garipov et al., 2018; Draxler et al., 2018; Freeman &
Bruna, 2016; Tatro et al., 2020; Jordan et al., 2022).

Weight-space networks. A growing area of research focuses on applying neural
networks to neural network weights. Early methods proposed using simple
architectures such as MLPs or transformers to predict test errors or the
hyperparameters that were used for training input networks (Unterthiner et
al., 2020; Andreis et al., 2023; Eilertsen et al., 2020). Recently, Navon et
al. (2023) presented the first neural architecture that is equivariant to the
natural permutation symmetries of (MLP) weight spaces and demonstrated
significant performance improvements over previous approaches. Zhou et al.
(2023a) proposed a similar approach and an extension to CNN architectures ,
which was later enhanced by the addition of attention mechanisms (Zhou et al.,
2023b). Finally, Zhang et al. (2023) proposes modeling the neural networks as
computational graphs and applying Graph Neural Networks to them, demonstrating
very good results on several weight space learning tasks.

Learning for combinatorial optimization. There exists a large body of research
on learning to solve hard combinatorial optimization problems such as TSP and
SAT (Cappart et al., 2021; Bengio et al., 2021; Khalil et al., 2017;
Vesselinova et al., 2020; Selsam et al., 2018). The key observation behind
these works, also shared by the current work, is that even though those
problems are computationally intractable in general, there may be efficient
methods to solve them for specific problem distributions. In these cases,
machine learning can be used to learn to predict these solutions. Specifically
relevant to this work are works that suggested learning to solve the graph
matching problem Fey et al. (2020); Zanfir & Sminchisescu (2018); Yan et al.
(2020); Yu et al. (2019); Nowak et al. (2017; 2018) which is a similar
alignment problem.

##  Appendix B Proofs for section 3

###### Proof of Proposition 1.

We write $\mathcal{G}(v,v^{\prime})=\text{argmin}_{k\in G}E(k,v,v^{\prime})$
with $E(k,v,v^{\prime})=\|v-k_{\\#}v^{\prime}\|^{2}_{2}$. First, we note that
the minimal value of the optimization problem $\mathcal{G}(gv,gv^{\prime})$ is
equal to the minimal value of $\mathcal{G}(v,v^{\prime})$, namely,

$\min_{k\in G}E(k,v,v^{\prime})=\min_{k\in
G}E(k,g_{\\#}v,g_{\\#}^{\prime}v^{\prime}).$

This is true since

$\min_{k\in G}E(k,g_{\\#}v,g_{\\#}^{\prime}v^{\prime})=\min_{k\in
G}\|g_{\\#}v-k_{\\#}g_{\\#}^{\prime}v^{\prime}\|^{2}_{2}=\min_{k\in
G}\|v-(g^{-1}kg^{\prime})_{\\#}v^{\prime}\|^{2}_{2}=\min_{k\in
G}\|k_{\\#}v-v^{\prime}\|^{2}_{2}=$

and using the fact that $k\mapsto g^{-1}kg^{\prime}$ is bijective.

Second, we show that if $k^{*}=\mathcal{G}(v,v^{\prime})$, or in other words,
is a minimizer of $E(k,v,v^{\prime})$, then $g\cdot k^{*}\cdot g^{\prime T}$
minimizes $E(k,gv,g^{\prime}v^{\prime})$ because:

$E(g\cdot k^{*}\cdot g^{\prime T},gv,g^{\prime}v^{\prime})=\|g_{\\#}v-(g\cdot
k^{*}\cdot g^{\prime
T})_{\\#}g_{\\#}^{\prime}v^{\prime}\|^{2}_{2}=\|v-(g^{T}\cdot g\cdot
k^{*}\cdot{g^{\prime}}^{T}\cdot
g^{\prime})_{\\#}v^{\prime}\|^{2}=E(k^{*},v,v^{\prime})$

∎

###### Proof of Proposition 2.

Similarly to the proof of the previous proposition, we have $\min_{k\in
G}E(k,v^{\prime},v)=\min_{k\in G}\|v^{\prime}-k_{\\#}v\|_{2}^{2}=\min_{k\in
G}\|k_{\\#}^{T}v^{\prime}-v\|_{2}^{2}=\min_{k\in G}E(k,v,v^{\prime})$.
Additionally, plugging in $k=\mathcal{G}(v,v^{\prime})^{T}$ achieves this
value:

$E(\mathcal{G}(v,v^{\prime})^{T},v^{\prime},v)=\|v^{\prime}-[\mathcal{G}(v,v^{\prime})^{T}]_{\\#}v\|_{2}^{2}=\|[\mathcal{G}(v,v^{\prime})]_{\\#}v^{\prime}-v\|_{2}^{2}=E(\mathcal{G}(v,v^{\prime}),v,v^{\prime})$

∎

###### Proof of generalization to other objectives.

We prove the generalization mentioned in the main text. Namely, that the
symmetries of the function $\mathcal{G}$, are shared by any function of the
form $\mathcal{G}(v,v^{\prime})=\text{argmin}_{k\in G}E(v,k_{\\#}v^{\prime})$
providing that the energy function $E$ satisfies

$E(gv,gv^{\prime})=E(v,v^{\prime})\text{ and
}E(v,v^{\prime})=E(v^{\prime},v),\quad\forall g\in
G,v,v^{\prime}\in\mathcal{V}.$

First, we note that the minimal value of the optimization problem
$\mathcal{G}(gv,gv^{\prime})$ is equal to the minimal value of
$\mathcal{G}(v,v^{\prime})$, namely,

$\min_{k\in G}E(v,k_{\\#}v^{\prime})=\min_{k\in
G}E(g_{\\#}v,k_{\\#}g_{\\#}^{\prime}v^{\prime}).$

This is true since

$\min_{k\in G}E(g_{\\#}v,k_{\\#}g_{\\#}^{\prime}v^{\prime})=\min_{k\in
G}E(v,g_{\\#}^{T}k_{\\#}g_{\\#}^{\prime}v^{\prime})=\min_{k\in
G}E(v,(g^{-1}kg^{\prime})_{\\#}v^{\prime})$

using the fact that $k\mapsto g^{-1}kg^{\prime}$ is bijective and the
$G$-invariance of $E$.

Next, we show that if $k^{*}=\mathcal{G}(v,v^{\prime})$, or in other words, is
a minimizer of $E(v,k_{\\#}v^{\prime})$, then $g\cdot k^{*}\cdot g^{\prime T}$
minimizes $E(g_{\\#}v,k_{\\#}g_{\\#}^{\prime}v^{\prime})$ This is because:
$E(g_{\\#}v,(g\cdot k^{*}\cdot g^{\prime
T})_{\\#}g_{\\#}^{\prime}v^{\prime})=E(v,k_{\\#}^{*}v^{\prime})$ using the
$G$-invariance of $E$.

We turn to proving a generalization of 2: similarly to the previous argument,
we have

$\displaystyle\min_{k\in G}E(v^{\prime},k_{\\#}v)$ $\displaystyle=\min_{k\in
G}E(k_{\\#}^{T}v^{\prime},v)=\min_{k\in
G}E(k_{\\#}^{T}v^{\prime},v)=\min_{k\in G}E(v,k_{\\#}^{T}v^{\prime})$
$\displaystyle=\min_{k\in G}E(v,k_{\\#}v^{\prime})$

where we used the fact that we can swap the inputs to $E$ and the invariance.

Additionally, plugging in $k=\mathcal{G}(v,v^{\prime})^{T}$ achieves this
value:

$E(v^{\prime},\mathcal{G}(v,v^{\prime})_{\\#}^{T}v)=E(\mathcal{G}(v,v^{\prime})_{\\#}v^{\prime},v)=E(v,\mathcal{G}(v,v^{\prime})_{\\#}v^{\prime})$

∎

##  Appendix C Proofs for section 5

#### Relation to activation matching.

Here we prove Proposition 4. Let the outer product matrix
$Z\in\mathbb{R}^{d_{m}\times d_{m}}$ be a cost matrix for the activation
matching algorithm. We say that Z is ${\epsilon}$-_friendly_ if it has the
following property: the exists some $i$ such that for all $j\neq i$ $\langle
Z,P_{i}\rangle>\langle Z,P_{j}\rangle+{\epsilon}$ where
$P_{i},P_{j}\in\Pi_{d_{m}}$. Intuitively this condition means that the cost
matrix is bounded away from the set of cost matrices for which the are
multiple optimal solutions. Let us now state Proposition 4 with full details:

######  proposition 6 (Full formulation of Proposition 4).

For any compact set $K\subset\mathcal{V}$,
$x_{1},\dots,x_{N}\in\mathbb{R}^{d_{0}}$ and for any ${\epsilon}>0$, there
exists an instance of our architecture $F$ and weights $\theta$ for that
architecture such that for any $v,v^{\prime}\in K$ for which all the cost
matrices used by the activation matching algorithm are ${\epsilon}$-friendly,
we have that $F(v,v^{\prime};\theta)$ returns exactly the same solution as the
activation matching algorithm.

###### Proof of Proposition 6.

First, to obtain weights for the DWS network that approximate the activations
of both networks on the set of inputs $x_{1},\dots,x_{n}$, we employ Lemma G.2
from Navon et al. (2023) (stated below for convenience). We note that this
lemma is proved by the approximation of all intermediate activations of the
input networks so the proof of this lemma trivially shows that there is a
DWSnet that outputs an approximation of all the activations.

Then, we set $F_{prod}$ to calculate outer products in order to generate the
outer-product-based cost matrix used in Ainsworth et al. (2022). It follows
that the composition of the DWS network above with the outer product function
approximates uniformly their limits, see Lim et al. (2022) (Lemma 6), and the
composition of their limits is exactly the function that takes two weight
vectors and computes the cost matrices from Ainsworth et al. (2022). We have
developed an architecture and weights for this architecture that can uniformly
approximate the cost matrices used by the activation matching algorithm to any
precision. As a final step, we apply our linear assignment projection.

To show that this architecture will return exactly the same solution $g$ of
the activation matching algorithm, we use our ${\epsilon}$-friendly assumption
and Lemma 8 which imply that there is an open ball of radius $\delta$ around
each cost matrix in which the linear assignment problem is _constant_. We can
now use the uniform approximation result from the previous paragraph to
approximate all the cost matrices up to $\delta$ and get that applying the
linear assignment problem to the approximated cost matrices yields the same
output as the original cost matrices.

∎

######  Lemma 7.

[Lemma G.2 in Navon et al. (2023)] Let ${\epsilon}>0$. For any
$x_{1},\dots,x_{N}\in\mathbb{R}^{d_{0}}$ there exist a DWSNet $D_{\epsilon}$
and weights $\theta_{{\epsilon}}$ such that for any $v\in K$ we have
$\|D(v;\theta_{\epsilon})_{i}-f_{v}(x_{i})\|_{2}<{\epsilon}$.

We note that Navon et al. (2023) restricted $x_{1},\dots,x_{N}$ to some
compact set but this assumption is not used in their proof.

######  Lemma 8.

Let ${\epsilon}>0$, $K\subset\mathbb{R}^{n}$ a compact domain and let
$f_{i}:K\rightarrow\mathbb{R}$, $i=1,\dots,m$ such that $f_{i}(x)$ are
continuous. Let

$S=\\{x\in K\mid\exists i~{}~{}\text{s.t}~{}\forall j\neq
i~{}f_{i}(x)>f_{j}(x)+{\epsilon}\\},$

define $g(x)=\text{argmax}_{j=1}^{m}f_{j}(x)$, then there exists some
$\delta>0$ such that for any $x\in S$, $g(x)$ is constant in an open ball of
radius $\delta$ around $x$.

###### Proof.

Since $f_{1},\dots,f_{n}$ are finite and continuous on a compact domain, there
exists some $\delta>0$ such that for any $x,y\in K$, $\|x-y\|<\delta$ implies
$|f_{i}(x)-f_{i}(y)|<\frac{{\epsilon}}{2}$ for every $i$. Then for any $x\in
S$ such that $\forall j\neq i~{}f_{i}(x)>f_{j}(x)+{\epsilon}$, if
$\|x-y\|<\delta$ , for every $j\neq i$ we have

$f_{i}(y)-f_{j}(y)=(f_{i}(y)-f_{i}(x))+(f_{i}(x)-f_{j}(x))+(f_{j}(x)-f_{j}(y))>0,$

which implies that $g(y)=i$. ∎

###### Proof of Proposition 5.

As mentioned in the statement of the proposition, we assume that $F$ uses
analytic non-constant activation functions, and $d\geq 2$ channels. In the
proof we consider our standard choice of $\phi$ as the normalized inner
product and set the scaling parameter $s$ to be $s=1$ for simplicity. We
consider the version of $F$ where the last projection layer $F_{proj}$
computes the permutations maximizing the correlation with the outputs
$Q_{1},\ldots,Q_{M}$ of the product layer (that is, we consider the version of
$F$ used in test time, rather than the version used in training which uses
differentiable Sinkhorn iterations). By equivariance of the model, it is
sufficient to show that for almost every $v,\theta$ the identity matrices are
the closest permutations to $(Q_{1},\ldots,Q_{M-1}).$

Next, recall that the indices of each $Q_{m}=Q_{m}(v,\theta)$ are given by

$Q_{m,i,j}=\frac{\langle a_{m,i},a_{m,j}\rangle}{\|a_{m,j}\|\|a_{m,j}\|}$

where $a_{m,i}$ is a $d$ dimensional vector. In particular, we have that
$|Q_{m,i,j}|\leq 1$, and the equality holds if $i=j$. It remains to show that
when $i\neq j$ we will have a strict inequality $|Q_{m,i,j}|<1$ for all $m$,
for almost every $v$ and $\theta$. Equivalently, we will need to show that for
all $m=1,\ldots,M$ and all $i\neq j$, the functions

$\phi_{m,i}(v,\theta)=\|a_{m,i}\|^{2}$

and

$\psi_{m,i,j}(v,\theta)=\|a_{m,j}\|^{2}+\|a_{m,i}\|^{2}-\langle
a_{m,i},a_{m,j}\rangle$

are non-zero for almost all $(v,\theta)$. We note that $\phi$ and $\psi$ are
analytic, and the zero set of a non-zero analytic function always has Lebesgue
measure zero (see Mityagin (2015)). Therefore it is sufficient to show that
there exists a single $(v,\theta)$ for which $\phi_{m,i}(v,\theta)\neq 0$, and
a single $(v,\theta)$ for which $\psi_{m,i,j}(v,\theta)\neq 0$.

Let us consider parameter vectors $\theta$ as follows: recall that the output
of each layer is a sequence of hidden weight matrices $W_{m}=W_{m,a,b,c}$ and
a sequence of hidden bias vectors $b_{m}=b_{m,i,c}$, where the last index $c$
runs over the channels. We choose $\theta$ so that each affine layer will map
all matrices $W_{m}$ to zero, and all bias vectors $b_{m}$ to a new value
$b_{m}^{\prime}$, where $b_{m,i,1}^{\prime}=b_{m,i,1}$ and
$b_{m,i,c}^{\prime}=1$ if $c\neq 1$. Since this describes an affine
equivariant mapping, and the linear layers in DWS can express all linear
equivariant function, the vector $\theta$ can indeed be defined to give this
function.

With this choice of $\theta$, we will obtain

$a_{m,i}=\left[\rho^{D}(b_{m,i}),1_{d-1}\right]\in\mathbb{R}^{d}$

where $\rho$ denotes the activation used in the DWS network, $D$ denotes the
depth of the DWS network, and $b_{m,i}$ is the $i$-th entry of the $m$-th
input bias vector. In particular, it is an entry of $v$. To conclude the proof
it is sufficient to show that we can choose a pair of $b_{i},b_{j}$ such that
$\rho^{D}(b_{m,i})\neq\rho^{D}(b_{m,j})$. If this is indeed the case we can
immediately deduce that $\psi_{m,i,j}(v,\theta)\neq 0$.

To prove the latter point all we need is to show that $\rho^{D}$ is not a
constant function. Indeed, since $\rho$ itself is non-constant its image
contains an interval $I$, and by analyticity $\rho$ cannot be constant on this
interval, so $\rho^{2}$ is non-constant. Continuing recursively in this way we
can show that $\rho^{D}$ will not be constant for any $D$. ∎

##  Appendix D Extending DWSNets to CNNs

In this work, we employed a DWSNet (Navon et al., 2023) as our $F_{DWS}$
block. Since the study in Navon et al. (2023) focused on MLPs, the original
implementation only supports input MLP networks. Here, we provide technical
details on the extension to input CNN networks under the DWSNets framework.
This requires only two simple adjustments. First, the kernel dimensions are
flattened into the feature dimension. Second, the first FC layer (after the
last convolution layer) in the input network generally requires special
attention. Specifically, denoting the output dimension of the last CNN with
$d_{0}$ and the dimensions of the first FC layer weight matrix by $d_{1}\times
d_{2}$, we reshape the weight matrix to $d_{0}\times
d_{2}\times(d_{1}/d_{0})$, i.e., folding the $d_{1}/d_{0}$ into the feature
dimension. This preserves the equivariance to the permutation symmetries of
the input network.

##  Appendix E Experimental Details

Table 4: CNN classifiers architecture. CNN Classifiers Arch.
---
3x3 Conv 16
3x3 Conv 32
3x3 Conv 32
2x2 MaxPool
3x3 Conv 64
3x3 Conv 64
2x2 MaxPool
3x3 Conv 128
3x3 Conv 128
2x2 MaxPool
Linear (2048, 10)

In all experiments, we use a 4-hidden layer Deep-Align network with a hidden
dimension of $64$ and an output dimension of $128$ from the $F_{DWS}$ block.
We optimize our method with a learning rate of $5e-4$ using the AdamW
(Loshchilov & Hutter, 2017) optimizer. For all experiments with image
classifiers, we train Deep-Align with all objectives as described in Section 4
($\ell_{\text{alignment}},\ell_{\text{LMC}}$ and $\ell_{\text{supervised}}$).
For the INR experiments, we drop the $\ell_{\text{alignment}}$ loss since we
found adding this loss significantly hurt the performance (see Appendix 7).

We use the entire dataset to estimate the activations for the Activation
Matching (AM) baseline. For the Sinkhorn baseline, we optimize the
permutations using learning rate $1e-1$ and for $1000$ iterations.

When using image datasets, we use the standard train-test split and allocate
$10$% of the training data for validation.

MLP classifiers. For this experiment, we generate two wight datasets,
consisting of MNIST and CIFAR10 classifiers. Each classifier is a 3-hidden
layer MLP with a hidden dimension of $128$. The input dimension is $784$ for
MNIST and $3072$ for CIFAR10. We train the classifiers for $5$ epochs with a
batch size of $128$ and learning rate $5e-3$. Both datasets consist of $10000$
networks, split into $8000$ for training and $1000$ each for validation and
testing.

We train Deep-Align for $25$K iterations. Since the $F_{DWS}$ block can grow
large when the input dimension to the input network is large, we employ the
method proposed in Navon et al. (2023) to control the number of parameters.
Thus, we linearly map the input dimension $784$ or $3072$ (MNIST or CIFAR10)
to $8$. See Navon et al. (2023) for details.

CNN classifiers. For this experiment, we generate two datasets by training
classifiers on CIFAR10 and STL10 datasets. Each network consists of 7 CNN
layers followed by a fully-connected layer. The full architecture is presented
in Table 4. For STL10, we apply a sequence of $3$ augmentations, first, we
random crop $64\times 64$ path, next we resize the patch to $32\times 32$, and
finally we apply random rotation drawn from $U(-20,20)$. We train the CIFAR10
and STL10 classifiers for $20$/$100$ epochs respectively using Adam optimizer
with $1e-4$ learning rate. We save the model’s checkpoint at the final epoch.
Both datasets consist of $5000$ networks, split into $4500$ for training and
$250$ each for validation and testing. We train Deep-Align for $300$ epochs.

Sine INRs. To generate the Sine wave dataset, we use the same procedure as in
Navon et al. (2023). Each INR is an MLP with $3$ layers, a hidden dimension of
$32$, and Sine activations. The dataset consists of $2000$ sine waves and two
INR copies (views) for each sine wave. We use $1800$ waves for training, and
$100$ for validation, and testing.

CIFAR10 INRs. The CIFAR10 dataset consists of $60K$ images. We split the
dataset to train, validation, and test with $45K$ / $5K$ / $10K$ samples
respectively. For each image, we create $5$ independent INR copies. Each INR
is a 5-layer MLP with a $32$ hidden dimension each followed by sine
activations. We optimize the INRs using Adam optimizer with a $1e-4$ learning
rate for $10K$ update steps. We train Deep-Align for $300$ epochs.

Generalization. To evaluate the generalization to STL10, we use the same STL10
CNN dataset from Section 6.1. Furthermore, we generate an additional test
dataset of $100$ CIFAR10 classifiers trained on images with random rotations.

Disjoint datasets. We use the same CNN network configuration as in the CNN
classifiers experiments, and train $2500$ networks for each split ($5000$ in
total). Each network is trained for $20$ epochs using the Adam optimizer with
learning rate $1e-4$. We allocate $100$ networks from each split for testing
and validation, and the remaining $4600$ networks for training. We train Deep-
Align for $50$K steps.

##  Appendix F Additional Experimental Results

(a) MNIST MLPs.

(b) CIFAR10 MLPs.

(c) CIFAR10 CNNs.

Figure 6: Deep-Align as initialization: Results for using Deep-Align as
initialization for the optimization-based approaches Sinkhorn re-basin and
weight matching.

Deep-Align as initialization. As discussed in the main text, our approach can
be used as initialization for optimization-based approaches, like the Sinkhorn
re-basin. Here, provide extended results on using the output of Deep-Align as
the initial value for the alignment problem. We evaluate two previously
proposed methods, weight-matching (WM) (Ainsworth et al., 2022) and Sinkhorn
re-basin (Peña et al., 2023). Initializing the Sinkhorn method significantly
improves the performance under all evaluated datasets. In addition, using
Deep-Align initialization greatly improves the convergence speed. Furthermore,
Deep-Align improves the barrier results of the weight-matching method.
Notably, using the Deep-Align initialization achieves on-par or improved
values for the WM objective.

Effect of sample size. We evaluate Deep-Align on the sine-wave INR experiment
from the main text, using a varying number of training examples. The results
are presented in Figure 7. Deep-Align achieves on par results w.r.t the
Sinkhorn method with random initialization using $1800$ training pairs. On the
other hand, initializing the Sinkhorn method with the Deep-Align alignment
shows significant improvement in the test barrier using only $100$ training
pairs. These results show the efficiency of Deep-Align both in producing model
alignments or initializing optimization-based approaches.

Figure 7: Effect of sample size: Deep-Align achieves on par results w.r.t the
Sinkhorn with $1800$ training pairs, while only $100$ pairs are sufficient to
significantly improve Sinkhorn by initializing the alg. with the Deep-Align
outputs.

Ablation on the Deep-Align objective. Here, we provide results for Deep-Align
trained with different objectives. Recall that we introduced three objectives
(losses) to train Deep-Align. The first is the supervised loss
$\ell_{\text{supervised}}$ computed using a model and its permuted version.
The second is $\ell_{\text{alignment}}$ which is the $L2$ loss between the
aligned weight vectors, and the third is $\ell_{\text{LMC}}$ which evaluates
the original task loss on the line segment between the aligned models. For
this ablation study, we use the MNIST and CIFAR10 MLP classifiers along with
the CIFAR10 INRs. The results are presented in Table 5. Using only the
supervised and alignment loss generally achieves insufficient results in terms
of the Barrier metric. Dropping the alignment loss and using the supervised
and LMC losses appears to have a minimal impact on the results in the
classifier experiments. However, interestingly, including the
$\ell_{\text{alignment}}$ in the INR experiment seems to have a detrimental
effect on the Barrier results, causing a significant drop in performance. This
suggests the alignment loss and the barrier metric are not always well
aligned. In these cases it is advised to drop the $\ell_{\text{alignment}}$
loss and optimize Deep-Align with the $\ell_{\text{supervised}}$ and
$\ell_{\text{LMC}}$ losses.

Table 5: Optimizing Deep-Align using different objectives: Test Barrier results averaged over 3 random seeds. | MNIST MLP CLS | CIFAR10 MLP CLS | CIFAR10 INR
---|---|---|---
$\ell_{\text{supervised}}+\ell_{\text{LMC}}$ | $0.007\pm 0.00$ | $\mathbf{0.070\pm 0.00}$ | $\mathbf{0.063\pm 0.00}$
$\ell_{\text{supervised}}+\ell_{\text{alignment}}$ | $0.061\pm 0.00$ | $0.343\pm 0.00$ | $0.127\pm 0.00$
$\ell_{\text{supervised}}+\ell_{\text{LMC}}+\ell_{\text{alignment}}$ | $\mathbf{0.005\pm 0.00}$ | $0.078\pm 0.00$ | $0.087\pm 0.00$

Visualization of predicted permutations. We visualize the predicted
permutation obtained using Deep-Align applied to three test sine-waves INRs.
Each network pair consists of an INR and its permuted and noisy version. For
clarity, we depict only the first permutation matrix, $P_{1}$. The rows of
Figure 8 correspond to the three test INRs. The left column represents the
output from the $F_{prod}$ layer, which then projected to the set of
permutations using $F_{prod}$ (middle column). Deep-Align is able to perfectly
predict the ground truth permutations (right column).

Figure 8: Predicted permutation matrices and ground truth permutations for
three test sine wave INRs and their permutated and noisy version. Deep-Align
outputs the exact ground truth permutations.