Dynamic Neural Graph: Facilitating Temporal Dynamics Learning in Deep Weight Space

Thank you for your valuable feedback and constructive comments. We address your concerns as below:

---

W1&2. Treating weight parameters as dynamic graphs may not be effective, as the neural graph structure between layers changes significantly, making it difficult to accurately capture their sequential and evolving characteristics. Are there any experiments that prove the improvement introduced by capturing the sequential nature between layers?

We thank for your question. In fact, we dedicated a significant portion of the text to explaining how our method captures the sequential nature of neural network processing, both theoretically and empirically.

• Theoretical analysis: We kindly refer the reviewer to Section 2.3 and Appendix D, where we provide a detailed theoretical analysis of static and dynamic neural graphs, demonstrating the expressiveness of neural networks in these contexts. These analyses demonstrate that our approach can effectively simulate the forward pass of input neural networks, capturing their inherent sequential processing nature. In contrast, the static counterpart may introduce challenging ill-posed problems, leading to a failure in approximating the functionality beyond the first MLP layer.

• Empirical analysis: We kindly refer the reviewer to Appendix E.1, where we include an experimental analysis to further compare the two approaches. These empirical results show that the method based on static neural graphs is primarily effective for simulating only the initial forward pass of the input neural network. In contrast, our proposed dynamic neural graph framework is capable of accurately simulating all forward pass steps of the input neural network.

W3. In the related work section, the paper lacks a deeper analysis to describe differences between the proposed method and static graph counterparts.

At the time of submission, the only work we know employing static graphs was Kofinas et. al. To highlight our primary contribution, we conducted an in-depth study, which is presented in detail in Section 2.3 of the main body. We kindly direct the reviewer’s attention to Section 2.3, where we thoroughly discuss the main limitations of static neural graphs.

W4. There lack experiments on the transformers architecture to validate the proposed method.

We would like to clarify that our primary contribution does not include applying the proposed method specifically to transformer architectures. While we acknowledge the potential of interpreting transformers as dynamic neural graphs (DNGs), this aspect was only discussed in the Appendix G.2 as a supplementary demonstration to illustrate the flexibility of our framework. It was intended as a conceptual reference rather than a core experimental validation.

To address your concern, we conducted an experiment to predict the generalization of transformers by processing their parameters, thereby validating the effectiveness of our proposed DNG framework in handling transformer architectures.

• Dataset: We follow Small CNN Zoo (Unterthiner et al.) to prepare a transformer dataset. Specifically, we trained 10,000 differently initialized Vision Transformer (ViT) models (used SimpleViT from ViT_pytorch library) to classify the CIFAR-10 dataset. Each model was trained up to a certain epoch before being stopped, and its parameters and test accuracy were saved. Among the trained models, 80% were used as the training set, 10% as the validation set, and 10% as the test set.

• Implementaion details: Similar to the settings for predicting generalization of CNNs discussed in Section 6.3, we use our proposed DNG-Encoder followed by an MLP to predict the test accuracy of a ViT model given its parameters as input. To construct the DNG, we follow the methods outlined in Appendix G to convert the ViT.

The table below summarizes the running time, memory usage, GFLOPs, and the Kendall rank correlation coefficient τ of our method and NG-GNN (Kofinas et al.) on the test set.

W5. Minor typos.

We have updated our paper to address and correct these typos.

Thank you for your valuable feedback and constructive comments. We address your concerns as below:

---

W1. The claim of invariance or equivariance to permutation symmetries requires further proof in the context of dynamic graphs.

To address your concern, we have revised our paper and added a detailed theoretical analysis, which can now be found in Appendices B and C.

W2. It would be valuable to explain how the proposed methods address these gradient vanishing and explosion problems.

In our method, we utilize Gated Recurrent Units (GRUs) as the recurrent memory updating function. GRUs are specifically designed to address the issues of gradient vanishing and explosion, which commonly affect standard RNNs, particularly when processing long sequences. These issues can be potentially mitigated by the gating mechanisms in GRUs, which effectively regulate the flow of information and gradients.

As per the reviewer fjzh’s request, we conducted additional experiments by replacing the GRUs with a standard RNN. We observed a performance difference, with GRUs outperforming RNNs. Naïve RNNs are well-known to gradient vanishing/explosion problems. However, we cannot conclusively attribute this performance discrepancy to gradient vanishing or explosion. During the training process of both networks, we monitored the gradient norms and found that, in both cases, the gradients remained within a reasonable range throughout the training.

We acknowledge that gradient vanishing or explosion may occur in larger networks. However, this field is still in its early stages of development, with relatively small datasets. In this paper, we have made further attempt by introducing the classification task on larger dataset, i.e., CIFAR-100-INR. Even in this classification dataset, most existing methods performed poorly. In contrast, our approach demonstrated a significant improvement, improving over the state-of-the-art by 10.75%. We hope that as the field progresses, gradient vanishing and explosion issues will become more evident, and addressing these challenges will be a key focus of our future work.

Q1. Can the proposed methods be applied to other types of neural networks?

Currently, our method supports processing structures such as MLPs, CNNs, and Transformers. For other types of networks, further research and adjustments would be needed to adapt our approach. We would also greatly appreciate it if the reviewer could suggest specific types of neural networks they are interested in exploring with our method.

W3. A more detailed theoretical analysis of time and space complexity should be included.

Thank you for your question. We chose not to provide a theoretical time and space complexity analysis primarily because such analyses often may not accurately reveal the deployment cost in practice. For example, in Kofinas et al., the authors may use the PNA backbone but significantly increase the parameters of $\phi_m$, the message function, to improve performance. In such cases, the overall computational complexity could be dominated by the larger message function. Similarly, many modern machine learning competitions prioritize real-world computational costs on specific datasets rather than theoretical estimates. This is why we believe the empirical results presented in Table 6 more effectively reflect the practical computational costs of different methods.

To further address the reviewer's concern, we conducted a computational complexity analysis of the mentioned papers, which is detailed below. Please note that the theoretical results for Kofinas et al. are derived by us, as we are unable to find any explicit theoretical analysis of time and space complexity in their paper. The computational complexity analysis of our method can now also be found in Appendix F.

Cont.

Theoretical Results: Computational Cost

To address the reviewer's concern, we present our theoretical results on the computational cost of the mentioned methods. We omit an analysis of [3] because reference [3] is the workshop version of reference [1]. Since neither our method nor [1,3] requires a decoder for inference, this analysis focuses on providing theoretical results for the computational costs within the encoder.

To determine the time complexity of applying an L-layer MPNN on a graph with the following characteristics:

• Number of nodes per MLP layer: $n$. For simplicity, we assume each MLP layer has $n$ neurons.
• Dimension of node/edge features: $d$. For simplicity, we assume the dimension of edge and node features are the same.
• Number of MLP layers: $L$.
• Number of MPNN layers: $L$. In [1], the number of MPNN layers is set to be the same as the number of MLP layers.

---

* Time Complexity of Our Method

1. Message Computation (Equation 4):

   • Per edge computation: $O(d^2)$ (including edge feature transformation).
   • Total edges computation: $O(n^2 \cdot L \cdot d^2)$. There are $n^2$ edges in an MLP layer and $L$ total layers.

2. Aggregation (Equation 4) :

   • Each node aggregation: $O(n \cdot d)$. Each node aggregates messages from its $n$ neighbors in the graph of previous time step.
   • Total nodes: $O(n^2 \cdot L \cdot d)$.

3. Recurrent Memory Updates (Equation 6): $O(n \cdot L \cdot d^2)$.

Total Computational Complexity: $O(n^2 \cdot L \cdot d^2 + n^2 \cdot L \cdot d + n \cdot L \cdot d^2)$

Analysis: For large input networks, $O(n^2 \cdot L \cdot d^2)$ dominates.

* Space Complexity of Our Method

1. Node and Edge Features: $O(n^2 \cdot d + n \cdot d)$. We only store the node and edge features of the graph from the previous time step, corresponding to the previous MLP layer.
2. Memory (Equation 6): $O(n \cdot d)$.

Total Space Complexity: $O(n^2 \cdot d + n \cdot d)$

---

*Time Complexity of [1,3]

1. Message Computation:

   • Per edge computation: $O(d^2)$ (including edge feature transformation).
   • Total edges computation: $O(n^2 \cdot L^2 \cdot d^2)$. There are $n^2$ edges in an MLP layer and $L$ total layers. The same computation is executed $L$ times since they have an $L$-layer MPNN.

2. Aggregation:

   • Each node aggregation: $O(n \cdot d)$.
   • Total nodes: $O(n^2 \cdot L^2 \cdot d)$.

3. Node Updates: $O(n \cdot L^2 \cdot d^2)$.
   Each node update costs $O(d^2)$, and there are $n \cdot L^2$ updates.

4. Edge Updates: $O(n^2 \cdot L^2 \cdot d^2)$.
   Each edge update costs $O(d^2)$, and there are $n^2 \cdot L^2$ updates.
   Note: We do not update edges.

Total Computational Complexity: $O(n^2 \cdot L^2 \cdot d^2 + n^2 \cdot L^2 \cdot d + n \cdot L^2 \cdot d^2)$

Analysis: For large input networks, $O(n^2 \cdot L^2 \cdot d^2)$ dominates.

*Space Complexity of [1,3]

1. Node and Edge Features: $O(L \cdot n^2 \cdot d + n \cdot d)$.
   NG needs to store the node and edge features of the whole static graph.

Total Space Complexity: $O(L \cdot n^2 \cdot d + n \cdot d)$

---

*Computational Complexity of [2]

In the algorithm of [2], both the weights and biases of the original input network are updated. Since the number of weights is significantly larger than biases, we calculate complexity based only on the operations on the network weights:

1. Self-Attention Layer Computation: $O(L \cdot n^3 \cdot d)$.
2. MLP Computation: $O(L \cdot n^2 \cdot d^2)$.

Total Computational Complexity:
$O(L \cdot n^3 \cdot d + L \cdot n^2 \cdot d^2)$

*Space Complexity of [2]: $O(n^2 + n \cdot d)$.

This is similar to the cost in transformer.

---

Summary:

In this analysis, we have demonstrated that our method achieves superior performance compared to the alternatives. Specifically, when compared with the static neural graph [1,3], the dominant term in our time complexity, $O(n^2 \times L \times d^2)$, is significantly lower than the corresponding term in [1,3], which scales as $O(n^2 \times L^2 \times d^2)$. Additionally, our space complexity, $O(n^2 \cdot d + n \cdot d)$, is more resource-efficient than the $O(L \times n^2 \cdot d + n \cdot d)$ requirement of [1,3].

Thank you for your valuable feedback and constructive comments. We address your concerns as below:

---

W1. The motivation for processing INR weights in the Intro is unconvincing. INR classification seems to have been used previously as a convenient testbed, while recent works add more practical use-cases.

The Motivation:

Thank you for raising this concern. We agree that the motivation of "uncovering information about the data they encode" aligns more closely with tasks such as classifying and editing INRs. In light of this, we have revised the motivation in our paper to reflect a broader perspective. Please refer to Lines 29-33 for the updated version.

The Value of INR Classification:

We argue that INR classification is not merely a testbed for evaluating methods, but rather has potential practical value in real-world applications. For example, recent work [1] demonstrates that INRs exhibit great potential as a compressed representation for images. Using INRs to store images can yield improved rate-distortion performance, outperforming commonly used formats like JPEG2000. This trend suggests that future image compression protocols may involve transmitting INRs between transmitters and receivers. In this context, at the receiver side, upon receiving the INR, an additional step is required to reconstruct the original image for human viewing. Interestingly, another line of research [2] points out that much of the captured visual content may not be intended for human perception, but rather for automated machine vision analytics. For example, given an image, machine may only be interested in determining whether it contains a dog or a cat, rather than reconstructing the original pixel-level content. This is essentially a classification problem. Inspired by this, when we receive an INR, engineers can adopt the method developed in our paper to directly classify the INR, eliminating the need to reconstruct the image.

Comparison with Kofinas et al. and Lim et al. [a] on Additional Two Tasks:

Thanks for the suggestion. We did not include experiments related to “learning to optimize” because the editing of INRs in our experiment demonstrates the capability to modify the weight space and transform the functionality of neural networks. In this task, our method outperforms Kofinas et al.'s approach by a significant margin, achieving an average improvement of over 50% in terms of MSE loss. Additionally, we apologize for not being able to complete the comparison within the rebuttal period, due to the unavailability of their code and dataset for the task. We plan to update the results in the final version of our paper.

Regarding the experiments in Lim et al. [a], although they investigate multiple architectures, their primary tasks are similar to those in our paper, including predicting neural network accuracy and editing INRs. While we intended to perform the comparison, we unfortunately found that both their code and the "diverse architecture" dataset are not publicly available. Given the time required to construct the datasets and reproduce their results, we will include an update on these experiments in the final version of our paper.

[1] Strümpler, Yannick, et al. "Implicit neural representations for image compression." European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.

[2] Choi, Hyomin, and Ivan V. Bajić. "Scalable image coding for humans and machines." IEEE Transactions on Image Processing 31 (2022): 2739-2754.

W2. Lim et al. [a] is not discussed, similar to neural graphs of Kofinas et al.

Thanks for point out this related work that we missed. We have included a discussion of this paper in both the main body and related work section.