Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models

## Abstract

Abstract Transformers have found extensive applications across various domains due to the 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 .

## 1 Introduction

Transformers (Vaswani et al., ) have revolutionized the field of deep learning, facilitating unprecedented advancements in natural language processing (Radford et al., ), computer vision (Dosovitskiy et al., ), and beyond (Dong et al., ; Arnab et al., ). Characterized by their attention 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 a crucial role in determining the output of each neuron within a neural network. Traditionally, simple nonlinearities such as Rectified Linear Unit (ReLU) (Nair & Hinton, ) and its variants (Hendrycks & Gimpel, ; So et al., ) 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 a novel category of polynomial composition activation functions (PolyCom), specifically engineered to 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 expressiveness endows the model with superior expressive capacity, enabling it to capture higher-order interactions that might otherwise be neglected. Unlike other forms of polynomials ((Hornik et al., ; Trefethen, )) that suffer from inadequate approximation, exploding values, and oscillatory behavior, we demonstrate that PolyCom possesses a more potent expressive capability than both ReLU and traditional polynomials and achieves optimal approximation within Sobolev space.

We posit that the integration of polynomial composition activations within transformer models can lead to enhanced performance in tasks requiring intricate data interpretation. To evaluate this hypothesis, 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 composition activations in comparison to those utilizing traditional activation functions. The results indicate that the proposed method not only improves model accuracy, but also accelerates convergence rates, thereby suggesting that polynomial composition activations provide a substantive advantage in deep learning applications.

The main contributions of this paper are summarized in the following.

- •
We propose a new activation function PolyCom which is a composition of the polynomial and other types of function. In particular, we introduce two instances of PolyCom: PolyReLU and PolyNorm, and details 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 a PolyReLU network of size $O(\epsilon^{-d/n})$ can approximate any function in Sobolev spaces with error tolerance $\epsilon$, achieving optimal approximation rates.
- •
Empirically, we validate the effectiveness of this new activation function on LLMs with both 1B dense models and MoE models with 1B 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 formulation of PolyCom and discuss its integration within transformer architectures. Section 3 delivers a comprehensive theoretical analysis of PolyCom, emphasizing its enhanced expressivity and effectiveness. In Section 4, we provide a detailed account of our experimental results involving large language models (LLMs). Section 5 provides an overview of related work in the field of activation functions and their applications in transformer models. Finally, we conclude the paper and outline potential directions for future research.

## 2 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.

PolyCom. The study of the polynomial activation function can be traced back to the seminal work of Hornik et al. (), which showed that neural networks with polynomial activation are 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, ). To overcome these limitations, we propose PolyCom, a novel composition of polynomial and other functions. Specifically, we explore two composition approaches:

$$ $\begin{cases}\text{Type I:}&x\mapsto\sum_{i=0}^{r}a_{i}\rho^{i}(x),\\ \text{Type II:}&x\mapsto\sum_{i=0}^{r}a_{i}\rho(x^{i}),\end{cases}\\ \quad a_{i}\in{\mathbb{R}},$ (1) $$

where $r\in{\mathbb{N}}$ denotes the order of PolyCom and $\rho$ represents an arbitrary functions 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. It can be theoretically shown that both approaches have equivalent expressivity, provided that $\rho$ is a non-linear function. This is because polynomial terms are symmetric with respect to composition, allowing both Type I and Type II to approximate similar function classes. In other words, rearranging the order of $\rho$ and the polynomial powers does not affect the ability to approximate complex non-linear functions. In practice, we use a third-order PolyCom ($r=3$) with trainable coefficients $a_{i}$. For initialization, we set $a_{i}=1/r$ for $i=1,2,\dots,r$ and $a_{0}=0$.

For Type I PolyCom, we specifically consider a composition involving the ReLU function due to its simplicity, which we term PolyReLU. An $r$-order PolyReLU is defined as:

$$ $\mathrm{PolyReLU}(x)=\sum_{i=0}^{r}a_{i}\mathrm{ReLU}^{i}(x),$ (2) $$

where$\mathrm{ReLU}^{i}(x)=\max\{x,0\}^{i}$. This formulation can be seen as an extension of both ReLU and square ReLU.

For Type II PolyCom, we introduce PolyNorm, which normalizes the powers to ensure consistent magnitudes across terms:

$$ $\mathrm{PolyNorm}({\bm{x}})=\sum_{i=0}^{r}a_{i}\frac{{\bm{x}}^{i}}{\|{\bm{x}}^ {i}\|_{2}},$ (3) $$

where ${\bm{x}}^{i}=[x_{1}^{i},x_{2}^{i},\cdots,x_{d}^{i}]^{\top}$ represents element-wise exponentiation, and $\|\cdot\|_{2}$ denotes the $L_{2}$ normalization.

Figure: Figure 2: Block diagrams of Transformer MLP blocks utilizing ReLU/GELU, SwiGLU, PolyReLU and PolyNorm. “FC” stands for Fully Connected layer. “$x^{i}$” represents the $i$-th power of the input tensor $x$, while “$a_{j}$” denotes the $j$-th element of the learnable weight vector $a$. “N” indicates a normalization operation.
Refer to caption: x2.png

Integration into Transformer.
The transformer architecture (Vaswani et al., ) consists of two alternating modules, Multi-Head Attention (MHA) and position-wise Feed-Forward Networks (FNN). Activation functions predominantly influence the performance of FFN layers. We begin by formalizing the common paradigm of FFN,

$$ $\mathrm{FFN}_{\rho}({\bm{x}})=\rho({\bm{x}}W_{1})W_{2}$ (4) $$

where $\rho$ represents the activation function such as ReLU, GeLU, PolyReLU, or PolyNorm. We replace the traditional activation function with our proposed PolyCom variants to enhance model capacity and performance, as illustrated in Figure 2.

## 3 Theoretical Analysis

As discussed in Section 2, PolyReLU and PolyNorm have equivalent expressivity. 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 us 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 by PolyReLU

In this subsection, we present theoretical results on approximating ReLU networks using PolyReLU networks. The following lemma shows that ReLU, ReLU^2, 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 .

$\mathrm{ReLU}$, $\mathrm{ReLU}^{2}$ and polynomial activation can be represented by $\mathrm{PolyReLU}$.

###### Proof of Lemma 1 .

For ReLU activation, set $a_{1}=1$, $a_{i}=0,\forall i\neq 1$, leading to $\mathrm{PolyReLU}(x)=\mathrm{ReLU}(x)$.

For ReLU^2 activation, set $a_{2}=1$, $a_{i}=0,\forall i\neq 2$, giving $\mathrm{PolyReLU}(x)=\mathrm{ReLU}^{2}(x)$.

For a general polynomial activation, observe that for $\forall x\in\mathbb{R}$ and $i\in\mathbb{N}$:

$$ $x^{i}=\mathrm{ReLU}^{i}(x)+(-1)^{i}\mathrm{ReLU}^{i}(-x),\quad\forall x\in{ \mathbb{R}},\forall i\in{\mathbb{N}}.$ (5) $$

Thus, for any polynomial activation of order $r$,

$$ $\mathrm{Poly}(x)=\mathrm{PolyReLU}_{1}(x)+\mathrm{PolyReLU}_{2}(-x),$ (6) $$

where $\mathrm{PolyReLU}_{1}(x)=\sum_{i=0}^{r}a_{i}\mathrm{ReLU}^{i}(x)$ and $\mathrm{PolyReLU}_{2}(x)=\sum_{i=1}^{r}(-1)^{i}a_{i}\mathrm{ReLU}^{i}(x)$.
∎

Building on Lemma 1, we can formally prove that any ReLU network can be exactly represented by a PolyReLU network of the same size, as stated in the following theorem.

###### Theorem 1 .

Let $f:[-1,1]^{d}\rightarrow[-1,1]$ be a ReLU network with depth $L$ and width $K$. Then, there exists a PolyReLU network $g:[-1,1]^{d}\rightarrow[-1,1]$ of size $O(LK)$
such that

$$ $f({\bm{x}})=g({\bm{x}}),\quad\text{for}\forall{\bm{x}}\in[-1,1]^{d}.$ (7) $$

This theorem, proved in Appendix A, shows that PolyReLU networks can exactly match the representational power of ReLU networks without increasing the model size.

###### Lemma 1 .

###### Proof of Lemma 1 .

###### Theorem 1 .

### 3.2 Approximating PolyReLU with ReLU networks

In this part, we give theoretical results on approximating PolyReLU networks using ReLU networks.
The following Lemma 2 demonstrates that the PolyReLU activation can be approximated by a ReLU network within a given error tolerance.

###### Lemma 2 .

For the activation $\mathrm{PolyReLU}(x)=\sum_{i=0}^{r}a_{i}\mathrm{ReLU}^{i}(x),x\in[-1,1]$ with $a_{i}\in[-1,1]$. Given any $\epsilon\in(0,1)$, there exists a ReLU network $f:[-1,1]\rightarrow[-1,1]$ with size $O(\ln^{2}(1/\epsilon))$, such that

$$ $\max_{x\in[-1,1]}|f(x)-\mathrm{PolyReLU}(x)|<\epsilon.$ (8) $$

The proof is provided in the Appendix A. Lemma 2 establishes an upper bound on the size of a ReLU network needed to approximate a PolyReLU activation function. This result highlights that while ReLU networks can approximate PolyReLU activations, they require a significantly larger number of parameters.

Building on Lemma 2, we derive the following theorem, which provides both upper and lower bounds for approximating PolyReLU networks with ReLU networks.

###### Theorem 2 .

Let $g:[-1,1]^{d}\rightarrow[-1,1]$ be a PolyReLU network with depth $L$ and width $K$, and PolyReLU activation with order $r$ and Lipschitz constant $\alpha$. Suppose each neuron computes $x\mapsto\mathrm{PolyReLU}(a^{\top}x+b)$ with the pair $(a,b)$ satisfies $\|a\|_{1}+b\leq 1$ and $\mathrm{PolyReLU}:[-1,1]\rightarrow[-1,1]$ (a, b, and PolyReLU are possibly distinct across neuron).
For any given $\epsilon\in(0,1)$, there exists a ReLU network $f:[-1,1]^{d}\rightarrow[-1,1]$ of size

$$ $O\left(LK\ln^{2}\left(\frac{L\alpha^{L}}{\epsilon}\right)\right),$ (9) $$

such that

$$ $\max_{{\bm{x}}\in[-1,1]^{d}}|f({\bm{x}})-g({\bm{x}})|<\epsilon.$ (10) $$

Conversely, the size of any ReLU network that approximates a PolyReLU network within tolerance $\epsilon$, must be at least

$$ $\Omega\left(KL\ln\left(\frac{1}{\epsilon}\right)\right).$ (11) $$

Theorem 2 tells us that the total number of trainable parameters required by ReLU networks to approximate a PolyReLU neural network within a tolerance of $\epsilon$ is $O(\ln^{2}(1/\epsilon))$. Conversely, there exists a PolyReLU network that can not be approximated by ReLU networks of size less than $\Omega(\ln(1/\epsilon)$. Combined with theorem 1, we conclude that PolyReLU networks are more efficient in terms of representational capacity than ReLU networks.

###### Lemma 2 .

###### Theorem 2 .

### 3.3 Approximation of General Smooth Function

Similar to Yarotsky (); Boullé et al. (), we also explore the universal approximation capabilities of PolyReLU networks in the context of Sobolev spaces (Adams & Fournier, ). Specifically, we show that PolyReLU networks achieve the optimal approximation rate within these spaces, meaning that PolyReLU networks require minimum parameters to approximate general smooth functions in Sobolev spaces, compared with networks with the other activation.

The definition of Sobolev space ${\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)$ is stated below. The set $[-1,1]^{d}$ can be replaced by any compact set in ${\mathbb{R}}^{d}$, we use it just for the sake of brevity.

###### Definition 1 (Sobolev Spaces) .

For $n,d\in{\mathbb{N}}$, Sobolev space ${\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)$ is defined as

$$ ${\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)=\left\{f\in L^{\infty}\left([- 1,1]^{d}\right)|\|f\|_{{\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)}<\infty \right\},$ (12) $$

with the norm which defined as the following

$$ $\|f\|_{{\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)}=\max_{{\bm{n}}:\|{\bm{ n}}\|_{1}\leq n}\operatornamewithlimits{ess\ sup}_{{\bm{x}}\in[-1,1]^{d}}\left \|D^{{\bm{n}}}f({\bm{x}})\right\|_{\infty},$ (13) $$

where ${\bm{n}}\in{\mathbb{N}}^{d}$ and $D^{{\bm{n}}}f$ is the respective weak derivative of $f$, and $\operatornamewithlimits{ess\ sup}$ means the essential supremum in functional analysis.

Intuitively, a Sobolev space is a space of functions endowed with a weaker notion of smoothness compared to differentiability and possessing generalized derivatives. The Sobolev space ${\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)$ contains functions from $C^{n-1}\left([-1,1]^{d}\right)$ which consists of functions whose derivatives of order $n-1$ are Lipschitz continous.
In the sequel, we mainly consider the unit ball within ${\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)$, which is defined as follows

$$ $F_{n,d}=\{f\in{\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)|\|f\|_{{\mathcal {W}}^{n,\infty}\left([-1,1]^{d}\right)}\leq 1\}.$ $$

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 $F_{n,d}$.

###### Theorem 3 .

Suppose that $d,n\in{\mathbb{N}}$ and $\epsilon\in(0,1)$. For any $f\in F_{d,n}$, there exists a PolyReLU network $g$ with size $O(\epsilon^{-d/n})$ that can approximate $f$ at a given error tolerance $\epsilon$, i.e.,

$$ $\max_{{\bm{x}}\in[-1,1]^{d}}\|f({\bm{x}})-g({\bm{x}})\|_{\infty}<\epsilon.$ (14) $$

Theorem 3 indicates that PolyReLU networks can achieve an *optimal approximation rate* of $O(\epsilon^{-d/n})$. In contrast, previous works by Yarotsky () demonstrated that ReLU networks require $O(\epsilon^{-d/n}\ln(1/\epsilon))$ parameters to achieve a similar approximation error. Similarly Boullé et al. () showed that rational neural networks need $O(\epsilon^{-d/n}\ln(\ln(1/\epsilon)))$ 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 DeVore et al. () shows that the total number of parameters required by neural networks to approximate functions in $F_{n,d}$ is $\Omega(\epsilon^{-d/n})$. Therefore, our PolyReLU networks achieve the *optimal approximation rate* in the context of Sobolev spaces.

###### Definition 1 (Sobolev Spaces) .

###### Theorem 3 .

## 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: a 1B dense model and a Mixture of Experts (MoE) model with 1B active and 7B total parameters. The 1B dense model contains approximately 1.3 billion parameters with an architecture similar to Llama 2 (Touvron et al., ). For the MoE model, we use the OLMoE framework (Muennighoff et al., ), which activates 1.3B parameters out of a 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 $d_{model}=2048$ and 16 attention heads. In the MoE model, the transformer is composed of 16 layers, with a hidden size of $d_{model}=2048$, 16 attention heads, and 64 experts. To maintain a consistent number of trainable parameters across all activation functions, we adjust the intermediate size accordingly. Specifically, for SwiGLU, the intermediate size is set to two-thirds that of the other activations in all experiments.

Datasets. The dense model is trained on the RedPajama-1T dataset (^1^11RedPajama-1T is available at [https://github.com/togethercomputer/RedPajama-Data](https://github.com/togethercomputer/RedPajama-Data).) (Computer, ), which was developed by the open-source AI community to enable competitive performance against proprietary models.
The MoE model is trained on the OLMoE Mix dataset (^2^22OLMoE Mix dataset is available at [https://huggingface.co/datasets/allenai/OLMoE-mix-0924](https://huggingface.co/datasets/allenai/OLMoE-mix-0924).) (Muennighoff et al., ).

Hyperparameters.
Unless otherwise specified, we use a third-order PolyCom by default and initialize the coefficients as $a_{i}=1/3$ for $i=1,2,3$ and set $a_{0}=0$. Model weights are randomly initialized. For optimization, we apply the AdamW optimizer with $\beta_{1}=0.9$ and $\beta_{2}=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 a cosine scheduler. The MoE model starts with a learning rate of 4e-4, also decaying according to a cosine schedule.

Evaluation To evaluate the performance of LLMs with PolyCom, we use a wide range of open benchmarks, including ARC-Easy (Clark et al., ), ARC-Challenge (Clark et al., ), HellaSwag (Zellers et al., ), PIQA (Bisk et al., ), SciQ (Welbl et al., ), CoQA (Reddy et al., ), Winogrande (Sakaguchi et al., ), MMLU (Hendrycks et al., ), BoolQ (Clark et al., ), COPA (Gordon et al., ), CSQA (Talmor et al., ), OBQA (Mihaylov et al., ), and SocialIQA (Sap et al., ).
We utilize the LM Eval Harness (Gao et al., ) for standardized performance evaluation.

### 4.2 Results on Dense Model

**Table 1: Overall results of the 1B dense model with different activation functions, reported in terms of training loss, validation perplexity, and downstream accuracy (%). ARC-E and ARC-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$\downarrow$ | PPL$\downarrow$ | ARC-E | ARC-C | HellaSwag | PIQA | SciQ | Winograde | Avg.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 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 |

Training Dynamics of 1B Dense Model. Figure 1 compares the training dynamics of the 1B 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 to 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.

Dowmstream Evaluation. Table 1 presents the training loss, validation perplexity, and downstream task accuracy (%) after processing 250 billion training tokens. The downstream tasks include ARC-Easy, ARC-Challenge, HellaSwag, PIQA, SciQ, and Winograde. More detailed results are provided in Appendix D. The results clearly demonstrate that the PolyCom family (PolyReLU and PolyNorm) outperforms the other activation functions. For instance, PolyNorm outperforms SwiGLU by an average margin of 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-1B-7B, which has 1 billion activate parameters and 7 billion total parameters (Muennighoff et al., ). Due to computational constraints, we compare only PolyNorm activation function, shown to perform best in dense models, with the widely used SwiGLU activation function, which is commonly employed in current LLM architectures.

Figure: Figure 3: Traning and validation loss on C4 and Wikipedia for MoE models with 200 billion training tokens. We compare models using SwiGLU and PolyNorm activation functions. PolyNorm demonstrates lower training and validation losses, indicating faster convergence.
Refer to caption: x3.png

Figure: Figure 4: Downstrean performance dynamics on HellaSwag, MMLU Var, ARC-Challenge, and SciQ for MoE models with 200 billion training tokens. Models with PolyNorm significantly outperform those with SwiGLU on downstream tasks.
Refer to caption: x4.png

Training dynamics of MoE model. In Figure 3, we report the training and validation loss of MoE models trained on 200 billion tokens. Models using PolyNorm consistently show lower losses compared to those using SwiGLU, indicating that PolyNorm enables faster learning. Figure 4 shows the downstream performance on HellaSwag, MMLU Var(^3^33MMLU Var is a variant of MMLU (Hendrycks et al., ) using varied few-shots (Muennighoff et al., ).), ARC-Challenge, and SciQ. PolyNorm outperforms SwiGLU on all tasks, with notable improvements, demonstrating superior generalization capabilities.

**Table 2: Validation losses of MoE models with different activation functions. CC denotes Common Crawl. Best results per column are bold.**
| Methods | C4 | Books | CC | peS2o | Reddit | Stack | Wiki-<br>pedia | ICE | M2D2 | Pile | Wiki-<br>text | Avg.$\downarrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 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 |

Dowmstream Evaluation. Table 2 presents the validation losses on 11 datasets. PolyNorm consistently achieves lower validation losses than SwiGLU across all datasets, with an average improvement of 0.02. In Table 3, we also observe that PolyNorm outperforms SwiGLU on 8 downstream tasks. These results highlight the superior performance of models using the PolyNorm activation function. Additional results can be found in Appendix E.

**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.**
| Tasks | MMLU<br>Var | Hella-<br>Swag | SciQ | ARC-C | ARC-E | PIQA | Wino-<br>Grande | OQA | COPA | Avg.$\uparrow$ |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| SwiGLU | 37.07 | 66.49 | 90.60 | 37.12 | 71.58 | 76.61 | 62.75 | 39.80 | 83.00 | 62.78 |
| PolyNorm | 37.27 | 67.63 | 92.40 | 38.46 | 70.70 | 77.04 | 62.19 | 40.60 | 84.00 | 63.37 |

### 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 set $\{2,3,4\}$ and plot the results in Figure 5(a). As seen, the convergence speed improves as the order increases. However, there is no noticeable difference between orders 3 and 4 in terms of convergence speed. Additionally, increasing the order can lead to computational overhead and overflow issues, particularly when using low-precision 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 function, 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 better outcomes.

Variants of ReLU. In Figure 5(c), we compare different variants of the ReLU activation function, including ReLU and ReLU^2. PolyReLU consistently outperforms both ReLU and ReLU^2 across all tasks, highlighting the benefits of using polynomial composition. This result reinforces the hypothesis 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 significantly increasing model size or complexity.

Rank of Weights. To understand how PolyCom enhances model performance, we analyze the rank of the weights in each FNN layer of the transformer. We use the effective rank (Roy & Vetterli, ) to measure the effective dimensionality of weights and its definition is in Appendix C.2. Figure 6 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 indicates greater capacity for representing complex patterns in the data. These findings suggest that PolyCom improves the expressibility of transformers by allowing the FNN 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, as deeper layers are not 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 downstream task performance. This increased diversity across layers enhances the model’s ability to capture complex relationships, thereby improving the overall effectiveness of LLMs.

Figure: ((a)) Different orders of PolyReLU
Refer to caption: x5.png

Figure: ((a)) Rank of $W_{\text{up}}$, dense
Refer to caption: x8.png

Figure: ((a)) SwiGLU, dense
Refer to caption: x12.png

## 5 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 (Goodfellow et al., ). 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, ) mitigated some of these issues by offering a simple, non-saturating nonlinearity, which has since become a standard in many deep learning applications. Variants of ReLU, such as Leaky ReLU (Maas et al., ) and Parametric ReLU (PReLU) (He et al., ), were developed to address the “dying ReLU” problem by allowing a small, non-zero gradient when the input is negative. Other functions, like the Exponential Linear Unit (ELU) (Clevert, ), aimed to provide smoother activation profiles, resulting in better generalization and faster convergence in certain tasks. Moreover, Manessi & Rozza () proposed a combination of weighted base activation functions for further enhancement.

Polynomial activation functions (Hornik et al., ; Oh et al., ), although less commonly used, have been studied in various contexts for their ability to model higher-order, complex relationships more effectively.
For instance, Lokhande et al. () introduced Hermite polynomial activations to improve pseudo-label accuracy, while Chrysos et al. () proposed polynomial networks, $\Pi$-nets, which apply to various domains such as image and audio processing. Building on this, Chrysos et al. () utilized regularization techniques to enhance the performance of polynomial networks. These works highlight the potential of polynomial functions to increase the expressiveness 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., ; Kidger & Lyons, ; Kubjas et al., ).

The choice of activation function in transformers has also become an important area of research. Originally developed for natural language processing, transformers (Vaswani et al., ) have been effectively adapted for diverse tasks, including image recognition, speech processing, and reinforcement learning. Despite their broad applicability, the activation functions predominantly utilized in transformers, ReLU and GELU, have seen minimal evolution. Recent studies, however, have begun to explore alternatives to these conventional activations. For example, the Swish activation (Ramachandran et al., ; Shazeer, ) and the Mish activation (Misra, ) 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. (), with SwiGLU (Shazeer, ), a prominent variant, being used in models such as LLaMA-Series (Touvron et al., ).

## 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, ablation studies indicate that PolyCom increases model expressivity by elevating weight rank and reducing redundancy across layers. These findings underscore the significant potential of polynomial-based activations to improve transformer models, thereby paving the way for future research endeavors.

## Appendix A Omitted Proofs

In this section, we provide the proofs that were omitted in the main body of the paper. The following proofs build upon the work of Yarotsky (); Telgarsky (); Boullé et al. ().

### A.1 Proof of Lemma 2

The proof of Lemma 2 leverages Lemma 3.4 from Telgarsky (), which we state below.

###### Lemma A.1 (Lemma 3.4 in Telgarsky ( 2017 ) ) .

Let $\epsilon\in(0,1)$ be given. Suppose $p:[0,1]^{d}\rightarrow[-1,1]$ be a $r$ order polynomial with $s$ monomials and coefficients within $[-1,1]$. Then there exists a ReLU network $f:[0,1]^{d}\rightarrow[-1,1]$ of size $O(\min\{sr\ln(sr/\epsilon),sd\ln^{2}(dsr/\epsilon)\})$ such that $\max_{{\bm{x}}\in[0,1]^{d}}|p({\bm{x}})-f({\bm{x}})|<\epsilon$.

Using this result, we now proceed with the proof of Lemma 2.

###### Proof of Lemma 2 .

First, we observe that $\mathrm{PolyReLU}(x)=\mathrm{Poly}(\mathrm{ReLU}(x)$, where $Poly(x)=\sum_{i=0}^{r}a_{i}x^{i}$ for $x\in[-1,1]$. By Lemma A.1, there exists a ReLU network $f_{1}:[0,1]\rightarrow[-1,1]$ of size $O(\ln^{2}(1/\epsilon))$ such that:

$$ $\max_{x\in[0,1]}|f_{1}(x)-\mathrm{Poly}(x)|<\epsilon.$ (15) $$

Thus, we construct $f=f_{1}\circ\mathrm{ReLU}$ for inputs $x\in[-1,1]$. This yields that:

$$ $\displaystyle\max_{x\in[-1,1]}|f(x)-\mathrm{PolyReLU}(x)|$ $\displaystyle=\max_{x\in[-1,1]}|f_{1}\circ\mathrm{ReLU}(x)-\mathrm{PolyReLU}(x)|$ (16) $\displaystyle=\max_{x\in[-1,1]}|f_{1}(\mathrm{ReLU}(x))-\mathrm{Poly}(\mathrm{ ReLU}(x))|$ $\displaystyle=\max_{x\in[0,1]}|f_{1}(x)-\mathrm{Poly}(x)|$ $\displaystyle<\epsilon.$ $$

Since $f_{1}$ is a ReLU network, the constructed function $f=f_{1}\circ\mathrm{ReLU}$ is also a ReLU network, completing the proof.
∎

###### Lemma A.1 (Lemma 3.4 in Telgarsky ( 2017 ) ) .

###### Proof of Lemma 2 .

### A.2 Proof of Theorem 1

The proof is an elementary extension of Lemma 1.

###### Proof of Theorem 1 .

Using Lemma 1, we can represent the ReLU activation on ${\mathbb{R}}$ using a 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 1 .

### A.3 Proof of Theorem 2

The lower bound of Theorem 2 follows directly from Theorem 11 in Liang & Srikant (), restated here for clarity:

###### Lemma A.2 (Theorem 11 in Liang & Srikant ( 2017 ) ) .

Suppose function $f:[0,1]^{d}\rightarrow{\mathbb{R}}$ is differentiable and strongly convex. Let $\epsilon\in(0,1)$ be given and $\tilde{f}$ be a ReLU network. If $\max_{{\bm{x}}\in[0,1]^{d}}|f({\bm{x}})-\tilde{f}({\bm{x}})|$, then the network size of $\tilde{f}$ is at least $\Omega(\ln(1/\epsilon))$.

Lemma A.2 shows that approximating the quadratic function $x^{2}$ with an error tolerance $\epsilon$ requires a network of size at least $\Omega(\ln(1/\epsilon))$. Since $x^{2}$ on $[0,1]^{d}$ is a degradation case of PolyReLU, any ReLU network approximating PolyReLU with error $\epsilon$ must also be at least $\Omega(\ln(1/\epsilon))$ in size. The upper bound is proved in the following.

###### Proof of Theorem 2 .

Denote $g_{i}$ as the $i$-th layer of PolyReLU neteeork $f$ for $1\leq i\leq L$, such that:

$$ $g=g_{L}\circ g_{L-1}\circ\cdots\circ g_{1}.$ $$

For each neuron, since $\|a\|_{1}+b\leq 1$, it follows that:

$$ $|a^{\top}{\bm{x}}+b|\leq|a^{\top}{\bm{x}}|+|b|\leq\|a\|_{1}\|{\bm{x}}\|_{ \infty}+|b|\leq 1,\forall{\bm{x}}\in\{{\bm{x}}|\|{\bm{x}}\|_{\infty}\leq 1\}.$ (17) $$

Additionally, note that the range of PolyReLU is $[-1,1]$. Hence, by induction, the output of each neuron remains within $[-1,1]$. For each subnetwork $g_{i}$, by applying Lemma 2, we can construct a corresponding ReLU network $f_{i}$ by replacing each PolyReLU activation $p_{i,j}$ in $g_{i}$ with a ReLU activation. Specifically, for any $i\in[L]$(^4^44We use notation $[L]$ to denote the set $\{1,2,\dots,L\}$.) and $j\in[K]$, there exists a ReLU network $f_{i,j}:[-1,1]\rightarrow[-1,1]$ that approximates the PolyReLU activation $p_{i,j}$ with given tolerance $\epsilon_{i}>0$.

Thus, the network $f_{i}$ is obtained by replacing each PolyReLU activation $p_{i,j}$ in $g_{i}$ with its ReLU approximation $f_{i,j}$. Obviously, $f_{i}$ is a ReLU network whose output dimensions are in range $[-1,1]$.

Next, we give the approximation error bound. Denote $h^{g}_{i}=g_{i}\circ\cdots\circ g_{1}$ and $h^{f}_{i}=f_{i}\circ\cdots\circ f_{1}$ for $i\in[L]$. For the sake of brevity, we assume $h^{g}_{0}=h^{f}_{0}$ as the identity map in $[-1,1]^{d}$. Hence, we have $h^{g}_{i}=g_{i}\circ\cdots\circ g_{0}$ and $h^{f}_{i}=f_{i}\circ\cdots\circ f_{0}$. Suppose $x\mapsto p_{i,j}(a_{i,j}^{\top}h^{g}_{i-1}+b_{i,j})$ be the output of $j$-th neuron of $g_{i}$. Denote the approximation between the PolyReLU network and the ReLU network at $i$-th layer and $j$-th neuron as $e_{i,j}$. And we use $e_{i}=\max_{j\in[K]}e_{i,j}$ to denote the approximation error between the PolyReLU network and the ReLU network at $i$-th layer. Then for any $i\in[L]$, we have that

$$ $\displaystyle e_{i,j}=$ $\displaystyle\max_{x\in[-1,1]^{d}}\left|h^{g}_{i,j}({\bm{x}})-h^{f}_{i,j}({\bm {x}})\right|$ (18) $\displaystyle=$ $\displaystyle\max_{x\in[-1,1]^{d}}\left|p_{i,j}(a_{i,j}^{\top}h^{g}_{i-1}({\bm {x}})+b_{i,j})-f_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})\right|$ $\displaystyle=$ $\displaystyle\max_{x\in[-1,1]^{d}}\biggl{|}p_{i,j}(a_{i,j}^{\top}h^{g}_{i-1}({ \bm{x}})+b_{i,j})-p_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})$ $\displaystyle+p_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})-f_{i,j}(a_{ i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})\biggr{|}$ $\displaystyle\leq$ $\displaystyle\max_{x\in[-1,1]^{d}}\left|p_{i,j}(a_{i,j}^{\top}h^{g}_{i-1}({\bm {x}})+b_{i,j})-p_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})\right|$ $\displaystyle+\max_{x\in[-1,1]^{d}}\left|p_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({ \bm{x}})+b_{i,j})-f_{i,j}(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})\right|$ $\displaystyle\leq$ $\displaystyle\max_{x\in[-1,1]^{d}}\alpha\left|(a_{i,j}^{\top}h^{g}_{i-1}({\bm{ x}})+b_{i,j})-(a_{i,j}^{\top}h^{f}_{i-1}({\bm{x}})+b_{i,j})\right|+\epsilon_{i}$ $\displaystyle\leq$ $\displaystyle\alpha\max_{x\in[-1,1]^{d}}\|a_{i,j}\|_{1}\left\|h^{g}_{i-1}({\bm {x}})-h^{f}_{i-1}({\bm{x}})\right\|_{\infty}+\epsilon_{i}$ $\displaystyle\leq$ $\displaystyle\alpha\max_{x\in[-1,1]^{d}}\left\|h^{g}_{i-1}({\bm{x}})-h^{f}_{i- 1}({\bm{x}})\right\|_{\infty}+\epsilon_{i}.$ $$

The first inequality is using the triangular inequality. The second inequality holds because the Lipschitz constant of $p_{i,j}$ is $\alpha$ and the ReLU subnetwork $f_{i,j}$ approximates $p_{i,j}$ with error $\epsilon_{i}$. In the fourth inequality, we used Hölder’s inequality.
Since $\|a_{i,j}\|_{1}\leq\|a_{i,j}\|_{1}+|b_{i,j}|\leq 1$, the fifth inequality holds.

Therefore, we derive the following approximation bound:

$$ $e_{i}=\max_{j\in[K]}e_{i,j}\leq\alpha\max_{x\in[-1,1]^{d}}\left\|h^{g}_{i-1}({ \bm{x}})-h^{f}_{i-1}({\bm{x}})\right\|_{\infty}+\epsilon_{i}=\alpha e_{i-1}+ \epsilon_{i},$ (19) $$

for $\forall i\in[L]$. Since $h^{g}_{0}=h^{f}_{0}$, we have $e_{0}=0$. Let $\epsilon_{i}=\epsilon/(L\alpha^{L-i})$ for $\forall i\in[L]$. It follows that:

$$ $e_{i}\leq\frac{i\epsilon}{L\alpha^{L-i}},\quad\forall i\in[L].$ (20) $$

Hence, the final error at the last layer is bounded by $e_{L}\leq\epsilon$.

Last, we need to estimate the size of the ReLU network $f$. By Lemma 2, the size of each ReLU subnetwork $f_{i,j}$ is $O(\ln^{2}(L\alpha^{L-i}/\epsilon))$. Therefore, the total size of the ReLU network $f$ is:

$$ $O\left(\sum_{i=1}^{L}K\ln^{2}\left(\frac{L\alpha^{L-i}}{\epsilon}\right)\right )=O\left(KL\ln^{2}\left(\frac{L\alpha^{L}}{\epsilon}\right)\right),$ (21) $$

where we use the fact that:

$$ $\sum_{i=1}^{L}\ln^{2}\left(\frac{L\alpha^{L-i}}{\epsilon}\right)=\sum_{i=1}^{L }\left(\ln\left(\frac{L\alpha^{L}}{\epsilon}\right)-i\ln\alpha\right)^{2}=O \left(L\ln^{2}\left(\frac{L\alpha^{L}}{\epsilon}\right)\right).$ (22) $$

This completes the proof.
∎

###### Lemma A.2 (Theorem 11 in Liang & Srikant ( 2017 ) ) .

###### Proof of Theorem 2 .

### A.4 Proof of Theorem 3

Before proving Theorem 3, we begin by introducing a few useful lemmas.

###### Lemma A.3 (Proposition 1 in Yarotsky ( 2017 ) ) .

Let $M\in{\mathbb{N}}$ and $\rho:{\mathbb{R}}\rightarrow{\mathbb{R}}$ be any continuous piece-wise linear function with $M$ breakpoints. Then the following two statements hold:

- •
For a network with activation $\rho$, depth $L$ and width $K$, there exists a ReLU network with the same depth $L$ and width $O(MK)$ that computes the same function as the original network.
- •
Conversely, if a ReLU network has depth $L$ and width $K$, there exists a network with activation $\rho$, depth $L$ and width $K$ that computes the same function on a bounded input domain ${\mathcal{D}}$.

This result, Combined with Lemma 1, directly leads to the following corollary, which demonstrates that PolyReLU networks can represent any piece-wise linear function exactly on ${\mathbb{R}}$.

###### Corollary A.1 .

Let $M\in{\mathbb{N}}$ and $\rho:{\mathbb{R}}\rightarrow{\mathbb{R}}$ be any continuous piece-wise linear function with $M$ breakpoints. Then there exists a PolyReLU network $g$ of size $O(M)$ such that:

$$ $\rho(x)=g(x),\quad\forall x\in{\mathbb{R}}.$ $$

In a similar manner to Proposition 10 in Boullé et al. (), we can show that PolyReLU networks can represent powers $x^{n}$ exactly for any $n\in{\mathbb{N}}$.

###### Lemma A.4 .

Suppose $n,r\in{\mathbb{N}}$ and $r\geq 2$. Then $x^{n}$ can be represented exactly by a PolyReLU network $g$ with an $r$-th order PolyReLU activation and size $O(\ln^{2}(n))$.

###### Proof of Lemma A.4 .

We first prove that $x^{n}$ can be represented exactly by a polynomial network $\hat{g}$ with $r$-th order polynomial activation and having size $O(\ln^{2}(n))$. Based on $\hat{g}$, we construct a PolyReLU network $g$ that satisfies the requirements.

By expressing $n$ in base $r$, we have that:

$$ $x^{n}=\prod_{i=0}^{k}x^{c_{i}r^{i}}=\prod_{i=0}^{k}\left(x^{c_{i}}\left(x^{r} \right)^{i}\right),$ (23) $$

where $k=\lfloor\log_{r}n\rfloor$, $n=\sum_{i=0}^{k}c_{i}r^{i}$, and $c_{i}\in\{0,1,2,\dots,r-1\}$. Each $x^{c_{i}r^{i}}$ can be represented by a polynomial network with $i+1$ layers and width 1. It follows that $x^{n}$ can be represented by a polynomial network of size

$$ $\sum_{i=0}^{k}(i+1)=O(k^{2})=O(\ln^{2}(n)).$ (24) $$

By Lemma 1, we know that a PolyReLU activation can represent a polynomial activation. Hence, there exists a PolyReLU network $g$ with an $r$-th order activation and size $O(\ln^{2}(n))$ such that

$$ $g(x)=x^{n}\quad\forall x\in{\mathbb{R}}.$ $$

∎

With the above lemmas, we can now prove Theorem 3.

###### Proof of Theorem 3 .

The proof is composed of two parts.
We first approximate $f$ by local Taylor polynomials and continuous piece-wise linear functions and then represent these functions using PolyReLU networks, following Yarotsky (); Boullé et al. ().

Part 1. Suppose $N$ is a positive integer. We begin by dividing $[-1,1]^{d}$ into a grid of $(2N+1)^{d}$ functions:

$$ $\sum_{{\bm{m}}}\phi_{{\bm{m}}}({\bm{x}})=1,\quad\phi_{{\bm{m}}}({\bm{x}})= \prod_{i=1}^{d}\varphi\left(3N\left(x_{k}-\frac{m_{k}}{N}\right)\right),\quad \forall{\bm{x}}=(x_{1},x_{2},\dots,x_{d})\in[-1,1]^{d},$ $$

where ${\bm{m}}=(m_{1},m_{2},\dots,m_{d})\in\{-N,-(N-1),\dots,0,\dots,N\}$, and $\varphi$ is defined as:

$$ $\varphi(x)=\begin{cases}1,&|x|<1,\\ 0,&2<|x|,\\ 2-|x|,&1\leq|x|\leq 2.\end{cases}$ $$

This function has the following properties:

$$ $\displaystyle\max_{x\in{\mathbb{R}}}|\varphi(x)|=1,\quad\max_{{\bm{x}}\in[-1,1 ]^{d}}\|\phi_{{\bm{m}}}({\bm{x}})\|_{\infty}=1,$ (25) $\displaystyle\operatorname{supp}\phi_{{\bm{m}}}=\left\{{\bm{x}}\biggl{|}\left \|{\bm{x}}-\frac{{\bm{m}}}{N}\right\|_{\infty}<\frac{2}{3N}\right\},\forall{ \bm{m}}\in\{-N,-(N-1),\dots,N\}^{d}.$ (26) $$

Part 2. We use a degree-$(n-1)$ local Taylor approximation of the function $f$, defined as

$$ $f_{N}({\bm{x}})=\sum_{{\bm{m}}\in\{-N,\dots,N\}^{d}}\phi_{{\bm{m}}}({\bm{x}})P _{{\bm{m}}}({\bm{x}}),$ (27) $$

where $P_{{\bm{m}}}$ is the degree-$(n-1)$ Taylor polynomial of $f$ at ${\bm{x}}={\bm{m}}/N$, i.e.,

$$ $P_{{\bm{m}}}({\bm{x}})=\sum_{{\bm{n}}:\|{\bm{n}}\|_{1}<n}\frac{1}{{\bm{n}}!}D^ {{\bm{n}}}f\left(\frac{{\bm{m}}}{N}\right)\left({\bm{x}}-\frac{{\bm{m}}}{N} \right)^{{\bm{n}}},$ (28) $$

with conventions ${\bm{n}}!=\prod_{i=1}^{d}n_{i}!$ and $\left({\bm{x}}-\frac{{\bm{m}}}{N}\right)^{{\bm{n}}}=\prod_{i=1}^{d}\left(x_{i}
-\frac{m_{i}}{N}\right)^{n_{i}}$.

The approximation error between $f$ and $f_{N}$ can be bounded as follows:

$$ $\displaystyle|f({\bm{x}})-f_{N}({\bm{x}})|$ $\displaystyle=\left|\sum_{{\bm{m}}\in\{-N,\dots,N\}^{d}}\phi_{{\bm{m}}}\left(f ({\bm{x}})-P_{{\bm{m}}}({\bm{x}})\right)\right|$ (29) $\displaystyle\leq\sum_{{\bm{m}}:\|x-\frac{{\bm{m}}}{N}\|_{\infty}<\frac{2}{3N} }\left|f({\bm{x}})-P_{{\bm{m}}}({\bm{x}})\right|$ $\displaystyle\leq 2^{d}\max_{{\bm{m}}:\|x-\frac{{\bm{m}}}{N}\|_{\infty}<\frac{ 2}{3N}}\left|f({\bm{x}})-P_{{\bm{m}}}({\bm{x}})\right|$ $\displaystyle\leq\frac{2^{d}}{n!}\left(\frac{2d}{3N}\right)^{n}\max_{{\bm{n}}: \|{\bm{n}}\|_{1}=n}\operatornamewithlimits{ess\ sup}_{{\bm{x}}\in[-1,1]^{d}} \left\|D^{{\bm{n}}}f({\bm{x}})\right\|_{\infty}$ $\displaystyle\leq\frac{2^{d}}{n!}\left(\frac{2d}{3N}\right)^{n}.$ $$

The first inequality is because of the triangular inequality and Eq. (25). In the second inequality, we used the fact that $\forall x\in[-1,1]^{d}$ belongs to the support of at most $2^{d}$ functions $\phi_{{\bm{m}}}$. The third inequality is a bound for the Taylor remainder and the fourth inequality uses the definition of $F_{n,d}$. Let

$$ $N=\left\lfloor\frac{2d}{3}\left(\frac{2^{d}}{n!\epsilon}\right)^{\frac{1}{n}} \right\rfloor+1,$ (30) $$

we have that

$$ $\max_{{\bm{x}}\in[-1,1]^{d}}|f({\bm{x}})-f_{N}({\bm{x}})|<\epsilon.$ (31) $$

Next, we construct a PolyReLU network $g_{N}$ to represent $f_{N}$ exactly. Let $a_{{\bm{m}},{\bm{n}}}=\frac{1}{{\bm{n}}!}D^{{\bm{n}}}f\left(\frac{{\bm{m}}}{N}\right)$. Since $\|f\|_{{\mathcal{W}}^{n,\infty}\left([-1,1]^{d}\right)}\leq 1$, $|a_{{\bm{m}},{\bm{n}}}|\leq 1$ for any ${\bm{m}},{\bm{n}}$, we rewrite $f_{N}$ as:

$$ $f_{N}({\bm{x}})=\sum_{{\bm{m}}\in\{-N,\dots,N\}^{d}}\sum_{{\bm{n}}:\|{\bm{n}} \|_{1}<n}a_{{\bm{m}},{\bm{n}}}\phi_{{\bm{m}}}({\bm{x}})\left({\bm{x}}-\frac{{ \bm{m}}}{N}\right)^{{\bm{n}}}.$ (32) $$

Therefore, $f_{N}$ is composed of at most $d^{n}(2N+1)^{d}$ functions $\phi_{{\bm{m}}}({\bm{x}})\left({\bm{x}}-\frac{{\bm{m}}}{N}\right)^{{\bm{n}}}$. Since $\phi_{{\bm{m}}}({\bm{x}})=\prod_{i=1}^{d}\varphi\left(3N\left(x_{k}-\frac{m_{k
}}{N}\right)\right)$ and each $\varphi\left(3N\left(x_{k}-\frac{m_{k}}{N}\right)\right)$ is a continuous piece-wise linear function, we can apply Corollary A.1, which guarantees that there exists a PolyReLU network $\hat{\phi}_{{\bm{m}}}$ of size $O(d)$ that can exactly represent $\phi_{{\bm{m}}}$ on ${\mathbb{R}}^{d}$, i.e., $\hat{\phi}_{{\bm{m}}}({\bm{x}})=\phi_{{\bm{m}}}({\bm{x}}),\forall{\bm{x}}\in{
\mathbb{R}}^{d}$. For $\left({\bm{x}}-\frac{{\bm{m}}}{N}\right)^{{\bm{n}}}=\prod_{i=1}^{d}\left(x_{i}
-\frac{m_{i}}{N}\right)^{n_{i}}$, by Lemma A.4, we know that there
exists a PolyReLU network $g_{{\bm{m}}}$ of size at most $O(d\ln^{2}(n))$ such that $g_{{\bm{m}}}({\bm{x}})=\left({\bm{x}}-\frac{{\bm{m}}}{N}\right)^{{\bm{n}}},
\forall{\bm{x}}{\mathbb{R}}^{d}$. Combining these results, we can now construct a larger PolyReLU network $g_{n}$ as follows:

$$ $g_{N}({\bm{x}})=\sum_{{\bm{m}}\in\{-N,\dots,N\}^{d}}\sum_{{\bm{n}}:\|{\bm{n}} \|_{1}<n}a_{{\bm{m}},{\bm{n}}}\hat{\phi}_{{\bm{m}}}({\bm{x}})g_{{\bm{m}}}({\bm {x}}),$ (33) $$

where the total size of the network is:

$$ $O\left(d^{n}(2N+1)^{d}(d+d\ln^{2}(n))\right)=O(\epsilon^{-\frac{d}{n}}).$ $$

Here, we use Eq. (30) to determine the size bound in terms of the error tolerance $\epsilon$. Clearly, we have:

$$ $f_{N}({\bm{x}})=g_{N}({\bm{x}}),\quad\forall{\bm{x}}\in{\mathbb{R}}^{d}.$ (34) $$

Hence, we conclude:

$$ $\max_{{\bm{x}}\in[-1,1]^{d}}|f({\bm{x}})-g_{N}({\bm{x}})=\max_{{\bm{x}}\in[-1, 1]^{d}}|f({\bm{x}})-f_{N}({\bm{x}})|<\epsilon.$ (35) $$

This completes the proof.
∎

###### Lemma A.3 (Proposition 1 in Yarotsky ( 2017 ) ) .

###### Corollary A.1 .

###### Lemma A.4 .

###### Proof of Lemma A.4 .

###### Proof of Theorem 3 .

## Appendix B Activation Functions

We provide definitions of several commonly used non-linear activation functions in Table 4.

**Table 4: Definition of activation functions.**
| Activation | Definition |
| --- | --- |
| ReLU (Nair & Hinton, ) | ReLU $(x)=\max\{x,0\}$ |
| ReLU^2 (So et al., ) | ReLU^2 $(x)=\max\{x,0\}^{2}$ |
| ReLU6 (Krizhevsky et al., ) | ReLU6 $(x)=\min(\max\{x,0\}$,6 ) |
| Leaky ReLU (Maas et al., ) | LeakyReLU $(x)=\begin{cases}x,&\text{ if }x\geq 0\\<br>ax,&\text{ otherwise }\end{cases}$, $a\in(0,1)$ is a constant |
| RReLU (Xu et al., ) | RReLU$(x)=\begin{cases}x&\text{if }x\geq 0\\<br>ax&\text{ otherwise,}\end{cases}$ |
|  | $a$ is randomly sampled from uniform distribution |
| Parametric ReLU<br>(PReLU) | PReLU $(x)=\begin{cases}x,&\text{ if }x\geq 0\\<br>ax,&\text{ otherwise ,}\end{cases}$ |
| (He et al., ) | $a$ is a learnable parameter |
| Tanh | Tanh$(x)=\frac{\exp(x)-\exp(-x)}{\exp(x)+\exp(-x)}$ |
| Softplus (Glorot et al., ) | Softplus $(x)=\frac{1}{a}*\log(1+\exp(ax))$, |
|  | $a$ is a constant (default 1.0) |
| Mish (Misra, ) | Mish $(x)=x*\text{Tanh}(\text{Softplus}(x))$ |
| Sigmoid | Sigmoid$(x)=\sigma(x)=\frac{1}{1+\exp(-x)}$ |
| SiLU(Swish) | SiLU$(x)=x*\sigma(x)$ |
| (Ramachandran et al., ) |  |
| ELU (Clevert, ) | ELU$(x)=\begin{cases}x,&\text{ if }x>0\\<br>a*(\exp(x)-1),&\text{ if }x\leq 0,\end{cases}$ |
|  | $a$ is a constant (default 1.0) |
| CELU (Barron, ) | CELU $(x)=\max(0,x)+\min(0,\alpha*(\exp(x/a)-1))$, |
|  | $a$ is a constant (default 1.0) |
| GELU (Hendrycks & Gimpel, ) | GELU$(x)=x*\Phi(x)$, |
|  | $\Phi(x)$ is CDF for Gaussian distribution |
| GLU (Dauphin et al., ) | GLU$(x)=\sigma(xW)\otimes(xV)$ |
| SwiGLU (Shazeer, ) | SwiGLU$(x)=\operatorname{SiLU}(xW)\otimes(xV)$, |
|  | $W,V$ are learnable parameters |
| Poly | $\mathrm{Poly}(x)=\sum_{i=0}^{r}a_{i}x^{i}$, $a_{i},i\in[r]$ are learnable parameters. |

## Appendix C Experimental Details

### C.1 Architecture

Table 5 outlines the model architecture used for the 1B dense model. To ensure comparable numbers of training parameters across different activation functions, we adjust the intermediate sizes accordingly. 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 1B dense model.**
| Params | Hidden size | Context Length | Intermediate size | Attention heads | Hidden Layers |
| --- | --- | --- | --- | --- | --- |
| 1.3B | 2048 | 4096 | 5504/8256 | 16 | 24 |

**Table 6: 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 | Context Length |  |  |
| 16 | 64 | 4096 |  |  |

### C.2 Definition of Effective Rank

We adopt the concept of effective rank from Roy & Vetterli () to measure the effective dimensionality of a matrix. Given a matrix $A$ with Singular Value Decomposition (SVD) $A=U\Sigma V^{\top}$, where $\Sigma$ is a diagonal matrix containing singular values $\sigma_{1}\geq\sigma_{2}\geq\dots\geq\sigma_{n}\geq 0$. we define the singular value distribution as $p_{i}=\sigma_{i}/\sum_{j=0}^{n}\sigma_{j},i\in[n]$. The effective rank of $A$ is then given by:

$$ $\operatorname{Erank}(A)=\exp{\left(-\sum_{i=0}^{n}p_{i}\ln p_{i}\right)}.$ (36) $$

## Appendix D Additional Results on Dense Model

More detailed results from our ablation studies are shown in Figures 8, 9, and 10. These figures illustrate the training loss, validation loss, and validation perplexity (PPL) for the 1B dense model under different configurations.

Figure: Figure 8: Traning loss, validation loss, and validation perplexity (PPL) for the 1B dense model with different orders of PolyReLU activation functions.
Refer to caption: x16.png

Figure: Figure 9: Traning loss, validation loss, and validation perplexity (PPL) for the 1B dense model with different polynomial compositions.
Refer to caption: x17.png

Figure: Figure 10: Traning loss, validation loss, and validation perplexity (PPL) for the 1B dense model with different variants of ReLU activation functions.
Refer to caption: x18.png

## Appendix E Additional Results on MoE model

More results for the MoE model are provided in Figure 11, , showcasing validation losses and downstream evaluations after 200 billion training tokens. The comparison highlights models with different activation functions, such as SwiGLU and PolyNorm. As shown, models with PolyNorm exhibit lower training and validation losses, along with superior downstream performance.

Figure: Figure 11: Validation loss and downstream evaluations for MoE models with 200 billion training tokens, comparing SwiGLU and PolyNorm activation functions. PolyNorm shows superior performance in terms of lower loss and better downstream results.
Refer to caption: x19.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: 1. Does theorem 2 make sense?
2. What is happening in the proof of lemma 2?
3. Do the PolyCom functions introduce additional learnable parameters?
4. Do the different models in the comparison have the same number of parameters?
5. How were the hyperparameters tuned?
6. How do the models compare when trained to convergence?
7. What does the performance look like as a function of training budget?
8. Does this method work also provide benefits when used outside of language modelling?
soundness: 2
strengths: - (clarity) Overall, the paper is clearly written and easy to follow.
 - (originality) The proposed activation functions are novel (to the best of my knowledge).
summary: This work presents a method to extend activation functions using polynomials.
Two examples of these extended activation functions are introduced: PolyReLU and PolyNorm.
PolyReLU is shown to achieve an optimal approximation rate in Sobolev spaces.
Language modelling experiments indicate that the proposed activation functions accelerate training.
weaknesses: - (quality) The proof of lemma 2 (and therefore also lemma 2) seems to be incorrect.
   It seems like the rank (and number of polynomials) from lemma 3.4 in (Telgarsky, 2017) have been ignored.
   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)$.
   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.
 - (quality/significance) The lower bound in theorem 2 does not make sense.
   If the PolyReLU network happens to have PolyReLU parameters $a_1 = 1$ and $\forall i \neq 1 : a_i = 0$,
   it should be possible to model it exactly with a ReLU network of the same size.
   Similarly, I do not see how the upper bound in theorem 2 would make sense.
   It should always be possible to make the network arbitrarily large by adding neurons for which all incoming and outgoing weights are zero.
 - (clarity/quality) The polynomial activation functions seem to introduce additional learnable parameters.
   However, the paper never explicitly states this anywhere.
   Futhermore, there is no discussion on whether/how the number of parameters of the different models was controlled.
   This possibly introduces a capacity advantage for the models with PolyCom functions compared to the baselines,
   leading to an unfair comparison.
 - (clarity/quality) There is no discussion on how the hyper-parameters were found.
   Furthermore, it seems like the same hyper-parameters were used for every model.
   If these hyper-parameters were tuned on the proposed models, this would be an unfair comparison.
 - (quality) The authors claim that the activation functions are able to capture higher-order interactions.
   However, based on the training curves (e.g. figure 3), none of the models were trained to convergence.
   As a result, it is not possible to conclude anything concerning the complexity of the interactions that can be learned.
   After all, it might be that the same interactions are captured for both models when converged, but one captures them "faster".
   As a result, these experiments do not confirm the stated claims.
 - (quality/significance) There is no discussion on the runtime overhead of the introduced functions.
   If the PolyCom functions introduce too much overhead, they might not be practically useful.
   Also, it would be useful to include a comparison performance for a given compute budget.
 - (clarity) It is not clear why the proposed activation functions should be especially suited for language modelling.
   I suspect that this should also work for models that do not require that much compute and allow for more extensive experiments.

###### Minor Comments
 - There seem to be some type-setting issues with equation (7) (Thereom 1).
 - Could it be that there is a superscript $d$ missing for the space of $\boldsymbol{m}$ on line 869-870?

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: - What is the memory and runtime overhead of switching from, e.g., squared ReLU to PolyCom?
- 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?
soundness: 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.
- The experiments are comprehensive, covering both pre-training and downstream evaluations.
- 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.
- The paper is well-presented.
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.
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. 
- 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.
- It's not clear that PolyNorm MoE has lower layer-wise cosine similarity in Figure 7 compared to SwiGLU MoE.
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
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 
- It’s stated that PolyReLU and PolyNorm have equivalent expressivity, e.g. on lines 164-165. Do you have a proof for this?
- 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?
soundness: 3
strengths: - The idea of activation functions with better expressivity, without adding more trainable parameters to the model, is significant, useful and interesting
- 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
- The combination of both theoretical and empirical results is stronger than including just either type of result alone
- Theoretical results are approachable and elegant
- Paper is well-structured and well-written: it’s easy to follow and there were no confusing typos or other writing issues
- Empirical evaluations are conducted on several different downstream datasets, adding strength to the authors’ claims
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.
PolyCom 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:
  - For Type 1, the authors use PolyReLU. This activation sets $\rho$ to ReLU and applies $\rho$ before the polynomial. 
  - For Type 2, the authors use PolyNorm. This activation has the L2 normalization for $\rho$ and applies the polynomial before $\rho$.
- The theory results focus on PolyReLU and proceed in three steps:
  - 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
  - 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.
  - 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}$.
- 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.
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
- 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
- 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.
- 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

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Check the weaknesses.
soundness: 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.

2) 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.

3) Overall, the paper is clearly written, with the description of polynomial activation easy to follow, and the results presented concisely.
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.
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?

2) How do the authors address stability or manage potential exploding gradient issues with the polynomial activation?

Typos: Line 274: Dowmstream Evaluation -> Downstream Evaluation

Meta-review: **Summary**

This work introduces a novel activation function called PolyCom, which is a composition of polynomials and other types of functions featuring specific instances like PolyReLU and PolyNorm. These are seamlessly integrated into transformer architectures, offering several key benefits. They maintain a comparable computational overhead and memory footprint to the widely used SwiGLU activation function while providing optimal theoretical approximation rates. PolyReLU, for example, is capable of achieving the optimal approximation rate in Sobolev spaces. Empirical tests on large language models (LLMs) and Mixture of Experts (MoE) models demonstrate that PolyCom can significantly enhance training efficiency, achieving up to 1.5 times faster convergence compared to SwiGLU.

**Strengths**


* The paper is well-structured and well-written, and the results are presented concisely.
* The core idea of the paper on devising activation functions with better expressivity without adding more trainable parameters to the model is significant, useful, and interesting.
* 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.
* Extensive experiments on downstream tasks demonstrate improved convergence rates in learning curves with fixed model parameter sizes. These results underscore the potential of the proposed activation function and substantiate the theoretical claims made in the paper.

**Weaknesses**

* The original submission lacked computational complexity analysis for the proposed activations.
* Without proper normalization layers, the proposed activations could lead to exploding gradients. However, this won't be an issue in transformer networks that utilize layer normalization. 

**Conclusion** 

The paper received positive feedback from all reviewers. The weaknesses highlighted by the reviewers were effectively addressed by the authors in their revisions, providing a more comprehensive understanding of the proposed work. After thoroughly reviewing the paper, the feedback from the reviewers, and the authors' rebuttal, I agree with the majority opinion and vote to accept this paper.

Decision: Accept (Poster)