Structure and Behavior in Weight Space Representation Learning

Authors: Léo Meynent, Ivan Melev, Konstantin Schürholt, Goeran Kauermann, Damian Borth
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: weight space learning, hyper-representations, deep weight spaces, representation learning, model reconstruction, model weights generation

Abstract:
The weights of neural networks (NNs) have recently gained prominence as a new data modality in machine learning, with applications ranging from accuracy and hyperparameter prediction to representation learning or weight generation. One approach to leverage NN weights involves training autoencoders (AEs) with contrastive and reconstruction losses. Indeed, such models can be applied to a wide variety of downstream tasks, and they demonstrate strong predictive performance and low reconstruction error. However, despite the low reconstruction error, these AEs reconconstruct NN models that fail to match the performance of the original ones. In this paper, we identify a limitation of weight-space AEs, specifically highlighting that structural weight reconstruction alone fails to capture some features critical for reconstructing high-performing models. To address this issue, we propose a behavioral loss for training AEs in weight space. This behavioral loss focuses on the features essential for reconstructing performant models, which are not adequately captured by structural reconstruction. We evaluate the capabilities of AE trained using this novel loss on three different model zoos: we demonstrate that when combining structural and behavioral losses, we can reconstruct and generate models that match the performance of the original models. With our exploration of representation learning in deep weight spaces, we show that a strong synergy exists between structural and behavioral features, and that combining them results in increased performance across all evaluated downstream tasks.

Decision: reject
Reviewer scores: 6, 3, 3, 5

Reviews and discussion:
<thread>
creation date:
1729709571406

thread_type:
official_review

<thread_object>
creator: reviewer_icwh

<content>
summary:
The paper discusses autoencoding of weight-spaces. In this field of research, 
weights of deep models are viewed as an object (input and output) of learning systems. 
Specifically, they look to improve the function of models encoded using weight-space AEs, 
by adding what they call, a "behavioral" loss, designed to improve network function. 
That loss is simply an L2 loss over the output of the model with the reconstructed weights. 

The paper runs meaningful experiments largely following the standards in the field. On three model populations "zoos", they show that adding the new loss improves the AE, in terms of generating and reconstructing more faithful weights.

My recommendation is borderline reject, because of the limited technical novelty and conceptual contribution, which means the audience for this paper would be very narrow. I am willing to be convinced otherwise.

strengths:
-- The paper tackles an important problem in the growing field of weight-space modality. 

-- The paper contains a large set of evaluation experiments for generation and reconstruction of weights in three model zoos.

weaknesses:
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. 

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.   

W3:  Every evaluation of the behavior loss involves 500 inference steps through the classifier model (250 samples x 2 models). How does this added compute compares with the total compute for training the AE? If it increases the overall compute required to train the AE, it should be explicitly stated.

W4: Definition of behavioral loss. For classifiers, L2 is not a great choice. One should use the same loss that the models from the zoo cis trained with, . e.g. cross entropy for classification etc. 


Specific and minor comments and suggestions 

-- Define key terms like "structural reconstruction error" and "model zoo" before using them. Add concrete definitions of all losses. 
    The paper should be made readable to researchers outside the weight-space community. 

-- I found Fig 3 totally confusing. Is there anything interesting insight to take from these histograms tat cannot be conveyed in a simpler way? Distributions are quite unimodal so a mean+std table would have been much simpler and easy to understand. 

-- Related work on generation of NN weights. It would make sense to mention Hypernetworks.

questions:
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%? 

Q2:  Not stated clearly how the loss weights were selected. Validation sets?
</content>
</thread_object>

</thread>
<thread>
creation date:
1729793138915

thread_type:
official_review

<thread_object>
creator: reviewer_kkrg

<content>
summary:
The paper considers the problem of constructing weight space encoders, i.e., learning an Auto-Encoder (AE) on the weights of trained NNs (for example, image classifiers). Previous works mainly considered two loss functions: a “structural” reconstruction loss, e.g., MSE between the reconstructed and input weights, and a contrastive loss. The authors propose augmenting the two losses with an additional loss, which they term “behavioral”: This loss is the MSE between the outputs of the original and reconstructed NN, evaluated on a set of inputs.

strengths:
1. The paper is generally well-structured and easy to follow.
1. The empirical evidence for the effectiveness of behavioral loss is strong.

weaknesses:
My main concern is the paper's limited novelty and contribution. The paper's main claimed novel contribution is the inclusion of behavioral loss.
1. 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.
1. 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].
1. Experiments are performed on small-scale NNs and datasets. It is not clear how this generalizes to real-world large-scale setups.
1. Please also include some measure of variability (e.g., standard derivation) in Table 1.
1. The results of Section 2 suggest that it is important to learn to match eigenvectors corresponding to both large and small eigenvalues, while that of mid-sized eigenvalues is less important. This provides a good motivation for augmenting the structural only loss, however:
    - It is not clear to me that the proposed loss achieves this specific goal.
    - Also, it would be beneficial to show some empirical evidence that including behavioral loss is improving the reconstruction of subspaces spanned by both the top and bottom eigenvectors.
1. As also noted by the authors, the proposed approach increases the computational cost. This increase may be too expensive for large-scale models (which are not included in this work).

References:

[1] Schürholt et al. "Towards Scalable and Versatile Weight Space Learning." ICML 2024.

[2] Luigi et al. "Deep Learning on Implicit Neural Representations of Shapes." ICLR 2023.

[3] Navon et al. “Equivariant architectures for learning in deep weight spaces,” ICML 2023. 

[4] Zhou et al. "Permutation Equivariant Neural Functionals." NeurIPS 2023. 

[5] Navon et al. “Equivariant Deep Weight Space Alignment,” ICML 2024.

[6] Kofinas et al. "Graph Neural Networks for Learning Equivariant Representations of Neural Networks." ICLR 2024.

[7] Herrmann et al. "Learning useful representations of recurrent neural network weight matrices.", 2024.

questions:
1. From Figure 3 and as pointed out by the authors, it appears that the behavior of reconstructed models does not necessarily reproduce that of the original model, rather, it is biased towards higher accuracies. Why do you think that is the case? These results suggest that the reconstruction of low performing models is not good, even with the added behavioral loss.
1. In Table 2, selecting the Max statistics seems too convenient. Using the mean or median or Min, for example, would have shown significant differences between the original and generated models.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730059915912

thread_type:
official_review

<thread_object>
creator: reviewer_fxsa

<content>
summary:
The authors introduced a loss function called *behavioral loss*, designed to reduce the performance discrepancy between original neural networks and those reconstructed through a weight-space autoencoder. This behavioral loss minimizes the difference between the outputs of the original and reconstructed networks by reducing the norm between their predictions.


The authors claim the following contributions:
- analyzing the error modes of weight-structure reconstruction.
- Introducing the behavioral loss.
- Demonstrating the improvement in loss performance across various experiments.

strengths:
- The paper is well written and easy to follow.
- The motivation for the proposed solution is well explained and investigated.
- Simple solution.

weaknesses:
- 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]?
- 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.
- It would be interesting to see how significant are $\beta$ and $\gamma$ hyperparameters in the AE optimization.
- The authors mentioned in the appendix that training with behavioral loss nearly doubles the training duration, which poses a challenge for large models.
-Rather than using a common model zoo [3,4], the authors chose to create their own, which raises some questions.
- The authors did not include confidence intervals in the experimental section. Moreover, in Tables 5 and 6 the authors provide additional results which are inclusive due to high std values. The same should be reported in the main text results.
- The only augmentation used for $\mathcal{L}_C$ is a random permutation, why the authors did not use augmentations from [5,6]? [6] should be cited in this work.

--------------
[1] Equivariant Deep Weight Space Alignment, Navon et al.

[2] Towards Scalable and Versatile Weight Space Learning, Schurholt et al.

[3] Model Zoo: A Dataset of Diverse Populations of Neural Network Models, Schurholt et al.

[4] Eurosat Model Zoo: A Dataset and Benchmark on Populations of Neural Networks and Its Sparsified Model Twins, Honegger et al.

[5] Equivariant Architectures for Learning in Deep Weight Spaces, Navon et al.

[6] Improved Generalization of Weight Space Networks via Augmentations, Shamsian et al.

questions:
- Based on the findings in Section 2, have the authors considered explicitly encouraging hybrid eigenvalue features? For example, they might incorporate a KL divergence loss relative to a predefined Beta distribution concentrated in the high and low values, effectively ignoring the intermediate range.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657594998

thread_type:
official_review

<thread_object>
creator: reviewer_6tzb

<content>
summary:
This work proposes a new behavioral loss to enhance the performance of hypernetwork autoencoders trained to reconstruct neural network weights. It demonstrates improved results in reconstructing convolutional neural network weights on vision datasets. The study provides extensive validation of reconstruction quality, combining contrastive, structural, and behavioral losses to achieve this improvement.

strengths:
* The paper is well-written and easy to follow, with clear, well-organized sections that ensure a smooth flow throughout. Additionally, the captions for tables and figures are self-contained, making it easier for readers to understand them in context.

* The results demonstrate a consistent improvement in both the performance of the reconstructed models and the richness of the information contained in the computed weight representations. The authors assess this information using linear probing to predict test accuracy and model generalization gaps, providing valuable insights into the learned representations.

* The authors have released the codebase for this work, which enhances the paper's reproducibility and allows others to verify and build upon their findings.

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

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

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

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.

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

* 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732128894931

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
We would like to thank all reviewers for the work and thought they put in their constructive reviews. With this common response, we aim to address concerns and answer interrogations raised by multiple reviewers.

## Contribution meaningfulness

Several reviewers mention that our submission lacks technical novelty. We respectfully disagree: our work proposes a principled approach in studying the relationship between structural and behavioural signals in the context of weight-space representation learning, and therefore goes beyond simply adding a behavioural loss to the SANE framework [1]. Inspired by similar, recent works in the computer vision domain by Balestriero and LeCun [2], we first identify that in weight space, learning by reconstruction similarly omits features that are essential to reconstruct models with high test-set performance. We then derive a behavioural loss to directly target those missing features, and show how that loss differs from the structural one. Finally, we perform a set of experiments on a variety of downstream tasks in order to empirically validate these findings. As noted by several reviewers, evidence shown in our experiments is strong. We acknowledge that our submission proposes a quite simple solution, and that probing-based losses have been used in previous works, such as HyperNetworks [3]. We nevertheless believe our work to be a relevant addition to the weight space literature because it highlights and explores the importance of the dual structural and behavioural signal in a way that has not been done before, which in itself is of interest for the development of more advanced methods in the future.

## Architectures and task variety

In our submission, we choose to only test our results on three different model zoos composed of smaller CNNs. Reviewers have noted that adding more architectures, tasks or even modalities could make our results more convincing. We would first like to highlight that with this submission, we mostly aim to show the existence of a synergy between structural and behavioural signals in weight space. As such, our experiments are aimed at validating these findings, and they show strong empirical evidence towards this. Nevertheless, we do agree with reviewers that adding more architectures would strengthen our paper.

Unfortunately, as discussed in the Limitations, scaling up our method to larger models is highly non-trivial. Indeed, while SANE [1] can work on arbitrarily large models, it can do so by processing “windows” of weights. Computing the $L^2$ distance between a window and its reconstruction is straightforward, but computing their behaviour is not since it relies on the model as a whole, not only the window of weights considered. It is possible to find ways to improve the SANE framework to circumvent this issue, but this is out of scope for this paper where we only aim to validate our findings that there exists a synergy between structural and behavioural signals in weight space. For this reason, we limit our experiments to model sizes that can be processed as a whole by the Transformer-based SANE. Since the size complexity of Transformers increases with the square of the size of the context window, this limits us to relatively small models.

We have nevertheless run experiments on the CIFAR-10 dataset with larger CNNs based on the LeNet-5 architecture, with around 6 times more parameters than those presented in the original submission: results are very similar and also highlight a strong synergy between structure and behaviour. We have added these new experiments to the manuscript as Appendix D.5.

## Choice of hyperparameters $\gamma$ and $\beta$

Regarding $\gamma$, we made the choice to retain the hyperparameter’s value used by Schürholt et al. [1]. Since it is not the focus of this work, we have not explored multiple values of this hyperparameter. We have nevertheless studied the opportunity to set it to $0$ to deactivate the contrastive element of the loss in Appendix D.1.

Regarding $\beta$, we have tested multiple possible values on our validation sets, and kept the one that showed the most consistent performance. We recognise that this experiment is currently absent from the paper. We have added the corresponding experiment as Appendix D.6, which justifies the choice of the value $0.1$ for this hyperparameter.

—

[1] Schürholt et al., "Towards Scalable and Versatile Weight Space Learning.", ICML 2024.

[2] Balestriero and LeCun, “Learning by reconstruction produces uninformative features for perception”, ICML 2024.

[3] Ha et al., “HyperNetworks”, ICLR 2017
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper focuses on the weight space of neural networks as a new data modality, and posits that simply weight reconstruction is not enough for learning. Rather, a 'behavior loss' is needed to tie the learning to specific downstream tasks. The paper mainly focus on experiments with convolutional neural networks. Reviewers share the common opinion that the paper does not clarify well enough the difference with SANE. The authors clarify that the behavior loss is not their main contribution, however,  I agree that this can be better written, since it is even claimed as such in the abstract (' To address this issue, we propose a behavioral loss for training AEs in weight space.'). Given that clear novelty is a fundamental requirement to be published in a conference like ICLR, I suggest the authors to revise their manuscript and resubmit.