Equivariant Neural Functional Networks for Transformers

Well done, all my concerns have been addressed. I still recommend the authors pay some attention to the larger and more practical scenario about the transformer in their future works. I will keep my score.

Q3. From my perspective, the performance of a checkpoint is related to the training data, the architecture of the network, the training algorithm, hyperparameter configuration, etc. So why does NFN only take the weights and the configuration into account? Is it a sufficient condition for performance prediction?

Answer. In our study, we evaluate Transformer-NFN on a test set comprising transformer weights trained on the same task as those in the training set. Our findings demonstrate that Transformer-NFN, by learning from a diverse set of weights trained under varying hyperparameter configurations and training algorithms, can effectively predict the performance of unseen transformer weights trained in different settings.

This highlights that transformer weights themselves contain substantial information, making them sufficient for predicting a model's generalization capability without relying on additional features. As shown in Table 3, the encoder weights hold the most relevant information, enabling accurate predictions of the generalization of pretrained weights. This aligns with intuition, as the transformer's encoder performs the majority of the computational work and is the most critical component driving the performance of the trained model.

Q4. The evaluated dataset is far away from the present scale of transformer. The weight of modern transformers is huge (bert-base 110M parameters), is Transformer-NFN able to predict the performance of it?

Answer. We agree that evaluating Transformer-NFN on larger models and datasets would provide a more comprehensive demonstration of its strengths. However, generating such datasets requires training thousands of large-scale transformers, which exceeds the computational resources available in our academic lab. For context, both our MNIST-Transformers and AGNews-Transformer datasets each consist of approximately 63k trained weights each. While our current experiments are limited to smaller-scale models due to these constraints, the core principles of Transformer-NFN are model-agnostic and specifically designed to scale with increasing model size.

We appreciate the suggestion to include small-scale experiments with larger models. We recognize its importance and will prioritize this direction in future iterations of our work.

Once again, we sincerely thank the reviewer for their feedback. Please let us know if there are any additional concerns or questions from the reviewer regarding our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

Q1+Q2. How about the performance on language modeling? Is Transformer-NFN able to predict?

How about the performance of the larger transformer?

Answer. We greatly appreciate your observation regarding the importance of evaluating Transformer-NFN on tasks such as language modeling. While this would enhance understanding of NFN's applicability, generating such a dataset within the rebuttal period is challenging due to computational demands. For context, both our MNIST-Transformers and AGNews-Transformer datasets each consist of approximately 63k trained weights each. Extending this effort to larger-scale models, such as training thousands of 110M-parameter transformers, exceeds the computational resources of our academic lab. Nevertheless, we recognize the importance of this direction and plan to address it in future work.

Currently, the closest evaluation we conducted to language modeling is text classification on the AGNews dataset. In this setting, Transformer-NFN demonstrated significant performance improvements over other baseline methods, as shown in Q2 of the General Response.

We recall the details of the experiment from Q2 in General Response as follows.

Experiment Setup. In this experiment, we evaluate the performance of Transformer-NFN on the AGNews-Transformers dataset augmented using the group action $\mathcal{G}_\mathcal{U}$. We perform a 2-fold augmentation for both the train and test sets. The original weights are retained, and additional weights are constructed by applying permutations and scaling transformations to transformer modules. The elements in $M$ and $N$ (as defined in Section 4.3) are uniformly sampled across $[-1, 1]$, $[-10, 10]$, and $[-100, 100]$.

Results. The results for all models are presented in Table 1. The findings demonstrate that Transformer-NFN maintains stable Kendall’s $\tau$ across different ranges of scale operators. Notably, as the weights are augmented, the performance of Transformer-NFN improves, whereas the performance of other baseline methods declines significantly. This performance disparity results in a widening gap between Transformer-NFN and the second-best model, increasing from 0.031 to 0.082.

The general decline observed in baseline models highlights their lack of symmetry. In contrast, Transformer-NFN's equivariance to $\mathcal{G}_\mathcal{U}$ ensures stability and even slight performance improvements under augmentation. This underscores the importance of defining the maximal symmetric group, to which Transformer-NFN is equivariant, in overcoming the limitations of baseline methods.

Table 1: Performance measured by Kendall's $\tau$ of all models on Augmented AGNews-Transformers dataset. Uncertainties indicate standard error over 5 runs.

W2+Q2. The architecture's intricate structure, requiring a specialized understanding of group actions and equivariant layers, could hinder adoption. The dense theoretical foundation, coupled with limited practical resources, may present obstacles to reproducibility and broader implementation in practical settings. Will the paper include code to facilitate reproducibility?

Answer. We recognize the challenges in understanding the implementation. To address this, our paper already includes an Implementation section in the Appendix H.2 and H.3 with pseudocode for clarity. The network consists of multiple components, represented by elegant einsum operators that are easy to read and adapt. Additionally, the supplementary materials provide code for dataset generation and Transformer-NFN layers, accompanied by detailed documentation to facilitate result reproduction.

Q1. Could the authors clarify the expected computational cost of Transformer-NFN compared to traditional NFNs, considering the added design complexity?

Answer. The computational complexity of the Transformer-NFN is derived from its invariant and equivariant layers as outlined in our pseudocode (Appendix H):

• Equivariant Layer Complexity: $\mathcal{O}(d \cdot e \cdot h \cdot D^2 \cdot \max(D_q, D_k, D_v))$.
• Invariant Layer Complexity: $\mathcal{O}(d \cdot e \cdot D' \cdot D \cdot \max(D \cdot h, D_A))$.

Where the parameters are:

• $d, e$: Input and output channels of the Transformer-NFN layer, respectively.
• $D$: Embedding dimension of the transformer.
• $D_q, D_k, D_v$: Dimensions for query, key, and value vectors of the transformer.
• $D_A$: Dimension of the linear projection of the transformer.
• $h$: Number of attention heads of the transformer.
• $D'$: Dimension of the output of the Transformer-NFN Invariant Layer.

The Transformer-NFN implementation leverages optimized tensor contraction operations (e.g., $`einsum`$), enabling efficient and highly parallelizable computations on modern GPUs. This ensures computational practicality while delivering significant performance improvements.

We sincerely thank the reviewer for the valuable feedback. If our responses adequately address all the concerns raised, we kindly hope the reviewer will consider raising the score of our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. While the use of equivariant layers in Transformer-NFN is theoretically compelling, empirical results suggest limited practical impact. Specifically, the marginal gains in predictive accuracy on the benchmark dataset raise questions about the method’s utility in real-world applications.

Answer. To further evaluate the performance of our Transformer-NFN, we conducted experiments testing all methods under variations in the maximal symmetric group (refer to Q2 in General Response). The dataset used was a set of Transformers trained on text classification for AGNews, representing a natural language application. Our model demonstrated consistent performance, while other baselines experienced a significant drop. This underscores our model's ability to process diverse inputs in the weight space, which could be advantageous for larger models trained on extensive datasets in real-world applications.

We recall the details of the experiment from Q2 in General Response as follows.

Experiment Setup. In this experiment, we evaluate the performance of Transformer-NFN on the AGNews-Transformers dataset augmented using the group action $\mathcal{G}_\mathcal{U}$. We perform a 2-fold augmentation for both the train and test sets. The original weights are retained, and additional weights are constructed by applying permutations and scaling transformations to transformer modules. The elements in $M$ and $N$ (as defined in Section 4.3) are uniformly sampled across $[-1, 1]$, $[-10, 10]$, and $[-100, 100]$.

Results. The results for all models are presented in Table 1. The findings demonstrate that Transformer-NFN maintains stable Kendall’s $\tau$ across different ranges of scale operators. Notably, as the weights are augmented, the performance of Transformer-NFN improves, whereas the performance of other baseline methods declines significantly. This performance disparity results in a widening gap between Transformer-NFN and the second-best model, increasing from 0.031 to 0.082.

The general decline observed in baseline models highlights their lack of symmetry. In contrast, Transformer-NFN's equivariance to $\mathcal{G}_\mathcal{U}$ ensures stability and even slight performance improvements under augmentation. This underscores the importance of defining the maximal symmetric group, to which Transformer-NFN is equivariant, in overcoming the limitations of baseline methods.

Table 1: Performance measured by Kendall's $\tau$ of all models on Augmented AGNews-Transformers dataset. Uncertainties indicate standard error over 5 runs.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. The paper is a bit dense to read and understand. Though this can be considered fine given the topic and detailed approach. There could've been a better balance at analyzing the different theorems presented vs covering the mathematical proof.

Answer. We appreciate the reviewer's suggestion and are in the process of restructuring the paper to enhance its flow. The revised version will be updated in a few days.

W2 + Q1. The experimentation is basic and doesn't necessarily capture the strengths of the method or even transformers. Transformers are known to generalize over large model sizes and large data and it is the biggest weakness of the paper on how this method would hold for larger models. Even if a small scale experiment was included, it would bolster the through mathematical approach much better.

How does this generalize to models with more parameters and trained on larger datasets?

Answer. We agree that evaluating Transformer-NFN on larger models and datasets would provide a more comprehensive demonstration of its strengths. However, generating such datasets requires training thousands of large-scale transformers, which exceeds the computational resources available in our academic lab. For context, both our MNIST-Transformers and AGNews-Transformer datasets each consist of approximately 63k trained weights each. While our current experiments are limited to smaller-scale models due to these constraints, the core principles of Transformer-NFN are model-agnostic and specifically designed to scale with increasing model size.

We appreciate the suggestion to include small-scale experiments with larger models. We recognize its importance and will prioritize this direction in future iterations of our work.

Once again, we sincerely thank the reviewer for their feedback. Please let us know if there are any additional concerns or questions from the reviewer regarding our paper.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Thanks for your response, and we appreciate your endorsement. We acknowledge the importance of applicability to larger models and consider this an intriguing challenge for future research.

Q2. I'm curious how sensitive is the maximal symmetric group of the weights in a MHA module, would affect the peformance of final peformance of NFN.

Answer. We show that Transformer-NFN’s equivariance to the maximal symmetric group action $\mathcal{G}_\mathcal{U}$ is crucial to its performance, by conducting an experiment on augmented dataset. The experiment is presented in Q2 of General Response as below.

Experiment Setup. In this experiment, we evaluate the performance of Transformer-NFN on the AGNews-Transformers dataset augmented using the group action $\mathcal{G}_\mathcal{U}$. We perform a 2-fold augmentation for both the train and test sets. The original weights are retained, and additional weights are constructed by applying permutations and scaling transformations to transformer modules. The elements in $M$ and $N$ (as defined in Section 4.3) are uniformly sampled across $[-1, 1]$, $[-10, 10]$, and $[-100, 100]$.

Results. The results for all models are presented in Table 1. The findings demonstrate that Transformer-NFN maintains stable Kendall’s $\tau$ across different ranges of scale operators. Notably, as the weights are augmented, the performance of Transformer-NFN improves, whereas the performance of other baseline methods declines significantly. This performance disparity results in a widening gap between Transformer-NFN and the second-best model, increasing from 0.031 to 0.082.

The general decline observed in baseline models highlights their lack of symmetry. In contrast, Transformer-NFN's equivariance to $\mathcal{G}_\mathcal{U}$ ensures stability and even slight performance improvements under augmentation. This underscores the importance of defining the maximal symmetric group, to which Transformer-NFN is equivariant, in overcoming the limitations of baseline methods.

Table 1: Performance measured by Kendall's $\tau$ of all models on Augmented AGNews-Transformers dataset. Uncertainties indicate standard error over 5 runs.

Unlike baseline methods, which lack such symmetry considerations and exhibit significant performance declines under augmentation, Transformer-NFN maintains stability and even demonstrates improved performance. This highlights the importance of leveraging equivariant group actions to achieve robustness and enhance final performance.

This is supported by existing research [1, 2, 3] showing that incorporating symmetric groups into NFN for MLP improves performance compared to those without symmetry consideration.

Reference.

[1] Zhou et al., Permutation Equivariant Neural Functionals. NeurIPS 2023.

[2] Kalogeropoulos et al., Scale Equivariant Graph Metanetworks. NeurIPS 2024.

[3] Tran et al., Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

We sincerely thank the reviewer for the valuable feedback. If our responses adequately address all the concerns raised, we kindly hope the reviewer will consider raising the score of our paper.

We appreciate the reviewer’s feedback and have provided the following responses to address the concerns raised about our paper. Below, we summarize the weaknesses and questions highlighted by the reviewer and provide our answers accordingly.

W1. The paper did not discuss the computational overhead associated with equivariant polynomial layers and whether this could affect scalability for larger transformer models.

Answer. The computational complexity of the Transformer-NFN is derived from its invariant and equivariant layers as outlined in our pseudocode (Appendix H):

• Equivariant Layer Complexity: $\mathcal{O}(d \cdot e \cdot h \cdot D^2 \cdot \max(D_q, D_k, D_v))$.
• Invariant Layer Complexity: $\mathcal{O}(d \cdot e \cdot D' \cdot D \cdot \max(D \cdot h, D_A))$.

Where the parameters are:

• $d, e$: Input and output channels of the Transformer-NFN layer, respectively.
• $D$: Embedding dimension of the transformer.
• $D_q, D_k, D_v$: Dimensions for query, key, and value vectors of the transformer.
• $D_A$: Dimension of the linear projection of the transformer.
• $h$: Number of attention heads of the transformer.
• $D'$: Dimension of the output of the Transformer-NFN Invariant Layer.

The Transformer-NFN implementation leverages optimized tensor contraction operations (e.g., $`einsum`$), enabling efficient and highly parallelizable computations on modern GPUs. This ensures computational practicality while delivering significant performance improvements.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Q2. Additional experiment to show the strength of Transformer-NFN.

Experiment Setup. In this experiment, we evaluate the performance of Transformer-NFN on the AGNews-Transformers dataset augmented using the group action $\mathcal{G}_\mathcal{U}$. We perform a 2-fold augmentation for both the train and test sets. The original weights are retained, and additional weights are constructed by applying permutations and scaling transformations to transformer modules. The elements in $M$ and $N$ (as defined in Section 4.3) are uniformly sampled across $[-1, 1]$, $[-10, 10]$, and $[-100, 100]$.

Results. The results for all models are presented in Table 1. The findings demonstrate that Transformer-NFN maintains stable Kendall’s $\tau$ across different ranges of scale operators. Notably, as the weights are augmented, the performance of Transformer-NFN improves, whereas the performance of other baseline methods declines significantly. This performance disparity results in a widening gap between Transformer-NFN and the second-best model, increasing from 0.031 to 0.082.

The general decline observed in baseline models highlights their lack of symmetry. In contrast, Transformer-NFN's equivariance to $\mathcal{G}_\mathcal{U}$ ensures stability and even slight performance improvements under augmentation. This underscores the importance of defining the maximal symmetric group, to which Transformer-NFN is equivariant, in overcoming the limitations of baseline methods.

Table 1: Performance measured by Kendall's $\tau$ of all models on Augmented AGNews-Transformers dataset. Uncertainties indicate standard error over 5 runs.

We hope that our rebuttal has helped to clear concerns about our work. We are glad to answer any further questions you have on our submission and we would appreciate it if we could get your further feedback at your earliest convenience.