Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models
Authors: Zhijian Zhuo
Ya Wang
Yutao Zeng
Xiaoqing Li
Xun Zhou
Jinwen Ma
Venue: ICLR 2025 Poster (2025)
Keywords: activation function
transformer
pre-training
large language models
Primary area: foundation or frontier models, including LLMs

Abstract: Transformers have found extensive applications across various domains due to their powerful fitting capabilities. This success can be partially attributed to their inherent nonlinearity. Thus, in addition to the ReLU function employed in the original transformer architecture, researchers have explored alternative modules such as GeLU and SwishGLU to enhance nonlinearity and thereby augment representational capacity. In this paper, we propose a novel category of polynomial composition activations (PolyCom), designed to optimize the dynamics of transformers. Theoretically, we provide a comprehensive mathematical analysis of PolyCom, highlighting its enhanced expressivity and efficacy relative to other activation functions. Notably, we demonstrate that networks incorporating PolyCom achieve the **optimal approximation rate**, indicating that PolyCom networks require minimal parameters to approximate general smooth functions in Sobolev spaces. We conduct empirical experiments on the pre-training configurations of large language models (LLMs), including both dense and sparse architectures. By substituting conventional activation functions with PolyCom, we enable LLMs to capture higher-order interactions within the data, thus improving performance metrics in terms of accuracy and convergence rates.  Extensive experimental results demonstrate the effectiveness of our method, showing substantial improvements over other activation functions. Code is available at https://github.com/BryceZhuo/PolyCom.

Published as a conference paper at ICLR 2025
POLYNOMIAL COMPOSITION ACTIVATIONS: UN- LEASHING THE DYNAMICS OF LARGE LANGUAGE MODELS
Zhijian Zhuo Ya Wang2* Yutao Zenga1 Xiaoqing Li
Xun Zhou²
Jinwen Ma'
School of Mathematical Sciences. Peking University
2 Cee--Fuunation-Modeee ByteDance
3 Capital University of Economics and Business
ABSTRACT
Transformers have found extensive applications across various domains due their powerful fitting capabilities. This success can be partially attributed to their inherent nonlinearity. Thus, addition to the ReLU function employed the original transformer architecture, researchers have explored alternative modules such as GELU and SwishGLU to enhance nonlinearity and thereby augment rep- resentational capacity. In this paper, we propose novel category of polyno- mial composition activations (PolyCom), designed to optimize the dynamics provide comprehensive mathematical analy-
transformers. Theoretically, sis of PolyCom, highlighting its enhanced expressivity and efficacy relative to other activation functions. Notably, we demonstrate that networks incorporat- ing PolyCom achieve the optimal approximation rate, indicating that PolyCom networks require minimal parameters to approximate general smooth functions in Sobolev spaces. We conduct empirical experiments on the pre-training con- figurations of large language models (LLMs), including both dense and sparse architcctures. By substituting conventional activation functions with PolyCom, we enable LLMs to capture higher-order interactions within the data, thus im- proving performance metrics in terms of accuracy and convergence rates. Exten- sive experimental results demonstrate the effectiveness our method, showing substantial improvements
other activation functions. Code is available https:/ /eg/gttuuuub/ com Brw.ceeucceeeeoeeeooo Poo.uuooooooCComm /olp://.....orrrce/o//////ooo//
Training
Validation PPL
HellaSwag ACC. (%)
ARC-Easy
2.40
3.70
2.35
3.50
PMMRU PelWarm
Mym
2.0
3,4
42
2.5
3.37
Sw
SWHLU
220
3.2
Mydall
40
PESTALU
Teyhoom
FobMorm
Takens (B)
Tokens (B)
d: of Takens (B)
Tokans (B)
Figure Training loss, validation perplexity (PPL), and downstream performance IB dense models. We compare models employing different activation functions, including SwiGLU, GELU, ReLU, PolyReLU, and PolyNorm. indicates that models using PolyReLU and PolyNorm exhibit lower training loss and validation PPL, alongside better downstream performance.
INTRODUCTION
Transformers Vaswani et al. 2017) have revolutionized the field of deep learning. facilitating un-
precedented advancements in natural language processing (Radford et al. 2019 Zeng et al. 2020 Li et al. 2024 Zuo et al. 2024), computer vision (Dosovitskiy al. 202 Wang et al. 2022 and
"Equal contribution.
Corresppnningg authors: Yutao Zeng www.aao.wwggouuoooooooooom and Jinwen Ma (jwma @iwma@math wwt.mmattthaaaae.mmmmmmmmmm pkw.edu.cn)

---

Published as a conference paper at ICLR 2025
beyond (Dong al. 2018 Arnab al., 2021 Wang et al. 2020). Characterized by their atten- tion mechanisms, transformers excel at capturing intricate relationships within data, making them indispensable in contemporary machine learning applications. However, despite their widespread success, there remain opportunities for further refinement, particularly concerning the selection of activation functions. The activation function plays crucial role in determining the output of each neuron within a neural network. Traditionally, simple nonlincarities such as Rectified Linear Unit (ReLU) (Nir Hinton, 2010) and its variants (Hendrycks & Gimpel 2016 Krotov & Hopfield, 2016 Li et al. 2019] So et al. (021) have been favored due to their computational efficiency and ease of implementation. Although effective, these activation functions are inherently limited in their ability to model complex higher-order relationships within data. This limitation can be particularly restrictive in transformer architectures, where the ability to capture subtle and complex dependencies is essential.
In this paper, we introduce novel category of polynomial composition activation functions (PolyCom), specifically engineered enhance the performance of transformer architectures. In contrast to conventional activation functions, which are predominantly linear or piecewise linear, polynomial composition activations facilitate the modeling of more complex patterns within data. This augmentation in the activation function's 's expressiveness endows the model with superior ex- pressive capacity, enabling to capture higher-order interactions that might otherwise be neglected. Unlike other forms of polynomials ((Hrrnik et al.) 1989 Trefethen 2019h) that suffer from inade- quate approximation, exploding values, and oscillatory behavior, we demonstrate that PolyCom pos- sesses more potent expressive capability than both ReLU and traditional polynomials and achieves optimal approximation within Sobolev space.
We posit that the integration polynomial composition activations within transformer models can lead to enhanced performance in tasks requiring intricate data interpretation. To evaluate this hypoth- esis, we conducted comprehensive experiments on the pre-training configurations of large language models (LLMs), including both dense and sparse architectures. These evaluations were performed across various benchmarks, assessing the performance of transformers employing polynomial com- position activations in comparison to those utilizing traditional activation functions. The results indicate that the proposed method not only improves the model accuracy but also accelerates con- vergence rates, thereby suggesting that polynomial composition activations provide substantive advantage in deep learning applications.
The main contributions of this paper are summarized in the following.
We propose new activation function PolyCom which is composition of the polyno- mial and other types of function. In particular, we introduce two instances of PolyCom: PolyReLU and PolyNorm, and detail its integration into the transformer architecture.
Theoretically, we derive bounds on the number of trainable parameters required for PolyReLU networks to approximate ReLU networks, and vice versa. Additionally, we show that PolyReLU network of size O(e-//n) can approximate any function in Sobolev spaces with error tolerance achieving optimal approximation rates.
Empirically, we validate the effectiveness of this new activation function on LLMs with both IB dense models and MoE models with IB active and 7B total parameters. The results of both models demonstrate that PolyCom can accelerate the converging speed and significantly outperform SwiGLU, GELU, and ReLU et al.
The outline of this paper is structured as follows: In Section 2 we present the mathematical formu- lation of PolyCom and discuss integration within transformer architectures. Section 3 delivers a comprehensive theoretical analysis of PolyCom, emphasizing its enhanced expressivity and effec- tiveness. In Section 4 we provide a detailed account of our experimental results involving large language models (LLMs). Section provides an overview of related work in the field of activation functions and their applications transformer models. Finally, we conclude the paper and outline potential directions for future research.
POLYNOMIAL COMPOSITION ACTIVATION FUNCTION
In this section, we present the mathematical formulation of the polynomial composition activation function (PolyCom) and detail its integration into the transformer architecture.
2

---

Published as a conference paper at ICLR 2025
SWGGLU
Palyoo
FC
FC
FC
X
Swish
ReLU GELU
x
--1
N
N
N
dr
0--
dg
8g
Figure 2: Block diagrams of Transformer MLP blocks utilizing ReLU/GELU, SwiGLU, PolyReLU and PolyNorm. "FC" stands for Fully Connected layer. "x"" represents the i-th power of the input tensor æ, "ay" denotes the j-th element of the learnable weight vector a, "N" indicates a normaliza- tion operation.
PolyCom. The study of the polynomial activation function can be traced back to the seminal work Hornik et al. (1989), which showed that neural networks with polynomial activation not dense within the space of continuous functions. Additionally, empirical evidence has shown that deep neural networks employing pure polynomial activations tend to underperform [Trefethen] 2019). To overcome these limitations, we propose PolyCom, a novel composition of polynomials and other functions. Specifically, we explore two composition approaches
Type I: x Type II: x x4 -iE0
ag E.
where r E denotes the order of PolyCom and p represents an arbitrary function such as ReLU, PReLU, Sigmoid, SiLU, or normalization. The key distinction between the two approaches lies in whether the function is composed before or after the power operation. The distinction between the two approaches lies in whether the composition implementation of PolyCom, we use 3-order PolyCom (r = 3) with trainable coefficients a₃ For initialization, set a; 1/r for i = ..... ; I' and dg 0. Our experiments large language models (LLMs) show that 3-order PolyCom can indeed achieve extraordinary performance.
the power is performed first. In the practical
For Type PolyCom, specifically consider composition involving the ReLU function due its simplicity, which we term PolyReL.U. An r-order PolyReL.U is defined as
PolyReLU(x) Σ (x),
where ReLU'(*) = max This formulation can be seen as an extension of both ReLU and square ReLU.
For Type PolyCom, we introduce PolyNorm, which normalizes the powers to ensure consistent magnitudes across terms
PolyNorm(z) Σ a. x
(3)
where x²' = [ .12. ""] represents element-wise exponentiation, and || |lz denotes the L2 normalization. PolyNorm incorporates normalization operators to rescale different powers into a manageable range. thereby preventing excessively large or small values. This makes the training procedures more stable.
Integration into Transformer. The transformer architecture (Vaswani et al.] 2017) consists of two alternating modules, Multi-Head Attention (MHA) and position-wise Feed-Forward Networks (FFN), Activation functions predominantly influence the performance of FFN layers. We begin by formalizing the common paradigm of FFN,
FFN,(*) = p(xWi)
where P represents the activation function such ReLU, GELU, PolyReLU, and PolyNorm. We replace the traditional activation function with our proposed PolyCom variants to enhance model capacity and performance, illustrated in Figure 2
3

---

Published as a conference paper at ICLR 2025
3 THEORETICAL ANALYSIS
From Figure I one can see that the expressivity of PolyNorm is greater than or equal to that of PolyReLU. To streamline the analysis, we focus solely on the theoretical properties of PolyReLU, specifically its expressivity and effectiveness. Additional, nonlinear activations such as GELU and SwiGLU can be locally approximated by Taylor polynomials around the origin, which allows to primarily compare PolyReLU with ReLU and polynomial activations. To avoid confusion, we refer to networks that use ReLU activations as ReLU networks, and those that use PolyReLU activations as PolyReLU networks.
3.1 APPROXIMATING RELU NETWORKS POLYRELU
In this subsection, we present theoretical results on approximating ReLU networks using PolyReLU networks. The following lemma shows that ReLU, ReLU², and polynomial activation are special cases of PolyReLU activation, highlighting the superior expressivity of PolyReLU, This implies that PolyReLU has stronger approximation abilities with fewer trainable parameters compared to ReLU and other polynomial activations.
Lemma 1. ReLU, ReLU2² and polynomial activation can be represented by PolyReT.U.
Building on Lemma ! we can formally prove that any ReLU network can be exactly represented by Poly network of the same size, as stated in the following theorem.
Theorem 1. Let f [-1, 1ja + [-1, 1] be a ReLU network with depth and width K. Then, there exists a PolyReLU network g [-1, 1, 1] + 1- 1] of size (LL)) such that
f(x) g(z),, for Va E [-
(5)
This theorem, proved in Appendix A shows that PolyReL.U networks can exactly match the repre- sentational power of ReLU networks without increasing the model size.
3.2 APPROXIMATING POLYRELU WITH RELU NETWORKS
In this part, we give theoretical results on approximating PolyReLU networks using ReLUU networks. The following network within given error tolerance.
demonstrates that the PolyReI.U activation be approximated by ReLU
Lemma 2. For the fixed positive integer r and the activation PolyReLU(x) a R, ( E [-1,1] with a, [-1, Given any E (0,1), there exists ReLU network f 111,,1] with size O (1/e)), such that
max If(z) PolyReLU(z)) <c -E1-
Lemma establishes an upper bound the size of ReLU network needed to approximate PolyReLU activation function. This result highlights that while ReLU networks can approximate PolyReLU activations, they require significantly larger number of parameters.
we derive the following theorem, which provides both upper and lower bounds for approximating PolyReLU networks with ReLU networks.
Building on Lemma
Theorem 2. Let [] + [-1,1] be PolyReLU network with depth L and width K, and PolyReLU activation with order r and Lipschitz constant a. Suppose each neuron computes x H> PolyReLU(a x (e b) with the pair (a,b) satisfies ||a]] + < 1 and PolyReLU 1| [-1, (a, b, and PolyReLU are possibly distinct across neurons). For any given E (0, 1), there exists
ReLU network f [-1, 1]at + [11, 1] of size
o LK (La))
such that
max ff(x) 9(x)) < =E|-1,11
(8)
Conversely, there exists PolyReLU networks cannot be approximated within tolerance by any ReLU network with a size less than
1 - )
Ω KLIn
(9)
4

---

Published as a conference paper at ICLR 2025
Theorem 2 tells us that the total number of trainable parameters required by ReLU networks to ap- proximate PolyReI.U neural network within tolerance of is (())(1/))) Conversely, there exists a PolyReLU network that can not be approximated by ReLU networks of size less than (lIn(//e). Combined with Theorem 1, we conclude that PolyReLU networks are more efficient in terms of representational capacity than ReLU networks.
3.3
APPROXIMATION OF GENERAL SMOOTH FUNCTION
Similar to Yarotsky 2 Boullé et al. (020), we also explore the universal approximation ca- pabilities of PolyReI.U networks in the context of Sobolev spaces (Adams & Fournier] [003). Specifically, we show that PolyReLU networks achieve the optimal approximation rate within these spaces, meaning that PolyReLU networks require minimum parameters to approximate gen- cral smooth functions in Sobolev spaces, compared with networks with the other activation.
The definition of Sobolev space W",.ce (I-1, 1] is stated below. The set [-1, 1]d can be replaced it just for the sake of brevity.
by any compact set in IRd Definition (Sobolev Spaces). For n., d E N, Sobolev space Wm, (1-1, is defined as
Wm.00 (( Loc (1-1, 119)
= {f E
< o0)
(10)
with the norm which is defined as the following
[Ifl ((-1))))
max ess sup | f(x
(11)
where E Nd and D is the respective weak derivative of f, and css sup means the essential supremum in functional analysis.
Intuitively, Sobolev space is space of functions endowed with weaker notion of smooth- ness compared to differentiabilityy and possessing generalized derivatives. The Sobolev space W7,00 1J) contains functions from ([-1, which consists of functions whose derivatives of order π - 1 are Lipschitz continous. the sequel, we mainly consider the unit ball within Win. (I-1, ) which is defined follows
Fy, = {f E Wm,00 ([-1, III (1 S 13.
With the above definitions established, we can present the following main results. We provide an upper bound on the size of PolyReLU networks required to approximate any function in Fi.d- Theorem 3. Suppose that d, n N and e E 1). For any J Fa,as there exists a PolyReLU network with size O(e-d/n) d/n) that can approximate f at a given error tolerance e, iee.
max # -111d
[1
-1())) g(x)l <e.
(12)
Theorem indicates that PolyReLU networks achieve optimal approximation rate of ((e-d/n) In contrast, previous works by Yarotsky (017)) demonstrated that ReLU networks re- quire d(e-d/n In(1/)) In(/c)) parameters to achieve a similar approximation error. Similarly Boullé et al. 2020) showed that rational neural networks need O In(In(1/e))) parameters for the same task. Therefore, the approximation ability of PolyReLU networks is superior to that of both ReLU networks and rational networks. Furthermore, Theorem 4.2 in Vore et al. (1989) shows that the to- tal number of parameters required by neural networks to approximate functions Fn.a is li (((e-d/m). Therefore, our PolyReLU networks achieve the optimal approximation rate in the context of Sobolev spaces. Additional disscution is included in Appendix
4 EXPERIMENTS
In this section, we demonstrate the expressivity and effectiveness of PolyCom within the transformer through experiments on LLMs.
4.1 SETUP
Baseline. We evaluate PolyCom across two series of models: 1B dense model and Mixture of Experts (MoE) model with IB active and 7B total parameters. The IB dense model contains approximately 1.3 billion parameters with architecture similar to Llama 2 (Touvron al. 2023
5

---

Published as a conference paper at ICLR 2025
Table 1: Overall results of the IB dense model with different activation functions, reported in terms of training loss, validation perplexity, and downstream accuracy (%). ARC-E and ARC-C C refer to ARC-Easy and ARC-Challenge. respectively. The best results in each column are highlighted in bold. "Avg." denotes the average accuracy of all downstream tasks.
Loss
PPL
ARC-E
ARC-C
HellaSwag
PIQA
SciQ
Winograde
Avg. 1
SwiGLU
2.19
3.22
56.61
27.47
49.23
68.61
86.10
56.83
57.47
GELU
2.20
3.24
55,43
27.73
48.42
68.12
87.40
54.78
56.98
ReLU
2.21
3.26
55.68
28.50
48.59
68.39
87.10
54.85
57.18
PolyReLU
2.17
3.18
57.53
27.99
50.19
70.29
87.60
55.72
58.22
PolyNorm
2.17
3.17
59.68
29.01
50.86
69.15
87.20
56.20
58.68
For the MoE model, we use the OLMoE framework [Muennighofi et al. 2024). which activates 1.3B parameters out total of 6.9B parameters. Both models are trained from scratch. We compare the performance of PolyCom with several activation functions, including ReLU, square ReLU, GELU, and SwiGLU, All experiments are conducted on NVIDIA A100-80G GPUs, 32 GPUs for the dense model, and 64 GPUs for the MoE model.
Model Configuration. For the dense model, the transformer consists of 24 layers with hidden size dmodel = 2048 and 16 attention heads. In the MoE model, the transformer is composed of 16 layers, with a hidden size of demodel = 2048, 16 attention heads, and 64 experts. To maintain a consistent number of trainable parameters across all activation functions, size accordingly. Specifically, for SwiGLU, the intermediate size is set to two-thirds that of the other activations in all experiments. More details can be found in Appendix
adjust the intermediate
Datasets. The dense model is trained on the RedPajama- I dataset Computer 2023 ). which was
developed by the open-source Al community enable competitive performance against proprietary models. The MoE model is trained on the OLMoE Mix dataset MMennnghofff et al. 2024).
perpparramesers. Unless otherwise specified, we use a 3-order PolyCom by default and initialize the coefficients as = 1/3 for = 1, 2, 3 and set ap = 0. Model weights are randomly initialized. For optimization, we apply the AdamW optimizer with 31 0.9 and B3 0.95. All models are trained on sequences of 4096 tokens. For the dense model, we set the initial learning rate to 3e-4, decaying to 1.5e-5 using cosine scheduler. The MoE model starts with a learning rate of 4e-4, also decaying according cosine schedule. We summarize the hyperparameters in Table 7
Evaluation. To evaluate the performance of LLMS with PolyCom, we use a wide range of open benchmarks, including ARC-Easy (Clark et al., 2018). ARC-Challenge (ARC-C) (Clark et al. 2018). HellaSwag (Zellers et al. (011)). PIQA (Bisk et al., 2002),, SciQ (Welbl et al. 2017). CoQA 2021). BoolQ
(Reddy et al. 2019). Winogrande (Sakaguchi al. 2021), MMLU (Hendrycks et al.
Clark et al.) 2019). COPA (Gordon et al. (012)) CSQA (Talmor et al. 2019), OBQA (Mihaylov al.]][]]000 2011), and SocialIQA (Sap et al. 2019). We utilize the LM Eval Harness (Gao et al.) 2023)
for standardized performance evaluation.
4.2 RESULTS ON DENSE MODEL.
Training Dynamics of 1B Dense Model. Figure compares the training dynamics of the IB dense model across different activation functions. As shown in the figure, models using PolyReLU and PolyNorm exhibit lower training loss and validation perplexity throughout the training process compared models utilizing other activation functions. This indicates that PolyCom accelerates the convergence of LLMs. The models with PolyReLU and PolyNorm also consistently outperform others in downstream tasks by large margins. highlighting the advantage of PolyCom in improving the overall expressivity and effectiveness of LLMs.
Downstream Evaluation. Table presents the training loss, validation perplexity, and downstream task accuracy (%) after processing 250 billion training tokens. The downstream tasks include ARC- Easy, ARC-Chall Chal-hhlee, HellaSwag. PIQA, SciQ, and Winograde, More detailed results are provided
"RedPajama-11 available http:s/ss...tiicccc..cc.c//ce/ee
PLLMoB Mix dataset is available at
OlMo 0924

---

Published as a conference paper at ICLR 2025
Training loss
Validation loss (C4)
Validation loss (Wikipedia)
2.55
SWiGLU Poly
2.88
SwiGLu PolyNorm
2.48
SwiGLU
2.53
2.45
PolyNorm
2.50
2.83
2.42
2.48
2.45
2.80
2:38
2.43
2.78
2.35
2.40
2.38
2.32
2.35
2.30
2.70
50
100
150
200
50
100
150 Tokens (B)
200
50
100
150
200
Tokens
# of Tokens (B)
Figure 3: Training and validation loss C4 and Wikipedia for MoE models with 200 billion train- ing tokens. We compare models using SwiGLU and PolyNorm activation functions. PolyNorm demonstrates lower training and validation losses, indicating faster convergence.
HellaSwag Acc. (%)
MMLU Var Acc. (%)
ARC-Challenge Acc. (%)
SciQ Acc. (%)
SAtGUU
SwIGLU
Palyforr
PolyWorm
Ihtyytarm
PolyNorm
40
54
34
35
36
34
344
33
of Takens (8)
# Tokens (8)
of Takans (8)
of Tokens (B)
Dynamics of downstream performance on HellaSwag, MMLU Var, ARC-Challenge. and SciQ for MoE models with 200 billion training tokens. Models with PolyNorm significantly outper- form those with SwiGLU on downstream tasks.
Figure 4:
in Appendix The results clearly demonstrate that the PolyCom family (PolyReLJ and PotyNorm) outperforms the other activation functions. For instance, PolyNorm outperforms SwiGLU by an av- erage margin of I. 1.21% across six downstream tasks. This underscores the expressivity and efficiency of PolyCom as an activation function in transformer models.
4.3 RESULTS ON MOE MODEL
Our experiments with MoE modes are based on OLMOE- B-7B, which has 1 billion activate param- eters and 7 billion total parameters (Muennighofi et al. 2024). Due to computational constraints, we compare only the PolyNorm activation function, shown to perform best in dense models, with the widely used SwiGLU activation function, which is commonly employed in current LLM archi- tectures.
Training dynamics of MoE model. In Figure 3 report the training and validation loss of MoE models trained on 200 billion tokens. Models using PolyNorm consistently show lower losses com- pared those using SwiGLU, indicating that PolyNorm enables faster learning. Figure 4 shows the downstream performance on HellaSwag. MMLU Va ARC-Challenge, and SciQ. PolyNorm out- performs SwiGLU on all tasks, with notable improvements, demonstrating superior generalization capabilities.
Dowmstream Evaluation. Table 2 presents the validation losses 11 datasets. Norm consis- tently achieves lower validation losses than SwiGLU across all datasets, with an average improve- ment of 0.02. In Table 3 we also observe that PolyNorm outperforms SwiGLU 8 downstream tasks. These results highlight the superior performance of models using the PolyNorm activation function. Additional results can be found Appendix
MML.U Var is variant of MMLU (Hendrycks et al. 2021 using varied few-shots Muennighoff et al.
2024
7

---

Published as a conference paper at ICLR 2025
Table Validation losses of MoE models with different activation functions. CC denotes Common Crawl. Best results per column are bold.
Wiki- pedia
Wiki- text
Methods
C4
Books
CC
peS2o
Reddit
Stack
ICE
M2D2
Pile
Avg.J
SwiGLU 2.72
2.59
2.79
2.16
2.93
1.01
2.30
2.50
3.07
2.07
2.37
2,41
PolyNorm 2.71
2.57
2.78
2.15
2.92
1.00
2.29
2.49
3.06
2.03
2.34
2.39
Table 3: Downstream evaluation results of MoE models with different activation functions. ARC-C, ARC-E, OQA denote ARC-Challenge, ARC-Easy, and OpenbookQA, respectively. Best results per column are in bold.
MMLU
Hella- Swag
Wino- Grande
Tasks
SciQ ARC-C ARC-E PIQA
OQA
COPA
Avg.t
Var
SwiGLU
37.07
66.49
90.60
37.12
71.58
76.61
62.75
39.80
83.00
62.78 63.37
PolyNorm
37.27
67.63
92.40
38.46
70.70
77.04
62.19
40.60
84.00
4.4 ABLATIONS AND ANALYSIS.
Order of PolyCom. We first investigate the effect of different orders of PolyCom. We vary the order r of PolyReLU in the range {2, 3, 4] and plot the results in Figure 5(a) As seen, the convergence speed improves the order increases. However, there no noticeable difference between orders and 4 in terms of convergence speed. Additionally, increasing the order can lead to computational overhead and overflow issues, particularly when using low- www-preiisoon arithmetic. Based on these observations, we select r = 3 as the default order for PolyCom in our experiments, balancing both performance and computational efficiency.
Different Polynomial Composition Functions. We evaluate the impact of different polynomial composition functions by comparing PolyReLU, PolyPReLU, PolyNorm, and PolyReLUNorm in Figure 5(b) Our results indicate that PolyNorm, which uses normalization as the composition func- tion, achieves the lowest training loss and best overall performance. This suggests that normalization plays a key role in stabilizing training and enhancing the model's ability to generalize. In contrast, combining ReLU with normalization (PolyReLUNorm) provides intermediate results, suggesting that more complex compositions do not always lead to hetter outcomes.
Variants ReLU. In Figure 5(c) we compare different variants the ReLU activation function, including ReLU and ReLU?. PolyReLU consistently outperforms both ReLU and ReLU² across all tasks, highlighting the henefits of using polynomial composition. This result reinforces the hy- pothesis that introducing higher-order terms through PolyCom enables the model to capture more complex data interactions, thus improving the expressivity of the activation function without signif- icantly increasing model size or complexity.
Rank of Weights. To understand how PolyCom enhances model performance, analyze the rank of the weights in each FFN layer of the transformer. We use the effective rank (Roy & Vetterli 2007) measure the effective dimensionality of weights and its definition is in Appendix [E.3] Figure [ shows that PolyReLU and PolyNorm result in higher weight ranks compared to other activation functions such as SwiGLU, GELU, and ReLU. A higher rank in the weight matrices usually indicates a greater capacity for representing complex patterns in the data. These findings suggest that PolyCom improves the expressibility transformers by allowing the FFN layers to better utilize their parameters, ultimately leading to better generalization on downstream tasks.
Layer-wise Similarity. We further analyze the layer-wise similarity of hidden states using cosine similarity, as illustrated in Figure 7 For both dense and MoE models, we compare SwiGLU with PolyNorm. The results reveal that PolyNorm consistently maintains lower layer-wise similarity compared to SwiGLU, indicating that PolyNorm promotes greater diversity between layers. This diversity likely enables the model to learn more complex representations, merely replicating the functionality of earlier ones. Notably, the gap in cosine similarity between PolyNorm and SwiGLU widens in the deeper layers, which are generally more crucial for down-
deeper layers not
8

---

Published as a conference paper at ICLR 2025
Training less
Training loss
Training loss
Zorder-RlyAel
PolyReLU
Rell
2.45
PolyAel Aurdee-Polylael
RolyPRelU
ReLUP
Polyftrrm
2.45
PolyBell
2.40
2.40
2,35
2.35
2.35
2,30
2.30
2.25
2.25
60
Tokens (B)
Tokens (B)
(B)
(a) Different orders of PolyReLU
(b) Different compositions
(c) ReLU variants
Figure 5: Training loss for IB dense models with different activation functions. 5(a) We compare
different orders PolyReLU. 5()] Comparison of PolyCom with different composition functions.
5(c) Comparison of different variants of ReLU activation function.
1500
Oleche 1ha
(a) Rank of Wap. dense
(b) Rank of W down. dense
(c) Rank of Wop. MoE
(d) Rank of Wdoan. MoE
Figure 6: Rank of weights in each FFN. 6(a) & 6(b) for the dense model, 6(c) & ( for the MoE model.
(a) SwiGLU, dense
(b) PolyNorm, dense
(c) SwiGLU, MoE
(d) PolyNorm, MoF
Figure 7: Layer-wise cosine similarity of hidden states.
((a)
( for 1B dense models with
SwiGLU and PolyNorm, respectively. 7(c) & 7(d) for MoB: models with SwiGLU and PolyNorm, respectively.
stream task performance. This increased diversity across layers enhances the model's ability to capture complex relationships, thereby improving the overall effectiveness of LLMs.
Training Stability. Through extensive experiments, we find both PolyReLU and PolyNorm main- tain a stable training process within transformer architectures. Our analysis indicates the combina- tion of normalization operators in transformers and standard gradient clipping strategies effectively stabilizes training dynamics. We specifically design PolyNorm to address potential instability asso- ciated with BF16/FP16 precision formats, incorporating normalization operators that rescale pow- ers to manageable range, thus preventing excessively large or small values. This is particularly beneficial for FP16 training, demonstrated in Appendix H In contrast, PolyReLU lacks this nor- malization feature, which may result in stability issues in -on-trassformer architectures like ResNet. As shown in Figure 5(a) 3rd order (our default setting) is sufficient. Based these findings,

---

Published as a conference paper at ICLR 2025
recommend the following configurations: (1) For transformer- models, either PolyNorm or PolyReLU. (2) For non- non-transformer models or those with lower stability, prefer PolyNorm.
Computational Overhead and Memory Footprint. We provide detailed analyses of the runtime and memory overhead for the proposed activation functions, including FLOPs ratios and memory consumption, which are included in Appendix F Overall, after applying the gradient checkpointing technique, the overhead and memory footprint are acceptable, and there is negligible difference in the training budget required compared to the widely used SwiGLU.
RELATED WORK
The design of activation functions has been a critical area of research in neural networks, directly influencing the performance and capabilities of deep learning models. Early activation functions like Sigmoid and Tanh were widely used due to their smooth nonlinear transformations {Goodfel- loo al. 2016). However, these functions faced challenges such as vanishing gradients, making it difficult to train deep networks effectively. The introduction of the Rectified Linear Unit (ReLU) (Nair& & Hinton 2010) mitigated some of these issues by offering simple, non-saturating nonlin- earity, which has since become a standard in many deep learning applications. Variants of ReLU, such as Leaky ReLU (Maas et al., 2013) and Parametric ReLU (PReLUU (He et al. 2015), were developed to address the "dyyng ReLU" problem by allowing a small, non-zero gradient when the input is negative. Other functions, like the Exponential Linear Unit (ELU) (Cleverl) 2015), aimed to provide smoother activation profiles, resulting in better generalization and faster convergence in cer-
tain tasks. Moreover, Manessi & Rozza (2018) proposed a combination of weighted base activation functions for further enhancement.
Polynomial activation functions (Homik et al. 1989 Oh et al. 2003). although less commonly used, have been studied in various contexts for their ability model higher-order, complex relation- ships more effectively, For instance, Lokhande et al. (2020) introduced Hermite polynomial activa- tions to improve pseudo-label accuracy, while Chrysos al, (2020) proposed polynomial networks, II-nets, which apply to various domains such as image and audio processing. Building on this, Chrysos et al. (2023) utilized regularization techniques to enhance the performance of polynomial networks. These works highlight the potential of polynomial functions to increase the expressive- ness of neural networks by capturing intricate, higher-order interactions. On the theoretical front, the expressivity and approximation power of polynomial functions have been rigorously explored (Kileel et al. 2019: Kidger & Lyons 2020 Kubjas et al. 2024). Additionally, (Li et al. 2019) in- vestigated the approximation capabilities of rectified power units (i.e., ReLU*), demonstrating that they achieve the same approximation rate as PolyReLU.
The choice of activation function in transformers has also become important area of research. Originally developed for natural language processing, transformers (Vaswani et al. (017)) have been effectively adapted for diverse tasks, including image recognition. speech processing, and reinforce- ment learning. Despite their broad applicability, the activation functions predominantly utilized in transformers, ReLU and GELU, have seen minimal evolution. Recent studies, however, have be- gun to explore alternatives to these conventional activations. For example, the Swish activation (Ramachandran et al., 2017 [hazzerr] 2020) and the Mish activation (Misral 2019) are smooth and non-monotonic functions that offer potential benefits in model performance and training stability. Additionally, Gated Linear Units (GLU) were proposed by Dauphin et al. (2017), with SwiGLU
Shazeer 2020), prominent variant, being used in models such as -LaMA-Seriec (Touvron al. 2023).
6 CONCLUSIONS
In this paper, we introduce the Polynomial Composition Activation (PolyCom) and demonstrate its effectiveness within transformer models. By enabling the capture of higher-order interactions, PolyCom enhances both the accuracy and convergence rates of these models. Our experiments, conducted across different large language model architectures and multiple benchmarking datasets, confirm that PolyCom consistently outperforms conventional activation functions. Furthermore, ab- lation studies indicate that PolyCom increases model expressivity by elevating weight rank and re- ducing redundancy across layers. These findings underscore the significant potential of polynomial- based activations to improve transformer models, thereby paving the way for future research en- deavors.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
Jinwen Ma was supported by the Natural Science Foundation of China under grant 62071171.
REFERENCES
Robert A Adams and John JF Foumiee. Sobolev spaces. Elsevier, 2003.
Anurag Amab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. Proceedings of the IEEE/CVE international conference on computer vision, pp. 6836-6846, 2021.
Jonathan T Barron. Continuously differentiable exponential linear units.
arXiv preprint
arXiv:1704.07483, 2017.
Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical com- monsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Nicolas Boullé, Yuji Nakatsukasa, and Alex Townsend. Rational neural networks. In Proceedings of the 34th International Conference on Neural Information Processing Systems, pp. 14243-14253, 2020.
Grigorios G Chrysos, Stylianos Moschoglou, Giorgos Bouritsas, Yannis Panagakis, Jiankang Deng, and Stefanos Zafeiriou. P-nets: Deep polynomial neural networks. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 7325-7335, 2020.
Grigorios G Chrysos, Bohan Wang, Jiankang Deng, and Volkan Cevher. Regularization of polyno- mial networks for image recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16123-16133 2023.
Christopher Clark, Kenton Lee, Ming-Wet Chang, Tom Kwiatkowski. Michael Collins, and Kristina Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 2924-29299999 2936, 2019.
Toutanova. Boolq:
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprini 1xXXvv8803.05555 2018.
Djork-Amé Clevert. Fast and accurate deep network learning by exponential linear units (clus). arXiv preprint arXiv:1511.0728) 2015.
Together Computer. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023.
URL httpss://github. con httpa/////oooooteeeeeeeoeerreeee /Ree///ePPPaajaa aaaed-aaaaa-aaaaa
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pp. 933-949--11 941.. PMLR, 2017.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scale hierarchical image database. In CVPR, 2009.
Ronald A Verer Ralph Howard, and Charles Micchelli. Optimal nonlinear approximation.
Manuscripta mathematica, 63:469-478, 1989.
Linhao Dong, Shuang Xu, and Bo Xu. Speech-transformme:: a no-recurrence sequence-to-sequeee model for speech recognition. In 2018 IEEE international conference acoustics, speech and signal processing (ICASSP), pp. 5884-5888, IEEE, 2018.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image worth 16x16 words: Transformers for image recognition at scale. In International Conference Learning Representations, 2021.
11

---

Published as a conference paper at ICLR 2025
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa. Jason Phang. Laria Reynolds. Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. Utttatt:: wrt.e//ooooooooreeeeeeeeeeeeeeee org/recrr////// 10256836
Xavier Glorot, Antoine Bordes, and Yoshua Bengio. Deep sparse rectifier neural networks. In Proceedings of the fourteenth international conference on artificial intelligence and statisties, pp. 315-323. JMLR Workshop and Conference Proceedings, 2011.
lan Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT Press, 2016.
Andrew Gordon, Zornitsa Kozareva, and Melissa Roemmele. Semeval-2012 task 7: Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In SEM 2012: The First Joint Conference on Lexical and Computational Semantics-Volume 1: Proceedings of the main conference and the shared task, and Volume 2: Proceedings of the Sixth International Workshop on Semantic Evaluation (SemEval 2012). pp. 394-398, 2012.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance imagenet classification. In Proceedings of the IEEE international conference on computer vision, pp. 1026-1034, 2015.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In CVPR, 2016.
Dan Hendrycks and Kevin Gimpel. Gaussian 10X.4:1606.08415, 2016.
error
linear units (gelus).
arXiv preprint
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob
Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021.
Kurt Hornik, Maxwell Stinchcombe, and Halbert White. Multilayer feedforward networks are uni- versal approximators. Neurul networks, 2(5):359-366, 1989.
Patrick Kidger and Terry Lyons. Universal approximation with deep narrow networks. In Conference on learning theory, pp. 2306-23663332 2327. PMLR, 2020.
Joe Kileel, Matthew Trager, and Joan Bruna. On the expressive power of deep polynomial neural networks. Advances in neural information processing systems, 32, 2019.
Alex Krizhevsky et al. Convolutional deep belief networks on cifar-10. 2010.
Dmitry Krotov and John J Hopfield. Dense associative memory for pattern recognition. Advances in neural information processing systems, 29, 2016.
Kaie Kubjas, Jiayi Li, and Maximilian Wiesmann. Geometry of polynomial neural networks. arXiv preprint (tt1::240200094444 2024.
Bo Li, Shanshan Tang, and Haijun Yu. Better approximations of high dimensional smooth functions by deep neural networks with rectified power units. arXiv preprint PrXiv:1903.053: 05855, 2019.
Zixuan Li, Yutao Zeng, Yuxin Zuo, Weicheng Ren, Wenxuan Liu, Miao Su, Yucan Guo, Yantao Liu, Lixiang Lixiang, Zhilei Hu, Long Bai, Wei Li, Yidan Liu, Pan Yang, Xiaolong Jin, Jiafeng Guo, and Xueqi Cheng. KnowCoder: Coding structured knowledge into LLMs for universal informa- tion extraction. In Lun-Wei Ku, Andre Martins, and Vivek Srikumar (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pp. 8758-8779, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 1/26533/112022acc-ooo
Shiyu Liang and R Srikant. Why deep neural networks for function approximation? In International Conference on Learning Representations, 2017.
12

---

Published as a conference paper at ICLR 2025
Vishnu Suresh Lokhande, Songwong Tasneeyapant, Abhay Venkatesh, Sathya N Ravi, and Vikas Singh. Generating accurate pseudo-labels in semi-supervised learning and avoiding overconfident predictions via hermite polynomial activations. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 1 11435-1 1443, 2020.
Andrew I. Maas, Awni Y Hannun, Andrew Ng, et Rectifier nonlinearities improve neural network acoustic models. In Proc. ICML. Atlanta, GA, 2013.
Franco Manessi and Alessandro Rozza. Learning combinations of activation functions. 2018 24th international conference on pattern recognition (ICPR), pp. 61-66. IEEE, 2018.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct elec- tricity? new dataset for open book question answering. In Proceedings of the 2018 Conference Empirical Methods in Natural Language Processing, pp. 2381-2391, 2018.
Diganta Misra. Mish: A self regularized non-monotonic activation function. arXiv preprint arXiv: 0868, 2019.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo. Jacob Morrison, Sewon Min, Weijia Shi, Pete Walsh, Oyvind Tafjord, Nathan Lambert, Yuling Gu, Shane Arora, Akshita Bhagia, Dustin Schwenk, David Wadden, Alexander Wettig, Binyuan Hui, Tim Dettmers, Douwe Kiela, Ali Farhadi, Noah A. Smith, Pang Wei Koh, Amanpreet Singh, and Hannanch Hajishirzi. Olmoe: Open mixture-of-experts language models, 2024.
Vinod Nair and Geoffrey E Hinton. Rectified linear units improve restricted boltzmann machines. In ICML, pp. 107-814,220 2010.
Kwun Oh, Witold Pedrycz, and Byoung-Jun Park. Polynomial neural networks architecture: analysis and design. Computers & Electrical Engineering, 29(6):70 2003.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Prajit Ramachandran, Barret Zoph, and Quoe Le. Searching for activation functions. arXiv preprint arXiv:1710.05990 2017.
Siva Reddy, Danqi Chen, and Christopher D Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-266, 2019.
Olivier Roy and Martin Vetterli. The effective rank: A measure of effective dimensionality. In EUSIPCO, 2007.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adver- sarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Common- sense reasoning about social interactions. arXiv preprint 4rX1v:1904.09728, 2019.
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05202 2020.
David So, Wojciech Manke, Hanxiao Liu, Zihang Dai, Noam Shazeer, and Quoc Le. Searching for efficient transformers for language modeling. Advances in neural information processing systems, 34:6010-6022 2021.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsenseqa: question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Lan- guage Technologies, Volume I (Long and Short Papers), pp. 4159- 2019.
Matus Telgarsky. Neural networks and rational functions. International Conference on Machine Learning, pp. 3387-3393. PMLR, 2017.
13

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Lloyd N Trefethen. Approximation theory and approximation practice, extended edition. SIAM, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Informa- tion Processing Systems, 2017.
Ya Wang, Dongliang He, Fu Li, Xiang Long, Zhichao Zhou, Jinwen Ma, and Shilei Wen. Multi- label classification with label graph superimposing. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 12265-12272, 2020.
Ya Wang, Xingwu Sun, Lian Fengzong, Zhanhui Kang, and Chengzhong Xu Xu. An anchor-based relative position embedding method for cross-modal tasks. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 5401-5413, 2022.
Johannes Welbl, Nelson F Liu, and Matt Gardner. Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User- Exeer-genrrate Text, pp. 10-, 2017.
Ross Wightman.
Pytorch image models.
https://githab.com com www.:/ow.wwwiigg.. http:///g...abb..ocmm///gag/a///
pytorch- menborch-imag--eee waaaaag-mmdeeee 2019.
Bing Xu, Naiyan Wang, Tianqi Chen, and Mu Li. Empirical evaluation of rectified activations in convolutional network (2015). arXiv preprint arXiv: 1oXiv:/505.0080) 505.00853, 2015.
Dmitry Yarotsky. Error bounds for approximations with deep relu networks. Neural networks, 94: 103-114, 2017.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a ma- chine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 4791-4800, 2019.
Yutao Zeng, Xiaolong Jin, Saiping Guan, Jiafeng Guo, and Xueqi Cheng. Event coreference res- olution with their paraphrases and argument-awart embeddings. In Donia Scott, Nuria Bel, and Chengqing Zong (eds.), Proceedings of the 28th International Conference on Computational Lin- guistics. pp. 3084-3084 3094, Barcelona, Spain (Online), December 2020. International Committee on Computational Linguistics. doi: 10. 10.18653/w1/2220.colll www.coooiingii...1111. main.275.
Yuxin Zuo, Wenxuan Jiang, Wenxuan Liu, Zixuan Li, Long Bai, Hanbin Wang, Yutao Zeng, Xiao- long Jin, Jiafeng Guo, and Xueqi Cheng. Alignxie: Improving multilingual information extraction by cross lingual alignment. arXiv preprint arXh::2411.04794, 2024.
14

---

Eizo a,xx for x [-1, By Lemma A.I there exists ReLU network fi [o, 1] + [ 1,11 of Published as a conference paper at ICLR 2025
A OMITTED PROOFS
In this section, provide the proofs that were omitted in the main body of the paper. The following proofs build upon the work of Yarotsky (007); Te Telgarsky (2017); Boullé et al. (2020).
A.I PROOF OF LEMMA I
Proof of Lemmal (e).U(x))
/
For ReLU activation, di 1, 04 0, Vi * 1, leading PolyRel.U(a)
For ReLU activation, set d2 1, a, = 0, + giving PolyReLU(z) (e)U3(x).
For general polynomial activation, observe that for Var IR and E N:
ReLLU (x) + (- (-1)ReLUU((a)) VEERVEN.
(13)
Thus, for any polynomial activation of order
Poly(*) PolyReLU, (x) + (olyRel.Ugg-1))
(14)
where (olyReLU, Ex &RRLUU (x) and (olyReLUg,s)) -E=s(-1)
L
A.2 PROOF OF THBOREM
The proof is elementary extension of Lemmal I
Using Lemma T we can represent the ReLU activation E using PolyReLU activation. Thus, we replace each ReLU activation in the ReLU network f with PolyReLU to construct a new network g. Obviously, such g satisfies the above requirements. Hence, the size and structure remain equivalent, and g serves as the PolyReLU network equivalent to the ReLU network.
Proof of Theorem !
A.3 PROOF OF LEMMA 2
The proof Lemma 2 leverages Lemma 3.4 from Telgarsky (2017), which we state below.
Lemma A.1 (Lemma 3.4 in Telgarsky (017 Let (0,1) be given. Suppose p : [0. 1]
[-1,1] be a r order polynomial with monomials and coefficients within [-1,1]. Then there ex-
ists ReLU network [0, 1] + [-1,1] of size O(min In(sr/e), sd sdln'
(dsr/e))) (dsr((err))))
such that
m0, (p(æ)) [(x)
Using this result,
proceed with the proof of Lemma 2
Proof of Lemma 2 First, we observe that PolyReLU(x)) Poly(ReLU(z), where Poly(x)
size (1/e) such that
max zE
(T() - <.
(15)
Thus, we construct f f1 ReL.U for inputs æ E [-1,1]] This yields that
max 1/(x) (())) = max |/1 ReLU(x)) (oly))LU(x)) =E(-1,1]
max |fa (ReeLU(())) - Poly(ReLU(z))
xE[-1,1|
(16)
max [f((x) - Poly(x))
=E[0,1]
Since fi is a ReLU network, the constructed function f = fi o ReLU is also a ReLU network, completing the proof.
15

---

Published as a conference paper at ICLR 2025
A.4 PROOF OF THEOREM! 2
The lower bound of Theorem 2 follows directly from Theorem 11 in Liang & Srikant (0017), restated here for clarity:
Lemma A.2 (Theorem 11 in Liang & Srikant (017)). Suppose function f : [0, 1] - R is differentiable and strongly convex. Let (0,1) be given and f be ReLU network. If lf()) - f()).. then the network size of is at least ()(lm(1/e)).
[0,
Lemma A.2 shows that approximating the quadratic function x² with error tolerance requires a network of size at least (lIm(//))). Since x² on [0, 1] is a degradation case of PolyReLU, any ReLU network approximating PolyReLU with error must also be at least ([In([/))) in size. The upper bound is proved in the following.
Proof of Theorem
Denote 9J the i-th layer of PolyReL.U neteeork / for 1 S i S I,, such that
9L 9L-1 o 91.
For each neuron, since ||a||] + b < 1, it follows that
la T x + bl S la x + |6| < |||||| + |6| S 1, Vo E {ø||<]]<> S 11..
(17)
Additionally, note that the range of PolyReLU is [-1, Hence, by induction, the output of each neuron remains within []. For each subnetwork gs. by applying Lemma 2 we can construct
corresponding ReLU network fa by replacing each PolyReLU activation Pi, 9, with ReLU activation. Specifically, for any i [L and j E [K]. there exists ReLU network ft.j [-1,1]
[-1, 1] that approximates the PolyReLU activation Pt,j with given tolerance 6 > 0.
Thus, the network fi is obtained by replacing each PolyReLU activation Pt,j in 9: with its ReLU approximation fag. Obviously, fü is a ReLU network whose output dimensions are in the range [],
Next, we give the approximation error bound. Denote hi = 9i ... 91 and hy = fi ....o fa for i E [T.). For the sake of brevity, we assume ng hg as the identity map in [-1, 1jd. Hence, we have go and h! = fi ....o fo- Suppose + (a. 4 b..;) be the output of j-th neuron of 9s- Denote the approximation between the PolyReLU network and the ReLU network at
hi 9:
i-th layer and j-th neuron etj. And we use e maxje[K] eij denote the approximation error
between the PolyReLU network and the ReLU network at i-th layer. Then for any i e [Z]], we have that
=E[-1,1 (Z)(z) ni, (2))
ei = max
pis(aT (x) + b(g) - Ses ((s(aT,(((112)) + brs))
max
=E[-1.1j
= max
PL. (a. 12 (z) + b1.)) (a (a 1
(x) + (s.s)
=E[-1,1]
+ + (a
T 11 (x) + b..)) Si (1 1/ (2)) + br.s)
< max
Pij (a A2 (z) b.n) P.o ( (aT, 1o6_ (x) + biu)
(18)
=f(-1,1d
max pay M/ (m) bb.)) fig (a [ (x) bas))
xE[-1.1]
(a 1/ (z)) b1.g) (aT, 1 (zz) + bio)) +
< max a
=1[-1,1
Sa max
||رره|| h/ (x) h (z) +
x1|-1,1]
Sa max =E[-1,1]
[ (x) 11 (x))
+
"We use the notation [T.] denote the set 21,2,.... T.).
16

---

Published as a conference paper at ICLR 2025
The first inequality is using the triangular inequality. The second inequality holds because the Lip- schitz constant of Ps.j is a and the ReLU subnetwork fig approximates Pi, with error 6,. In the fourth inequality, we used Hölder's inequality. Since ||a,,]|ı I ||atحl| + | I 1, the fifth inequality holds.
Therefore, derive the following approximation bound
max ei.j JE[K] aE(-1,1ja
S max
/1-1(x)) 11-1(10)
+ =-i=- EEs
(19)
for Vi E [L]. Since hg hg- we have eo O. Let c/(La for Vi [L]. It follows that
ic Lab-i
e
vi E [L].
(20)
Hence, the final error at the last layer is bounded by S.
Last, we need to estimate the size of the ReL.U network f. By Lemma 2 the size of each ReLU subnetwork fig is OInn ((att(Laa+/e))) Therefore, the total size of the ReLU network
Σ Knn (Lo)-)) =0 KI.
La" (Le))
(21)
where we use the fact that
Σ In La La
La
=0 Lin2 OLLnn (La²)
La"
Σ
In
i ilna
(22)
This completes the proof.
A.5 PROOF OF THEOREM 3
Before proving Theorem 3 we begin by introducing few useful lemmas.
Lemma A.3 (Proposition in Yarotsky (2017)). Let M N and p R + E be any continuous piece-wise linear function with M breakpoints. Then the following two statements hold:
For network with activation depth L and width K, there exists ReLU network with the same depth and width (M that computes the same function as the original network.
Conversely, ReLU network has depth L and width K, there exists network with activation p. depth L and width K that computes the same function on a bounded input domain DD
This result, Combined with Lemma D directly leads to the following corollary, which demonstrates that PolyReLU networks represent any piece-wise linear function exactly R. Corollary A.1. Let M E N and breakpoints. Then there exists PolyReLU network g of size (M) such that
R + R be any continuous piece-wise linear function with M
p(x) g(x),
Vr E IR.
In similar manner to Proposition 10 in Boullé et al. 2020), we can show that PolyReLUU networks can represent powers x" exactly for any n E N.
Lemma A.4. Suppose n,r N and > 2. Then x" can be represented exactly by PolyReLU network with an r-th order PolyReLU activation and size (ln?(n)).
Proof of Lemma A.4 We first prove that x" can be represented exactly by a polynomial network 9 with r-th order polynomial activation and having size (In2(n)). Based on 9, we construct PolyReLU network y that satisfies the requirements.
By expressing base r, we have that
k
orr
c (x")))
(23)
1=
1=0
17

---

Published as a conference paper at ICLR 2025
where & llog, n Ex= cor,, and E 1..1.2..... Each xn be represented
by a polynomial network with 2 + layers and width 1. It follows that x" can be represented by a polynomial network of size
Σ (0+1)= -(n)))((((()))))))
(24)
i=0
By Lemmal | we know that PolyReLU activation can represent a polynomial activation. Hence, there exists PolyReLU network g with an r-th order activation and size (In²(n)) such that g(x) = x". 4/2 E R.
With the above lemmas, we can now prove Theorem
Proof of Theorem 3 The proof is composed of two parts. We first approximate by local Taylor polynomials and continuous biece-wise linear functions and then represent these functions using PolyReLU networks, following Yarotsky (2017); Boullé al. (2020).
Part 1. Suppose N is a positive integer. We begin by dividing [-1, 1]d into grid of (2N + 1)d functions:
d II
Σ (m)(x)= 1, @m(
(3N
IR --N)) )) Ve (x1, 22, xa) 1-11,
i=1
where m = (m1, M2,,
ma) E E-NN,
-NN WN),,.... -))....
........... 0, 0,.... N) and xp is defined as
1,
|x|| <1 1
((z)) =
0,
2 < |x],
2 |x],
1 |x| 2. 2.
This function has the following properties:
max ZER
[p(x)] 1,
max [l0m())]l 1,
(25)
=E|-1,114
2
supp 6m x
<
Vm E (-N, -NN - NJJ.
(26)
3N
Part 2. We use degree- 1) local Taylor approximation of the function f. defined as
fx(x)= =
Σ
Pm))),
(27)
m
where Pm,. is the degree-(n (n 1) Taylor polynomial of f at m/N, i.e.,
Pm (x)
Σ
D
m
m
(28)
n!!
na: m:l/ <x
with conventions n! II nz! and (x m
IT i=
(z.
my )"
The approximation error between and fN can be bounded as follows
[f(x) (x) =
Σ
2m (f(x) - Pm
me{ N
Ny.
Σ
1s(x) Pmm(2))
m: [ 1ox *
2d
max
[f(x) - (m)
(29)
m: j-- Ioo < 2
us
2
2d
max
ess sup
17 fx=)lo
n!
3N
2
2d
n!
3.N
18

---

Published as a conference paper at ICLR 2025
The first inequality is because of the triangular inequality and Eq. we used the fact that Vx E [-1, 1] belongs to the support of at most 2d functions om. The third inequality is bound for the Taylor remainder and the fourth inequality uses the definition of Fn.d- Let
25 In the second inequality,
2d
2d
# +1.
N
(30)
nlE
have that
max -=111
If(x) - fN(()) <e.
(31)
Next, construct PolyReLU network ON represent SN exactly. Let Omn Since (1-11))) I 1, |mm.nl I for any m, we rewrite SN as
n! D" (M
fN(x)
Σ
E am Om mm) (x
(32)
mE{-N. N1 mellan
Therefore, JN is composed of at most (2N 1)d functions mm( (x m )² Since pm(a))
Ie (3N (rx- my and each yp (3N (Ix - The is continuous piece-wise linear function, we can apply Corollary A.1 which guarantees that there exists PolyReLU network @m of size (dd that can exactly represent 6m on RR, i.e., @m(x) = om 6m(x), Va E Rd. For (=- m )²
II (z. N )==, by Lemma A.4 we know that there exists PolyReLU network Im of size at most O(d (m)) such that Im (x) = (x - m
Vo Rd. Combining these results, we can now
N
construct larger PolyReLU network In follows
IN (a))=
Σ
Σ
am (2)mm(())))) 3 (2))gm (x))
(33)
mE(-N......N NI n: mnllis
where the total size of the network is
O (d + 1d + -(()))
Here, we use Eq. 30
determine the size bound terms of the error tolerance c. Clearly, we have
fx (x)),,
Va E Rd.
(34)
Hence, we conclude that
ff(x) - -xx(x)- = max 1f(x) fN(xx) <e. 1E1-1.1
max
(35)
xE1-1,11"
This completes the proof.
B DISCUSSION OF THE OPTIMAL APPROXIMATION RATE
For convenience, state Theorem 4.2 in Theorem (Theorem 4.2 in De Vore et al. F F. {f E X w IIJ1 & p g. N. then
Vore et al. (1989) in the following.
1989 ). Let 2 be a Banach space I.g on TR., s S oc.
sup inf llf - M())Il, 3 fer
(36)
BERm
where M be a mapping from Th' into X' which associate with each E R the element M(() E X, and C a constant.
Particularly, let tion error of the neural networks with m parameters to approximate p i.e., Fd,no larger than Cm #. Therefore, given error tolerance c, we have
do and X' = -,1- the above theorem tells us that the approxima-
6 < Cm
(37)
It follows that
> C C
(38)
Hence, the total number of parameters required by neural networks to approximate functions Fn.d ((E=()) #). Combining with Theorem 3 we have that our PolyReLU networks achieve the optimal approximation rate in the context of Sobolev spaces.
19

---

Published as a conference paper at ICLR 2025
C ACTIVATION FUNCTIONS
We provide definitions of several commonly used non-linear activation functions in Table
Table 4: Definition of activation functions.
Activation
Definition
ReLU (Narr & Hinton
2010
ReLU (a) = max 0}
ReLU
(So et al. 2021
ReLU² (x)) max 0}²
ReLU6 Krizhevsky et al.
2010
ReLU6 (x) = (2, 03.6)
LeakyReLU (x) = I, if 2 E (0, 1) is constant
Leaky ReLU (Maas et al.
2013)
az. otherwise
x:
if
RReLU (Xu al. 2015)
RReLU(())
o. otherwise,
is randomly sampled from uniform distribution
Ir, if f
Parametric ReLU (PReLU)
PReLU (x)
ax, otherwise,
(He al. 2015)
is learnable parameter
Tanh
exp(x)) Tanh(x) = exp(2)+oxp(-D
Softplus Glorot al. 2011
Softplus (x:) log(1 + exp(ax)),
a constant (default 1.0)
Mish Misraa 2019
Mish (ar) (aft(SSf)llussz))
Sigmoid
Sigmoid(z) = (x) T+exp(
SiLU(Swish) RRamachandran al.
2017)
SiLU(z) xx o(x)
if > 0
ELU Clevert
2015)
ELU(x) a (exp(x) 1), if S
is constant (default 1.0)
CELU Barron 2017)
CELU (x) = max x) + min(0, a (exp(x/a) 1)),
is constant (default 1.0)
GELU
Hendrycks & Gimpel 2016)
GELU(()) + ((x).
((z) is CDF for Gaussian distribution
GLU Dauphin et al.
2017)
GLU(x) d<+W) (=V)
SwiGLU Shazeer
2020
SwiGLU(x) = (iLU(>W) @ (zV),
W. V are learnable parameters
Poly
Poly(x) = Er= a,0" ag á E [r] are learnable parameters
20

---

Published as a conference paper at ICLR 2025
D PYTORCH IMPLEMENTATION OF POLYCOM
PyTorch implementations of PolyReLU and PolyNorm are provided in the following.
Algorithm PyTorch-Style Implementation of PolyReLU
impert terch
import checkpoint
import torch.nn. fun.tuccioona sa
(x,ly weight, bias, order=3);
return sum weeight [1] (x ** (1+1)) for 1 range (arder) bias
daf inii ini=__salf))
super (PolyReLU, selfi init
seif.weight Berch.nn.Pareeee onosch.ones 3) seif.bins
zcrea (1)
daf (self, checkpo' rting=trul:
(x)elu(x)
if checkpointing:
reture theckpoint self. weight,,, solé.bias, raa return _ooly(x (x, selt.weight, aelt.blas)
Algorithm PyTorch-Style Implementation of PolyNorm
_narn (x, epa=lc-6):
return x torchh (x. m-an keepdim=True) + epa)
_paly_nom (x, weight, bias, order=3) return sum whiight(1] _normix ** (1+1)1 for 1 in (anger)) blas
PolyNorm (wt.:ttooo...oooooo def ini: ini=__(eel)):
super (PolyNorm, _elf)._init__.. init__i |)
se.t.weight torch, torch.nn. Parameeter woorchoneesti 3:3)
seif.bise torch.nn.Paranee
worch.reres (1))
ooward (self,, checkpo' rtingetruel:
1: checkpointing:
return checkpoint (_paly_nomm, x, self, seighteight, self.bias, see_resntrant=Falee return _pely_norm self.weight, self.bios!
E EXPERIMENTAL DETAILS
E. ARCHITECTRRE
Table the model architecture used for the 1B dense model. To ensure comparable numbers of training parameters actoss different activation functions, we adjust the intermediate sizes accord- ingly. For SwiGLU, the intermediate size is set to 5504, while for other activation functions, it is set to 8256.
Table 6 outlines the model architecture used for the MoE models. Similarly, the intermediate size for SwiGLU is set to 1024, while for other activation functions, it is set to 1536.
Table 5: Model architecture of the IB dense model.
Params
Hidden Size
Context Length
Intermediate Size
Attention Heads
Hidden Layers
1.3B
2048
4096
5504/8256
16
24
E.2 HYPERPARAMETERS
In Table we list the hyperparameters that we use by default at training time for all our experiments for the 1B dense model and MoF- 1B-7B, unless stated otherwise.
21

---

Published as a conference paper at ICLR 2025
Table Model architecture of MoE model.
Activate Params
Total Params
Hidden Size
Intermediate Size
Attention Heads
1.3B
6.9B
2048
1024/1536
16
Hidden Layers
Exports
Active Exports
Context Length
Weight Tying
16
64
4096
no
Table 7: Pretraining hyperparameters for the 1B dense model and MoE- 1B-7B.
IB dense model
MoE-IB-7B
Optimizer
Adam
Adam W
Learning Rate (LR)
3E-4
4E-4
Minimum LR
3E-5
5E-5
LR Schedule
cosine
cosine
Weight Decay
0.1
0.1
81
0.9
0.9
82
0.95
0.95
Gradient Clipping
1
Warmup Tokens
620,000,000
Warmup Steps Init Distribution
2000
normal
trune normal
Init std
1/v2.5d
1/ 2.5d
Init Truncation
3x std
Load Balancing Loss Weight
0.01
Router z-loss Weight
0.001
E.3 DEFINITION OF EFFECTIVE RANK
We adopt the concept of effective rank from Roy & Vetterli (2077) to measure the effective dimen- sionality of a matrix. Given a matrix A with Singular Value Decomposition (SVD) A UVVT where D is diagonal matrix containing singular values σ 2 Ø2 singular value distribution as = 0 Σi=0 Jj. E [n)). The effective rank of A is then given by
o, O. we define the
Crank(A) = exp
- Inpi 1=0
(39)
F COMPUTATIONAL COMPLEXITY ANALYSIS
For the sake of simplicity, we only calculate the computational complexity in one-layer Feed- Forward Networks (FFN) since activation only. Support input tensor FFN is IR B where B, L, and H are the batch size, length of the sequence, and hidden size, respectively. Roughly, the relationship between computational FLOPs and model parameters can be regarded as proportional Therefore,
[ estimate the proportion of the computational cost incurred by the activation func- tion calculations within the total computational cost of the FFN matrix computations (2 The FLOPs ratio is calculated as
FLOPs ratio = FLOPS for activation
24BSH²
is important to note that the overhead and proportion often vary for different model sizes, so we provide the corresponding formulas directly and take // 1024, B 4 (each device), S = 4096, using BF16 precision as an example. For PolyReLU and PolyNorm, we use the 3-order default setting. The results are summarized in the following tables:
22

---

Published as a conference paper at ICLR 2025
Table Comparison of computational complexity for different activation functions without gradi- ent checkpointing.
Method
Intermediate Size
FLOPs for activation
FLOPs ratio
Memory Overhead
ReLU
4H
4BSH
1 0.016%
4BSH 128M
f2 0.29% Q
GELU
4H
72BSH
10BSH 320M B
H = 0.15%
SwiGLU
H
112 122 BSH
8BSH = 256M
ReLU2
4H
8BSH
37 0.032%
8BSH = B56MB
PolyNorm
4H
72BSH
Y 0.29%
12BSH 384M
PolyReLU
4H
40BSH
3H 0.16%
8BSH 2
Table 9: Comparison of computational complexity for different activation functions with gradient checkpointing.
Method
Intermediate Size
FLOPs for activation
FLOPs ratio
Memory Overhead
ReLU
4H
8BSH
0.033%
3
GELU
4H
144BSH
1 0.59%
SwiGLU
H
224 BSH
0.30%
0
Q 0.065%
ReLU²
4H
16BSH
O
3H 0.59%
PolyNorm
4H
144BSH
Hb = 0.33% 3H
PolyReLU
4H
80BSH
We assume that the scale of the input tensor set to [-1, In this case, the FLOPs for both tanh and exp are approximately 10 each. For a fair comparison, the intermediate size of models with SwiGLU activations set to 8/3H keep the overall numbers of parameters constant.
In practice, we utilized gradient checkpointing to reduce the additional memory overhead to O. While this may introduce certain computational overhead, given the overall modest computational cost of the activation functions, the overall increase in GPU memory and computational cost is quite small.
G SCALING CURVES
In Figure we present the training loss scaling curves for dense models utilizing the activation functions SwiGL.U, PolyReLU, and PolyNorm. As illustrated in the figure, both PolyReL and eters.
PolyNorm consistently outperform SwiGLU across model sizes ranging from 110M to 1.3B param-
The model sizes used for the scaling law experiments are detailed in Table and models employ the hyperparameters specified for IB dense models, as listed in Table 7 Models with 110M, 226M, and 502M parameters were trained on 200B tokens.
Table 10: Model sizes for scaling laws experiments.
Params
Hidden Size
Context Length
Intermediate Size
Attention Heads
Hidden Layers
110M
768
2048
2048/3072
16
12
226M
1024
2048
2560/3840
16
502M
1536
2048
4096/6144 144
16
16
1.3B
2048
4096
5504/8256
16
24
23

---

Published as a conference paper at ICLR 2025
SWiGLU
2.7
PolyReLu PolyNorm
2.6
Training Loss
2.5
2.4
2.3
2.2
Model Size
Figure Scaling curves of models with different activation functions.
H EXPERIMENTS ON VISION
To evaluate the effectiveness of PolyCom beyond language modeling, we trained ResNer50 (He al. (016)) on ImageNet-IK (Deng al. 2009) following the settings of timm (Wightman) 2019). For comparison, we replaced the ReLU activation in ResNet50 with PolyNorm and reported the training loss and top- lotp//5 accuracy on the evaluation set, shown in Figure 9 The results demonstrate that PolyNorm outperforms ReLU by a significant margin in terms of training loss, top-1 accuracy, and top-5 accuracy. Specifically, PolyNorm achieves a lower training loss of 2.026 compared to ReLU's 2. 2.121, and improves top-1 and top-5 accuracy to 75.117% and 92.099%. surpassing ReLU by +0.204 and +0.068, respectively.
Training Loss
Evaluation acc@1 (%)
Evaluation acc@5 (%)
ResNet50(ReLU)
6
ResNet50(PolyNorm)
80
60
60
4
40
20
ResNetSO(ReLU)
ResNet50(ReLU)
ResNet5O(PolyNorm)
20
ResNet50(PolyNorm)
2 0
50
100
150
200
0
50
100
150
200
50
100
150
200
Epoch
Epoch
Epoch
Figure 9: Training loss and evaluation accuracy ImageNet- KK for ResNet50 models with ReLU and PolyNorm activations. PolyNorm achieves lower training loss and higher top- 1pp-1/tp-5 accuracy, demonstrating improved performance.
24

---

Published as a conference paper at ICLR 2025
I ADDITIONAL RESULTS ON DENSE MODEL
More detailed results from our ablation studies are shown in Figures TO 11 and 12 These figures illustrate the training loss, validation loss, and validation perplexity (PPL.) for the TB dense model under different configurations.
The results of the IB dense models trained on 400 billion tokens are presented in Figure 13 As shown in the figure, models employing PolyReLU and PolyNorm consistently achieve significantly better performance compared to SwiGLU.
Training Loss
Validation Loss
Validation PPL
2order-Polylee
1.350
Ponder-FelyRall
Dorder-PelyRell
3order FolyReLU
PolyReLU
3order Poly
Polde--PlyRelLl
1.325
4order- 4order-PolyRall PolyRelU
3.8
4order-PulyBall
2.40
1.300
3.7
2.35
1.275
1.250
1.225
3.4
2.25
1.200
3.3
40
100
80
40
60
100
# of Tokens (B)
of Tokens (B)
Tokens (B)
Figure 10: training loss, validation loss, and validation perplexity (PPL) for the IB dense model with different orders of PolyReLU activation functions.
Training Loss
Validation Loss
Validation PPL
PolyReLU
1.350
PolyReLUU
3.9
PolyReLU
PolyPReLL
PolyFRwLU
PolyPReLU
2.45
PolyNorm
1.325
PolyNorm
3.8
PolyNorm
PolyReeLNNmmm
PolyfeeLLNormm
PolyfelLUNorm
1.300
3.7
2.35
1.275
1.250
2.30
1.225
3.4
2.25
1.200
3.3
20
40
60
80
100
20
60
80
100
40
60
100
# of Tokens (B)
of Tokens (B)
#
of Tokens (B)
Figure 11: training loss, validation loss, and validation perplexity (PPL) for the IB dense model with different polynomial compositions.
Training Loss
Validation Loss
Validation PPL
2.50
Relli
1.375
Rell ReLUF PolyReLU
ReUU
2.45
PolyReLU
Relu"
1.350
3.9
ReLy"
PolyRaLU
1.325
2.40
1.300
3.7
2.35
1.275
2.30
1.250
3.5
1.225
3.4
2.25
1.200
3.3
20
40
60
80
100
20
40
60
80
100
20
40
60
80
300
# of Tokens (B)
#
of Tokens (B)
#
of Tokens (B)
Figure 12: Training loss, validation loss, and validation perplexity (PPI.) for the 1B dense model with different variants of ReLU activation functions.
25

---

Published as a conference paper at ICLR 2025
Training Loss
Validation Loss
Validation PPL
3.70
PolyReLL Pory
PolyRall
PolyReLU
1.28
PolyWoorm
3,60
Polyorrm
2.30
1.25
1.23
3.40
1.20
3.30
2.20
1.18
1.15
2.15
1.12
3.10
100
200
300
400
# Tokens (B)
of Tokens (B)
Takens (B)
Figure 13: Training loss, validation loss, and validation perplexity (PPL) for the IB dense model with 400 billion training tokens.
J ADDITIONAL RESULTS ON MoE MODEL.
More results for the MoE model are provided in Figure 14 showcasing validation losses and down- stream evaluations after 200 billion training tokens. The comparison highlights models with differ- activation functions, such SwiGLU and PolyNorm. As shown, models with PolyNorm exhibit lower training and validation losses, along with superior downstream performance.
26

---

Published as a conference paper at ICLR 2025
Training
Validation
Validation
2.60
Nyliim
3.00
SwiU
Nlyiim
SWCL PolNam
Sm CLu
Nyluu
3.00
2.50
2.80
2.00
2.40
2.70
2.50
2.70
(peS2o)
Validation (Reddit)
Valldation (Stack)
Valldation (Wikipedia)
Saitill
SwGlll
SwiLil
2.40
PolyMorn
hwllem
Polloom
1.10
3.00
2.40
2.20
2.95
1.05
1.00
2.30
1cB
1c8
Validation loss (ICE)
Validation (M2D2)
Velidation loss (Pile)
Validation loss (Wikitext)
3.00
Polylom
PolyNom
Nory
SwitLU
Polylomm
2.80
3.20
2,30
2.40
1td
lxd
led
1sd
MMLUU stem Acc. (%)
MMLU humanities Acc. (%)
MMLU social sci Acc. (%)
MMLLU other Acc. (%)
47.5
32
37.5
45.0
30
42.5
32.5
28
SWKUU PolyHom
Swiku PolyKom
Se GLU
37.3
Polykam
25
1.60
160
160
1e0
MMLU Var Acc. (%)
HellaSSagg Acc. (%)
SclO Acc. (%)
ARC-Challenge Acc. (%)
61
50.0
$7.5
35.0
65.0
34
SW
Sw CLU
SWCW PolyNem
Sie GLU PolyWarm.
50
PoleNam
82.5
30.0
168
1c0
100
1.8
ARC-Easy Acc. (%)
PIQA Acc. (%)
WineGrande Acc. (%)
OpenbookOA Acc. (%)
65.0
72
Polyttomm
85
60.0
62
SantUi Payliom
SWIUUU PolyNorm
55.0
Se GLU PolyWomm
1xd
1xd
Boolq Acc. (%)
COPA Acc. (%)
CommonsenseOA Acc. (26)
SocialiQA Acc. (%)
Polyftum
60.0
45
57.5
35.0
SWUUU
GLU PolyWorm
Sie CLU PolyWorm
PetyWcmm
32
1.8
Tokenss
Figure 14: Validation loss and downstream evaluations for MoE models with 200 billion training tokens, comparing SwiGLU and PolyNorm activation functions. PolyNorm shows superior perfor- mance in terms of lower loss and better downstream results.
27


Reviews:
[{"review_id": "ncGvGTGiOW", "reviewer": "Reviewer_aNH6", "summary": "This work presents a method to extend activation functions using polynomials.\nTwo examples of these extended activation functions are introduced: PolyReLU and PolyNorm.\nPolyReLU is shown to achieve an optimal approximation rate in Sobolev spaces.\nLanguage modelling experiments indicate that the proposed activation functions accelerate training.", "questions": "1. Does theorem 2 make sense?\n2. What is happening in the proof of lemma 2?\n3. Do the PolyCom functions introduce additional learnable parameters?\n4. Do the different models in the comparison have the same number of parameters?\n5. How were the hyperparameters tuned?\n6. How do the models compare when trained to convergence?\n7. What does the performance look like as a function of training budget?\n8. Does this method work also provide benefits when used outside of language modelling?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- (clarity) Overall, the paper is clearly written and easy to follow.\n - (originality) The proposed activation functions are novel (to the best of my knowledge).", "weaknesses": "- (quality) The proof of lemma 2 (and therefore also lemma 2) seems to be incorrect.\n   It seems like the rank (and number of polynomials) from lemma 3.4 in (Telgarsky, 2017) have been ignored.\n   Furthermore, it is unclear why the minimum has disappeared because even $\\ln(1/x)^2 < \\ln(1/x)$ does **not** hold for $x \\in (0, 1)$.\n   Also, there is no argument for why the addition of the ReLU function in the proof does not change the size of the overall network.\n - (quality/significance) The lower bound in theorem 2 does not make sense.\n   If the PolyReLU network happens to have PolyReLU parameters $a_1 = 1$ and $\\forall i \\neq 1 : a_i = 0$,\n   it should be possible to model it exactly with a ReLU network of the same size.\n   Similarly, I do not see how the upper bound in theorem 2 would make sense.\n   It should always be possible to make the network arbitrarily large by adding neurons for which all incoming and outgoing weights are zero.\n - (clarity/quality) The polynomial activation functions seem to introduce additional learnable parameters.\n   However, the paper never explicitly states this anywhere.\n   Futhermore, there is no discussion on whether/how the number of parameters of the different models was controlled.\n   This possibly introduces a capacity advantage for the models with PolyCom functions compared to the baselines,\n   leading to an unfair comparison.\n - (clarity/quality) There is no discussion on how the hyper-parameters were found.\n   Furthermore, it seems like the same hyper-parameters were used for every model.\n   If these hyper-parameters were tuned on the proposed models, this would be an unfair comparison.\n - (quality) The authors claim that the activation functions are able to capture higher-order interactions.\n   However, based on the training curves (e.g. figure 3), none of the models were trained to convergence.\n   As a result, it is not possible to conclude anything concerning the complexity of the interactions that can be learned.\n   After all, it might be that the same interactions are captured for both models when converged, but one captures them \"faster\".\n   As a result, these experiments do not confirm the stated claims.\n - (quality/significance) There is no discussion on the runtime overhead of the introduced functions.\n   If the PolyCom functions introduce too much overhead, they might not be practically useful.\n   Also, it would be useful to include a comparison performance for a given compute budget.\n - (clarity) It is not clear why the proposed activation functions should be especially suited for language modelling.\n   I suspect that this should also work for models that do not require that much compute and allow for more extensive experiments.\n\n###### Minor Comments\n - There seem to be some type-setting issues with equation (7) (Thereom 1).\n - Could it be that there is a superscript $d$ missing for the space of $\\boldsymbol{m}$ on line 869-870?", "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": "zInXhRRGLj", "reviewer": "Reviewer_NayG", "summary": "This paper introduces polynomial composition activations (PolyCom), which is shown to be theoretically more expressive relative to common activation functions like ReLU and has an optimal approximation rate for general smooth functions in Sobolev spaces. Experiments show PolyCom, especially PolyNorm, achieves significantly better performance per token for training 1B parameters dense language models and 7B total parameters MoE model compared to SwiGLU, GeLU, and other ReLU variants.", "questions": "- What is the memory and runtime overhead of switching from, e.g., squared ReLU to PolyCom?\n- Can you show how PolyCom affects the scaling laws of loss v.s. model size or training computed by, for example, training additional smaller models?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- PolyReLU and PolyNorm show nontrivial performance gains for training language models with >1B parameters, even compared to strong baselines such as SwiGLU and squared ReLU.\n- The experiments are comprehensive, covering both pre-training and downstream evaluations.\n- PolyReLU has strong theoretical guarantees, showing it's more expressive than ReLU and has an optimal approximation rate for general smooth functions in Sobolev spaces.\n- The paper is well-presented.", "weaknesses": "- The paper does not discuss potential overhead in using PolyCom. For example, naively, using PolyCom would increase the activation memory by a factor of $r$ (set to 3 in the experiments) compared to ReLU. \n- The analysis section is not completely convincing. Why is a higher effective rank better? Both GeLU and SwiGLU have lower effective ranks than ReLU, but they achieve better performance.\n- It's not clear that PolyNorm MoE has lower layer-wise cosine similarity in Figure 7 compared to SwiGLU MoE.\n- An important unaddressed question is how the benefit of PolyCom scales to larger models. For example, is it most significant for smaller models and vanishing for larger models? I suspect this is not the case, but showing some evidence would be important.", "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": "hxDNyll3jq", "reviewer": "Reviewer_NvaU", "summary": "- The authors propose PolyCom, a set of polynomial composition activations that are tailored especially to the needs of transformer models. Theoretical and empirical evaluations compare the expressivity and other capabilities for PolyCom to common pre-existing activations, including traditional ReLU and SwiGLU.\nPolyCom applies a composition of a polynomial function and another function $\\rho$. Two variants of - PolyCom are focused on. Type 1 PolyCom applies $\\rho$ before the polynomial exponent, while Type 2 applies the polynomial exponent before $\\rho$ (Equation 1 in the paper). The paper’s evaluations concentrate on one specific substantiation for each type:\n  - For Type 1, the authors use PolyReLU. This activation sets $\\rho$ to ReLU and applies $\\rho$ before the polynomial. \n  - For Type 2, the authors use PolyNorm. This activation has the L2 normalization for $\\rho$ and applies the polynomial before $\\rho$.\n- The theory results focus on PolyReLU and proceed in three steps:\n  - In section 3.1, the authors show that the sets of all ReLU and ReLU^2 networks are subsets of the set of all PolyReLU networks, indicating that PolyReLU has stronger approximation abilities than ReLU or ReLU^2\n  - The main result of Section 3.2 (Theorem 2) is that the size of any ReLU network that approximates a PolyReLU network with depth L and width K within tolerance $\\epsilon$ must be at least $\\Omega(KL \\ln(\\epsilon^{-1}))$. By this, authors conclude that PolyReLU networks are more efficient than ReLU networks in terms of representational capacity.\n  - Theoretical results build to Section 3.3, which shows in Theorem 3 that PolyReLU networks achieve the optimal approximation rate: that is, there exists a PolyReLU network of size $O(\\epsilon^{-d/n})$ that can approximate arbitrary function $f$ in a unit ball within a Sobolev space $F_{d,n}$.\n- Empirical evaluations in section 4 apply PolyCom to one dense model with 1B parameters and one MoE model with 1B active and 7B total parameters. Comparison activations are ReLU, ReLU^2, GELU, and SwiGLU. Training datasets are RedPajama-1T for the dense and OLMoE Mix for the MoE model, and several other datasets are considered for downstream fine-tuning tasks. Authors state that lower training/validation loss and downstream task accuracy for PolyCom activations are indicative that PolyCom accelerates the convergence of LLMs and increases model expressivity, with PolyNorm performing generally better than PolyReLU as well. Ablations include the polynomial order of PolyCom and the choice of polynomial composition function $\\rho$, and the rank of model weights and layer-wise similarity are also compared favorably towards PolyNorm.", "questions": "- Would you mind including more results on the computational requirements required for PolyCom versus other activations, as I mention in the limitations section? It could also be interesting to include an experiment where PolyCom is compared to other activations, and model size/parameter count is somehow adjusted to ensure that models using PolyCom and models using other activations have comparable computational requirements \n- It’s stated that PolyReLU and PolyNorm have equivalent expressivity, e.g. on lines 164-165. Do you have a proof for this?\n- Have you tried any PolyCom activation in a non-transformer model? What causes you to think that they’re so good for transformers in particular, as opposed to also for other (deep-learning) architectures?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The idea of activation functions with better expressivity, without adding more trainable parameters to the model, is significant, useful and interesting\n- The construction/definition of PolyCom is very flexible and serves as a good base for future experimentation with other activation variants beyond the specific instantiations of PolyReLU and PolyNorm\n- The combination of both theoretical and empirical results is stronger than including just either type of result alone\n- Theoretical results are approachable and elegant\n- Paper is well-structured and well-written: it’s easy to follow and there were no confusing typos or other writing issues\n- Empirical evaluations are conducted on several different downstream datasets, adding strength to the authors’ claims", "weaknesses": "- No discussion of the differences in computational needs between PolyCom and prior activations: does it slow down training time a lot, or do you need much more memory? On lines 416-417 where you’re comparing different orders of PolyCom, you mention higher orders leading to overhead and overflow; I would want to hear more about that\n- Empirical results do not seem to be averaged across runs; I don’t see it stated anywhere. Results tables and graphs contain no error bars. Especially because the empirical results are fairly small in many areas, such as in Table 1, I would appreciate more evidence that these patterns hold across training seeds/different random initializations at the start of training\n- Other types of polynomial activations, such as those cited in Section 5’s paragraph on polynomial activation functions, are not compared against in the empirical results. Seemingly, no explanation is given for leaving these functions out.\n- Small comment: more discussion of the optimal approximation rate or a formal statement of DeVore et al.’s theorem 4.2 (cited on line 275) inside of your paper (even the appendix) could be nice for readers", "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": "ihRgLxAm9H", "reviewer": "Reviewer_TiU9", "summary": "This paper introduces PolyCom, a polynomial composition activation for transformers. Through theoritical analysis, PolyCom is shown to enhance expressivity and effectiveness over other activations. Empirical experiments on large language models (LLMs), both dense and sparse, demonstrate that replacing standard activations with PolyCom enables LLMs to capture higher-order data interactions, improving accuracy and convergence rates.", "questions": "Check the weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1) The authors demonstrate the expressivity of the proposed activation function both theoretically and empirically, using effective rank analysis and layer-wise similarity metrics. These methods provide insights into how the activation enhances the model’s ability to represent complex patterns and distinctions between layers, showcasing its potential advantages over traditional activation functions.\n\n2) Extensive experiments on downstream tasks, along with improved convergence rates in the learning curves with fixed parameter size model, highlight the potential of the proposed activation function.\n\n3) Overall, the paper is clearly written, with the description of polynomial activation easy to follow, and the results presented concisely.", "weaknesses": "1) Computational complexity analysis is not provided. Can the authors provide some analysis on the inference time throughput of the proposed activations with others?\n\n2) How do the authors address stability or manage potential exploding gradient issues with the polynomial activation?\n\nTypos: Line 274: Dowmstream Evaluation -> Downstream Evaluation", "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)