Equivariant Neural Functional Networks for Transformers

Authors: Hoang V. Tran, Thieu Vo, An Nguyen The, Tho Tran Huu, Minh-Khoi Nguyen-Nhat, Thanh Tran, Duy-Tung Pham, Tan Minh Nguyen
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: neural functional network, transformer, maximal symmetric group, equivariant model, dataset

Abstract:
This paper systematically explores neural functional networks (NFN) for transformer architectures. NFN are specialized neural networks that treat the weights, gradients, or sparsity patterns of a deep neural network (DNN) as input data and have proven valuable for tasks such as learnable optimizers, implicit data representations, and weight editing. While NFN have been extensively developed for MLP and CNN, no prior work has addressed their design for transformers, despite the importance of transformers in modern deep learning. This paper aims to address this gap by providing a systematic study of NFN for transformers. We first determine the maximal symmetric group of the weights in a multi-head attention module as well as a necessary and sufficient condition under which two sets of hyperparameters of the multi-head attention module define the same function. We then define the weight space of transformer architectures and its associated group action, which leads to the design principles for NFN in transformers. Based on these, we introduce Transformer-NFN, an NFN that is equivariant under this group action. Additionally, we release a dataset of more than 125,000 Transformers model checkpoints trained on two datasets with two different tasks, providing a benchmark for evaluating Transformer-NFN and encouraging further research on transformer training and performance.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 8

Reviews and discussion:
<thread>
creation date:
1730549414644

thread_type:
official_review

<thread_object>
creator: reviewer_ynvn

<content>
summary:
This submission provides a theoretical study of neural functional networks (NFNs) for transformers. By formally defining the weight space of the standard transformer block and the corresponding group action of such a space, the authors propose several design principles for NFNs and introduce Transformer-NFN to predict the test performance with a transformer checkpoint and the corresponding training configurations. Additionally, a dataset with 125k transformer checkpoints trained on MINST and AGNews will be released for reproduction and further analysis.

strengths:
- This topic is quite interesting and do have some foreground. In the foreseeable future, a plenty of pre-trained transformer checkpoints would be released and there is no golden rule to find out which checkpoint is appropriate for a specific task. And the NFN may be useful and draw some inspiration for future research.
- For all I know, this may be one of the first works to develop NFN for transformers. 
- The theoretical analysis is detailed enough.

weaknesses:
- According to the main presentation of this submission, no obvious weaknesses are found.

questions:
- How about the performance on language modeling? Is Transformer-NFN able to predict?
- How about the performance of the larger transformer 
- 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?
- 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730646778002

thread_type:
official_review

<thread_object>
creator: reviewer_yffw

<content>
summary:
This paper investigates Equivariant Neural Functional Networks for Transformer architectures by developing a specialized model, Transformer-NFN, designed to process the weights, gradients, or sparsity patterns of Transformers as input data. The paper defines a maximal symmetric group for multi-head attention modules and establishes a group action on the weight space of Transformer architectures, creating a systematic framework for NFNs with equivariance properties. Besides, Transformer-NFN, an NFN specifically tailored to Transformer structures, is evaluated on a new dataset containing 125,000 Transformer checkpoints. Experimental results indicate that Transformer-NFN outperforms other models in predicting Transformer model performance.

strengths:
1. The paper presents a novel theoretical framework for designing Neural Functional Networks that respects the inherent symmetries of Transformer architectures, representing an innovative approach previously unexplored in this domain.
2. The release of a dataset with 125,000 Transformer model checkpoints offers a valuable resource for future research on Transformer performance prediction and NFNs.
3. The paper provides a detailed methodological development, including thorough derivations and design principles, establishing the theoretical foundation for NFNs tailored to Transformer architectures, which may inspire future studies in this area.

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

questions:
1. Could the authors clarify the expected computational cost of Transformer-NFN compared to traditional NFNs, considering the added design complexity?
2. Will the paper include code to facilitate reproducibility?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730668341875

thread_type:
official_review

<thread_object>
creator: reviewer_tukj

<content>
summary:
The paper presents a formulation of NFNs for transformers beyond MLPs and CNNs. The paper provides thorough analysis of how to establish a weight space around multihead attention and a way to formalize NFNs through a parameter sharing strategy. The paper also shows the application on multiple transformer models across a basic vision and language task. The paper is very detailed in it's approach and provides for a good approach to understanding how the weights of transformer blocks are affected by hyperparameters and data

strengths:
The paper provides a very structured approach to a the solution of defining NFNs for transformers and is rigorous in it's approach.Some of the mathematical formulation around weight spaces of a multihead attention and the solution for parameter sharing of the NFN is good.

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

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.

questions:
How does this generalize to models with more parameters and trained on larger datasets
</content>
</thread_object>

</thread>
<thread>
creation date:
1731287711917

thread_type:
official_review

<thread_object>
creator: reviewer_wnmm

<content>
summary:
The paper introduces a novel concept of Transformer Neural Functional Networks (Transformer-NFN), designed as equivariant neural functional networks specifically for transformer architectures.The proposed Transformer-NFN addresses the challenge of applying neural functional networks (NFNs) to transformers by systematically defining weight spaces, symmetry groups, and group actions for multi-head attention and ReLU-MLP components. Experimental results demonstrate the effectiveness of Transformer-NFN in predicting transformer generalization and performance, surpassing traditional models across multiple metrics.

strengths:
- The is the first paper to my knowledge expore neural functional networks (NFN) in the context of transformer, where NFN treat the weights, gradients, or sparsity patterns of a deep neural network (DNN) as input data.
- The proposed Small Transformer Zoo dataset is a valuable resource for benchmarking and studying transformer-based networks, promoting reproducibility and further research.

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

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

</thread>
<thread>
creation date:
1731950924323

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response (1/2)

comment:
Dear AC and reviewers,

Thanks for your thoughtful reviews and valuable comments, which have helped us improve the paper significantly.

We sincerely thank the reviewers for their valuable feedback and constructive suggestions. We are encouraged by the endorsements that: 

1. Our Transformer-NFN is the first equivariant neural functional networks for Transformer models (Reviewer wnMM, YffW, ynVN), which may inspire and prove useful for future research (Reviewer YffW, ynVN).
2. The theoretical analysis is detailed (Reviewer ynVN), well-structured, and supported by rigorous mathematical derivations and design principles (Reviewer TUKj, YffW).
3. Our proposed Small Transformer Zoo dataset provides a valuable resource for advancing future research on Transformers (Reviewer wnMM, YffW).

---

Below, we address two common points raised in the reviews:

**Q1. Computational complexity of Transformer-NFN.**

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., $\texttt{einsum}$), enabling efficient and highly parallelizable computations on modern GPUs. This ensures computational practicality while delivering significant performance improvements.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732184200108

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Incoporating comments and suggestions from reviewers, as well as some further empirical studies we believe informative, we summarize here the main changes in the revised paper:
1. We add an experiment that shows Transformer-NFN’s equivariance to the maximal symmetric group action $\mathcal{G}\_\mathcal{U}$ is crucial for its consistent performance compared to other baselines (Table 9 in Appendix I). The results show that, unlike other non-equivariant baselines, Transformer-NFN maintains stable Kendall’s $\tau$ values across different ranges of scale operators, highlighting its robustness in when the input is transformed by maximal symmetric group action $\mathcal{G}\_\mathcal{U}$.

2. We add the computational complexity of our Equivariant Layer and Invariant Layer in Appendix J of the revised paper. The implementation uses the $\texttt{einsum}$ operation, which enables efficient and highly parallelized computations on modern GPUs. This design ensures computational feasibility while delivering significant performance gains.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper proposes an algorithm for equivariant transformer architectures of neural fields. Reviewers are unanimously in favor of accepting the paper, and initially had questions about the computational cost, as well as the significance of improvements. Authors answered with additional experiments, showing that the computational cost is reasonable, and improvements are also non-trivial. All in all, this is a good paper.