Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models

TL;DR

This work addresses the limited expressivity of conventional activation functions in transformers by introducing Polynomial Composition Activations (PolyCom), including PolyReLU and PolyNorm. The authors provide a rigorous theoretical analysis showing that PolyReLU can replicate ReLU networks with the same size and achieves optimal Sobolev-space approximation rates, while PolyNorm stabilizes training through normalization and can further improve expressivity. Empirically, PolyCom accelerates convergence and boosts downstream performance in both dense and MoE large language models, with PolyNorm frequently delivering the best results and maintaining training stability. The combination of strong theoretical guarantees and consistent empirical gains suggests PolyCom as a practical, scalable enhancement for transformer architectures across domains, with releasing code to facilitate adoption.

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 $\textbf{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.

Figures (14)

• Figure 1: Training loss, validation perplexity (PPL), and downstream performance of 1B dense models. We compare models employing different activation functions, including SwiGLU, GELU, ReLU, PolyReLU, and PolyNorm. It indicates that models using PolyReLU and PolyNorm exhibit lower training loss and validation PPL, alongside better downstream performance.
• 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$, "$a_j$" denotes the $j$-th element of the learnable weight vector $a$, "N" indicates a normalization operation.
• Figure 3: Training 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.
• Figure 4: Dynamics of downstream performance 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.
• Figure 5: Training loss for 1B dense models with different activation functions. \ref{['fig:ablation-different orders']}: We compare different orders of PolyReLU. \ref{['fig:ablation-different compositions']}: Comparison of PolyCom with different composition functions. \ref{['fig:ablation-ReLU variants']}: Comparison of different variants of ReLU activation function.

Theorems & Definitions (18)

• Lemma 1
• Theorem 1
• Lemma 2
• Theorem 2
• Definition 1: Sobolev Spaces
• Theorem 3
• proof : Proof of Lemma \ref{['lem:degradation cases']}
• proof : Proof of Theorem \ref{['thm:ployrelu2relu']}
• Lemma A.1: Lemma 3.4 in telgarsky2017neural
• proof : Proof of Lemma \ref{['lem:approximate ployrelu']}