Research map / Architectures and efficiency / Attention and transformers
Equivariant Neural Functional Networks for Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Characterizes the maximal symmetry group of multi-head attention, builds an equivariant polynomial neural functional network (Transformer-NFN) for transformer weights, and releases the Small Transformer Zoo of ~125k transformer checkpoints on MNIST and AGNews.
Problem
Neural functional networks (NFNs) that take transformer weights as input have not been designed; transformer weight spaces have symmetries (multi-head reparameterizations, permutations) that existing MLP/CNN NFNs do not handle.
Approach
Determine the maximal symmetric group of multi-head attention (S_h × GL_{D_k}^h × GL_{D_v}^h), define a transformer-block weight space with a group action including neuron permutations, and derive a G_U-invariant/equivariant polynomial NFN layer via parameter sharing that is applied to transformer checkpoints to predict generalization.
Key results
- Transformer-NFN reaches Kendall's tau 0.905 on MNIST-Transformers without threshold, outperforming MLP (0.866) and STATNN (0.881).
- On AGNews-Transformers, Transformer-NFN reaches 0.910 Kendall's tau without threshold, versus 0.879 for the best baseline (MLP).
- Using only transformer-block weights gives 0.902 vs 0.905 with all components on MNIST-Transformers.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Transformer-NFN instead of the strongest baseline | Generalization prediction · MNIST-Transformers, all five accuracy-threshold settings | Kendall's τ: performance gap ranging from $0.004$ to $0.026$ | Helped |
| Use Transformer-NFN instead of the strongest baseline | Generalization prediction · AGNews-Transformers, all five accuracy-threshold settings | Kendall's τ: Kendall's $\tau$ gaps ranging from $0.018$ to $0.033$ | Helped |
| Use encoder weights only (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.905 ± 0.002 vs. 0.902 ± 0.001 | Hurt |
| Use encoder weights only (ablation) | Transformer-NFN · AGNews-Transformers | Kendall's τ: 0.909 ± 0.001 vs. 0.910 ± 0.001 | Hurt |
| Use embedding weights only (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.424 ± 0.002 | Hurt |
| Use classifier weights only (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.847 ± 0.001 | Hurt |
| Use classifier weights only (ablation) | Transformer-NFN · AGNews-Transformers | Kendall's τ: 0.795 ± 0.008 | Hurt |
| Use embedding and classifier weights only, omitting the encoder (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.857 ± 0.003 | Hurt |
| Omit embedding weights (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.904 ± 0.001 vs. 0.905 ± 0.002 | No clear effect |
| Omit embedding weights (ablation) | Transformer-NFN · AGNews-Transformers | Kendall's τ: 0.910 ± 0.001 vs. 0.910 ± 0.001 | No clear effect |
| Omit classifier weights (ablation) | Transformer-NFN · MNIST-Transformers | Kendall's τ: 0.903 ± 0.001 vs. 0.905 ± 0.002 | No clear effect |
Surprising findings
- Encoder (transformer block) weights alone nearly match full-model prediction performance, while the embedding alone is weak (0.424 on MNIST).
Limitations
- Symmetry analysis covers only the multi-head attention and MLP permutations; scaling and layer-norm symmetries are left for future work.
- Experiments use a small fixed transformer (two 2-head blocks) on two classification tasks.
- Only Kendall's tau generalization prediction is evaluated; no downstream applications such as weight editing or optimization.
Experimental setup
Datasets: MNIST AGNews
Benchmarks: MNIST-Transformers AGNews-Transformers
Metrics: Kendall's tau
Baselines: MLP STATNN XGBoost LightGBM Random Forest
Builds on: Neural functional networks (permutation-equivariant NFNs) STATNN Kendall tau
Closest papers by meaning
- Equivariant Polynomial Functional Networks (ICML 2025 · similarity 0.89)
- Unraveling the Gradient Descent Dynamics of Transformers (NeurIPS 2024 · similarity 0.86)
- Rethinking Attention: Exploring Shallow Feed-Forward Neural Networks as an Alternative to Attention Layers in Transformers (Student Abstract) (AAAI 2024 · similarity 0.86)
- In-context Learning on Function Classes Unveiled for Transformers (ICML 2024 · similarity 0.86)
- What Does It Mean to Be a Transformer? Insights from a Theoretical Hessian Analysis (ICLR 2025 · similarity 0.86)
- The Need for Speed: Pruning Transformers with One Recipe (ICLR 2024 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.