Set-based Neural Network Encoding Without Weight Tying

We thank the Reviewer for taking the time to offer constructive feedback for improving the paper. We respond to the questions raised below.

Provide examples of the application of this research other than asserting that they must exist.

We add the following to elaborate on the potential applications of this research to complement lines 29-30: ‘The ability to infer such fundamental properties of trained neural networks using only the parameter values has the potential to open up new application and research paradigms such as learning in the latent space of neural network weights for tasks such as weight generation [1, 2], latent space transfer of weights across datasets allowing for transferring weights from one dataset to another as was recently demonstrated in [3] and latent space optimization using gradient descent where optimization is performed on the weight embeddings [4].’ We hope this addition enhances the motivation for this line of research.

Overall, the presentation is a weak spot, and should be revised. Some examples:

• We have added a citation to Implicit Neural Representations (INR) [1] and added the following description, “INRs are neural parameterizations of signals such as images using multi-layer perceptrons”.
• All occurrences of ‘modelzoo’ have been replaced with ‘model zoo’.
• We have changed the citation format to ‘Author et al., Year[num]’.
• We’ve included the following to captions on Tables 2/3 “Models are evaluated using Kendall’s Tau[5], a rank correlation metric.’ with a citation to the paper that introduces the metric.
• We have reduced the number of plots in Figure 2 to one per model and increased the size (see Figure 2 in the attached pdf). The remaining figures are moved to the appendix in enlarged form and the clipped legend fixed.
• The checklist has been placed after the references and supplementary as pointed out.

The evaluation is done on simpler datasets/networks, however, the limitations section highlights why.

Based on the suggestion of Reviewer nYWS we provide additional results on more complicated networks, vision transformers, in our response to nYWS and in the general response. We hope that these additional results on more complicated architecture demonstrates the utility of the proposed method.

Can the authors provide some commentary/comparison between these two umbrella fields (NAS and weight encoding)?

In NAS, the predictors mostly rely on the computation graph of generated architectures. This is quite different from the problem we consider where we require access to the parameter values themselves, optimized for some number of epochs. However, the general underlying problem has commonality between the two fields: namely predicting some properties of the networks in consideration. Perhaps a new line of exploration would consider combining both fields. For instance, a neural network encoder, like ours, could be trained to conditionally take as input either the weights, the computational graph or both. Such a pre-trained predictor could be applied to NAS where in test mode, only the computational graph is given to estimate the performance of newly generated architectures. This would require the encoder to be agnostic to architectural choices like our model. The discovered architectures when trained could then be used to update the property predictor using both weights and computation graphs. Such a combination would find application in both fields. To the best of our knowledge, we are unaware of any works that attempt such a combination and we leave this exploration as future work and will draw attention to such a merger of both fields in our discussion and limitation sections.

Similarly, is it possible to leverage the encoding mechanism (Eqs. 1-9) to derive which types of weights/layers/values are responsible for certain end-to-end attributes, e.g., performance of epochs, e.g., as [9] do?

Technically, such an analysis should be possible. For instance, we could train SNE to predict multiple properties for a given network and at test time, we can remove the encodings of different layers (in Eq 9) and see how it affects the predictions compared to the full network. This will provide a mechanism for determining which layers/weights are responsible for certain attributes. Currently we are unable to do such an analysis since all the model zoos used in this paper only come with a single attribute in the meta-data. Such analysis will require regenerating the model zoos, a task which we are unable to undertake during the short rebuttal period. However, we are inclined to agree with the Reviewer that such an analysis will be very useful.

References

[1] Hyper-Representations as Generative Models: Sampling Unseen Neural Network Weights, 2022

[2] Wang, Kai, et al. "Neural Network Diffusion." arXiv preprint arXiv:2402.13144 (2024).

[3] Soro, Bedionita, et al. "Diffusion-based Neural Network Weights Generation." arXiv preprint arXiv:2402.18153 (2024).

[4] Rusu, Andrei A., et al. "Meta-learning with latent embedding optimization." arXiv preprint arXiv:1807.05960 (2018).

[5] Kendall, M. G. A new measure of rank correlation. Biometrika, 30(1/2):81–93, 1938.

I thank the authors for their detailed rebuttal, which satisfactorily addressed my concerns and questions. After reading the other reviewer comments and corresponding rebuttals I find that the reasons to accept outweigh concerns to reject, and thus raise my score. If accepted, I encourage the authors to integrate rebuttal experimental results as well as revise related work and limitations to reflect these past discussions.

We thank the Reviewer for taking the time to offer constructive feedback for improving the paper. We respond to the questions raised below.

The writing is generally good but could be improved in certain areas. For example, the abstract contains a redundant sentence: “by utilizing a layer-wise encoding scheme that culminates to encoding all layer-wise encodings to obtain the neural network encoding vector,” which should be revised.

We have removed the sentence from the abstract and replaced it with ‘Furthermore, our $**S**$et-based $**N**$eural network $**E**$ncoder (SNE) takes into consideration the hierarchical computational structure of neural networks.’ We hope this resolves the redundancy pointed out.

Expanding Table 1 to include SNE's performance on additional outputs

Our request to [1] for the dataset of INRs for the ShapeNet dataset is currently pending and this task will be added to Table 1 when the dataset is received. However we note that fundamentally, there is no difference between the accuracy prediction task for INRs and the frequency prediction task that we present in Table 1 since these are all evaluations on MLPs.

We would also like to draw the Reviewers attention to the additional results we provide on generalization on vision transformers which we provide in our general response. We hope these additional evaluations further demonstrate the wide range applicability of our method.

References

[1] De Luigi, Luca, et al. "Deep learning on implicit neural representations of shapes." arXiv preprint arXiv:2302.05438 (2023).

We thank the Reviewer for taking the time to offer constructive feedback for improving the paper. We respond to the questions raised below.

could the authors provide results with a predictor trained on CNNs from arch1 and predict the performance of transformers?

We provide the requested results below. We generate a model zoo of transformer classifiers based on the [1] and test the transfer from arch1 to the transformer model zoo as requested. For this task, we are unable to benchmark against NeuralGraph[2] as was done in Table 2, since the model cannot process transformer weights when trained on Arch1. Hence we benchmark against the DeepSets baseline as in Table 2. We present the results in the table below:

From the results, the proposed SNE generalizes better to the unseen transformer architecture at test time than the baselines showing strong architectural transfer. Additionally, here, the model encodes an architecture with about 5 times the number of parameters in Arch1 demonstrating the scalability of our approach. The DeepSets baseline fails to generalize on this task.

could the author measure the importance of the training set size w.r.t. the predictor's performance, in order to estimate the cost of working on ImageNet-like tasks?

We provide the requested results in Figure 1 of the attached pdf. From this it can be seen that SNE is more data efficient compared to the baselines (DWSNet[2] and NFT[3]) across all percentages of the full training data demonstrating that the proposed method learns a good embedding even in the limited data setting. We thank the Reviewer for suggesting this ablation as it further demonstrates the strength of SNE.

References

[1] https://github.com/s-chh/PyTorch-Scratch-Vision-Transformer-ViT

[2] 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.

[3] 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, 2023.

We thank the Reviewer for taking the time to offer constructive feedback for improving the paper. We respond to the questions raised below.

The proposed method did not ensure weight-space equivariance by design but use a way that is more like an augmentation plus a consistency loss. The theoretic contribution is therefore a bit weak.

While we use a regularization approach to achieve approximate minimal equivariance in weight-space, our usage of the logit invariance regularizer [1] theoretically guarantees that we indeed learn the correct invariance property similar to the weight-typing approaches used by the baselines. Additionally, our formulation is what allows us to deal with arbitrary architectures using a single model, as opposed to the baselines, since strict enforcement of weight-space equivariance by design requires crafting a new model for every distinct architecture. In this sense, our approach provides a general encoder which in principle is applicable to any architecture, resolving the limitations of purely weight-tying approaches.

The padding and chunking designs look tricky and it is not clear for me how they help with the performance.

The padding and chunking operations are introduced in our pipeline to offer computational efficiency in cases where the parameters are large and hence infeasible to process all together. In Table 6 of the appendix we provide an ablation on the effect of the chunk size, which we reproduce below on the frequency prediction task:

Here we see that performance stays almost the same until the largest chunksize (which in this case is set to the largest size hence requiring no chunk/pad operations). The chunksize is selected to suit the gpu memory requirements allowing us to train our method even with limited gpu memory which in our case is a single GTX 1080 gpu with 11GB of memory.

Fig 1 is not really clear. (1) the legend is not sufficient to understand the computation diagram presented in the figure. (2) the input is shown as neurons but from the method I think it should be the weight matrix that connecting neurons in adjacent layers.

We have modified the figure to show the weight matrix instead of the layer in compressed form as shown in the figure in the updated paper. We thank the Reviewer for pointing this out as this improves the clarity of the presented pipeline. Additionally, we link the relevant approach subsections to each stage in the figure to aid clarification of the computation diagram.

In table 1 NeuralGraph is not reported. Is it because the method is not applicable there?

Yes. We were unable to apply NeuralGraph for this task. Before submission, correspondences with the authors of NeuralGraph could not resolve the issue we encountered when applying the method for this task and that of Table 4.

Does the method need more optimization steps than baseline methods that ensure weight-space symmetry by design?

No. In all experiments, we train our method for the same number of epochs as all the baselines. The addition of the logit invariance regularizer introduces no further optimization steps in all our experiments and we encountered no difficulties in minimizing the general objective together with the logit invariance loss.

References

[1] Artem Moskalev, Anna Sepliarskaia, Erik J Bekkers, and Arnold WM Smeulders. On genuine invariance learning without weight-tying. In Topological, Algebraic and Geometric Learning Workshops 2023, pages 218–227. PMLR, 2023

We would like to thank the Reviewer again and respond to the new queries below:

I wonder what the issue was with NeuralGraph on the experiments in Table 1 and Table 4? If the NNs for the INRs are MLP architecture, I do not expect any issue there since NeuralGraph is able to work on classification of MNIST / CIFAR in INRs.

• Firstly, we utilize the official implementation of NeuralGraphs provided at [1]. In Table 1 we are unable to get NeuralGraphs to converge on this task using multiple hyper-parameter runs.
  Secondly, the experiments of Table 4 are on ConvNets (the architecture is outlined in Tables 10 and 11 of the appendix. In Table 4, we evaluate in terms of cross-dataset performance prediction and the issue with NeuralGraph, under this evaluation setting, lies in the following: we are unable to obtain positive correlation coefficients when we evaluate it cross-dataset. As we elaborated in our response to Reviewer bv1s, in our correspondences with the authors of NeuralGraph, we were unable to find and resolve the issues with negative correlation values in Table 4. Interestingly, we do not observe these issues in our evaluation of NeuralGraph in Table 2. These explain the absence of NeuralGraphs in Tables 1 and 4.

This should be a more high-level and sophisticated task than sine wave frequencies, and direct comparison with performance reported by prior work would be more convincing.

• As we elaborated in our response to Reviewer EgRw, there is no fundamental difference between the task of INR frequency prediction and INR classification as these require the same guarantees in MLP weight space. Our request to [1] for the dataset of INRs for the ShapeNet dataset is currently pending and this task will be added to Table 1 when the dataset is received together with INR classification (a task which we are currently undertaking but do not expect to be completed during this discussion period. This is partly due to the large datasets required for for training INR classification for MNIST, FashionMNIST and CIFAR: all of DWSNet, NFT, NeuralGraph generate up to a million datapoints as opposed to the 60000 samples in the original MNIST classification task, resulting in long training iterations). Importantly, we point the Reviewers attention to the following: that all the experiments presented in Table 2, 3,4 and 7 of the appendix together with the newly introduced Transformer evaluation task suggested by Reviewer NYWS are much more difficult tasks compared to classifying INRs since these require respecting symmetric properties in both CNN (with MLP heads) and Transformer weight spaces. We hope this provides some clarity.

What is the NN architecture tested for the chunksize experiments?

• We utilize the architecture outlined in Table 15 of the appendix which consists of 3 linear layers each with hidden size of 32 and the first two followed by sinusoidal activation function. This is the architecture for generating the model zoo used in the experiments of Table 1.

Finally, we would like to thank the Reviewer for the continued engagement as both the discussions and our response to the original review concerns as well as the questions we’ve presently addressed greatly improves clarity of the paper and these will be incorporated in the paper.

References

[1]https://github.com/mkofinas/neural-graphs/tree/1f2b671ab4988ef212469363005a5b99eec16580

[2] De Luigi, Luca, et al. "Deep learning on implicit neural representations of shapes." arXiv preprint arXiv:2302.05438 (2023).