Set-based Neural Network Encoding

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

In Section 3.2, when applying flattening, padding, and chunking operations on a convolutional layer, they are applied on the kernel dimension. I'm wondering why is the case? I have the question because really we have $3 \times 3$ kernel size for Convs. Do we really need to chunk it anymore?

For convolutional kernels, they are small enough such that chunking is not necessary. We have made the presentation general since the set-to-set and set-to-vector functions are shared across all layer types(linear and convolutions). The chunking operation becomes more important for linear layers to avoid processing large tensors. We will clarify this more in the experiment sections.

Why Equation 2 will convert a vector in size of $c \times 1$ to $c \times h$? As mentioned in Section 3.6 $\Phi_{\theta_{1}}$ is SAB(SAB(x)), while SAB(X) is MAB(X,X) in Equation 10. It seems that MAB(X,X) should have the same size as X. How can we get from $c \times 1$ to $c \times h$. You may want to add the dimensions of Equation 10 to 13 for clarification.

SAB(X) correctly has the same size of X. However we introduce a projection MLP in the SAB block that projects its inputs to $h$. We will clarify this explicitly in the choice of set-to-set and set-to-vector section and add dimensions to Equations 10 and 13 as suggested.

To my understanding, the encoding process from Section 3.3 to 3.5 is like a nested process. Is it possible to provide pseudocode for it?

Yes Sections 3.3-5 involve a nested process. We will include pseudocode as suggested and also provide a reference implementation for easy reproducibility and readability.

About the experiments, could you have some results on popular CNN architectures like ResNet? Residual links are nowadays important part of the CNN models.

While we do not see any impediment to applying our method to networks where residual connections are important, we are currently unaware of any publicly available modelzoos with the appropriate meta-data to benchmark on. An important encoding step in such models will be incorporating the information of residual connections in the encoding pipeline. For this, we propose to introduce a special token(just as was done for layer types) when encoding the entire residual block. We agree with the Reviewer that adopting the proposed method to such architectures is important and we leave it as future work when such modelzoos become publicly available.

We thank the Reviewer for the constructive feedback. We respond to the raised questions below.

There is a lack of comparison against existing methods on established datasets. For example, there are no experiments on implicit neural representations, which DWSNet and NFN are evaluated on extensively. There is no reference to DWSNet in the text.

Firstly, we reference DWSNet in the text in Section 2(Neural Network Performance Prediction from Weights). With regards to experiments on INRs, we provide the following additional experimental results and benchmark against both DWSNet and the newly introduced NFN models.

Experimental setting: We utilize a modelzoo consisting of INRs fit to sine waves on [-$\pi$, $\pi$] and with frequencies sampled from U(0.5, 10). The goal is to predict the frequency of a given INR. All methods encode the INR to a 32 dimensional vector which is then fed into a classifier with 2 linear layers with hidden dimension of 512. All methods are trained for 100 epochs. We report the mean squared error and parameter counts for all models in the table below. We are unable to obtain a converged model for [2] hence we compare with [4] which is the second iteration of the method proposed in [2]. As can be seen, SNE performs better than the baselines and is parameter efficient compared to [3] and [4]. We note that increasing parameter counts for the MLP and [1] baselines results in overfitting and poor results.

The authors do not consider the case of architectures with branches and how the position encoding needs to be adapted for them

Our discussion ignores architectures with branches, such as ResNets, because we are currently unaware of any standard modelzoo of such architectures on which to benchmark on. However, we provide an outline, in the specific case of residual blocks, of how our method can be adapted to encode such networks.

Given that residual blocks are composed of convolutional and linear layers, each of these can be encoded independently as we already do. To account for residual connections, we propose to introduce a special token (just as was done for layer types) when we encode the entire block. While we do not see any impediment to applying our method to such networks, the unavailability of such modelzoos makes it impossible to test this out hence we leave it as future work when such modelzoos become publicly available.

We thank the Reviewer for the constructive feedback. We respond to the questions raised below.

Baseline accuracy for some of these NNs is very low (0.45 on MNIST?). Very marginal empirical improvements.

The evaluation metric in Table 2 is Kendall’s Correlation Coefficient (as stated in section 4.2. We will explicitly add this to the caption as in Table 1 to avoid further confusion) and not accuracy on the datasets. In table 2, we demonstrate the generality of the proposed method in the cross-architecture task and note that none of the baselines in Table 1 are applicable since the methods are developed for fixed architectures.

Across the tasks that we consider, our method performs better than the baselines as shown in tables 1 and 2. But more importantly, the flexibility of the model we propose allows far more capability such as cross-architectural transfer, a task which none of the relevant baselines[1-4] are capable of. Additionally, we provide additional experimental results on implicit neural representations (INR) below as suggested by Reviewer Yfht.

Experimental setting: We utilize a modelzoo consisting of INRs fit to sine waves on [-$\pi$, $\pi$] and with frequencies sampled from U(0.5, 10). The goal is to predict the frequency of a given INR. All methods encode the INR to a 32 dimensional vector which is then fed into a classifier with 2 linear layers with hidden dimension of 512. All methods are trained for 100 epochs. We report the mean squared error and parameter counts for all models in the table below. We are unable to obtain a converged model for [2] hence we compare with [4] which is the second iteration of the method proposed in [2]. As can be seen, SNE performs better than the baselines and is parameter efficient compared to [3] and [4]. We note that increasing parameter counts for the MLP and [1] baselines results in overfitting and poor results.

Why is the evaluation on a small validation set insufficient in this case? Also, I find the motivation for predicting post-training network performance quite lacking.

While in this work, and those of [1,2,3,4], we have focused on network performance prediction, the methods developed are not limited to this task only. For instance, there are properties of networks[1], such as the type of optimizer, the learning rate, the number of epochs the model was trained for, etc. which cannot be obtained by having access to the training/test/validation set. Such meta-data of pretrained networks may prove necessary in model selection for fine tuning and hence require a means for predicting them.

How does this method compare to things like zero-cost-proxies which try to predict network performance before training?

The formulation of the neural network weight encoding problem([1,2,3,4]) assumes access only to pretrained weights making the domain different from zero-cost-proxies used in NAS that predict performance using architectural information.

References

[1] Unterthiner, Thomas, et al. "Predicting neural network accuracy from weights." arXiv preprint arXiv:2002.11448 (2020).

[2] Navon, A., Shamsian, A., Achituve, I., Fetaya, E., Chechik, G., & Maron, H. (2023). Equivariant architectures for learning in deep weight spaces. arXiv preprint arXiv:2301.12780.

[3] Zhou, Allan, et al. "Permutation equivariant neural functionals." arXiv preprint arXiv:2302.14040 (2023).

[4] Zhou, Allan, et al. "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 motivation of the paper is not clear enough. For a large model, it will lead to large computation overhead, and for a small model you can just simply run the model on the test set. The necessity of the task should be well discussed.

While research on encoding neural network weights is relatively new [1,2,3,4], these methods have potential application in areas such as learnable optimizers[5,6], direct utilization of implicit neural representations[7], policy evaluation[8] in reinforcement learning, neural network editing[9,10] etc. However, to make progress in all these applications, encoding of network weights is the first step in this pipeline.

While in this work, and those of [1,2,3,4], we have focused on network performance prediction, the methods developed are not limited to this task only. For instance, there are properties of networks[1], such as the type of optimizer, the learning rate, the number of epochs the model was trained for, etc. which cannot be obtained by having access to the training/test/validation set. Such meta-data of pretrained networks may prove necessary in model selection for fine tuning and hence require a means for predicting them. To demonstrate the generality of the problem, we provide additional applications on implicit neural representations in our response to your second question.

The applications of the proposed SNE is still limited. The architectures and datasets are very simple.

SNE makes no assumptions on architectural choices. [1,2,3,4] all require that an architecture is known and fixed a priori. In section 4.2, we introduce the cross-architectural transfer task for neural network encoding to demonstrate the flexibility of SNE even at test time. The most relevant baselines [1,2,3,4] are all not applicable to this task given that they work only for a fixed architecture. Additionally, in Section 4.1, we have introduced the cross-dataset evaluation task, for the first time, for neural weight network encoding methods.

Secondly, with regards to architectures and datasets, we evaluate our method on the standard benchmarks for evaluating neural network encoding methods[1,2,3,4]. In section 4.1, we have generated a new modelzoo of pretrained networks to be able to demonstrate cross-architecture encoding. Currently, we are unaware of any large scale modelzoos of architectures such as ResNets or Vision Transformers with the relevant meta-data that can be used for this task. However, we note that the architecture agnostic formulation of SNE ensures that, when such modelzoos become publicly available, the proposed method will still be applicable. Finally, we provide additional results on implicit neural network(INR) encoding to demonstrate that our method is applicable beyond the performance prediction task.

Experimental setting: We utilize a modelzoo consisting of INRs fit to sine waves on [-$\pi$, $\pi$] and with frequencies sampled from U(0.5, 10). The goal is to predict the frequency of a given INR. All methods encode the INR to a 32 dimensional vector which is then fed into a classifier with 2 linear layers with hidden dimension of 512. All methods are trained for 100 epochs. We report the mean squared error and parameter counts for all models in the table below. We are unable to obtain a converged model for [2] hence we compare with [4] which is the second iteration of the method proposed in [2]. As can be seen, SNE performs better than the baselines and is parameter efficient compared to [3] and [4]. We note that increasing parameter counts for the MLP and [1] baselines results in overfitting and poor results.

Lack of ablation study of their methods, such as the effect of their hierarchical computational structure position encoding.

We remove the hierarchical compositional structure and the positional encoding mechanism from SNE and test it on the INR task introduced above. From the table below, we can see that these structures are essential to the performance of the method.