Structure and Behavior in Weight Space Representation Learning

We thank reviewer iCWh for their review, and will address the weaknesses they note and the questions they raise in this response.

W1: The idea is very straight forward, and the technical contribution is quite limited. Basically, they run experiments like those done previously in the field, while adding the new L2 functional loss. I am not aware that others have done this for weights autoencoders, but using "behavior" (out put of model parameterized by weights) to select weights has been in other context. For example a loss over the function parametrized by the weights is the standard way to train hypernetworks.

We agree with the reviewer that using a probing-based loss in itself is not a novel thing, and have added a relevant paragraph to our related work Section. As mentioned in the response to all reviewers, however, the scope of our contribution goes beyond that. By motivating why we use it, what issues it solves and how it works, we hope to bring more explanatory elements in understanding weight space self-supervised learning, and to highlight the importance of the structure-behaviour synergy for future works in the domain. Our experiments using the SANE [1] framework aim at validating these findings.

W2: The PCA analysis of Figure 1 has two issues. (1) It is suspicious (see questions below) and (2) it is not clear how it connects to the contribution of the paper, since the analysis is not used as far as I can tell to motivate or guide the method or the experiments.

Q1: Figure 1 looks highly suspicious, and many details are skipped. How come it is useful to add "Bottom" eigenvectors, which capture directions in space that have little to no information about the weights? How does the eigen spectrum look like? and what is the dependency of explained variance on the number of principal components? What is "reconstruction accuracy" in this figure? (the text said you measured test-set accuracy). What was the size of the training data and weight-vector dimension? same as in line 254? I suspect that the covariance matrix was low rank, and that the correct x-axis in the figure is "fraction of of eigen vectors taken", not "explained variance". That may explain the totally flat curves on the middle and right panels. And how comes the random accuracy in SVHN is 20%?

Figure 1 and the related experiment are highly related to that from Balestriero and LeCun [2]. In their work on autoencoders for computer vision, they show that the features most useful for classification of images lie with the lowest eigenvalues, although the features useful for reconstruction lie with the largest eigenvalues. In a similar fashion, we show that in weight space, one needs both features that lie with the “Top” and “Bottom” eigenvalues to reconstruct high-performing models.

Models are trained on the training sets of the corresponding datasets, and tested on the test sets. Weight vectors are flattened model weights, therefore length $10,853$. As the reviewer correctly guesses, the resulting covariance matrix is low rank. For this reason, we specifically target “Top” or “Bottom” eigenvalues that explain a certain amount of variance, and not a certain proportion of the number of eigenvalues. It follows that the bottom 10% explained variance contains many more eigenvalues than the top 10%, but they explain the same amount of variance. Finally, the first point in all plots sits at 10% explained variance, not 0%; the corresponding model is not random but reconstructed from eigenvalues that explain 10% of the variance, and therefore performs better than random guessing. We hope these details convince the reviewer of the soundness of Figure 1. We further add that the supplementary material contains the corresponding code and can be verified.

Regarding point (2), the goal with the Figure 1 experiment is to highlight that structural reconstruction, which naturally focuses on the features spanned by the top eigenvectors, will miss out some features that are essential for reconstructing high-performing models. As such, we motivate the need for an additional loss element that focuses on such features. Such a loss needs to be used in conjunction with the structural loss, since we also show that the features spanned by the top eigenvectors remain relevant. In Section 3, we derive our behavioural loss, and describe how it will focus on different features than the structural loss. In Section 4, we experimentally validate that with our new loss, we have indeed managed to target these missing features through the addition of the behavioural loss element. Whether that loss specifically targeted features spanned by the bottom eigenvectors remains however elusive, as the eigenvectors of reconstructed models are mostly misaligned with those of the original ones. Our results nevertheless empirically show that relevant features have been learned, since we do manage to reconstruct high-performing models.

We would like to thank reviewer KkRG for their review, and will try to address the weaknesses and questions they raise.

The proposed approach generally follows the SANE [1] paper, model arch, and experimental setup, with the only difference in the loss function. While the results presented show improvements over SANE, the contribution of the paper is limited.

As mentioned in the response to all reviewers, we do not believe our main contribution lies in adding a behavioural loss into the SANE framework. With our submission, we identify a particular issue in structural representation learning, derive and propose a solution, which highlights synergies between structural and behavioural signals in self-supervised weight space learning. The experiments we run using SANE [1] are used to validate our findings through the strength of the results when combining structural and behavioural loss elements.

Previous works for weight space learning used similar loss functions or probing-based features. At the very least, this deserves some reference and discussion by the authors. A similar loss function, although not identical (and not always in the context of AE), was proposed in e.g. [2, 3, 4, 5]. The idea is to evaluate the reconstructed model on a set of inputs and compute the loss w.r.t ground truth labels. This is very close to the proposed approach, with the difference being that the proposed approach uses pseudo labels instead of GT labels. However, this is almost identical in the context of INRs [2, 3]. Another relevant literature that is not discussed, is the inclusion of probing-based features, which evaluate the original function on (fixed or learned) inputs to provide additional information on the input net and functional behavior [6, 7].

We would like to thank the reviewer for the literature that they provide. We discuss how these works differ from ours here, and have also added a relevant paragraph in the related work Section of our manuscript.

De Luigi et al. [2] do use a probing-based loss in the context of autoencoders, but without incorporating a structural element. In addition, the queries are part of the inputs of the decoder, and they directly predict the true labels, whereas we favour reconstruction. Additionally, [2] performs reconstruction of the shapes represented by INRs, not INRs themselves, making direct comparison to our work difficult. To the best of our understanding, [3] mostly leverages structural signals in a supervised way, or contrastive signals for self-supervised learning. It uses a probing-based signal for domain adaptation using true labels on a corrupted dataset, but it seems limited to this use-case. [7] show high performance levels from using probing-based mechanisms to analyse recurrent neural networks, but we note that they use their decoders as an emulator of the function represented by the original model, whereas we directly reconstruct the original model. As mentioned in the response to reviewer fxSa, the learning setup in [5] is the closest to ours, as the authors leverage a mix of supervised, structural and behavioural signals, which shows great performance. Their probing-based loss, however, requires true labels and does not compare the behaviour of two models in the same way our behavioural loss does. With our contribution, we want to demonstrate that combining structural and behavioural signals plays a part in this, and to show the importance of that dual signal for the development of future weight-space methods.

Experiments are performed on small-scale NNs and datasets. It is not clear how this generalizes to real-world large-scale setups.

As explained above, in this work we aim to understand a fundamental issue in structural weight space learning. To that end, we evaluate failure modes and systematically develop an approach to avoid that failure mode. Our experiments are validations of that.

We acknowledge, however, that a more extensive set of architectures could further strengthen our claim. We address this point in more detail in the response to all reviewers, where we explain architectural challenges in scaling up the models we work on. In order to validate our findings on larger architectures, we have added a new Appendix D.5 where we run experiments on a new model zoo of LeNet-5 models trained on CIFAR-10, which are around 6 times larger that the models in our other zoos.

Please also include some measure of variability (e.g., standard derivation) in Table 1.

The results presented in Table 1 are computed as the $R^2$ score over the whole test set. It follows that there is no direct standard deviation measure to provide. Where applicable, such as in Table 5 and 6, standard deviation metrics are provided.

We would like to thank reviewer fxSa for their review, and for bringing three very interesting papers by Navon, Shamsian et al. to our attention [1, 5, 6]. We hope to address the weaknesses and questions raised by the reviewer with this response.

My main concern regarding this work is around its novelty the behavioral loss was already used in [1]. In [1] the authors calculated the loss based on the loss on the downstream task. In addition, there is no novelty (and it was not claimed) on the architecture side since the authors used SANE [2] for the weight space AE. How do the authors differentiate their work from [1]?

The work in [1] is very interesting and parallels can indeed be drawn with our work. To the best of our understanding, the authors use a combination of three losses to train their weight-space alignment model: a supervised loss $L_{supervised}$, an unsupervised loss that minimises weight distance $L_{alignment}$, and finally an unsupervised loss that minimises the model’s loss function on a linear combination of both models’ weights (i.e. the linear mode connectivity) $L_{LMC}$. $L_{LMC}$, similarly to our behavioural loss, is a probing based loss. Such losses have encountered success, for example in the context of HyperNetworks [7]. It however strongly differs from ours in its usage. In [1], authors use $L_{LMC}$ by computing the original model’s loss on point that lies linearly between the original neural network and the permuted neural network, so as to test for good linear mode connectivity. Ours differs in that, first, it does not need to compute the loss with regard to the original task, and therefore does not require any true labels. Second, its main interest lies in aligning the predictions made by the model and its reconstruction, and as such is closer to a soft-label distillation loss.

In addition, the main contribution of our paper lies beyond simply adding a behavioural loss to the SANE [2] framework. Our submission presents a principled exploration of the duality of structure-based and behaviour-based signals in weight space representation learning, exploring valuable insights in why they synergise well and providing strong experimental validation results. We are convinced that the good results in [1] further strengthen and help generalise our main claim that there exists a strong synergy between structural and behavioural signals. Where [1] succeeds at leveraging dual structural and behavioural losses to solve the weights alignment problem, our paper attempts to show why those dual signals synergise well, giving insights into why their choice of losses has been particularly effective.

We believe the discussion of the (di)similarities between [1] and our work to be a valuable addition to our manuscript. We have therefore incorporated a new related work paragraph that illustrates the differences between our approach and other probing-based ones that exist in the literature.

Limited experimental section. Given the fact that this work is mostly empirical, I expect the experimental section to be more extensive both in scale and learning setups. The authors focused on discriminative models and specifically only on small CNNs. There is no diversity in the reconstructed model architecture or the downstream tasks.

We respectfully disagree with the reviewer that our work is mostly empirical. In Sections 2 and 3, we provide insights into why there is a need to capture features not adequately captured by structural reconstruction alone, and why the behavioural loss is a good candidate. Then, in our Section 4, we validate these hypotheses through empirical experiments, which show strong performance improvements across all downstream tasks and datasets. In addition, we also disagree that there is no diversity in the downstream tasks studied: we evaluated 2 different discriminative downstream tasks, a reconstructive downstream task, as well as a generative downstream task.

We acknowledge, however, that a more extensive set of architectures could further strengthen our main claim. We address this point in more detail in the response to all reviewers, where we explain architectural challenges in scaling up the models we work on. In order to validate our findings on larger architectures, we have added a new Appendix D.5 where we run experiments on a new model zoo of LeNet-5 models trained on CIFAR-10, which are around 6 times larger that the models in our other zoos.

It would be interesting to see how significant are gamma and beta hyperparameters in the AE optimization.

We have also addressed this point in the response to all reviewers. $\gamma$ has been set to the same value as in SANE [2]. We have added a new Appendix with the experiments validating our choice of hyperparameter $\beta$ using our validation sets.

We would like to thank reviewer 6Tzb for their review, and will try to address the weaknesses and questions they raised with this response.

Table 2 shows that the proposed behavioral loss improves performance across all three datasets. However, it would be insightful to examine whether the representations produced by the original CNNs are similar to those generated by the reconstructed CNNs. This analysis could further strengthen the findings, clarifying whether the generated weights yield representations closely aligned with the originals and thereby contribute to the observed performance improvements.

In the Appendix, we have included more results that space limitations prevented us to include in the maintext. In particular, Table 5 shows measures of structural and behavioural similarity between models and their reconstructions. These results report both averages and standard deviations, and show that there is a substantial increase in model agreement when using a behavioural loss. Does that correspond to what the reviewer had in mind?

The reported results do not include standard deviations, which could offer insight into the consistency and robustness of the proposed method across different runs. Including this information would enhance the reliability of the findings.

We agree with the reviewer, and as such Table 5 includes both average and standard deviation for the structural and behavioural similarities of models and their reconstructions. For the discriminative downstream tasks, we compute the $R^2$ score over the whole test set, and therefore do not have any standard deviations to provide; results are however similar across all three datasets.

The experimental setup is primarily focused on convolutional neural networks on three vision datasets without exploring other types of architectures or data. Including a broader range of architectures and datasets could further validate the approach’s generalizability. (See questions).

It would be valuable to explore how the proposed approach performs when applied to architectures other than CNNs, such as MLPs or larger models. This could provide further insights into the generalizability and robustness of the proposed method.

We acknowledge that including more architectures would help demonstrate the generalisability of our methods. We discuss in more detail in the response to all reviewers why scaling up to models in the million parameters range or larger is challenging. As discussed there, our work focuses more on validating the synergy between structural and behavioural signals, and we aim our experiments at validating these findings. We have added a new experiment to the Appendix that validates our findings on a LeNet-5 based architecture, strengthening the validity of our results on models ~6 times larger than those in our original experiments.

It would be interesting to examine the effectiveness of the proposed method on other data modalities, such as text or graph data. Testing across diverse modalities could help determine the broader applicability and flexibility of the approach.

While we agree with the reviewer, we have chosen simpler models to facilitate experiment setup and comparison, and also because of limited compute resources.

It would be helpful to clarify whether all model zoos in the experiments were trained exclusively for image classification. Exploring the reconstruction of weights from models trained on other downstream tasks (e.g. image reconstruction) could provide insight into the versatility and task-specific performance of the method.

All models have indeed been trained for image classification. However, since the proposed behavioural loss is using the $L^2$ distance between outputs rather than cross-entropy, it remains applicable to regression or reconstruction models as is.

We indicate here all major changes we made to the revised rebuttal submission, compared to the original submission.

• Added an Appendix D.5 and Figure 10 in which we validate our methodology on a model zoo of larger LeNet-5 models trained on CIFAR-10
• Added an Appendix D.6 as well as Tables 7 and 8, in which we discuss our choice of hyperparameter $\beta$ to $0.1$. Also added a reference to that Appendix in Section 4.1.
• Added a new related work paragraph to incorporate other works using probing-based losses in the context of weight space learning
• Added a new reference to “Improved Generalization of Weight Space Networks via Augmentations” by Shamsian et al. to the related work, and the related discussion.
• Added a short definition of a “model zoo” at the beginning of Section 4.1