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
Venue: ICLR 2025 Poster (2025)
Keywords: neural functional network
transformer
maximal symmetric group
equivariant model
dataset
Primary area: other topics in machine learning (i.e., none of the above)
TL;DR: This paper systematically studies neural functional networks (NFNs) for Transformers, presenting a design principle, and an equivariant NFN called Transformer-NFN, along with a benchmark dataset for evaluation.

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.

Published as a conference paper at ICLR 2025
EQUIVARIANT NEURAL FUNCTIONAL NETWORKS FOR TRANSFORMERS
Viet-Hoang Tran!" Thieu N. Vo!
An Nguyen The Tho Tran Huu'
Minh-Khoi Nguyen-Nhat? Thanh Tran' Duy-Tung Pham? Tan M. Nguyen'
'National University of Singapore VVinUniversity, Vietnam
FPT Software AI Center, Vietnam
[hoang. tranviet, thotranhuu nus. edu, fthieuvo, fanmn]@nus.... (annt68, khoinnml, tungpd1 0 fpt. com, 21thanh, tq@vinuni. .en.i vn
ABSTRACT
This paper systematically explores neural functional networks (NFN) for trans- former architectures. NEN are specialized neural networks that treat the weights, gradients, or sparsity patterns of deep neural network (DNN) as input data and have proven valuable for tasks like learnable optimizers, implicit data represen- tations, and weight editing. While NFN have been extensively developed for MLP and CNN, no prior work has addressed their design for transformers, de- spite their importance in modern deep learning. This paper aims to address this gap by systematically studying NFN for transformers. We first determine the maximal symmetric group of the weights in a multi-head attention module and necessary and sufficient condition under which two sets of hyperparameters of the module define the same function. We then define the weight space of trans- former architectures and its associated group action, leading design principles for NFN in transformers. Based on these, we introduce Transformer-NFW an NFN equivariant under this group action. Additionally, release dataset of over 125,000 Transformers model checkpoints trained on two datasets with two tasks, providing benchmark for evaluating Transformer-NFN and encouraging further research on transformer training and performance. The code is publicly available http:/iiibbc.ooo/aaaaaaattclalla AAA.AAAAAAAAATSSSrreeeEEEEEEEE
INTRODUCTION
Deep neural networks (DNNs) have become highly flexible and adaptable modeling tools, widely employed in domains such as natural language processing (Rumelhart et al., 1986; Hochreiter & Schmidhuber, 1997; Vaswani et al., 2017; Devlin et al., 2019; Nielsen et al., 2025b; Gu & Dao, 2024; Vo et al., 2025), computer vision (He et al., 2015; Szegedy et al., 2015; Krizhevsky et al., 2012; Huang al., 2020; Nguyen al., 2023a), and natural science (Raissi al., 2019; Jumper et al., 2021; Portwood et al., 2019). Increasing attention is being given to constructing specialized neural networks that treat the weights, gradients, or sparsity patterns of DNNs as inputs. These specialized networks are referred to neural functional networks (NFNs) (Zhou et al., 2024b). NFNs have proven useful various applications, such as designing learnable optimizers for neural network training (Bengio et al., 2013; Runarsson & Jonsson, 2000; Andrychowiez et al., 2016; Metz al., 2022; Knyazev al., 2024). extracting information from implicit representations of data (Stanley, 2007; Mildenhall et al., 2021; Runarsson & Jonsson, 2000), performing targeted weight editing (Sinitsin et al., 2020; Cao et al., 2021; Mitchell et al., 2022), evaluating policies (Harb al., 2020), and applying Bayesian inference with neural networks (Sokota et al., 2021).
The construction of NENs presents significant challenges due to the complexity and high dimen- sionality of their structures. Earlier approaches sought to address this by constraining the training process, thereby restricting the weight space (Dupont et al., 2022h; Bauer et al., 2023; Luigi et al.). More recently, research has shifted towards the development of permutation-equiaarin semmutaiinn-euuuvarinnt NFNs, pable of operating on neural network weights without such constraints (Navon et al., 2023; Zhou al., 2024b; Kotinas et al., 2024; Zhou et al., 2024c). These advancements have resulted in NFNs
Equal contribution. Correspondence to: seang.tranviet@u... and tanmn @nus @nuu.cdu.sg

---

Published as a conference paper at ICLR 2025
that exhibit equivariance to weight permutations, which correspond to neuron rearrangements within hidden layers. Additionally, NENs that are equivariant to both permutations and operations such as scaling sign-flpppng have been proposed. These include graph-based message- passsingg meth- ods as described in (Kalogeropoulos et al., 2024) and parameter sharing techniques in (Tran et al., 2024a). However, most of the recent neural functional networks designed for Multilayer Percep- trons (MLPs) and Convolutional Neural Networks (CNNs).
Transformers have achieved remarkable success in the field of natural language processing (NLP) (Vaswani et al., 2017; Bahdanau et al.. 2015; Nguyen et al., 2023b; Nielsen et al., 2025a), powering large-scale language models. These models highlight the transformer architecture 's capacity to per- form a broad range of NLP-related tasks and its upplicability to real-world scenarios. The influence of Transformers extends beyond NI.,, finding applications areas such as computer vision through Vision Transformers (Dosovitskiy et al., 2020). reinforcement learning (Parisotto et al., 2020; Chen al., 2021), audio processing (Radford al., 2023), multi-modal data integration (Radford al., 2021), and robotics (Monastirsky al., 2022). Despite the significance of Transformers, tematic development of NFNs for Transformers has been realized to date. Addressing the design principle of NFNs for Transformers, well constructing equivariant NFNs capable of processing transformer models, has thus emerged as critical research problem.
sys-
1.1 BACKGROUND
Self-attention is the core component of transformer block (Cho et al.; Lin et al., 2017; Parikh al., 2016; Teo & Nguyen, 2025). Self-attention learns align tokens within input sequence by assessing the relative importance of each token relation to all other tokens. then transforms each token into a weighted average of the feature representations of other tokens, with the weights determined by importance scores between token pairs. These importance scores allow each token to attend others in the sequence, effectively capturing contextual representations (Bahdanau al., 2015; Kim et al., 2017; Vaswani et al., 2017; Nguyen et al., 2022; Abdullaev & Nguyen, 2025).
[X X2] T E R be an input sequence of L tokens with D features. Each token is vector X, E RR' The self-attention, denoted by Head, transforms X into the output sequence Head(X) RLDD, defined
Self-attention. Let X
(XW(O))(XWL00)) XWO,
Head(X; ((Q), w(k) w()) softmax
(1)
DDx
where the parameters ((Q), ((K) E Dx and W(V) E Dx are called the query, key and value matrices. Here, D& and D. are given positive integers.
Multihead attention. To jointly attend to information from different representation subspaces at different positions, multihead attention is used. Let h be positive integer which represents the number head. The multihead attention transforms the input sequence X E R4xD to an output sequence in L defined by:
MultiHead (X;W1O), (W(Q.) W(K.), w(Ki) (E
Head (x: w(Q,9, w(K.) w(v))) ) wro),
where each head is self-attention defined in equation 1, and ! is the concatenation operator. Here, the matrices WQQ.)) (K,K)) E DxD, W(V)) I DxD, and W(O) E DRDDD.D are learnable from input data for some postive integers D₁ and D...
1.2 CONTRIBUTION
This paper provides a systematic study on the development neural functional network (NFN) for transformer architectures. To achieve this, present three essential components for the study: (1) a design principle of NFNs for Transformers that incorporates the maximal symmetric group for the multi-head attention module, (2) equivariant NEN for Transformers, which we will call Transformer-NFi and (3) a benchmark dataset for testing the applicability and efficiency of Transformer-NFN. In particular, our contributions are four-fold:
2

---

Published as a conference paper at ICLR 2025
1. We determine the maximal symmetric group of the weights in multi-head attention module, establishing the necessary and sufficient conditions under which two sets of hyperparameters for multi-head attention module define the same function.
2. We formally define the weight space of a transformer architecture, along with group action on this weight space. In conjunction with the maximal symmetric group of the weights in multi-head attention modules, we characterize the design principles for NFNs in transformer architectures.
3. We introduce Transformer-NHH an NFN for transformer architectures that is equivariant un- der the specified group action. The main building block of Transformer-NFN an equivariant polynomial layer derived from parameter-sharing strategy.
4. We release Small Transformer Zoo dataset, which consists of more than 125,000 Transformers model checkpoints trained on two different tasks: digit image classification on MNIST and text topic classificaction on AGNews. To our knowledge, this the first dataset of its kind. This dataset serves as a benchmark for testing the applicability and efficiency of our Transformer- NFN, while also encouraging further research in this field to gain deeper understanding of transformer network training and performance.
We empirically demonstrate that Transformer-Ne NFNN consistently outperforms other baseline models on our constructed datasets. Through comprehensive ablation studies, we emphasize Transformer- NFN's ability to effectively capture information within the transformer block, establishing it as robust predictor of model generalization.
Organization. Following review of related works in Section 2, we derive the maximal symmetric group of multihead attention in Section dard transformer block and define corresponding group action. Section 5 introduces family of equivariant NFNs for Transformers, referred to as Transformer-NEW and details of the Small Transformer Zoo dataset in Section 6. Finally, in Section evaluate the applicability and efficiency of Transformer-NFW on this dataset.
Section 4, we construct the weight space of stan-
We then present the setting
2 RELATED WORK
Symmetries of weight spaces. The exploration of symmetries within the weight spaces of neu- ral networks, which relates to assessing the functional equivalence of these networks, is a well- established field of study (Allen-Zhu et al., 2019; Du et al., 2019; Frankle & Carbin; Belkin et al., 2019: Novak et al., 2018). This topic was first introduced by Hecht-Nielsen 1990). Following this, numerous studies have yielded insights for various network architectures, as detailed in (Chen et al., 1993; Fefferman & Markel, 1993; Kurkova & Kainen, 1994; Albertini & Sontag, 1993b;a; Bui Thi Mai & Lampert, 2020).
(Hecht-Niclsen,
Neural functional networks. Recent studies have aimed to develop representations for trained clas- sifiers evaluate their generalization performance and understand neural network behavior (Baker et al., 2018; Eilertsen et al.; Unterthiner et al., 2020; Schürholt et al., 2021; 2022a:b). Notably, low- dimensional encodings for Implicit Neural Representations (INRs) have been created for various downstream tasks (Dupont et al., 2022a; Luigi et al.). Schürholt et al. (2021) proposed neuron per- mutation augmentations, and other research has focused on encoding and decoding neural network parameters for reconstruction and generative purposes (Peebles et al., 2022; Ashkenazi et al., 2023; Knyazev et al., 2021: Erkoç et al., 2023).
Equivariant NFNs for MLPs and CNNs. Recent advancements have made considerable strides in addressing the limitations of permutation equivariant neural networks through the introduction of permutation equivariant layers. These layers employ intricate weight-sharing patterns (Navon et al., 2023; Zhou et al., 2024b; Kofinas et al., 2024; Zhou et al., 2024c), as well as set-based (Andreis et al., 2023) and graph-based structures of the network's weights (Lim et al., 2024; Kofinas et al., 2024: Zhou et al., 2024a), to maintain equivariance. Moreover, monomial equivariant NFNs, which equivariant to both permutations and scaling, have been proposed in (Kalogeropoulos et al., 2024) utilizing a graph-based message message-passing mechanism and in (Tran et al., 2024a; Vo et al.,
2024) through parameter sharing mechanism.
3

---

Published as a conference paper at ICLR 2025
3 MAXIMAL SYMMETRIC GROUP OF A MULTI-HEAD ATTENTION
As the first step in characterizing a principal design of NFNs for Transformers, we need to decide when two tuples of matrices with appropriate sizes, say ((w(e)) w(K.), w(v.)) i= W( and
(1 ( (P ( A
i=1 ( define the same Multihead map. We provide a complete
answer for this step in this section, thus characterizing the maximal symmetric group of the weights of Multihead.
Let consider the MultiHead map with h heads defined in Equation (2). We rewrite MultiHead as:
MultiHead (x:Wo) (wre.0, w(ka,  (oonn (iltaaaaaa )
h
Head (x:W(Q,), wiks) (eeean ( (oo)
+
i=1
Head Xx w(w,), w(k.)), w((v)) w(o..
soflmax (x. (wilad) (W((k)))
XT x. (WilK)) wlo.)), (
vDx
where W(0) (W10.1), ...., W(O.)) ((O.)) with each Wl0,)) E DoDD Here, the matrices (W(O))- (W(K)) /VD and W(V,). (00,0) have the same size D D. Based on this observa- tion, we define for each positive integer h and each sequence of matrices {A,, B in R D map
F (;
L
120
120
defined by
h
F(X; (A,, B.) =
softmax (X A,- -XT) X BB,
for each X RixD, is noted that, for each integer / O, the image of R via F is contained in Ri and F can be viewed as a generalization of MultiHead. With this setting at hand, we proved that:
Theorem 3.1 (Independence of heads in multi-head attention). Let D be a positive integer. Assume that for a positive integer h, matrices A,, Ag...., An E DXD and B1, B2...., E IR DxD we have
F (x: Xx: (,,,,.)... = 0,
(3)
for all positive integers L and X E LLDD Then, if A1, Ag...., An are pairwise distinct, then B = B₁ 0.
Remark 1. Roughly speaking, the above theorem says that, in the multi-head attention, each indi- vidual head plays its own unique role.
Based on the above theorem, we characterize the maximal symmetric group of the weights of MultiHead the following theorem.
Theorem 3.2 (Maximal symmetric group of multi-head attentions). Let h D, Dr. D, be positive in- tegers. Let (W(2.), W(K0), W(V.), w(O.)) and (W1)))) (K.i) (V.)) w wwoo.i) be arbitrary
elements of R. Dx x DD Dx DDxD with i 1,..., h. Assume that
(a) max(Ds, D.) D.
4

---

Published as a conference paper at ICLR 2025
(b) the matrices W(Q.). (w((.))) T, wll..0 are of full rank,
(wilk.)) T wwW. wlo.), and W .p -
(c) the matrices W(9.). (W(K.)) T
with i .i...,hh h are pairwise distinct,
(d) the matrices w ((9)) (Wi(ki))² with l,..., ,,, h are pairwise distinct.
Then the following equivalent:
For every positive integer L and every E RLXD, we always have (x x:((weed) w(K.), w(V)) wloo.) h i=1 (x: (x:(w<len) w(k.) w(v))) wlowin
MultiHead
= MultiHead
(4)
2. There exist matrices M
E GLD, (I) and N6) E GLD, (IR) for each 1,..., h, as
well permutation Sh, such that ( (w(0000) .. (K(Kr())) (.. w(00((0))
= =(wie.) (MO)T, W(K.) (M())- W(V.). NO, (NO))))) (N(.)) (5)
Remark 2 (Necessity of the assumptions (a), (b), (c), and (d)). The four assumptions in Theorem 3.2
are necessary for the implication (1.) = (2.). Indeed, the assumptions (c) and (d) are needed to utilize Theorem 3.1. In addition, it follows from Theorem 3.1 and item (1.) that
w(Q.0) (W(K.))
W (K(T())))
As, and
((((())) W (0.10)
B, for some matrices A,, B; E DXD, At this point, the assumptions (a) and (6) are necessary to utilize the rank decomposition of A; and B, obtain the symmetries of the factors (Piziak Odell, 1999). It worth noting that these assumptions are generically satisfied in practical implementations since the parameters MultiHead are randomly initialized with dimensions Dj & D and D,- < D. Remark 3 (Maximal symmetric group of multi-head attentions). Theorem 3.2 suggests that a multi- head attention in practice is characterized by the products W*(2.). (W(K.)) and W(Ka), w(0.)) each head. As a consequence, can consider the group S₁ GLD, (RR)* GLD, (RR) be the maximal symmetric group of the MultiHead map.
W(K9. W(O,)
WEIGHT SPACE OF A TRANSFORMER BLOCK AND GROUP ACTION
In this section, we construct the weight space of a standard transformer block and define a group action on it. Let us first recall necessary notations from permutation matrices before defining the weight space in Section 4.2 and the group action in Section 4.3.
4.1 PERMUTATION MATRICES
Definition 4.1. Let be a positive integer. A matrix of size n is called permutation matrix if it has exactly one entry equal to in each row and each column, and zeros elsewhere. We denote by Pn, the set of such all permutation matrices.
Remark 4 (Permutation matrix vs. permutation). For every permutation matrix P E Pn, there exists unique permutation æ S,, such that P is obtained by permuting the n columns of the identity matrix I, according to æ. In this case, we write P := Ps and call it the permutation matrix corre- sponding æ. Here, S, the group of all permutations the 21.2...., n). In particular, for permutation matrices P. E Pn, and P, E P... we have
Px x= (Zx-1(1): (1): (2), xx- (m))
(6)
and
(Pa A- P-) A+ 1())(1))
for all vector RR" and matrix Raxm.
5

---

Published as a conference paper at ICLR 2025
4.2 WEIGHT SPACE
standard transformer block, which will be denoted by Attn from now on, contains a multi-head attention and two-linear-layer ReLU MLP as well as a layer normalization. Formally, Attn trans- forms an input sequence X LxD to an output sequence Attn(X) R2DD defined as follows:
Xt) = LayerNorm (ReLU (x W(A) (X-W(A) + 1s .(4)) w(B) + +1 +11.b .L..)
(8)
where
x =
LayerNorm (MultiHead (x: (w(Q.), w(ki,, w(vi,, w(o.)) ))
with 1, [I....,,1]' RRx1 This transformer block is parametrized by the following matrices:
(wlQ.), W(K.), W(V)), w(o.)) i=1.... ) E (IR x R DxDe x DXD x R.D.D.D.
inside the multi-head attention, and
weights (W(A), W(B) DDDA x DAXD and biases (b(A), B(B)) RIXDA x RIXD
inside the two-linear-layers ReLU MLP.
The weight space of this transformer block, say U, is the vector space:
u (R Dx x T DxDx x DxDs x DDDDD (TR DXDA I DAXD) IRIIDA x R R!XD).
(10)
We write each element U E u in the form U = ((W))00)) [W](K)),,w(V)))) (W](o.3)
i=l.... (INjCA),
[w](Bi). (Iwj(A),[w1(BB) (1b(()), [6)(B)))
(11)
To emphasize the weights of each Attn map, we will write Attn U) instead of Attn(X). In particular, each element U E U will define map Attn(:;U): RixD R Rxx,, which maps sequence E R D to a sequence Attn(X; U) E RIxD defined by equation and equation In the next section, we will find sufficient condition for pair U and U in l such that Attn(:; Attn(;;U))
4.3
GROUP ACTION ON WEIGHT SPACE
Let us consider the weight space U standard transformer block defined in Equation (10) and Equation (11). We consider the symmetries of U/ which arise from the following two sources:
the maximal symmetric group of the multi-head attention, which are characterized in The- orem 3.2;
the permutation symmetries represented by permutation matrices, which arise from the unordered structure of neurons in the two-layer ReLU MLP.
Based on these observations, we consider the group
Gu S x GLD, (R)* GLD, (R)* x PD PDA
(12)
This group will serve as the symmetric group of the standard attention block. Each element g Gu has the form
(T. (Mi)) (N)) PRo, Pro, PrA):
(13)
where E Sn, TO E Sp. TA E Sp, are permutations and M,, N₁ are invertible matrices of appropriate sizes. The action of Gu on li is defined formally as follows:
Definition 4.2 (Group action). With notation above, the action of Gu on 4 is defined be map Gz x 44 - 24 which maps an element (g. U' E Su x U with U and given in equation 11 and equation 13 to the element
gU (19W)(9.0) (00])))),,[[[(((kk.. (o (,1](1,,,,ww(00..) i=I....h
(o0])a),Low/) (190(c4) lo0)(()))
(14)

---

Published as a conference paper at ICLR 2025
where
[9 [w](a.)) (M)(ar)))MMMtt))))))))
[ow](, [w]ckr(o) (Ar(r(o)) -1
[ow)(v.)) = (W](kr(0) N
(ow)(o,) (AvtrCo)) -1
-w](0,7(i)) Pro:
low)(a)= = PTO [W](d), PA lodi(A) [b](4). PAA*
[oW](B) = P (w)(B))
[ob] -6](P)
We conclude the construction of the group action on U in the following theorem which results in design principle for transformer blocks.
Theorem 4.3 (Invariance of Attn under the action of Gu). With notations as above, we have Attn(X:gU) Attn(X;U)
(15)
for all U E U. E G, and X E BLXD.
Remark (Other types of symmetries). The group Su defined above does not cover all symme- tries of the weight space U. Indeed, there are symmetries of scaling type arising from the internal architecture of the ReLU MLP (see (Godfrey al., 2022; Tran et al., 2024a)). In addition, layer normalization also offers additional scaling and sign-flipping symmetries for Lf (see Appendix B). We leave designing the maximal symmetric group of a transformer block for future study.
EQUIVARIANT POLYNOMIAL NENS FOR TRANSFORMERS
In this section, we introduce a family of NENs specifically designed for Transformers, referred to as Transformer-NENY which are equivariant to the group Gu as described in equation 12. The main building blocks of Transformer-NFW consist of -u-eqiivariant and Gu-invariant polynomial layers. We will outline the construction of a class of Gy-invariant polynomial layers below. The detailed construction of the Gu--equivariant polynomial layers, which follows similar approach, will be thoroughly derived in Appendix D.
In concrete, we define a polynomial map 1:44 - D' with maps each element U E l/ to the vector I(U) E D' of the form
D
D
D [WW](QK.)) P..
I(U)
& (QQK,))
(QK.s)
ΣΣΣ
(vo..s)
[WW](VO.a) P..
p=1
[W]]9
[W](K*)
D
D.
h 3
D.
D
(V,
[W][K
+
4 (0,))pp
[W]
8=1
1==1
DA
DA D
2 (A)sp
[W]
ΣΣ
(BB)p [W](3]
p=19=1 1=1
$ (B) (()ai) [B](4) Σ
[b](8) 11
(16)
4=1
where
[WW](QKM) (2)() (I7((K.)) and [v (V)(V.)) (W](0.)) (17)
and the coefficients d_s are matrices of size x 1. Remark (I(U) as a quadratic polynomial). Intuitively speaking, I(U) is a linear combination of all entries of the matrices [w](Q.s) [W](K.ss, [W](V.s) [W](0.s), [W] [W](B), (4) and
[B](B) in U, as well as all entries of the additional matrices [WW] and [WW](VO.>) defined in Equation (17). Since the entries the matrices [WW](QK,a) and [WW](VO.) are polynomials of degree 2, the map I(U) is indeed quadratic polynomial in the entries of U. These additional quadratic terms help us incorporate more relations between weights inside the input multi-head attention block, thus allowing NFN to maintain its expressivity.
7

---

Published as a conference paper at ICLR 2025
The above formula for [U) is irredundant in the following sense: Proposition 5.1. With notation above, if (U)) for all U U, then Φ 0 for all coefficients $_
To make I to be Gu-invariant, the parameters + must satisfy system of constraints (usually called parameter sharing), which are induced from the condition T(gti) = T(U) for all g E Gu and U E U. We show in details what are these constraints and how derive the concrete formula of I in Appendix E. The formula of / is then determined by
D
D
h
D Σ
h Ewwnroon)
I(U)
Σ
$ PQQ
):paa"
[WW)OKK)) Pin
(QK,,)
(vo.
p=
1=1
p=1 g=1
S (A).,
[W)4
(B)
[W]
DA
1 (A):
(B((4]
+ Σ 1 (B) (b) + 1
(18)
In the above formula, the bullet indicates that the value of the corresponding coefficient is inde- pendent of the index at the bullet.
Theorem S.2. With notation as above, the polynomial map I: U RD' defined by Equation (18) is Gu iiviiaariant. Moreover, if a map given in Equation (16) is Gu-invariant, then it has the form given Equation (18).
The concrete formula for the polynomial Gu Gu-equivariant layer is presented in detail in Appendix D.
THE SMALL TRANSFORMER ZoO DATASET
Large-scale empirical studies have produced datasets of trained classification models with varied hyperparameters (Eilertsen et al.; Unterthiner et al., 2020), enabling data-driven approaches to mod- eling generalization. However, dataset of Transformer-bet based classification models is absent. Mo- tivated by this shortage, we introduce the Small Transformer Zoo dataset to experimentally demon- strate the efficacy of our proposed method. This dataset contains the weights of a fixed Transformer architecture trained two distinct datasets, spanning both vision and language tasks. Each en- try in the dataset includes a checkpoint weight along with its corresponding accuracy metrics and hyperparameter configurations.
General settings. We focus on two prevalent deep learning tasks: digit image classification using the MNIST dataset (LeCun & Cortes, 2005) for vision, and text topic classification using AGNews (Zhang et al., 2015) for natural language processing. This selection covers two primary data modal- ities in deep learning-imen most common and fundamental the field. Our model architecture contains three components: an embedding layer, encoder, and a classifier. The embedding layer processes raw input data to produce initial token embeddings, which the encoder then transforms through self-attention mech- anisms capture contextual relationships. Finally, the classifier generates the classification output based on these enriched embeddings. While we adapt the embedding and classifier components to each specific task, maintain consistent encoder architecture across both tasks, consisting of two stacked two-head transformer blocks as defined in Equation (8).
and text-while addressing classification tasks, which are among the
Our resulting datasets, named MrIST-Transsormers and AGNews- Transformers, consist of 62756 and 63796 model checkpoints, respectively. These models were generated by varying key hyper- parameters such as the optimizer, learning rate, weight regularization, weight initialization, and dropout. as detailed in Appendix G. By making these datasets available, we aim to facilitate and inspire further research into the inner workings and behavior of Transformer models.
AGNews-Transformeen
EXPERIMENTAL RESULTS
We empirically evaluate the performance of the proposed Transformer- model two datasets: Transfrrrasorrmrrs and AGNews-Transformer Additionally, we conduct ablation studies to ex- amine the contribution of each component within the Transformer architecture in predicting network
8

---

Published as a conference paper at ICLR 2025
Table 1: Performance measured by Kendall's 's of all models on MNIST-Transformers dataset. Un-
certainties indicate standard error over 5 runs.
Accuracy threshold
No threshold
20%
40%
00%
80%
ML.P
0.86G + 0.002
0.873 + 0.001
0.874
0.003
0.874 + 0.006
0.873
0.007
STATNN (Unsterthiner al., 2020)
0,881 0.001
0.872 4: 0.001
0.868
0.001
0,86 #00001
0.856 0.001
XGBoostt (Chen & Guestrin, 2016)
0.803 0.002
0.839 + 0.004
0.860 0.003
0.846 + 0.001
0.884 0.001
LightGiBM (Ke ct.,. 2017)
0,858 1: 0.002
0.835 =: 0.001
0.817 = 0.001
0.822 + 0.000
0.830 2 0.001
Random Forest (Breiman, 2001)
0.772 0.002
0.758 0.004
0.769
0.001
0.752
0.769 0.001
Transformr--NH (ours)
0.905 + 0.002
0.899 + 0.001
0.895 0.001
0. 0.895 + 0.002
0 0.888 + 0.002
Table 2: Performance measured by Kendall's T of all models on AGNews-Transformers dataset. Uncertainties indicate standard error over runs.
Accuracy threshold 40%
No threshold
20%
60%
80%
MLP
0.879 0.006
0.875 4: 0.001
0,841 0.012
0.812 :: 0.001
0.862 0.006
STATNN (Usterthiner al., 2020)
0.841 0.002
0.839 0.003
0.812
0.003
0.813 + 0.001
0.812 0.001
XcBBoost (Chen & Guestrin, 2016)
0.859 1: 0.001
0.852 :: 0.102
0.872
0.002
0.874 4: 0.001
0.872 +: 0.001
LightGBM ctal., 2017)
0,833
0.815 0.001
0.837
0.001
0.835 +: 0.001
Random Forest (Breiman, 2001)
0.774 1: 0.003
0.801 1: 0.001
0.797
0.001
0.798 1: 0.002
0.773 1: 0.001
Transformer NFN (ours)
0.910 0.001
0.908 1 0.001
0.897 0.001 0.896 1: 0.001 0.590 0.001
generalization, and investigate the impact of varying the Transformer-NFN dimension and the num- ber of layers on the overall performance. Our analysis yields three key findings: (1) Transformer- NFN, with its enhanced layers for processing transformer block parameters, outperforms existing baselines in both vision and NLP Transformers datasets, maintains consistent performance across different accuracy thresholds of the data (2) the information embedded in the weights of Trans- former blocks provides a strong predictor for the performance of transformer model, and (3) good performance for Transformer-NFN can be obtained with compact setting.
We use Kendall's rank correlation (Kendall, 1938), ranging from [-1, 1]. as the evaluation metric to assess how closely predicted accuracy rankings align with ground truth accuracy rankings. A value near I indicates strong agreement, as shown in the scatterplot in Figure 2 (in Appendix H). All results in this section are averaged over 5 runs with different random seeds, with details on hyperparameters and training settings provided in Appendix H.
7.1 PREDICTING VISION TRANSFORMERS GENERALIZATION FROM PRETRAINED WEIGHTS
Experiment Setup. this experiment, we focus on predicting the test accuracy of pretrained Vision Transformer models using only their weights, without access to the test set. To perform this task, we utilize our TNIST--rransfommers dataset. We evaluate our model against 5 models: MLP, STATNN (Unterthiner et al., 2020), XGBoost (Chen & Guestrin, 2016), LightGBM (Ke al., 2017), and Random Forest (Breiman, 2001). As shown in Figure 1 (in Appendix H), the accuracy distribution of the MrisseTranssormerrs dataset is highly skewed (notice the log scale on the y-axis). Therefore, we evaluate each model's 's prediction performance not only on the entire dataset but also on four smaller subsets, each filtered by accuracy thresholds of 20%, 40%, 60%, and 80%. As a significant portion of pretrained models in the dataset exhibit higher accuracy, achieving strong Kendall's correlation becomes increasingly challenging as the accuracy thresholds increase.
Results. Table 1 illustrates the results of all models when predicting generalization of Vision Trans- former networks trained on MNIST-Transformer dataset. As expected, Kendall's 7 generally de- creases as the accuracy threshold increases. In addition, our Transformer-NFN consistently outper- forms all four baseline models across all dataset settings with performance gap ranging from 0.004 to 0.026, demonstrating the effectiveness of our model's design in capturing the information within each transformer block.
7.2 PREDICTING TEXT CLASSIFICATION TRANSFORMERS GENERALIZATION
Experiment Setup. In this experiment, we utilize the AGNews-Transfomin Transformers dataset to predict the performance of pretrained transformer models in text classification. The goal is to evaluate the effectiveness of Transformer-NFN in predicting the performance of pretrained models specifically trained on language tasks. Similar to Experiment 7.1, we assess model'ss 's capabilities across dif- ferent dataset configurations by using five subsets: the entire dataset without any accuracy threshold, and four subsets with accuracy thresholds of 20%, 40%, 60%, and 80%, respectively.

---

Published as a conference paper at ICLR 2025
Table 3: Ablation study the important of each component of the input networks on predicting gen- eralization of the input network, the metric being used is Kendall's 7.
No. of components
MNIST-
AGNews-
Components
Transformers
Transformers
Encoder
0.902 + 0.001
0.909 0.001
Embedding
0.424 + 0.002
Classifier
0.847 + 0.001
0.795 0.008
Embedding + Classifier
0.857 + 0.003
2
Encoder Classifier
0.904 + 0.001
0.910
Encoder Embedding
0.903 + 0.001
Encoder Embedding Classifier
0.905 + 0.002
Results. Table 2 shows that our model consistently outperforms all baselines across all text clas- sification dataset configurations. Compared to MNIST-Transformerr, the performance gain even more pronounced, with Kendall's gaps ranging from 0.018 to 0.033. As language tasks involve more complex syntax, the transformer encoder captures richer information, allowing Transformer- NFN to predict NLP performance more effectively than in vision tasks. Figure 2 (Appendix H) further illustrates its superiority, particularly in generalization for low-accuracy networks.
7.3 IMPORTANCE OF ENCODER IN PREDICTING THE GENERALIZATION
interesting question arises is how much information about the network generalization ability embedded in each component of the transformer model. To investigate this, restrict our Transformer-NEF NFNS to access only certain subsets of the Transformer's components and train the model on Trasforranssormrss dataset and (GNews-Transformet dataset. Our goal is deter- mine the importance of each component, both individually and in combination, for predicting gen- eralization of the input network.
Table shows that for both alone provide strong performance predictions. The classifier is the second most important compo- nent, followed by the embedding. Even with only transformer block weights, our model achieves Kendall's 's 7 score nearly identical to using all components: 0.902 vs. 0.905 for and 0.909 vs. 0.91 for AGNews-Transformen Transformers.
and Tranessr-Tanssorerrr.. transformer blocks
7.4 ABALATION STUDY ON VARYING THE DIMENSION AND NUMBER OF LAYERS
In this section, we conduct an ablation study to explore the impact of varying the hidden dimension and the number of equivariant layers in Transformer-NFF of hidden dimensions and layer counts on the AGNews-Transfome Transformer dataset, with dimensions e [3, 5, 10, 15), and number of layer E [1, 2]. These configurations allow us assess how changes in model size affect performance and efficiency.
We evaluate different configurations
Table (in Appendix H) demonstrates that strong performance can be achieved with relatively small dimension and few parameters. For example, with dimension of 15 and single layer, the model reaches a Kendall's 7 score of 0.913, matching the best performance across all settings. These results suggest that good performance for Transformer-NFr be obtained with compact setting.
8 CONCLUSION
In this work, we made significant contributions to the understanding and application of NFNs in transformer architectures. We determined the maximal symmetric group of the weights in a multi- head attention module. We also formally defined the weight space of a transformer architecture and introduced a group action on this weight space, thereby characterizing the design principles for NFNs. Additionally, we presented Transformer-NNN, tures that is equivariant under the specified group action. Finally, we released a dataset more than 125,000 transformers model checkpoints trained on two datasets with two different tasks, mark- ing a significant resource for benchmarking the applicability and efficiency of Transformer-NFN and promoting further research enhance understanding of transformer network training and performance.
NFN designed for transformer architec-
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This research project is supported by the National Research Foundation Singapore under the Al Singapore Programme (AISG Award No: ALSGG-TC-2223-01--S000 supported by the Ministry of Education, Singapore, under the Academic Research Fund Tier (FY2023) (A-8002040-00--0) A-8002039-00-000 This research / project is supported by the NUS Presidential Young Professorship Award (A-0009980-001-00) and the NUS Artificial Intelligence Institute-Seed Funding (A-8003062-00-00),
This research project is
Thieu N. Vo is supported by the Singapore National Academy of Science under the SASEA Fellow- ship Programme (Award No: NRF-MP-2025-0001)
Thanh Tran acknowledges support from the Application Driven Mathematics Program funded and organized by the Vingroup Innovation Fund and VinBigData. Thanh Tran acknowledges support from the VinUni's Student Research Grant Program AY24-25.
Ethics Statement. Given the nature of the work, do not foresee any negative societal and ethical impacts of our work.
Reproducibility Statement. Source codes for our experiments are provided in the supplementary materials of the paper. The details of our experimental settings are given in Section 7 and the Ap- pendix All datasets used this paper are publicly available through an anonymous link provided in the README file of the supplementary material.
REFERENCES
Laziz Abdullaev and Tan Minh Nguyen. Transformer meets twicing: Hamessing unattended residual information. The Thirleenth International Conference on Learning Representations, 2025.
URL https:tp//// /openreviiii nnt./ foet/foooe forum? i d= 6kG5
Francesca Albertini and Eduardo D Sontag. Identifiability of discrete-time neural networks. In Proc. European Control Conference, pp. 460-465. Springer Berlin, 1993a.
Francesca Albertini and Eduardo D Sontag. For neural networks, function determines form. Neural networks, 6(7):975-9)) 990, 1993b.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. convergence theory for deep learning via over- parameterization. In International conference on machine learning, pp. 242-252. PMLR, 2019.
Bruno Andreis, Soro Bedionita, and Sung Ju Hwang. Set-based neural network encoding. arXiv preprint wxXiv:2305.16625 2023.
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
Maor Ashkenazi, Zohar Rimon, Ron Vainshtein, Shir Levi, Elad Richardson, Pinchas Mintz, and Eran Treister. Nern: Learning neural representations for neural networks. In The Eleventh Inter- national Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. 3rd International Conference on Learning Representations, 2015.
Bowen Baker, Otkrist Gupta, Ramesh Raskar, and Nikhil Naik. Accelerating neural architecture search using performance prediction. In 6th International Conference on Learning Representa- tions, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Workshop Track Proceedings. OpenReview.net, 2018. URL https: https:://ppezeeiieee//c
net www.ormmmmngg.1111111
Matthias Bauer, Emilien Dupont, Andy Brock, Dan Rosenbaum, Jonathan Richard Schwarz, and Hyunjik Kim. Spatial functa: Scaling functa to imagenet classification and generation. arXiv preprint arXiv:2302.03/ 2023.
11

---

Published as a conference paper at ICLR 2025
Mikhail Belkin, Daniel Hsu, Siyuan Ma, and Soumik Mandal. Reconciling modern machine- learning practice and the classical bias-variance trade-off. Proceedings of the National Academy of Sciences, 116632):15849-1555 15854, 2019.
Samy Bengio, Yoshua Bengio, Jocelyn Cloutier, and Jan Gecsei. On the optimization of synaptic learning rule. In Optimality in Biological and Artificial Networks? pp. 265-287. Routledge, 2013.
Leo Breiman. Random forests. Machine learning, 45:5-32, 2001.
Phuong Bui Thi Mai and Christoph Lampert. Functional vs. parametric equivalence of relu net- works. In 8th International Conference on Learning Representations, 2020.
Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge language models. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, pp. 6491-6506. Association for Computational Linguistics, 2021.
An Mei Chen, Haw-minn Lu, and Robert Hecht-Nielsen. On the geometry of feedforward neural network enror surfaces. Neural computation, 5(6):910-927, 1993.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Adiyya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 14:15084-15099 2021.
Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, volume 11 of KDD '16, pp. 785-794. ACM, August 2016. doi: 10.1145/293997222333555
Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pp. 1734.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapplis, MN, USA, June 2-7, 2019, Volume I (Long and Short Pupers), pp. 4171-4186, Association for Computational Linguistics, 2019. doi: 10.18653/ / http////nn1111///11
V1/N19-1423. URL https: /ww.i http:/ii.og.g.g//oo
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 19929, 2020.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. In International conference machine learning, pp. 1675- 1685. PMLR, 2019.
Emilien Dupont, Hyunjik Kim, S. M. Ali Eslami, Danilo Jimenez Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat it like one. In International Conference on Machine Learning. ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 5694-5725 PMLR, 2022a. URL https:: / http://pooeeeiceeee//o m ml r. prrsss//// / v1 62/ duddottt222 html.
Emilien Dupont, Yee Whye Teh, and Arnaud Doucet. Generative models as distributions of func- tions. In International Conference Artificial Intelligence and Statistics, AISTATS 2022, 28-30 March 2022, Virtual Event, volume 151 of Proceedings of Machine Learning Research, pp. 2989- 3015. PMLR, 2022b. URL https: proceedoceeiieooooooooooo ngs mm www.r.press...m /vaass/vvvl v1 /dupont222.. html.
Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman. Classifying the classifier: Dissecting the weight space of neural networks. In ECAI 2020 24th European Conference on Artificial Intelligence, volume 325 of Frontiers in Artificial Intelligence and Ap- plications, pp. 1119-1126. IOS Press. URL https: /ttp///.../// .wrooooorr//oo 3233//ATA2002220220
12

---

Published as a conference paper at ICLR 2025
Ziya Erkoç, Fangchang Ma, Qi Shan, Matthias Nießner, and Angela Dai. Hyperdiffusion: Generat- ing implicit neural fields with weight-space diffusion. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 44300-14310, 2023.
Charles Fefferman and Scott Markel. Recovering feed-forward net from its output. In Advances in Neural Information Processing Systems 6, [7th NIPS Conference, Denver, Colorado, USA, 1,93], pp. 335-342 Morgan Kaufmann, 1993. URL http: /ptpprapaaa.a//////e nips. cc/paper/ re8-vecovering-aa oaw.:vvvrrroao--------------rrre
from- its -tsttttut....... output.
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. URL https: /topppppevvvieeeeeeeeeeeeeeeeeee ew net. forummrrmmm/mm ? rJ1-b3RcF7.
Charles Godfrey, Davis Brown, Tegan Emerson, and Henry Kvinge. On the symmetries of deep learning models and their internal representations. Advances in Neural Information Processing Systems, 35:1 11893-1995,, 2022.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference Language Modeling, 2024. URL https: http::://pteeeeeeeeeee /wtp//ppeeeveeeeeeeeeeeeeeeeeeee eew net tere/fforuee forum?id tEYskwiVY2.
Jean Harb, Tom Schaul, Doina Precup, and Pierre-Luc Bacon. Policy evaluation networks. arXiv preprint arXiv:2002.11833, 2020.
Kaiming He, X. Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition.
2016 IEEE Conference on Computer Vision and Puttern Recognition (CVPR), pp. 770-778, 2015.
URL https::////o /api......// semante nt cscholar org/ Corpus ID: 20 206594692
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Ad- vanced Neural Computers, pp. 129-135. Elsevier, 1990.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8): 1735-1780, 1997. doi: 10. 10.116220 1.11162/NEO..1997.9.88 neco. 1 1997.9999 9. 8. 1735.
hRLhttts:///dt orr.ii..rgglio /10 1162/
Yujia Huang, James Gornet, Sihui Dai, Zhiding Yu, Tan Nguyen, Doris Tsao, and Anima Anand- kumar. Neural networks with recurrent generative feedback. Advances in Neural Information Processing Systems, 33:535-545, 2020.
John M. Jumper, Richard O. Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ron- neberger, Kathryn Tunyasuvunakool, Russell Bates, Augustin Židek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A. A. Kohl, Andrew J. Ballard, Andrew Cowie. Bernardino Romera-Paredes, Stanislav Nikolov, R. D. Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David L. Silver, Oriol Vinyals, Andrew W. Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis. Highly accurate protein structure prediction with alphafold. Naturr,,5667855 596(7873))5883-5888 July 2021. doi: https:/ / ///ee//n...//////////////eeeee.. www.//gg/--/....... 21:0002-11-200 1 983-988388888 826-886660 945.
URL
Ioannis Kalogeropoulos, Giorgos Bouritsas, and Yannis Panagakis. Scale equivari- ant graph metanetworks.
Advances in Neural Information Processing Systems https//cccppppeeees/pppeeeeee files/paper/2222 4/22h//
38, 2024. URL http:://apers..ipps wwtpsrssiis...ccpppppoe.o c13d5a 10200006888111 httrsaaeteeeeatatttt-s--aaatt-o-
Guolin Ke, Qi Meng, Thomas Finley, Taifeng Wang, Wei Chen, Weidong Ma, Qiwei Ye, and Tie- Yan Liu. Lightgbm: highly efficient gradient boosting decision tree. Advances neural information processing systems, 30, 2017.
M. G. Kendall. A NEW MEASURE OF RANK CORRELATION. Biometrika, 30(1-2):81-93, 06 1938. ISSN 0006-3444. doi: 10.1099/bioee/33..122..1.. URL https: //oo: org/ ///////000 10. 1093/br omet//ee//3
13

---

Published as a conference paper at ICLR 2025
Yoon Kim, Carl Denton, Luong Hoang, and Alexander M. Rush. Structured attention networks. In International Conference on Learning Representations, 2017. URL WKLLtttps///ppeeeviiiee net/ wttr//orrnnnneeeeeeee
Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero Soriano. Parameter prediction for unseen deep architectures. Advances Neural Information Processing Systems, 34:29433-294483 2021.
Boris Knyazev, Abhinav Moudgil, Guillaume Lajoie, Eugene Belilovsky, and Simon Lacoste- Julien. Accelerating training with neuron interaction and nowcasting networks. arXiv preprint arXiv:2409.04433 04434, 2024.
Miltiadis Kofinas, Boris Knyazev, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, Efstratios Gavves, Cees G. Snoek, and David W. Zhang. Graph neural networks for learning equivariant represen- tations of neural networks. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttp//pppeeeeeeeeeeeeeeeeeeeeeep net foreefoouuuee forum? wtw.eeeoooonmndoooeoFree www.o00FFsMyDyBS
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep con- volutional neural networks. In Advances in Neural Information Processing Systems 25, pp. 1106-1114, 2012. URL https: wttp://ccweeennn.eeeiieeeeeeeeee
wttps//cppppppccccc/papaappa c399862d3b9d 6677668433300 62t.eee6606664446444444444466330 Abstract. html.
Vera Kurkova and Paul C Kainen. Functionally equivalent feedforward neural networks. Neural Computation, 6(3):543-558, 1994.
Yann LeCun and Corinna Cortes. The mnist database of handwritten digits. 2005. URL https:
//api. semanticscholaa. org. /Crr//rrpssID: 60282629 629..
Derek Lim, Haggai Maron, Mare T. Law, Jonathan Lorraine, and James Lucas. Graph metanetworks for processing diverse neural architectures. In The Twelfth International Conference on Learning Representations, 2024. URL https:ttt......cc..tt /ttp///pppeeeeeeeeeeeeeeeeeeeeee net/ / jerrumii--jjKhhyxeeee
Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A STRUCTURED SELF-ATTENTIVE SENTENCE EMBEDDING. In Interna- tional Conference on Learning Representations, 2017. URL htttps://eeenreview net/
forum? (ae-__d__BBU__ee.
Luca De Luigi, Adriano Cardace, Riccardo Spezialetti, Pierluigi 7ama Ramirez, Samuele Salti, and Luigi Di Stefano. Deep learning on implicit neural representations of shapes. In The Eleventh International Conference on Learning Representations, ICLR 2023. URL https:
// opoeneevieww net wwt//eet//oooe woooummidd-oooo 3uadi.
Luke Metz, James Harrison, Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:2211.09760, 2022.
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. Nerf: Representing scenes neural radiance fields for view synthesis. Communications of the ACM, 65(1):99-106, 2021.
Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D. Manning. Fast model editing scale. In The Tenth International Conference on Learning Representationss fttte//fommmmmmdddceeeeeee
2022. URL https: http////ppeeeviieeeeeeeee
Maxim Monastirsky, Osher Azulay, and Avishai Sintov. Learning to throw with a handful of samples using decision transformers. IEEE Robotics and Automation Letters, (2):576-583 2022.
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant architectures for learning in deep weight spaces. In International Conference on Machine Learning, pp. 25790-255 25816. PMLR, 2023.
Tam Nguyen, Tan Nguyen, and Richard Baraniuk. Mitigating ovor-shgo in transformers via regularized nonlocal functionals, Advances in Neural Information Processing Systems, 36:80233- 80256, 2023a.
14

---

Published as a conference paper at ICLR 2025
Tam Minh Nguyen, Tan Minh Nguyen, Dung DD Le, Duy Khuong Nguyen, Viet-Anh Tran, Richard Baraniuk, Nhat Ho, and Stanley Osher. Improving transformers with probabilistic attention keys. In International Conference on Machine Learning, pp. 16595-16621. PMLR, 2022.
Tan Minh Nguyen, Tam Minh Nguyen, Nhat Ho, Andrea L.. Bertozzi, Richard Baraniuk, and Stanley Osher. A primal-dual framework for transformers and neural networks, In The Eleventh Interna- tional Conference on Learning Representations, 2023b. URL https::///// https::/cpenneeiers forummidd- 5hC1V.
Stefan Nielsen, Laziz Abdullaev, Rachel SY Teo, and Tan Nguyen. Elliptical attention. Advances in Neural Information Processing Systems, 37: 19974881999978, 2025a.
Stefan Nielsen, Rachel Teo, Laziz Abdullaev, and Tan Minh Nguyen. Tight clusters make specialized experts. The Thirteenth International Conference on Learning Representations, 2025b. URL https://cpeereet
net. Eorum?1d=Pu3c02205
Roman Novak, Yasaman Bahri, Daniel A Abolafia, Jeffrey Pennington, and Jascha Sohl- Dickstein. Sensitivity and generalization in neural networks: an empirical study. arXiv preprint arXiv:1802.08760, 2018.
Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable atten- tion model for natural language inference. In Proceedingy of the 2016 Conference on Em- pirical Methods Natural Language Processing, pp. 2249-2255, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/11/D16-- 1244. URL https: / /aclanthhlogyy. wwg//oo-g200 12444.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning, pp. 7487-7498. PMLR, 2020.
William Peebles, Ilija Radosavovic, Tim Brooks, Alexei A Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209. 12892, 2022.
Robert Piziak and Patrick L Odell. Full rank factorization of matrices. Mathematics magazine, 72 1911193201, 1999.
Gavin Portwood, Peetak P Mitra, Mateus Dias Ribeiro, Tan Minh Nguyen, Balasubramanya T Nadiga, Juan A Saenz, Michael Chertkov, Animesh Garg, Anima Anandkumar, Andreas Dengel, et Turbulence forecasting via neural ode. Second Workshop on Machine Learning and the Physical Sciences (NeurIPS 2019), 2019.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning, pp. PmLR, 2021.
Alec Radford, Jong Wook Kim, Tao Xu, Greg Brockman, Christine McLeavey, and Ilya Sutskever. Robust speech recognition via large-scale weak supervision. In International conference on chine learning, pp. 28492-28518. PMLR, 2023.
M. Raissi, Perdikaris, and G.E. Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equa- tions. Journal of Computational Physics, 378:686-707, 2019. ISSN 0021-9991. doi: https://doi. 10.045. URL https: /wwww encedireeteee com/ science//
David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams, Learning internal representations by error propagation. 1986. URL https: api. semanticschen .rg/ orgg/CrppssD: 62245742.
15

---

Published as a conference paper at ICLR 2025
Thomas Philip Runarsson and Magnus Thor Jonsson. Evolution and design of distributed learning rules. In 2000 IEEE Symposium on Combinations of Evolutionary Computation and Neural Net- works. Proceedings of the First IEEE Symposium Combinations of Evolutionary Computation and Neural Networks (Cat. No. 00, pp. 59-63. IEFE, 2000.
Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-supervised representation learn- ing on neural network weights for model characteristic prediction. Advances in Neurul Informa- lion Processing Systems, 34: 111116141199999 2021.
Konstantin Schürholt, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Hyper- representations as generative models: Sampling unseen neural network weights. Advances in Neural Information Processing Systems, 35:27906-27920, 2022a.
Konstantin Schürholt, Diyar Taskiran, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Model 2oos: A dataset of diverse populations of neural network models. Advances in Neural Information Processing Systems, 35:38134 1188813--1444 188, 2022b.
Anton Sinitsin, Vsevolod Plokhotnyuk, Dmitry V. Pyrkin, Sergei Popov, and Artem Babenko. Ed- itable neural networks. In 8th International Conference on Learning Representations, 2020. URL https: htttps:///peenreieww met/ wttp://wwwwweeefouteeeeeee
Alicja Smoktunowicz, Ryszard Kozera, and Gianluca Oderda. Efficient numerical algorithms for constructing orthogonal generalized doubly stochastic matrices, Applied Numerical Mathematics, 142:16-27, 2019.
Samuel Sokota, Hengyuan Hu, David Wu, Zico Kolter, Jakob Nicolaus Foerster, and Noam Brown. A fine-tuning approach to belief state modeling. In International Conference on Learning Representations, 2021.
Kenneth O Stanley,
Compositional pattern producing networks: A novel abstraction of development. Genetic programming and evolvable machines, 8:131-162, 2007.
Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott F. Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. Going deeper with convolutions. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2015, Boston, MA, USA, June 7-12, 2015, pp. 1-9. IEEE Computer Society, 2015. doi: 10. 10.11000/VPR. URL httph:tp////// / doii wtt.ooo......../.......
www.111199..... 2015. 7298594.
Rachel SY Teo and Tan Nguyen. Unveiling the hidden structure of self-attention via kernel principal component analysis. Advances in Neural Information Processing Systems, 111100193-101441 2025.
Hoang Tran, Thicu Vo, Tho Huu, An Nguyen The, and Tan Nguyen. Monomial matrix group equivariant neural functional networks. In Advances in Neural In- formation Processing Systems, volume 37, pp. 48628 48665. Curran Associates, Inc., 2024a. URL https: http:://///oedieeeeeeee weweeinnn...nnniiee neurips cc wttppppprfllsspppeee//ee/ae F11e/5777d5866ecT (eteBBeeeeeeeeee8811344444 09bbee9933aae 936ae Paper-Cofferercce pdf.
Hoang V Tran, Khoi NM Nguyen, Trang Pham, Thanh T Chu, Tam Le, and Tan M Nguyen. Distance-based tree-sliced wasserstein distance. arXiv preprint 2rXiv:2503.11050, 2025a.
Thanh Tran, Viet-Hoang Tran, Thanh Chu, Trang Pham, Laurent El Ghaoui, Tam Le, and Tan M Nguyen. Tree-sliced wasserstein distance with nonlinear projection. arXiv preprint arXiv:2505.00% 00968, 2025b.
Viet-Hoang Tran, Trang Pham, Tho Tran, Tam Le, and Tan M Nguyen. Tree-sliced wasserstein distance on a system of lines. arXiv preprint orXiv:2406.13725 2024b.
Viet-Hoang Tran, Thieu N Vo, Tho Tran Huu, and Tan Minh Nguyen. A clifford algebraic approach to (n)-equivar high-order graph neural networks. arXiv preprint arXiv:2410.0466 046692 2024c.
Viet-Hoang Tran, Thanh T Chu, Khoi NM Nguyen, Trang Pham, Tam Le, and Tan M Nguyen. Spherical tree-sliced wasserstein distance. arXiv preprint arXiv:2503.1 /249, 2025e.
16

---

Published as a conference paper at ICLR 2025
Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and llya Tolstikhin. Predict- ing neural network accuracy from weights. arXiv preprint arXiv:2002.11448, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30, pp. 5998-6008, 2017. URL https: /pwoceedingseseseeeeeeeppp.s..se hew.oeeeeddinnn...nii........ neurips. httppp//00.//ass/ffffee222//e// html.
Thieu Vo, Duy-Tung Pham, Xin T. Tong, and Tan Minh Nguyen. Demystifying the token dynamics of deep selective state space models. The Thirteenth International Conference on Learning Representations, 2025. URL https: opp..evieeeeeeeeeeee net Sttee//orummmi--gTPPPPGG
Thieu N. Vo, Viet-Hoang Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khoi Nguyen- Nhat, Duy-Tung Pham, and Tan Minh Nguyen. Equivariant polynomial functional networks, 2024. URL https:://///. wtt.vv.iiivv//////////aa////aav. org/ abs 24100.0213
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
Allan Zhou, Chelsea Finn, and James Harrison. Universal neural functionals. Advances in neural information processing systems, 2024a.
Allan Zhou, Kaien Yang, Kaylee Burns, Adriano Cardace, Yiding Jiang, Samuel Sokota, Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. Advances in Neural Infor- mation Processing Systems, 36, 2024b.
Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Neural functional transformers. Advances in Neural Information Processing Systems, 36, 2024c.
17

---

Published as a conference paper at ICLR 2025
Supplement to "Equivariant Neural Functional Networks for Transformers"
Table of Contents
Maximal Symmetric Group MultiHead
19
A.1 Multi-head Attention
19
A.2 Maximal symmetric group of multi-head attention
20
Matrix Group Preserved by LayerNorm
25
C Weight Space and Group Action on Weight Space
27
C.I Weight space
27
C.2 Group action on weight space.
27
D An Equivariant Polynomial Layer
30
D.1 A general form with unknown coefficients
30
D.2 Auxiliary lemmas
32
D.3 Finding the constraints for the unknown coefficients
:o..............................
33
D.3.1
Computing E(gU)
33
D.3.2 Computing gE(U).
43
D.3.3
Coefficients comparison in the equation E(gU)=
47
D.4
Final formula for the equivariant polynomial layer
55
E
An Invariant Polynomial Layer
58
E.1
A general form with unknown coefficients
................................
58
E.2
Compute [ggU)
................................
59
E.3
Compare [(UU) and I(U)
................................
59
E,4
Final formula for the invariant polynomial layer
60
Computation Complexity of Equivariant and Invariant Layers
60
Additional Dataset Details
60
H Additional Experimental Details
62
H.1 General details
62
H.2 Archiecture and hyperparameters of Transformer-NFY
H.3 Architecture and hyperparameters for other baselines
63
H.4 Experiment on augmented AGnsfseTranssrrmerr dataset
64
I Implementation of Equivariant and Invariant Layers
64
LI Summary of Equivariant and Invariant Layers
64
L.11 Equivariant Layers with bullet notation
64
18

---

Published as a conference paper at ICLR 2025
11..2
Invariant Layers with bullet notation
66
1.2
Equivariant Layers Pscudocode
67
1.2.1
[E(W)]
Pseudocode
67
1.2.2
[E(W)] (Ka)
Pseudocode
67
j.k
[E(W)] (V=)
1.2.3
3.1k
Pseudocode
67
1,2.4
[E(W)]5,
3.k
Pseudocode
67
1.2.5
[E(W)] Pseudocode
68
[E()]] (A) Pseudocode
1,2.6
69
1.2.7
[E(W)] (B)
Pseudocode
70
1.2.8
[E(b]]}
(B) Pseudocode
I.3
Invariant Layers Pseudocode
72
A
MAXIMAL SYMMETRIC GROUP OF MultiHead
This section aims provide a complete answer to the question of when two of parameters of a multi-head attention mechanism will define the same function. The answer to this question will serve the core of the design principle NFNs for Transformers. We recall the formulation of Head and MultiHead attention maps in Section A.1. Then, the maximal symmetric group of Head derived in Section A.2.
A.1 MULTI-HEAD ATTENTION
Let D be positive integer. Recall the notion of the parameterized map Head as follows: For positive integer L and X E BLXD we have
(XXW())) (xww)o)))
(xwii)).
Head X; (x:W(9) W(K), w(n) sofumax
vDx
where W(Q), w(K) E DDD, and W(V) E DxD. By definition, we have
Head (:: (11()), ((R), w())
R
Ix
50
150
and for all > 0, the image of R is contained in Ri (W(K))² IvDa DDD we can rewrite the map Head as follows
By combining (W(Q))
Head (X:W(), w(K), w(n) softmax (x. (W(D)) (W(K)) T xx)) (XWII)))
VD.
By this observation, we define a class of parameterized maps as follows: For X E RLXD we have S(X; A) = softmax (XAXT): X,
where A R DDD, Similarly, for the MultiHead map, have
X; ((0), {w(ew, W(K.), w(v,)) )
MultiHead
+ Head (x; w(Q.1) W(ks)) ((kk)) w(o)
Head (x:W(Q.)) w(ki)) w(vs)) w(o,))
19

---

Published as a conference paper at ICLR 2025
(xW(Q..)) (xxwixa))) x (wrISs) wrlo.)).
softmax
D
where h is a positive integer, W(O) (w(o.),,... w(ooh) with each W(O#)) E P.xDD Con- sidering W(V.), w(O,)) matrix B; E DDD we define new class of parameterized maps as
follows: For E RRLDD we have
Σ ((X;A,)
Bi,
t=1
where h is positive integer and A., Br ERDDD Note that,
rank ((w1e)) (W(K) T /VDx & minn(D,D)) <
rank (w(V,)) (y(V))-((o.))
(min(D,Do) &
So, in general, the new class of maps contains the class MultiHead maps. Note that, F simply weighted summation of some maps and is linear with respect to (B,)
Remark We can see, in S. the matrix A plays the role as a core matrix that defines f. Similarly, in F, each A; defines each fi, and B; is the weight of each component contributes to F.
A.2 MAXIMAL SYMMETRIC GROUP MULTI- HEAD ATTENTION
In this section, we present a theoretical result that shows the following: Roughly speaking, in the multi-head scenario, each individual head plays its own unique role. The main results are Theo- rems A. (which corresponds to Theorem 3.1 in the main text) on the unique role of each head and Theorem A.7 (which corresponds Theorem 3.2 in the main text) the maximal symmetric group of multi-head attention.
Theorem A.1. Let D be a positive integer. Assume that for a positive integer h, matrices A1, A2.... ..... A₁ E RRDxD and B1, B2,..., B₁ E RDDxD have
h Σ f(X:A.) = i=1
(X,
(19)
for all positive integers L and X R LxD Then, if A1, A1,,A2,..... An pairwise distinct, then
Bi = Bv. 0.
Proof. From Equation (19), we have
h
Σ softmax (XA:XT) X B; =
(20)
Consider > 1. We write
(z,y, y. ..... .....) T) LLDD where (I Tp) and y (yi...., MD) in RIxD (So in X, y appears L -- times). We will consider the entry in the first row and first column of both sides of Equation (20). Let bi, bn,,,, E Dx1 be the first column of matrices B1, B2,..,, Br. From Equation (20), have
LE
Σ
e + T vV)
b, 00
(21)
ezAi + Le
By substituting y (1,0...., 0) R to Equation (21), we have the sum of all first entries of by...., by is equal to O. Similarly, for every j = =,, 1,..., equal O. It shows that
1...,D, D. the sum of 7th entries of by...., by is
b + +bb2 ....
(22)
From Equation (21), have
h 0=
T
1 eRAix T + LeRA v/ v)) b,
LET
Σ
(23)
G2A.07
LeRANT
t=1
20

---

Published as a conference paper at ICLR 2025
caAa" + Le
LeaAxy
2 + caA,a' + LeaniNt (y 2) b,
(24)
C2A,
Le
+
@2A -LnAAy (y -2) .b.
(25)
h
Le
x
@tA.w
+ LozAN
(y x1) 12)))- bi
(26)
h
h + xx)
LEZAND
Σ
bi
Σ
LeadiNt bi
(27)
exAix
h
LE
(y 2)
bj
(28)
Axx
Lea
L
(y x)
(29)
A,x -mAANt +
L
= (y x)
bi
(30)
exAilzzv)) + Li
By let 2 = x -- y. from the above equations, we have
Σ @
L
bi
0,
(31)
+ T,
or
h
Σ
2 @ + T,
bi = 0,
(32)
i=1
for all (I1,.... ..p)). and (Z1,..., 2D) in R.D Now, each xA,² can be viewed
polynomial in 2D indeterminates 21...., .D and 21.....
2D as follows
(Ad) p.g Tp24.
(33)
p=1g=1
AA T are pairwise distinet polynomials. By Lemma A.3, there exists u E RD and non-empty open set V E R.D, such that for all V, we have uAiz" ...., uA are pairwise distinct real numbers. Obviously 0 & V. Now fix an v e V and let t for t R. Denote uA,v'
Since AI,..., An are pairwise distinct, so xA12
St. from Equation (33), we have
t.v Σ cta +
by 0
(34)
for all / E R, or
b, 00
(35)
clio:
for all E R \ (0). By continuity, this still holds for 0. So we have
c:.. I b, 0
(36)
for all E R. Now, consider the set
{
S 1.081 6 L' ct
TR² LE R
(37)
pt.sh + I
= From Equation (36), we have the linear span of S, i.c. span(S), satisfies that: For all s ($1, .,,..., sh)) sSan(S), we have
h
V
Σ
si bi = 0.
(38)
21

---

Published as a conference paper at ICLR 2025
In other words, and Σا= 8, b, are orthogonal to each other as two vectors in R Since $1:...,8h R are pairwise distinet, by Lemma A.4, there exist fi..... ta E R such that their h corresponding vectors in S form basis of R This implies span(S) R We have the set
span (Ab1...., bb,))
Σ
8, by s= ($1...., 8h) E R
(39)
b, == ($1..... 9h) span(s)
(40)
By the previous observation, it implies that is orthogonal to every vectors in span (161, b% }) In other words, the orthogonal complement of span (fbi...., H), which is denoted by span ((b1,.... bnt)+. contains This holds for every vectors in the non-empty open set V. so
(41) Since span (Lbi,..., bb})) linear subspace of R2D that contains non-empty open set of Rd so by Lemma A.2, we have span ( ((b1..... bn)) = RD. This implics span (fbi...., bn = O, which means 61 ... by 0. So the first column of Bı..... Bh are all equal to 0, and a similar proof is applied for every other columns of B,...., B₁, So B = =... = B, 0.
V span
b,,))
Lemma A.2. Let D be a positive integer, and V is non-empty open set of R with the usual topology. U is a linear subspace of RDD that contains V, then U = R..
Proof. Since V is non-empty, let x E V.. Since V is open, there exists > 0 such that the closed ball
Bx(x) = ty RR.
||- vll <r.cv
(42)
Then for all y E 2 that y # O, have
1 )
y
r
+
r.
=
(43)
which means
++7
y
E (x)x) CVCU.
(44)
lvl
Since is also in V c. U, and U is a linear subspace, then E U. So, for E RD that y # 0, we have y E U, and clearly 0 U, so R.D
Remark 8. Lemma A.2 still holds if we replace RD by normed vector space.
Lemma A.3. Let h be two positive integers, and fi....,Jh Sn E R[ sn)] be h pairwise distinct real polynomials in n variables. Then there exists a non-empty open subset U of IR" such that fi(),...., fx( f(((),...,f((()) are pairwise distinct for all a E U.
Proof. Consider the polynomial
II ( J5)
(45)
1eijjh
Since f1,..., fn are pairwise distinct, then fi - fj is non-zero for all <i<j < h. It well- known that the real polynomial ring an integral domain, so f is non-zero. There exists E E such that (a) A O. Since Hausdorff, we can choose a open set V IR that contains a and does not contain O. Let U = fV), and since f is continuous, U is open. We have f(x) 0 for all are pairwise distinct for all x E U
x E U, which means f1(
Lemma A.4. Let h be positive integer, and h distinct real numbers $ Sh. Then there exists h real numbers fI....,f, tn and positive integer L such that the matrix
A (a:))) 1gjgh E Rixh where aij-
(46)
et::
is full rank. In other words, there exists h real numbers ti...., ty and positive integer L such that h vectors
E E
(47)
efris +L chi +
cti + L
for 1...., br form basis of Rbb
22

---

Published as a conference paper at ICLR 2025
Proof. We first have observations. Let ti,...,th th be fixed and be variable. Consider the matrix
1
A(x) (at)) where aij
(48)
ch + x
Then the determinant of A(z), denoted by det(A(z)) can be viewed as a real rational function, i.c. function that can be written the ratio of two real polynomials. So, in the case that this rational function is zero, (A(x)) = 0 for all x E R, and in the case that this rational function is non-zero, there are finite number of x E R such that det(A(=)) =
(ex))) is not defined. In other words, A (ax is not full rank for all æ, or for only a finite number of x. So, in Lemma A.4, if there exists real numbers .1,..., la and a real number L that makes A becomes full rank, then L can be made to be a positive integer.
Back the problem, we will prove by mathematical induction. We will show that for every h, it is possible to choose fi...., iy and L to make det(A) becomes non-zero. For h = 1, then the matrix A is full rank since its single entry is always positive for f = holds for a positive integer h 1, we will show it holds for h. For j 1,...,h,
1. Assume that the result
h, let Bj is the (h 1) x (h --1 1) matrix obtained by removing the first row and the 5th column of matrix A. By computing the determinant of A via the Laplace expansion along the first row, we have
dct(A)
Σ (-1)++ 0n ddt( j=1
111
e +2 det(B,)
(49)
Denote Cj = (-1) det( (B)), and note that ej depends on the choice of ......, ty Without loss of generality, assume 81 * 0. Since .2,...;* 8h are induction hypothesis, there exists ....... la such that e₁ is non-zero for at least one T. E R. So with this choice f2..... the there exists a E R such that
pairwise distinct real numbers, by the
Cj is non-zero for all L < a.
(50)
Since .1,.... 84 are pairwise distinct and 81 # 0, we can choose = 4 E R such that:
1. ehrs 1 - and,
2. 1e c6 3; for all (i,j) + (1,
With this choice of t1, let A = 1-1- , eh Then for L A
1. We have ghre + L E [-1,]]
2. For (i,j) + (1, since le m ctoss] > 3, then etray + $ [-1,
3. We have L T 1 t1- < (1 a)
We show that with this choice of t1, 12, the there exists E such that det(A) is non-zero. Assume the contrary that
ch. + Cj =0
(51)
for all L E R. This implies that
h
1
C1
Cjj
(52)
e
+
L
e +
j=2
so
h
ct +
1aa
Σ j=2
chies +
Cj
cli,, + L
|cyl
(53)
j=2
Considering CI...., C₁ as functions L. we have these functions are well-defined on the closed interval A, since they are determinants of matrices where their entries are
(54)
cles + L
23

---

Published as a conference paper at ICLR 2025
for 1 < & h and & j & h, and these entries defined on A by the choice of t f....... t/- Moreover, Ci..... Ch. are continuous on A. Since a continuous function on compact set is bounded, so there exists 51 such that
lcil...., leal < 51, for all A. Moreover, since L E A implies I. < a, then |cil > 0 for all L. < a, which means there exists 82 such that
(55)
1a 62, for all LEA
(56)
Similarly, for 1 <j & h,
(57)
etr* L' considered as functions in L, is well-defined and continuous on A, so there exist d3 > 0 such that
c L < 83 for all I. E A.
(58)
f
From Equation (53), since we have Equations (55), (56), (58), then for all I. E A (-
82 et1-81 L (h 1) 51 63-
(59)
As T. - -1 the LHS of Equation (59) goes to but the RHS is constant, which is contradiction. So with the choice of 11,...,4,, 4, there exists E T& such that (et )) * 0. The result holds for h. By mathematical induction, it holds for every positive integers h. The proof is done.
Remark If fix positive integer L, there might not exist (1,--- ln satisfy the condition. For instance, if h > 4, and +
0, then the matrix A is not full rank for all ti....,ty f/p
We have two direct corollaries of Theorem A.I.
Corollary A.5. Let D be a positive integer. Assume that, for two positive integers h N, collections (A (AAE (Bi) (Bi of matrices in IR xx we have
F (X: ( [A),
=
(60)
for all positive integers L and X E BLXD, Then, for all A E RDDD, we have
Σ B Σ B4.
(61)
A:A.=A
1: A:AA=A
Proof. From Equation (60), we have
Σ f(XAX)) X. D
Σ A:AmA
B EB-
1 A:AA
B'
=0,
(62)
then the result directly from Corollary A.5.
Corollary A.6. Let D. Dis Dr and h, h' be positive integers. For & h and < h², let
(W(Q.1) W(K,), ((.)), w(O.))
(63)
(Q3) (KJ) (VJ) ((s))
(64)
be elements of Dx x R Dx x DDD, x D.DD Assume that, the two corresponding Multi- Head's 's identical, i.e.
X: w(ee,, ((k))), ((()), w(oo)) "
MultiHead
1=1
MMultiHoad x: (x: (wrlew) ((ks), (VVi) mp(oas)) /
(65)
)=11
for all positive integers L and X BLLDD Then, all A E DDD Σ
we have
W(V)) w(0.)
Σ
WiV)) (p(0.)
(66)
i: (.I..)
j w (wii
=a
24

---

Published as a conference paper at ICLR 2025
We characterize the symmetries of the weights of MultiHead in the following theorem.
Theorem A.7. Let h, D.D,D. Dk: Dy be positive integers. Let (W(2.), W(KA), W(V,), w(o.)) and
W ( W ((K.)) (Vi)) .w(Oii) be arbitrary elements of Dx Rx R DxD. R.xD
with i = I,... h. Assume that
(a) ma, D.) & D.
(b) the matrices ((Q.) (w((k))) w(a.0) (wilka)) w(V.s) and (V "........... are of full rank,
(c) the matrices W(Q.). (W(X.))) T with l.... h are pairwise distinct,
(d) the matrices W(Q.) (wlk.))) with i I,..,, h are pairwise distinct.
Then the following are equivalent:
1. For every positive integer I. and every E RLLD, we always have
x; (W100) w(ki), wikk, w(on)) s
MultiHead
MMlttHeadd (x:(w w(K.)) ((.))) w "
2. There exist matrices M E GL.,, (R) and NO E GLD, (TR) for each 1, 1,.... h, as well permutation Sj, such that
(w(rr()) w(Kk((t) (VV1)) wlorr()))
= =wile.) (M) w(K.)) (M w(V)) N (N(0)))1 .
Proof. The implication (2.) + (1.) is clear. Let us consider the implication (1.) + (2.). For each h, we set
(*) = ((Q.)). (WrKK.s)
7 (e.)) (TP(K.N))
B(*) = ((Va) w(O.s),
B(a) (Vs)), (0,0.s)
which are matrices in DDD, By applying Corollary A.6 for A(s), we see that B(a)= Σ
(67)
A A(is
However, since the matrices ( with j = 1,... h are pairwise distinct, there exist a unique index js E (1,...,h) h} such that A (5x) = A(s), The correspondence s + j. is a permutation of {1...., h).- Therefore, we can write ja = T(s) for some Sp. Hence, (r())) A and thus equation 67 becomes
B (T()) B(s).
The theorem is then followed by the rank factorization of the matrices (*) and B(*) (Piziak & Odell, 1999).
B MATRIX GROUP PRESERVED BY LayerNorm
In our setting, layer normalization is row-wise operator. In particular, for row vector "i....,;p! E RRP, the standard layer normalization of x, denoted by LayerNorm is deter- mined
Layer.Norm(z) = VD.
I f 1p
(68)
| T 1p
25

---

Published as a conference paper at ICLR 2025
where E 3 (2) + .... zD) is the mean of the coordinates of x and 1p is row vector in R whose coordinates are all equal to 1. Geometrically, we can view LayerNorm as the composition LayerNorm = 7 P.
(69)
of two transformations and p on R2 defined belows.
The perpendicular projection map p: IR (1p) 1 which maps z r -- 1p for each row vector æ E RD. Here, the hyperplane (1p) 1 is the orthogonal complement of (1p) in IRD with respect to the standard dot product. It is noted that p is linear map with the kernel Ker(p) = (1n), which is the one-dimensional vector subspace of RRD generated by the row vector D-
The scaling map r: R \ {0) - IR which maps + VD TRTTE The image of is the sphere centered at the origin of radius vD.
noted that, the LayerNorm operator not defined on the whole RR"" Indeed, since T is not defined at the origin, it is necesarily that p(z) * O. This means that z should not lie in the kernel of p. Therefore, the LayerNorm operator actually defines a map from the set I P \ (1p) to RD.
Under the view of the LayerNorm as composition above, it is natural to ask which matrix (or linear map represented by this matrix) will commute with both the projection p and the scaling т. The following theorem gives a complete answer to this question.
Theorem B.1. Let be a positive integer. Let p and I be the projection and scaling map defined above. The following are equivalent for an arbitrary matrix M E (RR::
1. M commutes with p and t, i.e. p(xM) = p(x) and T (x:M = M for all row vector
x E IRD.
2. M is an orthogonal matrix such that its row sums and column sums are all equal.
Proof. One can direct verify (2.) implies (1.). To show (1.) implies (2.), let x = (I1, ID) E RRD
and M (m.j)ısiJED E R DDD. From p(zM) M( we have
xM zM- 1p (x - I 1p)M,
(70)
which means
x. 1n = = 1n M.
(71)
So
))
D
D
D
D
D
Σ
Ft
Σ
mij
1n
Σ
Pi
mal
ma2
Σ
M&D
(72)
i=1
J=1
implies that
D
D Σ
D Σ x;
D Σ im
D Σ
D
Σ
mij
mm1 =
x
Σ
m.n
(73)
The above equation holds for all feasible x (as p and 7 are not defined the whole Rd). But since the set of feasible x is dense in IRd, by continuity, the equation holds for all E Idd. It implies row sums and column sums of M are all equal.
In addition, since (I) M())M we have |aM = ||0||2. By the same argument, it holds for all Rdd. Hence, M is orthogonal.
Remark 10 (generalized doubly stochastic matrix). A matrix whose row sums and column sums are all equal to one is called a generalized doubly stochastic matrix. Smoktunowicz et al. (2019) characterized all orthogonal generalized doubly stochastic matrices. In particular, let O(D) is the set of D x D orthogonal matrices, and Up is the subset of O(D) consists of all orthogonal matrices with the first column is (1/VD)-(1n) then every orthogonal generalized doubly stochastic matrix can be writen in the form
()
0 X
(74)
for some (D (D-1) (D- 1) orthogonal matrix X.
26

---

Published as a conference paper at ICLR 2025
There invertible matrices that do not permute & and but still have nice behavior with LayerNorm, as we see in the following theorem.
Theorem B.2. Let D be positive integer. Then for arbitrary permutation matrix P PD and real mumber À + O, we have
LayerNorm(AxP) = (ighn)))))))))))aaa))) LayeNNorr))))
(75)
for every row vector x E RDD
Proof. We have
LAye = =TO ((AxP)
(76)
= AzP 1D)
(77)
T(AEP AI 1D P)
= sign(A)T(x - W 11))P
(78)
sign(A) Layer
(79)
We finish the proof.
C WEIGHT SPACE AND GROUP ACTION ON WEIGHT SPACE
In this section, we recall the weight space of a transformer block and the group action on it. Then, we prove Theorem C.5 (which corresponds to Theorem 4.3 in the main text). Group actions on models extensively studied in Machine Learning Tran et al. (2025a;;;;; 2024c:b).
C.I WEIGHT SPACE
Recall that a standard transformer block, which is denoted by Attn, is defined as follows: for each X LXD, we have
XX)n(X)= = LayerNorm (ReLUU( w(1) + 1r, b(A)). w(B) + 1, .(())
(80)
where
X LayerNorm ( w(K,0) w(lk w(o.)))/1))
(81)
with 1, = [1....,1]] 1]_ RR
The weight space 21 of the above transformer block is the vector space:
u [R DxDx x DxDe RDDD BDDDD DxDA Baad) (RIxDa x R1XD).
(82)
An element U E is of the form
U ((W)(Q.0) (W)(K, [w]vs,) (w])o.))
(Iwjaa), (II(ca,[[w](B)) (ll(la) (18(4),, (8)(8)))
i=l....
(83)
To emphasize the weights of Attn, will write Attn(X;U) instead of Attn
C.2 GROUP ACTION ON WEIGHT SPACE
Consider the weight space U defined in Equation (82) and Equation (83), set
Gu S,, GLD, (IR)* GLn, (R)* x PD
PD
(84)
Each element E Gue has the form
(T, (Me))m (Ne)i=1,... (1))i=1,.... P Proy Pr.). Pro)
(85)
where no, are permutations and Mr, N; are invertible matrices of appropriate sizes. Definition C.1 (Group action on Gv) The action of Gu on 1 is defined to map 9y x - U
which maps element (g. U) Gu x U with U and given in Equation (83) and Equation (85) to the element
gU- (loweee,, (oW)(K.)) (gW]((.)) low](o.)) ;-1....
(86)
[ow]((,loww(y) (1.laeg) (00)(18))
(87)
27

---

Published as a conference paper at ICLR 2025
where
low](00) [w](e.)))
(Artr(n))
[ow)(K.)) [W](k.00) (AZ(r((3))
(o [w](2)()) N(T())
low)oww (AetrCa) [W](0.(()) Pro:
(aw)(()) PMd [W](A). PRA.
[ow](m PTá [W]
[gb](A) (B](4). PrAa
(o6](B) [6] (B)
In addition to the above definition, we will also need the following ones for the construction of the equivariant and invariant maps later.
Definition C.2. With notations as above, for each i = 1,..., h, we denote:
[WW](OK.) == [w](2.). (IW7(K:))" and [WW] := [W](Va) (w]] (88)
Remark 11.
The terms [W (QK.)...)) and [ [W]](V..) are equivariant under the action of g. since [sW wWW)() == (law)ee,) (o
(Iijaa (2r)(1)7) (IW7(K+()) -(x)(())-)) T
[W7(0.000) (W](K.+(0))
= [WW]]QK.>(0)
and
[gW (WgW](vo.)
= [oW](K (gW](O..) (IN)(V . M(()))
(.)tr (w](0.(0)) Pro)
[W](.=(0)) (1]0 Pro
[WW](VOr()). Pro.
Therefore, can intuitively say that these terms compatible under the action of Cu. Proposition C.3. With notations above, have
[gW (WWgW1( (QK,i)
[WW]] (QK.+()) jik
lo gW (VO.)) (VO.)) wwwjyoord
j.k
(.To(k)
(Q.)) j.k
[1)().)0)) (MIT))))
j.
[w](K.00)
j.k
[W] N(T(O))
[oW](o.))
(MTTTO)) -1 [[]
j.)o(k)
[gW](1) 1 [W7
(B))
(g)
(B)
[W](R)
j.k
28

---

Published as a conference paper at ICLR 2025
lobb) ( Lob)/ (B)
for all appropriate indices j and k.
Proof. This proposition follows immediately from the definition of the group action on U4 Lemma C.4. Let Tu be a positive integer. The following are equivalent for matrix M GLL (IR):
1. ReLU(x M) ReLU(()) M for all row vector x E R
2. M is a monomial matrix whose nonzero entries are positive numbers.
Proof. See (Godfrey et al., 2022).
Theorem C.5 (Invariance of Attn under the action of (u). With notations as above, we have
Attn(X; (an(X;gle gū)) Attn(X;U)
(89)
for all U E U, E G, and X E RLXD
Proof. According to Equation (80) and Equation (81), we have Attn(X:gU)
(90)
LayerNorm (ReLU(X [oW](4) + 1r. (g [ggW( +,, [nb]
(91)
LayerNorm [ReT.U(X Pad1 [W] + 1 (b].()(1) (b](A) PFA) PaT [W] + 11 (B]
(92)
1r,bb) 1r.bb) LayerNorm 2,
1r (B)
(93)
LayerNorm [ReI.U(XP:
[b] (1. (w) 1r)
(94)
In the above equalities, Equation (92) following from the definition of the group action of Gu on
U, while Equation (93) follows from Lemma C.4. In addition, we proceed the term X P-1 inside Equation (94) as:
P-1 30
(95)
(( LayerNorm MultiHHead ( [ (oa](vo), [om ((r))) P-1
Head(X; (lgW] (kW) Mo wo  )o..ooooooo).. (oralm
(96)
LayerNorm
P-11
(97)
) Mo
(
LayerNorm
Σ Head (WW(((0))))))) ( ((-t))) (lleklla]]
(98)
Head (x: (11 (w) (iitaaaa (aollm)
LayerNorm
(99)
X
(100)
In the above equalities, we used the notations
W (Qr(())
(T((2-00) (M/1T(0))T,
W (Kr(()) [](().(0) ((()))
(V.r(0) (V.())() ((Vr((i)
W (0.r (2/(V.+(0))11 (W](0.00)
In the above equalities, Equation (98) follows from Theorem A. and Corollary B.2. In addition, Equation (99) obtained by permuting terms in the sum by r. From Equation (94) and Equation (100), we see that
Attn(X: ((; ) = Attn(X;U),
for all X, g and U. The theorem is then proved.
29

---

Published as a conference paper at ICLR 2025
D AN EQUIVARIANT POLYNOMIAL LAYER
We now proceed to construct a Gu-equivariant polynomial layer, denoted as E. These layers serve as the fundamental building blocks for our Transformer-NFYy Our method is based on parameter sharing mechanism.
D.1 A GENERAL FORM WITH UNKNOWN COEFFICIENTS
We want to build map E: 4 - U such that E is Gu-equivariant. We select E(U) to be linear combination of entries of the matrices [W](Q.+) [W](K,)) [W](V.), [W](0.s), [W](A), [W](8), [b](A), [b](M) as well as the matrices [WW](KK;> and [WW](VO.s) ie
EEU) (1F(W))(Q,, [F(W)] (F( [E()))), i=1.....h"
(E(W)((0))) (18(6)(4), (E()))))))
(101)
where
h
D
D
h
D
D
[E(W)]
(Q.):
1Q.05.
[WW]
p.
(VO..):p.
p.
1=p=-1= 9=1
$
1Q.0%..
[W] (Q.)):..
(K..):p.. (w
1Q.05
D
D.
1Q.00
W]( (Vaa):p.
1Q.0:
[W] (0.s):p.
DA
DA D
1Q.0300
(A)pp (WbA))]
(2.0):
( (B)
1..
(B)p.
pM
DA
$ 19.00. MD)) (A)a
Σ
M + (QQ.):
(102)
(B)g
9=1
D
( (K.)
T
(Ki)) (QK.)):p.
(QK,))
(K.0akk (VO.))
(W]
(Q,)) (K.)):Lk
(K.)) (K.)).na
(Q.)):p.1
D
D.
D
(
( (V.)) (V.a)ina
(K.)):
T (K.b)s.k (0.s):ng P..
P.
DA ( (A):p.
DA D
$
P.g
(B):p.
p.u
DA
Φ (Ki): (B)g
(b) + + (K.)ss
(103)
(A)o0
h
D
D
h
D
D
B
$
(V.i).jk [ [WW](QKs)) (QK.a):p.g
(Vi).j.
[WW]
(VO.slppgg
D
D2
A
D
D2
(VJ).kkk (Q.)):..
(Q,))
+
(K.) (Vi)) (K.)).p.q .pg
9=1
D D.
D. ΣΣΣ
D
ΣΣΣ.
(Vi)ij.
W (V.))
$
(Va))s.k (0.s):.g 1..9
[W](0.a)
(V.e):p.q p.u
p=1111
30

---

Published as a conference paper at ICLR 2025
D DA
Da D
[W] (B)
(A)Pa
+
(B)p.
p.g
g=1
DA
$ (
(6%,
(4)
+
Σ
(V
o (V.0s.
(104)
(A)xq
(B)xg
q=1
q=1
[E(W)]
T (0.0
(W])OKKs)
$ (0.00 (VO,.):p.g
1p-0
s=1p=1
Dx ME
(1) (2.))..4
10,0 [T](ka)
(K.8)p..
s=1
4
D
h D. ΣΣΣ
D
(V.)):p (O.0
10,0 (0.s):p.g
(0.a)
p=
g=1
g=1
D
DA
DA D
$ 1O
(A)
$
(0 (B)ra
(B)
(A)pa
p.
1=1
DA
$
()(A)
Σ
M (0..i.i
(105)
(1)0
(Bja
D
[E(W)]
$
(A):j. (QK.)):p.
(W](KK))
(A)sk
WW)
(VO.9):p.
D2 (A)ikk (Q:)):p.g
D
D2
(Q
(1)..
(K.s)
p.y
(K.)):n0
p.
D.
D (4)g.
(A) (V.s):pag p.
[W]
$
[W](0.)) 1p.9
(O.s):p.g
DA (1)3 (1)=ma [W
(A)
(W] BB)pp
DA
D
(1)sk (A)o
+
Φ (A)s (1 + (A):.
(106)
T=1
h
D
D
A D ΣΣΣ
BB (B)
(QK,) (QK.
(B)
(B)ikk
p.y
(VO.e):p.g
p.q.
D
DR
( (Q.)):.44
(Bibkk
(B)jkk
(K.s)
...
(K.)) p.
D
D
D. ΣΣΣ
D
(B)
$ BBiAk
[W](0,a)
(V.elip.q
1.
(..)).p.
p= g=1
s=1 p=
D DA
DA D
(A):..
(B)ssk
1p:
(4)
+
(B)
(B)
w1 (B)
p.
p= p=lg= g=1
p=
DA
T (B)) ((A))
+
② (Bss A [ + (BB} (B)g
(107)
(A)xq
q=1
D
[(D)((A)
T (A)ik (QK..):p.
(W](QK,a)
$ (A) ΣΣ
[W](VO.s)
1p-0
(VO..):p.
11
31

---

Published as a conference paper at ICLR 2025
D
DR
h
D
DR
(A)ak
(1
(A)ck
[W]
(Q.s):..g
+
(K.
8=1
D
D2
h
D2
D
(V.o):p.q) [Wpp
(A):
(A)sk
(0.));..
q=1
DA
DA -Σ>
(A):
$ (4):k
(B):p.q [W]6
(A):p.g
p=1q=1
AA
(1))
(1)))
11)
(1))
(B) + T(A):t
(B)
(108)
(A)g
(B):
q=1
h
D
D
h
D
D
[E(b){
(B)
$
(B)bk (Q...):p.g
(WK(O p.
(B)B
(VO.)
(VO.e):p.g
p=1
D
h
D2
$
BB)b (Q.)):p.
(B)bk
(K,
P.
(K.
p.y
p= 4=1
D
D.
D
$ (V.s):pu
(B)b
(V.
$
(Bikk
[W]00
p-0
(O.s):p.g
DA
DA
$ (A)=ri
(B)bk
[W(A))
+
(B)b
(
(B)sp.
p=1q=1
DA
$
(B)
(4)
+
(B)* (B)g
| + (B)b
(B)
(109)
(A)o
9=1
and the constants I are parameters.
In the following subsections, we will determine the constraints on the coefficients Φ such that E(gU)- = -E(U) for all U e ll and E Gu.
D.2 AUXILIARY LEMMAS
We will need the following auxiliary lemmas in order to determine the constraints of the coefficients of E later.
Lemma D.1. Assume that E: - U is a function defined as in Equation for some coefficients $_. Ef. = for all U E U, then all coefficients are equal to zero.
Proof. We view the entries of U as indeterminates over R. Under this view, since E(U) = for all U E Ut we can view E(U) as zero polynomial. Moreover, the set including all entries of the matrices [W](Q.e), [W](K.*), [W](K.), [W](O.e), [W](4), [W](B), [6](4), [B] well as the matrices [QK]( and [WW](VO..) is linear independent set over IR Therefore, the coefficients of E, which are Φ_, must be equal to zero.
Lemma D.2. Let h and D be positive integers. Let 122)) _!22) DxD - R be R--inear functions for each .,.... h. Assume that there exists constant X F such that
s!(2) ((arl)))) X,
Σ )!!)) (Mr()))
(110)
for all (M))..... M() GLD GLD(R) Then
(( (M) (12) (M) X 0
for all s 1,...,h h and M (R).
32

---

Published as a conference paper at ICLR 2025
Proof. Fix arbitrary element (M.)
....,.,))).. M of GLD(R)*. Then for arbitrary nonzero num-
bers A1..... An, the tuple (A1 ..... M( is also an element of (LD( Since f f!?) are linear, it follows from Equation (110) that
Exeld" (Mr())
+ 1 X
1!2) (()rl))))) A.
(111)
==1
Consider the function
1e II
Σ taf!!! (
(( )(x)))
to - 5500
P(t tn)
(112)
as polynomial the variables f1...., la with coefficients IR. Then, according to Equation (11
have P(A1,-..,4 >n) = 0 for every nonzero numbers x1.---; An.- This happens only when P is
zero polynomial. Thus all coefficients of P must be equal to zero. In particular, we have À = 0 and
E GT. (TR)* ( (M()) n!²) ((M())) ")) = for all s 1,...,l , and (M(1)..... (()))
The lemma is then proved.
Corollary D.3. Let h and D be positive integers. Let fa f:² g.t gs functions for each s = 1..... h. Assume that there are constants d1,72 E R such that:
(1)
(2), DXD - R be linear
1! (M(()) + 1!2) ((ar(o)) + 11
Σ
Do (1)
(Mr(e)
+ )
((arlo))
+ A2,
(113)
for all (M(()..,,M E (R). Then 1 = ys (2) (M) and 12a = 1/² (M) and
A1 = for all = 1....,h h and M E (R).
Proof. Apply Lemma D.2 with f!d)) is replace by f! go )!(2) is replace by f!²) gb²² and replaced by xg- Å1.
D.3 FINDING THE CONSTRAINTS FOR THE UNKNOWN COEFFICIENTS
In the following, we will find necessary and sufficient conditions for the coefficients $ such that E(gU)= (EU for all U e U and Gu. We follow the parameter-sharing strategy. In particular, we first determine the entries of E'(gU) and gLU)). Then, we compare the corresponding entries to determine the exact constraints the coefficients.
D.3.1 COMPUTING E(gU)
Following Equation (86), we have E(gU)
(BGW))((
([[TW))",(((WW))"))))) ((E(b)))) [E(G))]()))
(114)
where
(QK.kkn 19 (WWW((Q.))))
19,0
D
P
(QA)
LoW (VO,a)
(VO.s)
4
D
Dr
h
Dr
[g (Q.s) (Q.4):p.
p.g
+
(Ke) l.
9=1
D.
D
ΣΣΣ
(V.))
$
(0.)) (0.0)).. p.g (0.)
(Q,
(V.))
33

---

Published as a conference paper at ICLR 2025
D
DA
D4 P
LoW (A)igg
(R.)i
[gW
p.g
+
(B)p.
p=
1
(Q.):: (A)y
Lob)!
+
(Q.0): (B)g
[gb] + 199,00
(115)
q=1
q=1
[E(gW))(K
(K.))
(K.) (QK
(WWgW]
j.k
D
(K.))
[9W WWgWhyo (VO.s)pp.
D
Dr
h
D
D2
(9,s) (Q.e):p.g P.
(K.0
g
(K.)):p. 9 (W,)
D.
$
(.K.)
( ( (V.s):Pig
T (K.))
gW (0.a)
pi
(0.0):p.g p.
9=1
(K.i)
$ (K.i)
[g
(B)
(B):p. P.
(K.i)sik (A)or
ggb((2)
+
$ (Ki): (A):
[ab]
(B) + (KK.)i (116)
D
D
(VJ):j.kk (QK,s):naa
[g ( gWWgW
(v.0 W gW] 12.9
(VO.s)=p.g
h
(. (Q (0.)):p.p ( ...
(Vi)
(g
(K.s)ong
D.
D
(.iiz (V.)). (V.)p.glg (WW
(0.))... (0..),..19 (9W]
DA ( (A)x..
Da D
[gW](44
Σ
(V.))
LoW]3
...9
(B):p.g
..g
p= q=1
p=1
DA
$ A
[ob](
+ Σ
$ (Vi)iskk (B)x
[ab
(B) + (V.))...
(117)
(A)ag
h
[E(gW)1(00
(0 (QK.))p
[gW (WgW)(QK,,)
D
D
(O.)iskk
Wo gW] (VO.s)
(VO.s)
P.9
Dt
h
Dr
(oDik
(Q.0) (Q.0)
(g l..
p.1
(K.s):
9=1
D.
D
ΣΣΣ
$
( (V
(V.9))7-
$
(0.=):p.gggh) (0.
(0.))
34

---

Published as a conference paper at ICLR 2025
D
DA
D4 P
LoW (A)ig
(..)i
(B)p. [gW
p.1
+
p=1==
1
(0.) (A)a
Lob)!
+
(
[ob] + 10.00,
(118)
(B)g
9=1
q=1
[B(gW)|
(A)
(WggW(0kk.)
j.k
(Q.
D
(A)
g9 @WgWTyo.
(VO.s):p.
D
Dr
h
D
D2
$ (A)isi
[gW (Q.=):p.g P.
(A)J.k
(K.)):p. (
D.
D
$ (()) (V:)) 1.
$
(A)3
1o
(0..):p:
9=1
( (A)sp.g
(A)ikk
(9W)( p.
pMu
BB)sp
D
[ob](4) (A)o
+
$
(Arsk
gabl gbb(() + ((A):i
(B)
(119)
(M)g
D
D
[E(gW))!
(QK,s):na [g (g (WWW))
(B)sj.k
(B)jj. (VO.s)=p.g .2.9
h
(B)jk [oW]) (Q.s):p. (
(K.s)op. ( (gW)(K
(B)jjk
...
q=1
D.
D
(V.)), [gW]
(BBjkk
(0. (oW)
(B)jj.k
DA (B)
DA
(4)
(B)Jk
1y
(B)
(A):p.g
1..
(B)
p.
p=
q=1
p=1
DA
Σ
$ (B)J
(gb)(4 (A)y
+ Σ
$
(Bjjk
(gb((B) $ (Bijk
(120)
(B)g
h
[E(gb)](a)
(QK [gW (WWgW)QQkk)
(A)ak
D
D
T (Vo.s) loW AWgWiivo. P.9
(A)ik
De
h
D2
$
(9 (Q.)) (Q.s):p.
(A)&
$
(A):
[gW](K,))
P.1
(K.s):p.g
9=1
D. Σ
D
ΣΣΣ
$ (A):k
( (W
$ (A):k
(0.s):p.glgh) (9
(V:))
35

---

Published as a conference paper at ICLR 2025
D
DA
DA
D
(AA)pp.g (A
(A)ik
[gW
o (A)k
(B):p.g (B)..g [ (B) ...
1..
Da
(A)sk
(A)i 186b/1)
+
( (Big) + ((1)x,
(121)
q=1
[B(gb)}k
BB)b
(WgW)(kk.
P)
D
(B)b
( WWgWhyo
(VO.s)p.
D
Dr
h
D
D2
(Q.=):p.g [gW
(B)kk
$
(BBb
[g [
(K.)):p.q
D.
D
$ BBBb
$ (B):k
(0,s)... (0..)p..1997>>. p.0
(V:)):..0 1..
T
BB)ck
[oW](a)
BB)pp |oW p.
BB)a
(A)sp.g
B)ak (A)og
[o0]
(4)
+
(B):k (B)g
(B) + (B):k
(122)
Using Proposition C.3, we can proceed the right hand side of the above equation further as:
ADD D
[F(gW)]
12.0
$
(Q05k
(W(QQ.F0)) P4
(QK.3)pp4
9=1
(Q.) (VO.s)
WW) (VO+(+)) p.no(
(Q.5 (Q.a):n.a
[1((27(0)) (Mr()()))
p.g
(Q.):Ak (K.s):pa
[w) (ar()(2))
-1]
p.g
D. $
(Q.)):jk (V.a)
Nx((()) [w]
p.q
①
(Q.)::AA (0.9):m.
(NTr(ct))
[W]((.)())
p.mo(l)
DA Q.:)
Da D
( (1)
Q.:5 (B)
(A):pa
BBpp.
Da (
(A)a Ta(a)
(1)
(4)
+
+ (Q.0)
(123)
4=1
[(GW))
ΣΣΣ
Σ
$
(Kigjkk
(QK,7()))) [QK.3)
-..g
36

---

Published as a conference paper at ICLR 2025
A
D
D
(K.)issk
[WW]
(
(VO.))
(K.)).,k
[w](Q.-())) (Ag
(Q.+):p.q
(K. (K..):ppg
[W](K (as(+(a)) -1
P.9
$
(K. (V.)):pia
[W/(V55())) .ao...... -N((()))
P.0
D (KK.) (0.0):p.
(Mitrol)) -1 (W](0
p.To(a)
D
Da
Da D Σ
T
(KKi)) (A):p.q
Whoplpp To(p)
[ (B)
(B):p.
ma(p).
p= p=1 q=1
$
(K.)): (A)9
(K + ((Ki):
(124)
(m)
D
E 1.k
(V.)
(QK (Q
(V.i)..k
(W (.)
(VO..)pp.
$ (V.i).ik (Q.+):n-
[W](0r(0))
(M1
p.a
4 (V.)-..k (K,)):p. [W](K (ae(r(o7)
-1
p.9
D2
$
(V.o.j.k (V:s):pu
p.g
D
$
(V.)s..k
(0.x):p- (NNTTOO)) -1 [W](0.r(a))
1==1
p.50())
DA
Da
(A) (V.i.j.
(B)pq [W
(A):p.
(V.)ssi (4)
Φ (B:g
(1 + (V.):
(125)
(A):
D
D
[E(gW) (0,)
10.00k
(QK.a) WW (QK,(())
P.4
h
1 (VO..):p.g
(VO.V0 p.)
pm
q=1
ΣΣΣ
$ (.
(((2.7(0)) (Mr(rioD)
] P.
(Q.s):p.
37

---

Published as a conference paper at ICLR 2025
A
D
[w](Ki((o)) (artr(69)
(K.)):p.q
P.9
De
[IN](CTO>>> (x((((a)) (
(Vis):pig
(0.): (0.0):P-1
(NITTOO))
-1 [W](0
p=1q=1
p.)o(a)
D
DA
D
(
(W7 (o(2))aag))
(A)
$
(0.0
(1)(8) (B)
(A):p.q
(B)pa
xx(P)4
==
DA
D
(A)g
M (A)
Talq)
+
(B)g
1
+ $ (O.).i.
(126)
$
(A)uk
(KK-())))) (QK
P4
(A)ik
WW] (W] p.oo())
(VO,
(Vo.s)
(A)S.A
[ (ar(rto))
(Q.s):n-a
p..
(A):ik
[W](K)(O)) (ar(r()))
-17
(K,5):p-1
2.g
$
(A):j.k (V:s):p.
[W](V
-x(((()))
(A):k.k (O.s):n-
(Nir(ON))
[W]((.)(0))
p.no(1)
DA (A):jA (A)sp.
DA D
(1)s (B)
(oO())
+
(B)sp
(a(())
p=1
DA
$
(A)
(A)3g (TTa()))
+
Σ
② (
(B)y [0, + (())
(127)
0=1
ΣV
(WK,QK PQ. -2.
(B)
h
(B
(W)(
D
(B)ikk (Q.e):p.q
[w](9-(0)) (NI()1O7)
(128)
p.g
(B) (K.e):p.q
-17 [w](K.(0)) (2)()(22)
q=1
p.g
ΣΣΣ
T
(B)
[W](V.7(0)) N((((1
(V.s):p.
P.q
38

---

Published as a conference paper at ICLR 2025
D
D
((Nitco))] [W](0
(B))
(0.0):p.
q=
p.To(1)
D
D4
D4
(B)),,
(A)
BB
W (B)
(A):p.q
[B]pp.
p=1==1
DA
T (B)s (A):4
(b)d]
Σ
$ 1
1
B)
(129)
[B)a
[E(bb)] (4)
$ (A)a
(Q. (KK,(K))))
h
$ (A):k
WW] (v),
(VO,s)
p.ro(1)
$
(A)a (Q.+):n-a
[W(Q.0()))
(Arbtio7)
$
(A)ik (K;+):p
(W](K.(())
-17 (as(r(a)))
P.g
De
(A)ck (V.s):p00
p.g
D
(A)ck (O.)):na (NITTOO)
(w()),)(s))
p.vo(g)
DA
DA
(A): (A):paq
[W](
(d)ck (B)
[W](8)
(p)
(B)p. mx(())
g=1
(A)= (A)=7 11
(M)g
(B() + ((A:,K
(B)
(130)
D
D
LB):k
(Q. WW (QK,())
(B)b
WW] (VO,(0)))) p.)o(g)
(Vo,
q=1
(B)b (Q.5):p.q
[1](0.(0)) (MV(+(9)))
(B)bk (K.))p.
-1] [W](K.(()) (Mr1)667)
p.g
D
De
(B):k
[WN N(G(
(V.)):P.
p.g
D
(NNTTCe)) -1 (0])).)())
$
(B)b
(B)bk
(O.)):pu
p=
(.)
Da
DA
ΣΣ (A)
$ (B)bk
T B)a
W (B)
(A):p.q
(B).p.g
(p)
p=1q=l
39

---

Published as a conference paper at ICLR 2025
DA
D
B (6) (A):g Talg)
(4)
+
Σ
BB)
B
+ (B)k
(131)
Using the symmetries of the indices, we can rearrange the right hand sides of the above equations as:
h
[E(gW)] (Q.))
Σ
$ (Q.)
[WW](QK.)) P.9
(QK, "())p,
$ (Q.0:1
(V..)) P.
(V0,r-(())):ppp (4)
1Q.0%
[w)len)) - (Ar(s))
T7
1
(Q.r (x)):pp
P-9
Dr
& 1Q.0
[W]((Kss) (Nc(c))
(K,7 (0))):p
(Q.0:
[In -V())]
(V.r 16)))p
p.9
D.
[(NI.0) [0,
(Q:0
(4):p.00
P-q
DA
D
(Q.1::jkk
(m/6 +
(Q.)::tkk
(B)
(A)= (1))=7'(11) (1)
(B)=m* (p).g p.g
1=
DA
D
Σ
$ 1Q.0001k (1)o0 (n)
(B(A)
Σ
$
(B) + (QQ.0:
(132)
(B)
1
h
(F(WW)
(k.)
(s)):ppl (WT(QK.>)
T
(K.):
WW](VO,)
(5)) (1)
D
(k.))
[mija.))(M11) (Mr(.)
T
(s))
P.g
(K.)) (o)::a
-17 [W](Koo) (Mr(s)
P.o
(Ka)j (Vir (6))s
[In](Vi)) -v()]]
p.g
A
D.
[ANvee) (0,)
(K.j)siA
(0,z-1()):pp-
P.9
D
Dr
DA D
(K.)::skk
[W](A)
+
(K.))..k
[B]
(A)=To (1)
P.g
(B) (p).g
..g
1=
1=
DA
Σ
1 (K.)sskk
(A) (g) (b((41
(K.)),.k (B)
4 (K.00..
(133)
40

---

Published as a conference paper at ICLR 2025
4
D
D
[ jik
(V.)s
W (QK,))
(QK,T-1())))))
D
D
$
(V.):...
(VO,r-(())::p..a (g)
p.o
$ (V.))...k
[W(a, (Ar(a))
(+)):pp
Dr
(V.))..kk
[w](K.)) . (Ar(ss)
(K, (+)):p.)
h
D
D2
[wi](is)s N(o0
p.9
A
D.
[Arlo)) -1 [w]
(V.))...k
(0,- (2))x_0))))
P-
T
(V.):
(W(())
$ (V.) (W(B) ΣΣ
(A)=T (p):)=())) (1)
(B)=
p=1
DA
(Ka):jk
[1(4)= +
(V.) p B + $ (V.).
(134)
(1):T (4)
h D
D
[E(gW)]
$ (O.)::Ak
(QKKr-((s))pp (WW](QK,))
j.k
...
$ (0,0):,k
[WW]VO,,) p.
(VO.r-(0))): (00)) (0)
ΣΣΣ
$ 10.0..k
[w]ce (MI(1) (2)())
(Q. (+))
(s))=p
p.y
ΣΣΣ
Dr
0.00. (s)):pp
-17 [w](K (ar(s))
P...
A
D (.i5s
ΣΣΣ'
[w]]( .(oo]
(c)
ng
D.
[Avicc) -1 (0s)
(0,
1=1
P.g
DA
DA D
Σ (O
(I(() +
ΣΣ
(..)sskk
(3]
(A)imo (p)x) (1)
BB) (B) (p).g
..g
DA
D
(ONNU (A)aT (e)
(1)
(((A)
+ Σ
1 + Φ (0)}
(135)
(B)g
9=1
q=l
h
D
$
(A)ifA
(QKs) p.
p=
D
D
Σ
$ (A)iiA
(Vo.)
(VO,z (60)))077 (4)
p.
41

---

Published as a conference paper at ICLR 2025
A
D
DR
T7 [wpsa (AI(T))
(Akik
(Q. (w)):pi
p.g
D
(A):
-17 [W]jikso] (Ae(c))
(k,77((s)):pm
p.g
$ (A):
[In](Vss] -V(ss]
(V.T- (0)))pp0
p.9
A
D.
(A)
(0, (0))spp00
p= q=1
P-q
D
DA
DA D
(A)a.
(W/6)) +
ΣΣ
T
(A):si
(A)oso (p),x²'(0) (1)
(B)= (p),g p.9
DA
D
$
(A)U,A
[B((A)
Σ
$
(A)
(A))AA
(136)
(1)TTa (4)
(B)
1
q=1
q=1
h
(E(gW)!
(B).&k
(W](QK.))
(QK,r- (0)):p.)
P.9
$ (B)ij.k (VO,-- (1)) (4)
(VO,s
(B)jjkk
(mije (Ar(s))
T
(137)
(s))=r
P.g
$
(B)ikk
-17 [W]KKo). (Mr(c))
(K,T- (s)):pa
p.g.
ΣΣΣ
(B)j..
[WT] N(0o
(W)-so
A.D ΣΣΣ
[Avco))
(B)jkk
- (O,)
(0,- (w))pp
P.9
DA
D4 D
Φ (B):jk (A): (p), (9)
+
B):jk
[W]
1.g
BB) (p)
1=
1=1
Σ
$ (B)iAA
(1(441
Σ
(
x (B:sk
(138)
(A) (4)
BB)
/
4 D
[B(gb)](A)
(A):k
(1))) (W [WW](QK.>)
(QK,
D
D
(A):k
[WW]VO.>)
(VO,r (6))) (0)
A
D
D
[w] (M)()) T
(A):
(Q,7- (s)))pp
p= g=l
Dz
(A)ak
[wyjkko). (Mr(.) -17 P.g.
(K. (s)):p.
42

---

Published as a conference paper at ICLR 2025
h
D
De
(A)ak
[In](vss] N(0) _N(0)
(V.r (+)):.p
q=1
p.9
A
D
[(Avcc)) -1 [w]]
(A):
(0,r (e)):p.To''
p= q=1
DA
T (A)i
(W((4)
T (A):
(7
(A)sTO (p),x
(B)= (p).g
DA
D
(A):
(1)=T (4)
[b](4) + Σ
(A):
(B) 4
D B + $ (A):k
(139)
1=1
h
D
D
[E(gb)] (B)
(B)*
(QK,s)
(QK, (0)):p.1
P.9
D
$ (B)sk
(W)VV P.g
(VO,7 (5)) (1)
(B):k
[Wice (ar(s))
$
(Q. (s))
p.g
$ (B):k
[w](Ko)) (Mr(s)
(K,r (.)):p-7
P.q.
(B)b
[WT(Via) NT(0]
(V.r (3))sp07
q=1
P.9
[Avic)) - [W]
(B)*
(0,- (w))ppTT''(1)
P.g
DA D
(B):k
(A)oso (p).x (1) (W](A)
BBik
+
(B) (p).g
DA
D
+
(B)c
(A)
(g)
[b(())
+
(B)* (B)
BB)g 1
+ (B)b
(140)
1=1
D.3.2 COMPUTING gE.(U)
According to Equation (86) and Equation (184), we have:
gE(U) (O (OE(W)) ( (OEW)))()
i=1.....
(LFF())((A) (OF(W))(8))) (16F())))))((((())())))) (GGF(()))))) (F(())(B))
where
(
(E(W))(R (ar()(01)
j.k
(SL(W)) [K(W)](K>T0)) (ar(t(o)
j.k
(BE(W)) [E(W)] N((((1)] [N(())
j.k'
( (NNTO0)) -1 (B(W))((0-(1)
(k)
43

---

Published as a conference paper at ICLR 2025
(OE(W) (4) [B(W)) (OO)
(OE(W)) [E(W)) Ta()).
(o)())) (A) - [
[(b)) [F(C)) (
We calculate these entries in detail below.
[B(W)](R-6))
" i.k
Σ
((W) (Q.(0,)
(MS(r(O)) 1,k
1=1
D
(Q.1t))
[WW] (KK.sipp
(Q
WW] (vo.)) P.
(VO.s):pp
p=
(0-1 (Q.))
Σ
(K.a)pp. (Q. (
(2.).:p
P.
P:
11pp=1e- q=1
Dy
4
D
D2 ΣΣΣ
M((())
[W
(V.s)ap.
(0.s):p. p.
sml
D
DA
Da D (07
[W]bd] (A)PPa
(Q.T()):
(Byp.
1Q) (D(A)) (A)a
10000 + (.7 (B)
(BB(W))(
(Kx
[LE(W))(A+T))) (Nr1~OOD) -17
j,k
((ar61o0)) -)
1.
(K-(0) ((ar(re0))1 [EW))
(141)
1,
(K.(l))ddd (QK.)) (QKa))m
(Kr(i)) (vo,ahm
Ww (vo.)) p.
p=
(K. (Q.)):.
+
Σ
(K.r [W](K
( (K.o)pp
Di
h
D.
Σ
( (Viskaia
(W
1KK-0a p $
1=1
1.k
4=1
D
DA
(K-700) (A)T-g
(W](a)
(K-(0)) (w)
(B)sp4
pml qm1
(K.70) (A)T
+ Σ
+ (K, (K)ra
(B)a
7=1
(BE(W)) = (V ((((
(142)
Dv
Σ
(E(W)) N((T(0)
(143)
44

---

Published as a conference paper at ICLR 2025
D
4 (vo.s)
[W
@KKo)pa
== p=
(V (K.))
(W]
(Q.5):p
(K.)):p
p=1
q=1
Do
Σ
N1T(1) 1.K
(Krro) (Vi))pp.
(Kr(0) (O.)::
[W]
P Da
DA
Φ (V (Alipa
[W]
(V (W] (
(B)aaa
D
$ (
(A)s
(B)a
2=1
q=1
[wvbrco))
(B
(E()))(
(144)
(k)(()
=
=Σ
Dr
((vir10))))
NE(W))
(145)
(WW)LOK.N) (KK.) 1.1
(0)(
[W J.9
(QK,skip
(VO.s)ipa
q=1
Dv
(M.r0))a
[W] (Q.0)p4
(W)]/
(W](ka)
(K.a)pp
D.
-)
ΣΣΣ
h
D
D. ΣΣΣ
h
$
(0) W]]
(Vaa)
(..))pa
sm1 p=1 1=
sm1 p-p1- 1
(0
(W)(A)
(k)
(
(A)ma
+
(B)r
p= g=1
(0.
(()00))))))
(A)s
(B)
(0E(W)) [E(W) (A) 1
(.)))
(146)
h
D
D
(aa)
(KK.)
(QK.e)ppg
D
(vO.s
(VO.s):p.
p.
q=1
(
[W]
(1)
(Q.9):0.4
(K.))pp
pam
D
De
h
D=
(a)
(W) (Ve))pp
$
(A)
(0.s):p.
sm1 pm
D
DA
Da
D (A)
(W](A)
[W]
(B)
(A):p.
p-j
(B):p.
p.q
1=1
p=1 1=1
DA
(A)
T
T
(1)=a
(B):
8=1
9=1
(1)
(147)
45

---

Published as a conference paper at ICLR 2025
(gE(W)
(E(WW)
(148)
(B)
WW (QK.)) P.
(QK.y)aa
D
(B) (VO,s)
[W]]
Dr
h
D
Dr (B) (W))
(B)
(Q.2) P.
(Q.)):m
(Ka))n
P.
D.
$
(). (0,s):P0
(0.))
(V.s)ap.
p.2
1=1
DA
(
(WTTA))
(B)
[W]
(A)sr.
(B):p.
g=1
(B)
(6)(A)
(A)a
+
(B):
9=1
+ $ (
(149)
(gE(())(A) (A)
[E(b)] (4)
(150)
(K))
(R) [WW](OK.a)
(QK
(k) (vO.s)
[WW]
D
Dr
(k)
[W] (Q.1k:pgg
(A))
(W)
p.
(K,0):p
p.
9=1
h D.
(4)) (W]((()))))) (Vo)p..
(A):Ta(k)
(0,)):p.
(A) (W)())
a (R)
(B):p.
p=1 9=1
D
(A)= (1)ca
+ Σ
(A): (B):
[b]( + (1))
(151)
(A)x
9=1
( = (E(6)((0)
(152)
D
D
(Bbak
(QK,s)
(p.
D
D
(B)b
(VO.))
(VO.s):p.g
P.g
DR
ΣΣΣ
BBikk
(Q.s)
ΣΣΣ
(B)c
([](K,s)
(Q.)):p.
p.
(K.)):p ( l..
p=
1=1
46

---

Published as a conference paper at ICLR 2025
D
D.
h
D.
D
& (V.a):p.g
(B)*
ΣΣΣ
$
(..)p. W
BBbk
9=1
D DA
DA
(B):
[W]
+
(B)p.. [1 (B)
(B)
(A):p.q
p=1 q=1
p=1q=1
(B) ()(4) (A):
$ (B):
q (B)B
(153)
D.3.3 COFFFICIENTS COMPARISON IN THE EQUATION F(OU) = F(U)
Since E(gU)= FE(),, we must have:
[E(gW) = (OE(W)
[[(W))((K.) = (O((W)
[E(gW)((Va) = (
[E(g))](o.) = (EE
[E(gW)](4) = (E)
[F(W)) = ()())
[E(G6)](4) = (9E(6)
[B(gb) = (gE(b)) (
and these equalities hold for all 1,...,,,, h. U and g E Gu.
In the following, we will solve these equalities, one by one, to find the constraints of the parameters p_s of E(U).
Case 1. Solve the equation [F(gW))(Q)) (g For every J, k. we have [F(gW)) j.k = (o (Q.))
or equivalently,
ΣΣΣN
$ (s) (W)(0K.))
(Q.)
$
1Q.0sss
(W)(VO,,)
(10,.-1())) (0)
D D2 10.ck
T [W](Q.0) (ar(ct)
())=m
p.g
-1] [W](K)). (Mr(e)
$
p.g
(c)) [W](V NV(0]]
p.g
D.
[(A16))
11.0
(0,)
(+))(=))) ((4) (0,--
P..
D Dr (Qii
DA D
[W](A))
+
@U))..kk
[W]
(A)=To (p),
p.g
(B) (p).g
p=1
$
[b]((A
$
(Q.)) (B) Φ (
(1
(A) (q)
(B)4
9=1
47

---

Published as a conference paper at ICLR 2025
De -EEW (Mrtrio))
(154)
1.6
Observe that the right hand side of the above equation is just linear function M(7()) By applying Corollary D.3, we see that all terms of the left hand side are equal to zero, except those containing M(T(() will be identically equal to the right hand side, i.e.
ΣΣΣ
(QA) (oKKr-(())))
[ W (QK,*) -P-9
D
1Q.0i
[WW]sVOs)
(0)
p=
h
D
Dr
Σ
ΣΣ
$ (.))
[W](Q..]]
(Mi(i))
(Q. (c)):pp
p= p=
P-g
(
A D
D2
Qoak
[w]]k..) (ar(c))
(K. (+))
p qmi
P.q
A
D.
[W]cvio] .N(e)]
(V.r (+))sp.
D.
(0, (1))pp (a) (Ar(.)) [W](0.s)
1==1q==
P--
DA
D
(Q.0:j
(W((A) +
$
Q.iiss
[W](3) (
(A)= (p)=5 ( (1) (1)
(B) (p).g
p.9
p=1
p=19=1
Da
D
+
(Q.0:31
(A)=T (1))
[b((4)
+
Σ
(Q.)) ( + (Q.i)
(B):
[b]
q=1
= 0,
(155)
and
D Dr +
De
ΣΣ< pml
@ (Q.))
[w](e-00)
(Mr
= -Ie(W)
(0.0)
(MI(TO))
(156)
1.
We choose to be the identity of Gu (i.c. TO, WA, Ms) and (() are all identities) in Equa- tion (155), then substituting the obtained result into Equation (102), we obtain a more compact formula for [E(W)](R follows:
D Dv
[E
(157)
(Q:):
Then by substituting this compact formula of [F(W)) constraints:
(Q.0) j.k
to Equation (156), we obtain the first
ΣΣ 1=
1Q.0kkk (Q.1):.g
[
(M/+()))
P.a
Dc
D Dv
$ (1) (1.(0) (Q.((0):
M(7(1)
1=1
or equivalently,
Ds (0(()))
T [W] (AA(:)
ΣΣ
(Q.(1))::
p=19=1
p.q
48

---

Published as a conference paper at ICLR 2025
De
D
D
Σ E=
ΣΣ
(Q (Q.0p..
12.
( W
-M
(158)
k.t
We view the left hand side of Equation (158) as a sum of linear functions on the q-th rows of M with y = 1,..., DR. while the right hand side just linear function on the k-th row of M The by applying Corollary D.3, we see that all of the coefficients (Q.r())) of the left side are equal :Q:r
to zero, except those with the indices h: = q. Therefore, we can reduce the formula for [E(W)](Q.)
and the constraint further as follows:
[E(W) (9.0 (Q.0
Σ
Φ (Q.1 [W] (Qil: p.k
(159)
p=1
To find the constraints for the coefficients, we observe that, the left hand side of Equation (158) depends T, while the right hand side not. Therefore, by varying т, the left hand side does not change. As consequence, we must have
(0)()))
+b (Q.0).. (Q.(l)))p.k (Q.5)-p.k
(160)
for i, P, k and r.
Next, by substituting Equation (159) and Equation (160) to Equation (156), we obtain:
② Σ
u. (Q.0)..k
Σn
M (
Σ
Σ
+
( p..
(Q.)
[W]
(161)
1=1
By comparing the corresponding coefficients both sides, finally obtain the constraints
Φ ((Q,7 Φ 1Q.05. (Q:((0))pk (Q.i)pp
(162)
for all i, k. and т.
(K.)) One verify directly that (E(W)) defined in the form
[E() (Q.0
1,
Σ
( (Q.0 (Q.ip.kk p.k
(163)
with the constraints
(Qa) = + (Q) (Q.):..
(164)
())
for all i.j. k. T, satisfies the constraint Equation (156).
Case 2. Solve the equation [E(gW)](KK) (g)(W))(K.))
By using the same argument, also
obtain [F(W)) (K.)) in the form
((W)) (K,1)
T (K.0s.
(W(Kk))) (Ki))ppk p.k
(165)
with the constraints
$ (K.)):fkk (K,i):p.k
क6 (K.
(166)
(K,(()):p.1)
for all i.j.p, k. k
Case 3. Solve the equation
By using the same argument, we also
obtain [E(W)) in the form
[(W)) ji
Σ
Φ (V.0) (
(W]
(167)
(Vo)),,
p=1
with the constraints
(V.))..k (V.))p..
(V.
(168)
(V.
for all i.j. P, K
49

---

Published as a conference paper at ICLR 2025
Case 4. Solve the equation (E(gW)](0)) [gE(W))(O.)) Since
[E(gW)] = (gE(
follows from Equation (105) that
T
(0
[W](KK,) (QK,s)
(0))=
P-
D
WW)(VO.)
(VO,r 1())) (0)
D
D2
1000
[1(19 (ar(.i)
(+)::pa
P-g
D2
-1] [w]cK..) (Ar(co)
(K,r- (o)):pu
p-9
$
(0,1:
[wn((is) _N(0]]
D.
[(Nr6)) [W7
(0,-
(+)) (e)
P.
DA
DA
D
(ON)
(W
+
Σ
(O.i...k
[W]63]
(A)=T (1)
p.
(B) (p).g
p.g
1=
p=1
D
(A) (0) 10/(11
10.05
$
1
x (0,0
(B)a
((NIrOO))
Σ 1=1
-IE(W))(O
(0,(()) (1.
(169)
The right hand side of the above equation is just linear function on (NV())1) -! Therefore, it follows from Corollary D.3 that all terms of the left hand side are equal to zero, except those containing (NV(((0) -1 will be identically equal to the right hand side, i.c.
$
(O.)..k
WW (QK.s) .p.
(A)
10, (W]](V.s)
(Vo.r 2(0))==5 (0)
T (ar(e))
[W](2.0).
p.g
D
D2
-17 [W](k.s) (ar(c))
(>))a
P.g
D
[wn]ais) Nr(e]]
(s))
p.a
(0.r 130)) [avi6)) [W](..s)
Σ
(O.))..A
p.g
50

---

Published as a conference paper at ICLR 2025
D
DA
DA D
(0,):
[W +
ΣΣΣ
B [W]63
(A)=o (p),x) (A) (1)
p=
Da
D
Φ
(0.0:1 (A)=T (0)
(b(((1]
10.) (B)g
(B) (O.i):ik
9=1
q=1
0.
(170)
By choosing y be the identity element of Gu in Equation (170), and substituting the obtained result to Equation (184), we obtain a shorter formula for [E(W))
j.k as follow:
[E())
(0.)
(0.) (W] (0.i);p.
(171)
Thus, the constraint becomes
D.
10.0 (0.):: (e)
(AvtrI0)) -1
" -[W](O.(0))
P-
Dv
((NTTO))
Σ
(E( (0.r(0))
(1))
equivalently,
ΣΣ
D (0)
(MV())
[W]
(0)(())))
P.
Σ (NO))
D.
(E( 1
(172)
1=1
5.
Observe that, the left hand side of Equation (172) is just a linear function on the j-th row of (NO)) -1 while the left hand side is sum of linear combination of p-th rows of (N(1) -1 with
,,....,, D.. Therefore, only coefficients of [E(W)] with = are nonzero, i.e
[E()) Σ
$ 10.0g.. (0.0): j.
w (0.1)
(173)
In addition, since the right hand side of Equation (172) is independent of and To, the left hand side must remain unchange if we vary T and тo. As a consequence, we have
$ (0,+())))mo())) (o)
क (O.)):..k
(174)
(0.9):
for all i, k, T and TO. Next, by substituting Equation (173) and Equation (174) to Equa- tion (172), we obtain
Σ -=1
D
(0.00)))000))) (0.(0));:000)) (MNM) -1
-Im](o.))
39
Σ ((N1))
(
$ (o.)
b.g
or equivalently,
((N1)))
ΣΣ
1 (1
[W](0.))
D.
((vio) (o.i)
O.i.aL
[W](0.))
(175)
L= 1=1
By comparing corresponding coefficients both sides, we have
Φ (O) o (0.i)..
(176)
(O)(())))) (O.i)d.q*
51

---

Published as a conference paper at ICLR 2025
Finally, one verify directly that [B(W)]{( defined [E(W)) (1))
with p = are nonzero, i.e
(0.1) [E(W)]
11
10.0 (0.1) (0.0::4 1/
(177)
with the constraints
$ (0)
(178)
(0.0):44 (0.r(0)):' mo())
satisfy the constraint in Equation (172) for all j.j.'', k 9 and To.
Case 5. Solve the equation [E(gW]] = ()) Since [E(gW)] = (1E((W)(()) we
have
D
T
(A)AAA (QK (6))..
(WK)
(A)skk
(VO.r-1())) (00))) (n)
(1). [w](Q..) (Ar(ct) 77
77
(Q.-- (2
P.9
(
-1] [w]CKis) (Ar(e)
(K,r- (1)):n
$
(
[IN](Vis), NT6)
(t))
P-2
D
[(aves))
(1)ikk
(0,)
(0,--1 (0)) (0)
P.2
Da
DA
D
(1)iikk
(A)=T (p).mz (1)
[W][A]]
p.
+
(A)ii
(B) (0).g
[W]p
p.g
Da
Σ
(1)si&
(0((()
+ Σ
(1)
+ $ (A).
(A)=T (4)
(B)xg
q=1
D
D
(A)
[WW](OK.>)
(M (VO..):p
|
(VO.s) (VO,s) p.
D D
(1) (Q.s):p.g
(Q.0)
(A)
( pa
(K,s)
(K.a)pp.
p=
T
(A)) (M (V.+):p.
T
(Aa (0.00.)) (O.s):p.q
D Da
Da
(a)
(W](d)
+
(a)
p. (W (B)pp.
(A):p.q
p= g=1
p=1
Da
(4))
(a) (1))) + - ((a2)) (179)
+
(A)a
(B)g
1=1
1=1
52

---

Published as a conference paper at ICLR 2025
We observe that the left hand side a sum of linear functions on the matrices M (M)s) -1
.N
and (Nv+)) -1 while the right hand side is independent of these matrices. Therefore, according to Corollary D.3, the terms containing these matrices in the left hand side must be equal to zero. For the remaining terms, we can use Lemma D. to compare the corresponding ones with those in the right hand side and obtain the constraints:
$ (A):kkk
$ (QK,)):p
(QK,T-1()))p) (6)))pp
(A):..k
(A)TT())TTa())
(VO,r
(s) (q)
(VO.s):p.
$ (Q.s):p-
(A) (lk) 0,
(A)= (k) 0,
(K,)):pi
1 (A) (k) O, (V.a):p-
$ (0,0):pa
(A)
0,
(A):iA
$ (A..aa)
(A): (p) (a)
(A):p.
o (B): (p),9
(A):i.
(B)iia
(A):i.
= (A) (i).T(()
(A) (a)
(A)aa
4 (A))
1 (i).((()
(B)a
(B)1
$ (1)s = Φ (1)
Therefore,
(1)jkk
$ (A)
(QK,s)
$ (A)ijk
$ (A)55(())=2(8)
(VO,s):p.
(VO.r(e))n,mo((1)
o (A)st
O,
(Q,a)sp.
(A)sj.
O, (K,s):pp
& (A):j,
(V.s):p O,
(1)g.k 0,
(0,s)sp.
(A):f.A
Tb (A)
(A)rp.q
(A)=
$ (A):j.k
T (A)
(B)p..
(B)
(A):j.k
+ (A)
(A):q
(A):TA(1)
(A)g..
Tb (A):
(B)4
(B):9
$ (A):.k
s (A): (K).
One can verify directly that [E(W)]
(4) defined by
[E(W)
(QK,.): [WW](QK.))
(A)j&
[V (vo..)
(ADS&A
DA
DA D
(1)sk [W]
(A):2
...
(Bi:p.g
p= 1=1
Da
$ (A) (1(
(A)j.A (B)
(1
(B) 4 (A):AA
(180)
(A):
9=1
53

---

Published as a conference paper at ICLR 2025
with the constraints
(A)jkk
(A
(Q
(QK,
$ (A):i.
$ (A)
(VO,)):p.u
(VO:((e))p.m0()))
$ (A):j.
4 (A)
(A):p.
(A)
$ (A)g.k
& (A)
(B)p.g
(B)aa(p)a
$ (A):g
(A):j.k
T (A)
(A)ira(1)
$ (A):j.k
$ (A)
(B)-q
(B)a
$ (A):..
for all j.3', J.kKKKK K,, 8, s', P, p', 9.9. o. no, R satisfies the condition [E(gW)] = (oE(W)
Case 6: Solve the equation BB = (B (gE(W)) we also see that:
By using the similar argument as Case
[B(W) (B)
$
(Bbjk (QK..):p.q
(QK)
1p.4
p
(B)ik& (VO.s):p.g
[WW](VOs)) 1--2
pml
DA
DA
D
BBUA (A)ip.
(A)
p.g
+
$
(B)p.
(B)ik
(] (B)
D
(A)
+
(B)a + + (B)jA
(B)
(181)
(A)i9
with the constraints
(B)jk
$ (B)
PQK
(QK
$ J.
T (B):RA().k
(VO.s):P:
(VO,+(s)))p
$ (A):p.4
(B)jfk
T (B):
(A)
(B)j..
(B):
(B)PP4
(B):
(B):j.A
T (B)=FA().
(A):q
(A)=aA(T)
(B)j,k
+ (U).
(B)s
(B):
$ (B)skk
(BB):AAA)..
for j,k, 8, 4, TO, nA satisfies the constraint (F(gW)) (gF(W)) (
Case 7. Solve the equation [ = (18()))) By using the same argument as Case 5, we
also see that [E( (4)
defined by
[E()] (A)
ppp [W]](KK.))
(A)s
D
(1)ak
(VO.s) p.g
(VO.el:p.g
DA
DA
ΣΣ
(A)s
(A)
4 (A)sk
[W]
(A).p.
.p.g
[B):p
p=19=1
p=1y=1
54

---

Published as a conference paper at ICLR 2025
DA
(A)a
(4)
(B) 1i + ((A);
(182)
with the constraints
4 (4)ak
1 (k)
(QK,T())):p-
(A)a
$ (A)
(VO,s)pp
(V)==(s))=
$ (A):k
(k)
(A)ppo
(A)
(A):k
+ (k)
(B)>P-
(B)=FA()).
(A):
$ (A):
(A):g
(A)=a(()
$ (A):k (A)
(B)- (B)a
4 (A):k ((A)=a(()
for all s, TO, TA, satisfies the constraint (E(gW)] loE(W))
Case 8. Solve the equation
[ [gE(()) By using the same argument as in Case 5,
see that [E(b)]{ defined by
[bb) (
(B)a (Q.
[WW]
D
D
$
(B)+k
[WW]
(VO.s):p.g
DA
DA D
$
(B):k
(W](1) (A):p.g
ΣΣ
$ BBik
[W]
..g.
(B):p.
p.
1=
p=1y=1
$ (Bb4
+ (B)
(183)
(A)sg
BB)
with the constraints
$ (B)s
(B):k
(QK,*)
(QK,r()))pp
② (B)ak
(B):
(VO,s):
(V)
(B):k
+ (B)sk
(A):p.4
(A)=
(B):k
+ BB)k
(B):p.q
(B)
(B):k + (B):k
(A)xg
(A):Ta())
(B)ak + (B)ck
(B):
(B)a (B)ak
for all j.k, 8,P, 1, T, AA, To, satisfies the constraint [F(gW)) (OE(W)
D.4 FINAL FORMULA FOR THE EQUIVARIANT POLYNOMIAL LAYER
We gather the results of the computations of the equivariant layer in the previous section here for further implementation.
Theorem D.4. The map E: u U with
E(U) = (1EBW)))Q0) (B())))(,,,,,W)(((.),.,), ( (E())(O.))
(184)
55

---

Published as a conference paper at ICLR 2025
is Gu-equivariant.
Here, the components of EUU) are given below.
1. [E(W)](Q8) is defined as
D
[B(W)
Σ
(Q.0 ( (Q.).
(Q.)., p.k
p=1
with the constraints
$ (Q.0):t T
(Q7(0)): (Q)(k)))p.))
(Q.)):p.k
2. (B(W)](K)) is defined as
[B()
(K.
$ (k.jjjj
(W)sk
(K.i))p.kk p.k
with the constraints
$ (K.i):p.k
(K.)).kk
D (KTU)))
(0)2k))
3. [V()) is defined by
(F(W))
(V.) = Σ
T (V.0.s. (W))
(V.0p..k
with the constraints
(V
$ (V.rl)))
(V.)):..k
(Vir()))p.))
4. [F(W)](O.) is defined by
D
[(()
$ (O.)) (0.0)
j.k
344
q=1
with the construints
T (0.):..kk
$ (o)
(0.1):4
(0,
5. [E(W)](A is defined by
h
$ (A)i&
[W]](K.e))
(Q.
1p4
D
(
[WW)(VO.a) 1.g
(Vo.
DA
DA D
(A)
(A)ino
(A)
+
(A).j.
(B):p.q
DA
+
(A).
(A)
+
(A) (B) + ((A):,. (B)x
(A)e
1=1
9=1
with the constraints
T (Abikk
$ (A)
(QK,5):p
(QK,r(+)):0g
(A):.k
4 (A)
(VO.s):P.
(VO.r(s))2.x000))
4 (A)j..
$ (A)
(A):p.
(A):Fo(P).TA()))
$ (4):j.
Φ (A)=
(B):p.q
(B)sTA(P).
$ (A):
(1)
(A):
(A)IxA(1)
$ (A)
& (1)
(B)4
(B))
$ (A)..k & (A)))
56

---

Published as a conference paper at ICLR 2025
6. [E(W)]() is defined by
h [F(W)
D
$
(B).k (QK,a):p.g [W])(kks) P.
D
D
Σ
(B)
(VO,.):... WW],VO,s
DA
Da D
(
(A)
(A)ipa
(B)iaa
Da
Σ
(B)s
(b((( (A)a
+
(B)sk
(B)a D
B) + 2 (B)akk
9-1
with the constraints
(B)ikk
(B):ma))..
(QK
((K)
(B):j.
(B):TA())
(VO,
(VO.7(e)):pP000))
(BB):A(()).
(A):pa
(A)
(B)s.
(B).TA))..
(B)p-
(B):TA(P);
(B)jj.
$ (B)=AA)).
(A)-g
(A)iTala)
$ BB
$ (B):Tali).k
(B)a
(B):9
$ (B)jkk
(BB:
7. [E(b)](A) is defined by
D
[E()){( A)
$ (QK..):p.
(A):k
(KK)Q PM
D
D
$
(4)sk
WW] (VO.s)
(VO.s):p.
p.
DA
(1(4) (A)=p.
(A)ck
p.
+
$ (A)sk
(B)=p.0
[W]68]
p=1g=1
DA
D
$ (1An (4)
11 (8) + の (A):
(A):
(B):g
-=
with the constraints
1 (A)c
(K)
(QK,))
(QK,r(s))sp
T (A):k
T (K)
(VO.s):P.
(A)csk
2 (A)
(A)=p4
(A)
(A):k
$ (A)
(B)P-4
(B):
$ (A)ag
(4)ck
T (A)
(A):Fa(O)
$ (A)sk
$ (B)-
(A)
(B):a
$ (A):
(())
8. [E(b)(()) is defined by
[E66]] (B)
ΣΣΣ
$ (B):
[WW](QK.)) (QK.))
(QK,s) p.
57

---

Published as a conference paper at ICLR 2025
D
D
Σ
&
BBak
[W]]
(VO.s):p.g
DA
DA D
(B)* (Al:p.g
[W]
+
ΣΣ
(B): (B):p.q [W]s
q=1
p=1==1
+
( (1)
$
11 $ (B)b
(A):q
with the constraints
① (B):k
$ (B)ck
(QK
(QK,7(5)):
$ (B).k
(B):
(VO,))
(VO.r())))
$ (B):
T (B)b
(A):p.4
(A)
$ (B)b
$ (B)b
(B)p.4
(B)=
$ (B)
$ (B):
(A)=g
(A):
$ (B)a
(B)ak
T (B):k
(B)a7
(BB)k 4 (B):k
for all 3.J', , Ji, ke, 8, s', ,..''s 4444 and T; no: T.A
E AN INVARIANT POLYNOMIAL LAYER
In this section, we will construct a map 1: u - D such that is equivariant with respect to Gus ie.
1(gU)= = I(U):
(185)
for all U E and g E Gu-
E. A GENERAL FORM WITH UNKNOWN COEFFICIENTS
We select I(U) to be linear combination of entries of the matrices [W](Q.>), [W](K,)) [W](v.),
[W](O.0), (W](A), [W] (b)4), (B)() as well as the matrices (WW](Q,,>) and [WW](vo.s) ie
D
I(U) Σ
ΣΣX
4 ((2 [WK](QK.) P-9
p=
(W]]VV,,) p.2
((VO,s)
D
D2
[W.) +
4 (KC,)))p (K,s)
(Q
D
D.
D. ΣΣΣD
D
D (V..))p
$ (0.)):
[W]
p.q
l..
1=1
==1 p=1=1
D
DA
DA
((A):p.
[W]
Σ
$ DBB).pg (W](B)
p=1
p=1ę=1
D
(
+
Σ
$ (B)
(b) + +1.
(186)
9=1
where the coefficients O are matrices of size D' 1.
In the following, we will determine constraints on the coefficients such that will be Gt invariant under these constraints.
58

---

Published as a conference paper at ICLR 2025
E.2
COMPUTE [g)))
By using the same arguments in the computation of ElgU), we have
D D
I(gh) = Σ
((ok,- (s))=m) (KK.))
P.
p=1
[WW] (V))
(4)
(Q. (c)) [w](a..) (Ar(s))
[W](K.)) (MS(1))
(
P.9
D
[wn] N(0)
P(V.--
p.
(o0r-"0)):00' (NN.)) (w)(
P-
DA
$ (A):=o (p)an (A))r='(())T (v) (W +
(B)x (p). pp
(B)
DA
+
+ (A) (4) (1) (1)))
(B) + $. + Σ $ 11
(187)
q=1
E.3 COMPARE T(gU) AND T(U)
According to Lemma D.1, in order to have [ggU) T(U), the corresponding coefficients of (gUU) and IUU) must be equal. As a consequence, we have
(QK, 16)))pp 4 ()
$ ((vor-14)) T (VO.s):
((0
((x,r-(0)):p
(():r-1()))-
00, (4)))/0)) (4)
4 (p), 4 (A):pa
( (p).g (BB)p
① (A))xT (g) ((A)a
((B).g 4(B):
Φ $1
or equivalently,
$ (QKK))) = & (QK,T(H))pP
(VO.s)p (VO,r(e))pp004))
(QQ.))= 0
(KK,)) 0
(V..))
59

---

Published as a conference paper at ICLR 2025
(0,)):p 0
0 (A):p. 0 (4)
& (B)pr. = (B)EFA())
((4)² (A))TA(4)
BB) (B):v
$ 11-
E.4 FINAL FORMULA FOR THE INVARIANT POLYNOMIAL LAYER
We summarize the above computation here for implementation.
Theorem E.1. The map I: U - D defined by
h
D
1(U) = Σ
((KK)) [WW]
D
P
[WW](VO.)) P.
q=1
DA
DA
(A):p.
[W]
Σ
+ (B)pp
[W]6
p:
+
Σ
(A)= (b](4) +
Σ
4 (B)g
[b)( +1
(188)
q=1
with the constraints
$ (QKK))pp = (QKKT())):p4
(VV..)):P = ([VO,=(s))sp
& (A):P.4 = (1):
((B)pP.4 ((B):AA())
( (AA:=A(())
(B) (B)):
is Gu-equivariant.
F COMPUTATION COMPLEXITY OF EQUIVARIANT AND INVARIANT LAYERS
The computational complexity of the Transformer-NFN is derived from its invariant and equivariant layers as outlined in their pseudocode (Appendix 1.2, I.3):
Equivariant Layer Complexity: O(d e h D2 max(Dy. Dx. D.)) Invariant Layer Complexity: O(d D' D max(D h, DA))
Where the parameters follow our notation in Table
The Transformer-NFT implementation leverages optimized tensor contraction operations (c.g., einsum), enabling efficient and highly parallelizable computations on modern GPUs. This ensures computational practicality while delivering significant performance improvements.
G ADDITIONAL DATASET DETAILS
To create wide range of transformer model, we opt to vary six hyperparameters in our experiments: train fraction, optimizer (SGD, SGDm, Adam, or RMSprop), learning rate, L2 regularization coeffi- cient, weight initialization standard deviation, and dropout probability. The train fraction determines the proportion the original training dataset used, while the optimizer dictates the algorithm for parameter updates. Learning rate, I.2 regularization, and weight initialization standard deviation
60

---

Published as a conference paper at ICLR 2025
Epach
Epech
10h
109
109
10:
1.0
Figure 1: Accuracy histogram of MNIST task and AGNews task in the Small Transformer Zoo. The number of samples is showed in log scale for improved visibilityy
control various aspects of the training process, and dropout probability helps in preventing overfit- ting.
We select a range of typical values for each hyperparameter independently, then generate all pos- sible combinations to form our set of hyperparameter configurations. During our preliminary periments, we observed that the optimal range of hyperparameters varies significantly between op- timizer types. Consequently,
settings into categories: Adam-RMSprop and SGD-SGDm. Table 4 provides a detailed overview of our hyperparameter configurations. Overall, there are 8000 configurations for each category, resulting in 16000 configurations in total. These configurations are consistently applied across both tasks to ensure comparability. All models are trained for 100 epochs. with checkpoints and accuracy measurements recorded at epochs 50, 75, 100, and at the epoch with the best accuracy. During the process, we eliminate any runs that crash.
divided
Table 4: Hyperparameter configurations of the Small Transformer Zoo Dataset
Hyperparameter
SGD-SGDm
Adam-RMSprop
Train Fraction
[1.0, 0.9, 0.8, 0.7]
[1.0, 0.9, 0.8, 0.7]
Dropout
[0.2, 0.15, 0.1, 0.05, 0]
[0.2, 0.15, 0.1, 0.05, 0]
7e-3,
[3e-4, 5e-4, le-3, 3e-3,
Learning Rate
3e-2, Se-2, 7c-2]
5e-3, le-2, 3e-2, 5e-2]
Weight Init Standard Deviation 1.2 Regularization
0..1,..15...2 0.25, 0.3]
[0.1, 0.2, 0.3, 0.4, 0.5]
[1e-8, 1e-7, 1e-6, le-4, 1e-2]
[1e-8, 1e-7, 1e-6, 1e-4, 1e-2]
NNISST-Transfommrrs. The MNIST dataset (LeCun Cortes, 2005) is widely-used benchmark in computer vision, consisting of 28 x 28 pixel grayscale images of handwritten digits (0-9). For this vision task, the model objective to perform digit classification. The embedding component first applies 2D convolution to encode image patches and then adds fixed positional encoding to provide spatial information. The encoder, comprising two stacked transformer blocks, processes these embeddings to capture complex relationships between different parts of the image. The clas- sifier applies global average pooling to the encoder's output and then passes it through two fully connected layers with ReLU activation in between, finally outputting probabilities for ten classes corresponding to the digits 0-9. Using our hyperparameter configurations, we generate a total of 62756 model samples for the MNIST task, including 15689 checkpoints epochs 50. 75, 100, and the best persfrrerorming epoch, each trained with a distinct combination of hyperparameters. The accu- racy distribution for MNIST in Figure displays strong concentration in the 80% to 100% range
61

---

Published as a conference paper at ICLR 2025
Table Ablation study on varying the hidden dimension and number of layers in Transformer-NFNY trained on the AGNews- TaaNess-raanfrrmen dataset. Dimensions of all equivariant layers are indicated inside square brackets.
Transfomer-NFF
[10]
[15]
[10,
Kendall's r Num. params
0.007 0.002
0.900 :: 0.002
0.900 =: 0.001
0.913 4 0.001
0,W05 4 0.003
0.911 0.002
0.913 0.00
0.013 0.002
0.491M
0-840M
1.700M
2.857M
1.901M
4.757M
17.467M
38. 38.101M
Table 6: Number of parameters for all models
Model
MNIST
AGNews
Transformer-NFN
1.812M
1.804M
MLP
0.933M
0.896M
STATNN
0.203M
0.168M
while remaining models distribute almost uniformly across lower accuracies, with a slight increase around 10%.
The AG's News Topic Classification Dataset (Zhang et al., 2015) is a collection of news articles from the AG's corpus of news articles on the web, categorized into four classes: World, Sports, Business, and Sci/Tech. We take the description of the articles input and train models predict its corresponding topic. Our transformer-based model adapted for this task as follows: The embedding component uses a pre-trained Word2 model to map each token an embedding vector. These embeddings are combined with fixed positional encodings to retain sequential information. The encoder, consisting of two stacked transformer blocks, processes these embeddings to capture contextual relationships within the text. The classifier applies global average pooling to the encoder's output, then passes it through two fully connected layers with a ReLU ac- tivation in between, finally outputting probabilities for the four categories. Our experiments on this task yield diverse set of 63796 model checkpoints, derived from 15949 unique model configura- tions. These checkpoints are collected at four key points during training: epochs 50, 75, 100, and at the epoch of peak performance. The accuracy distribution for AGNews in Figure shows notable concentration in the 50% to 90% range, with a peak around 80% while also exhibiting a smaller cluster of models around 25%.
CGGNew--Transforee
H ADDITIONAL EXPERIMENTAL DETAILS
H.1 GENERAL DETAILS
Training details The models were trained for total of 50 epochs, using a batch size of 16. We employed the Adam optimizer with a maximum learning rate of 10-3, A linear warmup strategy was applied the learning rate, spanning the initial 10 epochs for gradual warmup. We utilize Binary Cross Entropy for the loss function.
Number of parameters Table 6 summarizes the total parameter count for all models. The archi- tectural details and hyperparameter configurations for each model are provided in H.2 and H.3. Importantly, we optimized the baseline models to their best configurations, and further increasing the number of parameters would likely result in overfitting.
H.2 ARCHIECTURE AND HYPERPARAMETERS OF TRANSFORMER-NER
The Transformer-Nui model is composed of three main components designed to handle the input weights of transformer network. For the embedding and classifier components, both of which are MLP-based, we employ standard MLP with ReLU activation to process individual component. The transformer block itself handled by an invariant architecture, incorporating several equivariant polynomial layers of Transformer-NFN. These layers utilize ReLU activation, which specifically operates on the two MLP components within the transformer block. Following this, the output is passed through invariant polynomial layer of Transforme-NNFN. The outputs of each of these components are represented as vectors, which are concatenated and passed through a final MLP head with Sigmoid activation for prediction.
62

---

Published as a conference paper at ICLR 2025
1.0
0.91
0.84
r = 0.88
0.8
0.6
0.4
Predicted accuracy
0.2
0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
0.2 0.4 0.6
0.8 1.0
Transformer-NFN
STATNN
MLP
1.0
T = 0.83
0.77
T 0.86
0.8
0.6
0.4
0.2
0
0.6 XGBoost
LightGBM
Random Forest
Test accuracy
Figure 2: Visualization of all models on test set of AGNews-Transtomerrs dataset.
In our experimental setup. the embedding component is modeled using a single-layer MLP with 10 hidden neurons, while the classifier component two-layer MIP, each layer containing 10 hidden neurons. For the invariant architecture of Transformer-NHY. we apply an equivariant polynomial layer of Transformer-NFN with 10 hidden channels to process the weights, followed by invari- ant polynomial layer of Transformer-Nei NFN and an MLP that outputs 10-dimensional vector. The resulting vectors from these components are concatenated and passed through single classification head to generate the final prediction.
H.3 ARCHITECTURE AND HYPERPARAMETERS FOR OTHER BASELINES
Here we describe the architecture of all baselines:
MLP the MLP model, the weights of all components are first flattened and processed through separate MLP for each component. Specifically, the MLP handling the trans- former block and embedding components consists of a single layer with 50 hidden neurons, while the classifier component is modeled using two-layer MLP, each layer containing 50 neurons. The outputs from all components are concatenated and passed through a final MLP to produce the prediction.
STATNN (Unterthiner et al., 2020) For the STATNN model, we adapt the original approach to work with the transformer block. In this case, we compute statistical features from the weights of the query, key, value, output, as well as weights and biases of the two linear layers. These features are concatenated and passed through single-layer MI. with 256 hidden neurons. The classifier component uses the original STATNN model to extract features, which are then processed through an MLL with 256 hidden neurons. For the embedding component, a single-layer MLP with 64 hidden neurons is employed. The outputs from all components are concatenated and passed through single-layer MLP for the final prediction.
XGBoost (Chen & Guestrin, 2016), LightGBM (Ke et al., 2017), Random Forest (Breiman, 2001): We flattened the weights of all components and directly employed popu- lar gradient boosting libraries for regression. The hyperparameters for all three tree-based models were maximum depth of 10, minimum child weight of 50, and maximum of 256 leaves.
63

---

Published as a conference paper at ICLR 2025
H.4 EXPERIMENT ON AUGMENTED AGNEWS- TRANSFORRERS DATASET
Experiment Setup In this experiment, we evaluate the performance of Transformer-Ne on the AGNews-Transformer dataset augmented using the group action Gt. We perform a 2-fold augmen- tation for both the train and test sets. The original weights are retained, and additional weights are constructed by applying permutations and scaling transformations transformer modules. The elements M and (see Section 4.3) are uniformly sampled across [-1,1], [-10, 10), and [-100, 100].
The results for all models are presented in Table The findings demonstrate that Transformer-FFy maintains stable Kendall's across different ranges of scale operators. Notably, as the weights are augmented, the performance of Transformer-NFN improves. whereas the per- formance 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.
Results
Table Performance measured by Kendall's 7 of all models on augmented AGNews-Transionn dataset using the group action Gu. Uncertainties indicate standard error over 5 runs.
Original
1, 1]
[-10, 10]
[-100, 100]
XGBoost
0.859 + 0.001
0.799 + 0.003
0.800 + 0.001
0.802 + 0.003
LightGBM
0.835 0.001
0.785 + 0.003
0.784 0.003
0.786 0.004
Random Forest
0.774 = 0.003
0.714
0.715 : 0.002
0.716 + 0.002
MLP
0.879 0.006
0.830 + 0.002
0.833 0.002
0.833 + 0.005
STATNN
0.841 = 0.002
0.793 = 0.003
0.791 = 0.003
0.771 0.013
Transformer-NEr
0.910 - 0.001
0.912 + 0.001
0.912 = 0.002
0.913 - 0.001
Gap
0.031
0.082
0.079
0.080
The general decline observed in baseline models highlights their lack of symmetry. In contrast, Transformer-NFN. equivariance to Gu ensures stability and even slight performance improvements under augmentation. This underscores the importance of defining the maximal symmetric group, to which Transformer-NEr equivariant, overcoming the limitations of baseline methods.
IMPLEMENTATION OF EQUIVARIANT AND INVARIANT LAYERS
We present the multi-channel implementations of the Gu-equivariant map E: 2d - 2 and the Gu-invariant map I: 2d - ex We summarize the Equivariant and Invariant Layers with the bullet notation and adopt einops-like pseudocode to maintain consistency and standardization in transformer weight space manipulations.
To facilitate understanding of the implementation, we summarize the key dimensions involved in Table 8 and define the shapes of the input terms in Table 9.
I.1 SUMMARY OF EQUIVARIANT AND INVARIANT LAYERS
We summarize our derived Equivariant and Invariant Layers from Appendix D and Appendix E. We use the bullet notation simlify the notation, making it easier implementation. Roughly speaking, the bullet stands for "the parameter is unchanged when the bullet varies across all possible index values"..
I.1.1 EQUIVARIANT LAYERS WITH BULLET NOTATION
[E(W]]
[W] (Q+)p
j.k
[E(W)]
(K:t)
(K:))
(K:)
(K:.):p .p.9
[E(W)]
(V:0)
(V:
(V:)
(V:o)p..
[E(W)
(0:0):..
(0:0)...) Σ
[w) + (0:)):.)))
k=
64

---

Published as a conference paper at ICLR 2025
Table 8: Summary of key dimensions involved in the implementation
Symbol
Description
Number of input channels for the equivariant and invariant layer Number of output channels for the equivariant and invariant layer Embedding dimension of the input and output sequences of the transformer block Embedding dimension for key and query vectors in the transformer block Embedding dimension for value vectors the transformer block Embedding dimension for the linear projection step in the transformer block Number of attention heads in the transformer block Batch size
D
DR
Da
Do
DA
b
Embedding dimension of the invariant layer's output
Table 9: Shapes of imput terms used in the implementation
Term
Shape
[ (Q.0
[6, d. h, D, D
le-
WILK (K.))
(b,d h, D, DE
[b, D, De
[b, h, Du, D
[W] (KK() (QK,))
(,,d D, D
p.4
WW] (VO.))
b, D, DI
PP
W
D.
P-
bI
[6,d, DA]
W
16,d, DA. [6, DJ
b
(E(W) (4)
Σ
+ (A):
Σ
Dwwwjokk. p.
()p.
(A):
Σ W [WW](OOs))
(VO..):p.))
h Σ (ww)yVo.s)
c (A)
(VO.)):......
T (1))):))))) (A):o.-)
(W](A)
Σ
(())::.))) (A):0.)
(W7(2)
P.
9=1
(1):o.))
(w7(1)
(W7,02)
(A):o.))
DA
(B):::)
(W)B)
( (B).
(W)b)) (B)
kk.g
② (1)) ) [b](4) + (A): )₂ [b](21
(1):.)))))
Σ
$ (A):
(6 (B) (A):...
(Bi:g
65

---

Published as a conference paper at ICLR 2025
D
D
h
LE
(B8)
BB
(IW]]OKk.))
(QK..):p.
p-l==1== D
a=
D
h
$ (B))
(WW)/O0.)) P.0
(VO.-):p.*
1=
T (B)) (A):0..)
[W]b)]
T (8):
[
1..9
(A):
pjj
p=l
DA D
(B)))
T (B)
[W]/6
(B))
(B](4) +
(B)
(A):
/
Φ (B) (B)3g
[61 (B))
[E(6))44)
1 (A):
Σ
WW QK.s) P.
D
h
(4):
Σ
[WW [wwWJvoO. (VO.))
(VO.o)pp.
(1)):a.)
pig
+
(A):
pp.
p=1
Φ (A):e
[W]m
Σ
1 (B)):a))) [ (B)
(B):..4)
q=1
b (A)::) (A):
(b](4) + (d)):
[b)(2)
(A):
(A): (B) + (A):0 (B):g ln
+
1
D
D
(E())
(B)
$
BB)*
(IWW p.o
(Q
11
D Σ
$
:B)&
(WW] l.
(VO..)...
DA
Da
(B)ck
[W](A)
(B)kk (B):.g
[W][2]
(A)
..g
p=1 g=1
D
(B)* (A)
(A))
(B) + 4 Byak
(Byg
L.1.2 INVARIANT LAYERS WITH BULLET NOTATION
D ((V) ΣΣ p=19=1
D
h
2 (QKKNpp
66

---

Published as a conference paper at ICLR 2025
A
D
D
Σ [WW]vo p.g
& (vO..):p
g=
s=1
D
D D
2 (A)):.
[W]S
Σ
2 (B)
[W]ba
pl
p=1
q=1
(A):
(b]
Σ (B)a
[b][8] 4
I.2 EQUIVARIANT LAYERS PSEUDOCODE
1.2.1
[E(W))
PSEUDOCODE
From the formula:
[E(W)
(Q:) 4b (Q:)):
[W]
(Q:)).p.
We define the pseudocode for each term:
For + (Q:.):p.*
(Q:
[W][Q:) with [W][2-)) of shape [b, d, h, D, Dy] and o (Q:a):
(Q:+):p of shape [e, de D,
Corresponding pseudocode: e einsum(bdlpk;; edjy - behik)
1.2.2
[B(W)]
(K:))
PSEUDOCODE
From the formula:
[E(W)) (K:)).... [W]KKs j.k (K:.):p.e
1.9
We define the pseudocode for each term:
For o (K:e)...
[W][Ksa)
with (K*( of shape [6, c, h, D, Du] and & (K:e)s -p.g
of shape [e, d, D, D]
(K:.):p.
p.g
(K:.):p
Corresponding pseudocode: e einsum(bdlpk,, edjp behjk)
1,2.3
[E(W)
(V:0) PSEUDOCODE
From the formula:
[E(W)) & (: (V.p.
We define the pseudocode for each term:
For t (V.
[W](V:), with [W][V of shape [b, d, h, D. Do] and o6 (V:))i
(Vio) of shape [e, d, D, D]
(V:o).p.
Corresponding pseudocode: e einsum(bdhpk; edjp behjk)
1.2.4
[E(W)]
PSEUDOCODE
From the formula:
Eumbo 1i + + (0)))): [W1
[E(W)] = (0:)):))))
k=
We define the pseudocode for each term:
Eiw (Oi) with (W](00) of shape (b, d, Dr, Dj and (o) of shape [e,
For
Corresponding pseudocode: einsum(bdhjj,,d ed behj). unsqueeze(-1)
For ② (0:)))))) [W]]0:) with (W]00 of shape [b, d, h. D.. D) and p (00)))))))) of shape [c, d]
Corresponding pseudocode: e nsum(bdhjjk, ed behjk)
67

---

Published as a conference paper at ICLR 2025
[B(W) (A)
1.2.5
PSEUDOCODE
From the formula:
h
(F(W)
$ (A)e
)Q...
9=1
T (A):o.
ww [ww
(vO.0)
4 Aiio )VO..):p.o)
[wwwrrro pitt
D DA
DA
① (A)io:))
[W]La]] P.g
+
$ (Ax0::)
[W](1]]
9=1
db (A)iee))
[W]bat]
(b) (A):
[W111t]
(A)):)))))
1
DA
(B))0.0)
[W]60]]
+ Σ
1 (A):0..)
(B): [W]ke] (B)
(A):o:)))
(A): ), (M][4) + ( (A)=
[b)(4)
+
q=1
$ (A):
()()) + ((A):,.
+
(B)g
q=1
We define the pseudocode for each term:
(A) For +
Σ[WWJOK.s) with [W]](Kke)) of shape [b, h, D, DJ and (A)
(QK
8=1
of shape [e, d, DJ Corresponding pseudocode: einsum(bdhpp, edpg - be). uesqueeeezze))) (-nuuuuuueeueeeeee
( (1)=... (VV...)..)
[www]ro..) p.g with [WW]]VO. of shape [b, d, h, D, DJ and
For
p.g
( ( (V of shape (e, d, D]
Corresponding pseudocode: e cinsum(bdhppp edp - be) unsquuczee- uinuuuuucee
( (A) (VO.)))))
ZwwT (Vo.s) with [WW] pp
(VO,s) of shape [b, d, h, D, D) and
For
==1 pj
② (A): of shape D)
(vo..)).))
Corresponding pseudocode: einsum(bdhpjj, edp bew))..nsqueezeee1111
For
@ (o))):)))
[W]64] with [W](4) of shape [6,d, D, DA) and
(4) (o(d)::)) of shape [e, df
Corresponding pseudocode: cinsumbbdpp ed - be). (onsuueeee--1..unssee
④ (A)))))) (A):...)
[W]]6]] with [W](4) of shape d, D, DAI and
( (2)):)) of shape (e, df
For
Corresponding pseudocode: einsum(bdjg, ed bej) unssuucee((-1)
For T (A))))))
(W/(2)
pp.k with (W)(4) of shape db D. DAj and
(T (d))):))
of shape (e, d]
(A):..))
3
Corresponding pseudocode: e innummbbppl ed Juns)uuusueeeee
68

---

Published as a conference paper at ICLR 2025
( (A)))))
[WIgh]] with [W1kk shape [6, D, DAl and (o (A))):)) of shape (e, df
For
Corresponding pseudocode: einsum(bdjj, ed - bejk)
① (4))))))) (B)
[W] with [Wipa]] of shape [b, d, A DJ and (o (B())ie))) of shape [e, d, DJ
For
p.g
.ig
Corresponding pseudocode: einsummbdpy edq - be) junsqueezeeeeeee uum.uuuueee
For
pb (A))))))) (B):..)
[Wike] with [W]ke] of shape [b. d, A D) and ( (B):0.0)
of shape [e, d, D
kig
Corresponding pseudocode: e rinsummbdkie edq - bek) uusqueeze((2)
(A))..)) (A):
[b] with (b([A) of shape 1b, d, DA] and (A)))) of shape [e, d]
For
Corresponding pseudocode: einsum(bdg ed - be) www.uuueeeeee)i.eemmooo (usqueezee-
For
(p (A)))
(A):
[b][[1 with [b] of shape (b, d. DA] and
(A)=+) <
of shape [e, d]
Corresponding pseudocode: einsum(bdl: ed - bek) unsqueeze(-[]
For 4
(A):..
(B)= [6], with (b)(7) of shape [b, d, Dj and 4b (B)3 shape [e, d. D]
Corresponding pseudocode: e einsum(bdlg edq - be). wttp://eeeee--l..uueceeeee
For +( with shape [e]
Corresponding pseudocode: einsum(e - e).
1,2.6 [E()){( PSEUDOCODE
From the formula:
D ΣV
D
h Σ Eww P.9
[E(b)]
(A)
(QK kpp.
1=
(A):
WWW (VO,s) 10.9
(VO.*) Jop.
D
((())ee))
(W(a) +
0 (A):
[w](1)
(A):0.0)
T (AA): (A))e..))
[W](B)
+
T (A):e.))) [W]ke]
P.
-=1
(A)=) [b((4)
(A):
(1)
(A):
(A)= B)ng B8( + ((A):
We define the pseudocode for each term:
For T
(A): of shape [e, d, D] Corresponding pseudocode: ei einsum(bdhpp, edpq - be) unsqueeze(-1)
Σ (WW]Skk)) with (WW)(OK.)) of shape [b. d, D, D) and o6 (A): p.
(QK.
(QK.)):
h Iiwn b-g
For o (A): (p.
with [W of shape (b, d. h. D, DJ and (A)
(Vo.a)p.
of shape [e, d, D] Corresponding pseudocode: e nsum(bdhpp, edp be) unsqueeze(-1)
69

---

Published as a conference paper at ICLR 2025
For
(((a)ee)); [W]]
p.g' with [WTpg of shape d, D, DA] and (o(4)) (o(a)o..), of shape [e, d]
lig
Corresponding pseudocode: einsum(bdpg, - be) unsqueeze(-I)
(A()))))) [W1/hh with [WT, of shape [b, d, DAl and (o(4)o..) of shape [e, d]
For
p.k
Corresponding pseudocode: einsum(bdple ed - bek)
For
(b))) Corresponding pseudocode: e nsum(bdpp, edq - be) unsqueezee-1)
(B))).))) [W]be), with [W]]9] of shape [b. d, DA, DJ and (B) of shape [c, d, D]
p.y
ld)e (B):..)) [W]k.]] with [W]ke] of shape [b, d, DA, DJ and (B)+, of shape (e, d, DJ
For
Corresponding pseudocode: einsum(bdkg, edq - bek)
(gld)e (A)c [6141, with [b][4) of shape [6, d. DA] and
( (A))=) of shape (e, d]
For
Corresponding pseudocode: einsum(bdg, ed - be)
unsqueeze(-1)
(
(1)= ( with (b)[")) of shape [b, d. DA] and (Φ) (A):a)
For
(A):
of shape (e, d]
(A):
Corresponding pseudocode: einsum(bel,, ed - bek)
For 4
(4)::
(B):a [6] with (b](8) of shape [b, d, D) and Φ (A): of shape [e, d, D]
(B):
Corresponding pseudocode: einsum edq be) unsqueeze(-1)
For +(A)) of shape (c).
Corresponding pseudocode: e einsume e) [nssueerr((.uussqueeea(111
1,2.7
[E(W),
(B) PSEUDOCODE
From the formula:
D
D
h
[EWW
11
ΣΣΣ
(B) (wWW)Ok..))
(QKK.):p.0
1=
(B):o.k (VO..)pp..
① (d)t..)) (W]64]
1 (B):,k)
(A):0,0) [W](1) (4)
+
(B):,.))
[W]{R]]
( (B): [W]
(B):
$ (B:ckk
(A): ) [b](4)
( (B): .k)
(A): ), [1)54]
q=1
Σ
Th (B) (B) + (B):o.k (B):
We define the pseudocode for each term: For Tb B (QK,))sp.) [WW(QQKa)) with (QK,) of shape [b, d, D, DJ and o (B) p.y
p-1
(Q
of shape [e, d, D. D. D. D]
Corresponding pseudocode: einsum(bdhpp, edkpy - bek) .unsqueeze(-)
For 4
( of shape [e, D.D.D)
(vo..)pp.* [VW])VO.>) with [WW](VO.s) of shape [6, d. h, D, DJ and o (B)
(vO,,))
Corresponding pseudocode: e nsuu(bdhhpg edkp bek) uusqueezze(22
70

---

Published as a conference paper at ICLR 2025
@ (d)x,,); [W]p4] with [W]p4] of shape [6,d, D, D. and (o (B(B))))) of shape (e, d, D]
For
l.g
Corresponding pseudocode: einsum(bdpg, edk bek) unsqueeze(-2)
@ (B)))) (A)x.,) 2
[W],3, with [W16] of shape [6, d. D, Da] and (M)))))))
For
of shape (e, d, DJ
Corresponding pseudocode: einsum(bdpjj edk - bejk)
For
( (B)c))) (B):
[W]pe] with [W]s
p.g of shape [b, d. AA Dj and (B):
of shape (e, D,
Corresponding pseudocode: e nsuu(bdpp, edkq - bek) unsqueeze(22)
( For
(B):..g [W](e]] with [WI[e of shape [b, d. A DJ and (o (B(
of shape [e,d,
Corresponding pseudocode: einsum(bdjg, edkq - bejk)
( ); [b],") with [b][4) of shape [b, d. DA] and
( (
) of shape [e, d, D]
For
Corresponding pseudocode: einsum(bdg, edk - bek) (unsqueeze(22) (B)))) ( with [b)[A) of shape [b, d, Da) and
For
( (B)a))))
(A):
of shape (e, d, D)
Corresponding pseudocode: einsum(bejj edk - bejk)
For 4 (B):a
[6] with (b) shape [b, d, D] and o (B shape [e, d. D.
(B)g
Corresponding pseudocode: einsum(bdg, edkg bek) unsqueeze(-2)
For of shape [e, D),
Corresponding pseudocode: e einsum(ek - ek) .unsqueeze())
1.2.8
(B) [E(D)]
PSEUDOCODE
From the formula:
E(D) (B)
$
(B)b (QK
(WW])eKK.)
(B)b
[WW]
(VO..)pp.
DA
DA D
(B):k
(1)
$ (B)k
[W]6 (B):0.g
(A):
..
9=1
4=1
(B)
We define the pseudocode for each term: For T (B)* (QK, . )p.- (WW)(KK.) with (QK)QKK.) of shape [b, d, h, D, DJ and कb (B)sk p.g
(B)a
p..
(QK
of shape [e, d, D,
Corresponding pseudocode: einsum(bdhpp, edkpq - bek)
For & (B) pp..")pp.* [WW]
with [WW](VO.>> of shape (b. d, h D, DJ and & (B)*
(vO.))..
of shape [e,d, d, D, D)
Corresponding pseudocode: einsum(bdhpp, edkp bek)
For ④
(R)sk
. [W][A),
p.g with [W](A) of shape [b, d, D, A and a (B)ck
(A) of shape je, d, D]
Corresponding pseudocode; cinsum(bbpp, edk bek)
For & (B)sk
[W]{8) with [W](R) of shape [b, d, DA, D] and ( (B)sk (B)e.g
of shape [e, d, D.D]
(BB)..g
Corresponding pseudocode; e einnaummbddp edkg bek)
71

---

Published as a conference paper at ICLR 2025
For $ (a)=
(A): [0]44, with (b]/4) of shape [b,d, DA] and $ (A): of shape [e, de,d, DJ
Corresponding pscudocode: einsum edk - bek)
For + (B)= [ with [6],8) of shape [b, d, D] and + [B)ck (B)=
of shape [e, d, D,
Corresponding pseudocode: e einsum edkq - bek) For ((B): of shape [e, D),
Corresponding pseudocode: e einsum(ck - ek) unsqueeze(()
1.3 INVARIANT LAYERS PSEUDOCODE
From the formula:
D I(U)
D P(OK))
WWLOK.s)) (QK.s) p.g
p=
g=1
D
W (WW)yoos)
)p..
p=1g=1
D
S (A):
[W](4]
+ (B)
[W]68]
q=1
p=1q=1
DA
D
(A):
(b((4)
+
4 (B)
[b +1
0=1
We define the pseudocode for each term:
For ((QKK))).." [WW],OK.) with (WW](OKs)) of shape (b, h, D, D) and ((QK,))P
piq
of shape [e, d, D, D, D"]
Corresponding pseudocode: einsum(bdhpg, edpgk bek)
For (vv..))). [WW] with [WW] of shape (b, d, h D, DJ and (vv..))
of shape [e, d, D']
Corresponding pseudocode: einsum(bdhpp, edpk - bek) .p.y
For ...)).. [W][A] with (4) of shape [b, d, D, Da) and (A)):, of shape [e, d, D*)
Corresponding pseudocode: nsuumbdpg edk: bek)
For ((B): [W] with [W](B) of shape [b, d, D D] and (BB)a of shape [e, d, D. D
Corresponding pscudocode: einsum(bdpy, edqk - bek)
For ((A)) [b](a), with [b](-4) of shape [b, d, Da] and (A): of shape (e, D')
Corresponding pseudocode: einsum(bdg, edk: bek)
For P( [b) with [b]()) shape [b, d, D] and & (B); of shape [e, d, D, D')
Corresponding pscudocode: sinsum(bdg, edgk - bek)
For & of shape [e, D'].
Corresponding pseudocode: e nsum(ek - ck) .unsqueeze())
72


Reviews:
[{"review_id": "14HGfs6tlo", "reviewer": "Reviewer_wnMM", "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.", "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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9J1OFDSpEd", "reviewer": "Reviewer_TUKj", "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", "questions": "How does this generalize to models with more parameters and trained on larger datasets", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5T4ruZs3oP", "reviewer": "Reviewer_YffW", "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.", "questions": "1. Could the authors clarify the expected computational cost of Transformer-NFN compared to traditional NFNs, considering the added design complexity?\n2. Will the paper include code to facilitate reproducibility?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2. The release of a dataset with 125,000 Transformer model checkpoints offers a valuable resource for future research on Transformer performance prediction and NFNs.\n3. 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.\n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "r76qsWtIts", "reviewer": "Reviewer_ynVN", "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.", "questions": "- How about the performance on language modeling? Is Transformer-NFN able to predict?\n- How about the performance of the larger transformer \n- 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?\n- 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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n- For all I know, this may be one of the first works to develop NFN for transformers. \n- The theoretical analysis is detailed enough.", "weaknesses": "- According to the main presentation of this submission, no obvious weaknesses are found.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)