Equivariant Deep Weight Space Alignment

The proposed framework requires new training for every new input neural network architecture.

DEEP-ALIGN is not restricted to a single input network architecture and can be used to align a variety of input network architectures using the same DEEP-ALIGN model. This can be done by replacing the DWSNet encoder, with an architecture-agnostic model like the concurrent submissions of https://openreview.net/forum?id=ijK5hyxs0n and https://openreview.net/forum?id=oO6FsMyDBt. There, the input neural networks are modeled as weighted graphs, and the weight space encoder is implemented using a graph neural network (GNN). That type of GNN architecture can be trained and applied to a wide range of neural architectures, with early experiments showing very promising generalization abilities to OOD architectures (see subsection 5.1 here https://openreview.net/forum?id=ijK5hyxs0n).

---

How to select the weights for the linear combination of the three loss functions?; How do the three losses interact with each other?

Our framework enables the practitioner to adapt the loss to the task at hand. In our experiments, we observed that combining all losses with equal weighting works well for most learning setups, except in the INRs experiment where removing the L2 loss was beneficial.

---

Demonstrate the framework’s effectiveness in applications that require weight alignment.

As the reviewer mentioned, one important application is Federated Learning (FL) where it was shown that alignment-based averaging, instead of the standard federated averaging, leads to better performance and reduced communication cost (see e.g., https://arxiv.org/pdf/2002.06440.pdf). Following the reviewers' comments and motivated by applications from FL, we added a new experiment in which we align models trained using disjoint datasets. To that end, we split the CIFAR10 dataset into two disjoint splits with different class distributions. We train our DEEP-ALIGN model to align CNN networks trained using the different datasets. The merged models achieve lower classification loss compared to the original models. The results are provided here and are added to the revised version of the paper (Subsection 6.1, page 9, “Aligning networks trained on disjoint datasets”).

Link to results

Other possible applications are alignment-based mixup as data augmentation for learning in deep weight spaces (see, e.g., http://arxiv.org/abs/2311.08851) and weight-space clustering, for example for analyzing large INR collections.

---

Does the exactness property (proposition 5) apply to test data as well?

Yes. The result is very general and holds for all weights and all input vectors aside from a set of measure zero (which implies zero probability for any continuous distribution).

---

Why does the loss decrease on the interpolation around lambda=0.1 or 0.9?

This is a common phenomenon in model merging, as demonstrated in the git rebasin (https://arxiv.org/abs/2209.04836) and Sinkhorn matching (https://arxiv.org/abs/2212.12042) papers.

---

Why is incorporating equivariance beneficial to the model's performance?

Equivariance is well known to be beneficial for learning data with symmetries (see, for instance, the geometric Deep Learning proto book https://arxiv.org/abs/2104.13478). Among the main benefits of this approach is that the model does not have to "memorize" different equivalent input representations of the same object, resulting in better generalization. Moreover, many equivariant architectures, such as the architecture used in this paper, have significantly fewer parameters and are considerably more efficient than fully connected networks. This is a result of the inherent parameter-sharing schemes that these equivariant models employ. We refer the reviewer to Figure 3 in the DWSNet paper, which compares equivariant and non-equivariant models. The performance gap as well as the improved sample complexity are substantial even for very simple tasks. We are happy to include a similar experiment for the weight alignment task if required.

---

Missing definition of rho_1 and rho_2; Similar colors in Figure 5.

Thank you. We fixed these issues in the revised version of the paper.

The network's generalization capabilities are limited; The method has to be trained using the same network architecture, solving the same task as during inference.

As a learning-based method, our approach indeed requires training on the task of interest given data from the targeted data distribution. The upside, however, is that at inference time our learning-based approach results in high-quality alignments at a fraction of the computation time compared to non-learning-based approaches. This opens up the ability to use alignment in applications that require solving many alignment problems in real-time, such as federated learning and weight space clustering. Nevertheless, as the reviewer mentioned, our OOD performance is very reasonable and we believe it can be even further improved by incorporating recent GNN-based deep-weight space encoders that can process different architectures (such as these concurrent submissions: https://openreview.net/forum?id=ijK5hyxs0n and https://openreview.net/forum?id=oO6FsMyDBt). That type of GNN architecture can be trained and applied to a wide range of neural architectures, with early experiments showing very promising generalization abilities to OOD architectures (see subsection 5.1 here https://openreview.net/forum?id=ijK5hyxs0n).

---

Provide more explanations on mapping from weight to activation space; Why feed the weights W to the network at all and not just use the bias?

Thank you. We understand that this part was not adequately explained in the original paper. Following your comment and since weight space networks are new and relatively unknown, we expanded our previous work section to include more information about DWSNets in the revised version of the paper.

In a nutshell, weight space networks are composed of weight space layers that map weights and biases $(W_1, b_1,...,W_M,b_M)$ to new weights and biases $(W'_1,b’_1,...,W'_M,b'_M)$. Importantly, these layers combine information from both biases and weights to output the new set of weight and bias representations. Thus, the output biases, and consequently the predicted permutations are also a function of the input weights, and not only the input biases.

The reason that using the output bias representation makes sense is that activations and biases reside in the same space. In fact, we use the terms bias space and activation space interchangeably in the paper, which we believe is what caused the confusion.

How does DEEP-ALIGN scale to larger networks (e.g. VGG models)?

Thank you. Following this important comment, we added experiments with deeper architectures, VGG11 with 9M parameters and VGG16 with 15M parameters. We used the same amount of data, 4500 training examples, similar to the CNN experiments in the paper. The experiments show that: (1) DEEP-ALIGN scales well and can process large architectures; and (2) DEEP-ALIGN generalizes well when applied to larger networks without increasing the data size. This demonstrates that DEEP-ALIGN is suitable for processing standard deep architectures. The results are provided here and are added to the revised version of the paper (Appendix F, page 18, “Scaling to large networks”).

Link to results

---

How does weight-matching + Sinkhorn compare with Deep-Align + Sinkhorn? Does Deep-Align reliably beats activation matching?

Comparison to Weight matching + Sinkhorn: Following your comment and feedback, we added comparisons to the weight-matching (WM) + Sinkhorn baseline in Tables 1 and 2, and Figure 3 in the revised version of the paper. Both DEEP-ALIGN and DEEP-ALIGN+Sinkhorn constantly outperform the WM + Sinkhorn baseline in all datasets and model architectures. We provide here the main results:

Results for CNNs:

Discussion on results w.r.t. activation matching: DEEP-ALIGN outperforms activation matching in $3/4$ of the datasets and performs comparably on the remaining dataset. This is achieved while being almost $\times 20$ faster at inference time.

Is there any possibility that one pertained DEEP-ALIGN model can work for networks of slightly different architectures?

Yes. DEEP-ALIGN is not restricted to a single input network architecture and can be used to align a variety of input network architectures using the same DEEP-ALIGN model. This can be done by replacing the DWSNet encoder, with an architecture-agnostic model like the concurrent submissions of https://openreview.net/forum?id=ijK5hyxs0n and https://openreview.net/forum?id=oO6FsMyDBt. There, the input neural networks are modeled as weighted graphs, and the weight space encoder is implemented using a graph neural network (GNN). That type of GNN architecture can be trained and applied to a wide range of neural architectures, with early experiments showing very promising generalization abilities to OOD architectures (see subsection 5.1 here https://openreview.net/forum?id=ijK5hyxs0n).

---

Scaling to deeper networks and the effect on the required amount of training data.

Thank you. Following this important comment, we added experiments with deeper architectures, VGG11 with 9M parameters and VGG16 with 15M parameters. We used the same amount of data, 4500 training examples, similar to the CNN experiments in the paper. The experiments show that: (1) DEEP-ALIGN scales well and can process large architectures; and (2) DEEP-ALIGN generalizes well when applied to larger networks without increasing the data size. This demonstrates that DEEP-ALIGN is suitable for processing standard deep architectures. The results are provided here and are added to the revised version of the paper (Appendix F, page 18, “Scaling to large networks”).

Link to results

---

The real-world applications for this method are unclear due to long pre-train time.

Like other learning-based methods for solving combinatorial optimization problems, our method is particularly well-suited to applications that require solving many problems repeatedly. One key application is Federated Learning (FL). There, several models trained by clients need to be combined into a single model at the hub (model merging). This is often achieved by FedAvg, which merges models by averaging their weights elementwise in each round of communication. It has been shown (e.g., https://arxiv.org/pdf/2002.06440.pdf) that averaging aligned versions of these weights improves the performance of FedAvg, reduces the number of communication rounds, and improves training convergence.

Motivated by applications from FL, we added a new experiment in which we align models trained using disjoint datasets. To that end, we split the CIFAR10 dataset into two disjoint splits with different class distributions. We train our DEEP-ALIGN model to align CNN networks trained using the different datasets. The merged models achieve lower classification loss compared to the original models. The results are provided here and are added to the revised version of the paper (Subsection 6.1, page 9, “Aligning networks trained on disjoint datasets”).

Link to results

Other possible applications are alignment-based mixup as data augmentation for learning in deep weight spaces (see, e.g., http://arxiv.org/abs/2311.08851) and weight-space clustering, for example for analyzing large INR collections.

---

What is the definition rho_1 and rho_2?

Thanks for spotting this. $\rho_1$ and $\rho_2$ are representations. We fixed the issue in the revised version of the paper.