Paper: Equivariant Deep Weight Space Alignment (ICLR-2024)

Reviewer comment (Experiments):
The proposed framework requires new training for every new input neural network architecture, even for minor changes such as adding or modifying the size of a layer.

Author response:
***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).

Reviewer comment (Experiments):
It is not clear how to choose the weights for the linear combination of the three loss functions in Section 4.2. Also, as shown in Appendix F, the two unsupervised loss functions are not always aligned, and it is not clear which loss to include before training. Finally, the evaluation metrics, barrier and AUC, are based on model merging tasks, but the supervised loss does not seem to help on these metrics directly.

Author response:
***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.

Reviewer comment (Experiments):
The practical contribution would be more convincing if the authors can demonstrate the framework’s effectiveness in applications that require weight alignment.

Author response:
***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*](https://github.com/DeepAlignPaper/deep-align-results/blob/main/cifar10_disjoint_datasets.png)

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.

Reviewer comment (Writing):
$\rho_1$ and $\rho_2$ in the first paragraph of Section 2 are not defined. Are they representations?

Author response:
***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.

Reviewer comment (Writing):
Section 4.1: might be better to state what $\theta$ is somewhere.

Author response:
***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.

Reviewer comment (Presentation):
Figure 5: The “Weight Matching” and “DEEP ALIGN + Weight Matching” have very similar colors.

Author response:
***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.

Reviewer comment (Theory):
How do the three losses interact with each other? Does the two unsupervised loss provably help improve the supervised loss?

Author response:
***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.

Reviewer comment (Evaluation):
Figure 3: Why does the loss decrease on the interpolation around $\lambda=$ 0.1 or 0.9, especially in CIFAR10 CNNs?

Author response:
***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.

Reviewer comment (Writing):
The network that maps from weight embedding to activation space and its connection to stage 1 is unclear (see below for questions).

Author response:
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.

Reviewer comment (Theory):
I don't fully understand the network mapping from weight embeddings to activation space. If the network just maps onto the bias vectors, the input weights do not have any influence on the estimated permutation anymore, would that be correct? This seems to be unintuitive to me and I would like the authors to clarify.

Author response:
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.

Reviewer comment (Experiments):
The main limitation is that results are not given for sufficiently large/complex models/tasks. For the barrier results, it is known (e.g. in Ainsworth et al. or Jordan et al.) that deeper networks are harder to align so that they are linearly connected, whereas easier networks (e.g. 3-layer MNIST networks) are relatively trivial to align.

Author response:
***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”).


|                       | CIFAR10 (VGG11) |                | CIFAR10 (VGG16) |                |         |
|:---------------------:|:---------------:|:--------------:|:---------------:|:--------------:|:----------------------:|
|                       |    Barrier ↓    |      AUC ↓     |    Barrier ↓    |      AUC ↓     |       Runtime (Sec) ↓                  |
|         Naive         |    1.273±0.04   |   0.738±0.02   |    1.131±0.04   |   0.777±0.03   |           ---          |
|        Sinkhorn       |    0.368±0.00   |   0.041±0.00   |    0.921±0.05   |   0.212±0.02   |         158.40         |
|       DEEP-ALIGN      |    0.310±0.00   |   0.059±0.00   |    0.779±0.01   |   0.356±0.01   |          0.98          |
| DEEP-ALIGN + Sinkhorn |  **0.099±0.01** | **0.000±0.00** |  **0.349±0.01** | **0.027±0.01** | 159.38 = 0.98 + 158.40 |

[*Link to results*](https://github.com/DeepAlignPaper/deep-align-results/blob/main/cifar10_VGG.png)

Reviewer comment (Experiments):
It is not clear if Deep-Align reliably beats activation matching (which is still relatively fast and scalable), or if Deep-Align + Sinkhorn can beat weight matching + Sinkhorn.

Author response:
***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:

|                       |    MNIST (MLP)   |                  |   CIFAR10 (MLP)  |                  |
|:---------------------:|:----------------:|:----------------:|:----------------:|:----------------:|
|                       |     Barrier ↓    |       AUC ↓      |     Barrier ↓    |       AUC ↓      |
|        Sinkhorn       |   0.027 ± 0.00   |   0.002 ± 0.00   |   0.183 ± 0.00   |   0.072 ± 0.00   |
|     WM + Sinkhorn     |   0.012 ± 0.00   | **0.000 ± 0.00** |   0.137 ± 0.00   |   0.050 ± 0.00   |
|       DEEP-ALIGN      |   0.005 ± 0.00   | **0.000 ± 0.00** |   0.078 ± 0.01   |   0.029 ± 0.00   |
| DEEP-ALIGN + Sinkhorn | **0.000 ± 0.00** | **0.000 ± 0.00** | **0.037 ± 0.00** | **0.004 ± 0.00** |
|

Results for CNNs:

|                       |   CIFAR10 (CNN)  |                  |    STL10 (CNN)   |                  |                      |
|:---------------------:|:----------------:|:----------------:|:----------------:|:----------------:|:--------------------:|
|                       |     Barrier ↓    |       AUC ↓      |     Barrier ↓    |       AUC ↓      |    Runtime (Sec) ↓   |
|        Sinkhorn       |   0.313 ± 0.01   | **0.000 ± 0.00** |   0.366 ± 0.00   |   0.163 ± 0.000  |         79.81        |
|     WM + Sinkhorn     |   0.333 ± 0.01   | **0.000 ± 0.00** |   0.371 ± 0.00   |   0.165 ± 0.00   | 80.02 = 0.21 + 79.81 |
|       DEEP-ALIGN      |   0.237 ± 0.01   | **0.000 ± 0.00** |   0.382 ± 0.01   |   0.182 ± 0.00   |         0.44         |
| DEEP-ALIGN + Sinkhorn | **0.081 ± 0.00** | **0.000 ± 0.00** | **0.232 ± 0.00** | **0.097 ± 0.00** | 80.25 = 0.44 + 79.81 |
|

**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.

Reviewer comment (Theory):
The Siamese structure and why it guarantees equivariance to transposition should in particular be fully described.

Author response:
Q: ”The major methodological innovations could use much more explanation”.
A:  Thank you for this important comment. Weight space networks are indeed new and relatively unknown. Thus, we expanded our previous work section to include more information about DWSnets in the revised version of the paper. In addition, we will attempt to explain the principle behind the equivariance to transposition. In short, this equivariance property means that when the input networks are swapped, i.e. we would like to map from $v_2$ to $v_1$ and not from $v_1$ to $v_2$, the output permutations should map in the opposite direction, and the mapping itself should be the inverse mapping. This is a fundamental principle in matching/alignment problems.
In our case, since these permutations are represented as permutation matrices, the inverse mapping is exactly the transposed permutation matrix. Using Siamese architecture is a natural way to implement this idea. For example, when we use a simple outer product layer as $F_{prod}$ we get the following for the $m$-th output:

$F(v_1,v_2)_m = F_E (v_1)_m * F_E (v_2)_m^T$

Where $F_E$ stands for DWS encoder. Note how changing the order of the inputs generates the transpose exactly as desired.

$F(v_2,v_1)_m = F_E(v_2)_m*F_E(v_1)_m^T = F(v_1,v_2)_m^T$


Which means that we obtained the equivariant property we wanted.

Reviewer comment (Experiments):
How does this method scale (regarding accuracy and computation time) to larger networks (e.g. VGG models) and residual networks?

Author response:
***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”).


|                       | CIFAR10 (VGG11) |                | CIFAR10 (VGG16) |                |         |
|:---------------------:|:---------------:|:--------------:|:---------------:|:--------------:|:----------------------:|
|                       |    Barrier ↓    |      AUC ↓     |    Barrier ↓    |      AUC ↓     |       Runtime (Sec) ↓                  |
|         Naive         |    1.273±0.04   |   0.738±0.02   |    1.131±0.04   |   0.777±0.03   |           ---          |
|        Sinkhorn       |    0.368±0.00   |   0.041±0.00   |    0.921±0.05   |   0.212±0.02   |         158.40         |
|       DEEP-ALIGN      |    0.310±0.00   |   0.059±0.00   |    0.779±0.01   |   0.356±0.01   |          0.98          |
| DEEP-ALIGN + Sinkhorn |  **0.099±0.01** | **0.000±0.00** |  **0.349±0.01** | **0.027±0.01** | 159.38 = 0.98 + 158.40 |

[*Link to results*](https://github.com/DeepAlignPaper/deep-align-results/blob/main/cifar10_VGG.png)

Reviewer comment (Experiments):
How does weight-matching + Sinkhorn compare with Deep-Align + Sinkhorn?

Author response:
***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:

|                       |    MNIST (MLP)   |                  |   CIFAR10 (MLP)  |                  |
|:---------------------:|:----------------:|:----------------:|:----------------:|:----------------:|
|                       |     Barrier ↓    |       AUC ↓      |     Barrier ↓    |       AUC ↓      |
|        Sinkhorn       |   0.027 ± 0.00   |   0.002 ± 0.00   |   0.183 ± 0.00   |   0.072 ± 0.00   |
|     WM + Sinkhorn     |   0.012 ± 0.00   | **0.000 ± 0.00** |   0.137 ± 0.00   |   0.050 ± 0.00   |
|       DEEP-ALIGN      |   0.005 ± 0.00   | **0.000 ± 0.00** |   0.078 ± 0.01   |   0.029 ± 0.00   |
| DEEP-ALIGN + Sinkhorn | **0.000 ± 0.00** | **0.000 ± 0.00** | **0.037 ± 0.00** | **0.004 ± 0.00** |
|

Results for CNNs:

|                       |   CIFAR10 (CNN)  |                  |    STL10 (CNN)   |                  |                      |
|:---------------------:|:----------------:|:----------------:|:----------------:|:----------------:|:--------------------:|
|                       |     Barrier ↓    |       AUC ↓      |     Barrier ↓    |       AUC ↓      |    Runtime (Sec) ↓   |
|        Sinkhorn       |   0.313 ± 0.01   | **0.000 ± 0.00** |   0.366 ± 0.00   |   0.163 ± 0.000  |         79.81        |
|     WM + Sinkhorn     |   0.333 ± 0.01   | **0.000 ± 0.00** |   0.371 ± 0.00   |   0.165 ± 0.00   | 80.02 = 0.21 + 79.81 |
|       DEEP-ALIGN      |   0.237 ± 0.01   | **0.000 ± 0.00** |   0.382 ± 0.01   |   0.182 ± 0.00   |         0.44         |
| DEEP-ALIGN + Sinkhorn | **0.081 ± 0.00** | **0.000 ± 0.00** | **0.232 ± 0.00** | **0.097 ± 0.00** | 80.25 = 0.44 + 79.81 |
|

**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.

Reviewer comment (Experiments):
This method has to be pre-trained in advance. For experiments conducted on MNIST and CIFAR10, 8000 shallow networks are used to train the model to achieve the claimed performance. For deeper networks that are more common in nowadays deep learning task, this number may grow fast and the training cost can be unacceptable. Also, this paper does not provide experimental results about the relationship between the complexity of weight vectors and the scale of needed training data.

Author response:
***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”).

|                       | CIFAR10 (VGG11) |                | CIFAR10 (VGG16) |                |          |
|:---------------------:|:---------------:|:--------------:|:---------------:|:--------------:|:----------------------:|
|                       |    Barrier ↓    |      AUC ↓     |    Barrier ↓    |      AUC ↓     |      Runtime (Sec) ↓                  |
|         Naive         |    1.273±0.04   |   0.738±0.02   |    1.131±0.04   |   0.777±0.03   |           ---          |
|        Sinkhorn       |    0.368±0.00   |   0.041±0.00   |    0.921±0.05   |   0.212±0.02   |         158.40         |
|       DEEP-ALIGN      |    0.310±0.00   |   0.059±0.00   |    0.779±0.01   |   0.356±0.01   |          0.98          |
| DEEP-ALIGN + Sinkhorn |  **0.099±0.01** | **0.000±0.00** |  **0.349±0.01** | **0.027±0.01** | 159.38 = 0.98 + 158.40 |

[*Link to results*](https://github.com/DeepAlignPaper/deep-align-results/blob/main/cifar10_VGG.png)

Reviewer comment (Evaluation):
The real world applications for this method are unclear due to long the pre-train time. It seems that the traditional methods can compute the weight alignment much more effectively unless a ton of weight vectors need to be aligned so that the inference phase of DEEP-ALIGN dominates. What scenarios feature such characteristics?

Author response:
***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*](https://github.com/DeepAlignPaper/deep-align-results/blob/main/cifar10_disjoint_datasets.png)

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.

Reviewer comment (Writing):
What is the definitions of $\rho_1$ and $\rho_2$ in the section 2 Preliminaries?

Author response:
***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.

Reviewer comment (Theory):
Is the assumption that the minimizer $k$ in equation (2) is unique reasonable? Any argument that such cases would be rare in practice?

Author response:
***Is the assumption that the minimizer k in equation (2) is unique reasonable? Any argument that such cases would be rare in practice?***

Thank you for raising this question. Our equivariance results do not depend on the single minimizer assumption - this assumption was used to simplify the discussion. Equivariance holds even when multiple minimizers are present. This was stated in our submission (page 4, just after proposition 2), but we made efforts to make this clearer in the revised version of the paper.

Reviewer comment (Experiments):
I'm curious about the generalization capability of the model. As mentioned in the paper, the model will be first trained on a dataset of weight vectors and then applied to unseen weight vectors. Is there any possibility that one pertained DEEP-ALIGN model can work for networks of slightly different architectures? Due to the expensive cost of training phase, it can be helpful if one pertained model can work for more scenarios.

Author response:
***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).