Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models

Thank you for your thoughtful comments and time. We have tried our best to address your concerns and revised our paper accordingly.

Q1. Computational complexity analysis is not provided.

A1. We analyze the runtime overhead introduced by the activation functions using a typical feedforward network (FFN) with input tensor $x\in \mathbb{R}^{B \times S \times H}$, where $B$，$S$ and $H$ represent the batch size, sequence length, and hidden size, respectively. The relationship between computational FLOPs and model parameters can generally be regarded as proportional (as discussed in Eleuther AI’s transformer math). Below, we estimate the proportion of the computational cost incurred by activation function calculations within the total computational cost of the FFN matrix computations ($24BSH^2$). The FLOPs ratio is calculated as:

$\text{FLOPs ratio} = \frac{\text{FLOPs for activation}}{24BSH^2}$

The results are summarized in the following table:

Note:

• We assume that the scale of the input tensor is set to [-1, 1]. In this case, the FLOPs for both tanh and exp are approximately 10 each.
• To reduce memory overhead during large language model pretraining, we typically employ gradient checkpointing (refer to PyTorch Docs). Although this approach incurs some additional computational cost, its overall impact on GPU memory and runtime is minimal.

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

A2. Stability or potential exploding gradient issues were not observed in our training process. The normalization operators in the transformer block effectively stabilized the training, as evidenced by the consistent trends shown in Figure 1 and Figure 3 in the paper. Additionally, for PolyNorm, we hypothesize that its integrated normalization mechanisms contribute further to stabilizing the training process. These results suggest that the activation functions operate reliably within the proposed framework without requiring additional stability interventions.

Thank you for highlighting the typos; we have corrected them accordingly.

Thank you for the follow-up questions about training stability, which is indeed an important point. We appreciate the opportunity to further elaborate on this matter.

Inherent Reasons for Stability: Currently, the reasons why PolyReLU and PolyNorm can maintain a stable training process within the transformer architecture lack a definitive theoretical explanation. A plausible hypothesis is that the normalization operators within transformers, combined with commonly used gradient clipping strategies, effectively stabilize the training process. And we will explore this topic in future research.

Key Experimental Observations: From extensive experiments, we observed that:

1. Both PolyReLU and PolyNorm exhibit stable behavior in transformer architectures.
2. In ResNet-50 training, while PolyNorm remains stable, PolyReLU encounters stability issues. For ResNet-50 experiments, we simply replaced ReLU with PolyNorm while preserving the original sequence of BatchNorm and activation.

Recommendations for Practical Usage: Actually, when designing these activation functions, we have taken into account the widespread adoption of BF16/FP16 precisions in contemporary model training procedures. To mitigate potential instability, we specifically introduced PolyNorm, which includes normalization operators to rescale different powers to a manageable range, thereby avoiding excessively large or small values. It is particularly beneficial for FP16 training, as in ResNet-50. PolyReLU doesn't include this normalization property, which may introduce some potential stability issues in other non-transformer-based architectures such as ResNet.

Also, our experiments reveal that for orders ranging from 2 to 4, the choice of order has minimal impact on stability. Indeed, extremely high orders may pose inherent stability risks. As shown in Figure 5(a), the utilization of 3-order (our default setting) is sufficient.

Based on these findings, we recommend the following settings:

1. For transformer-based models, use PolyNorm or PolyReLU .
2. For non-transformer-based or less stable models, use PolyNorm.

As highlighted in the title and abstract of our paper, in the current stage, our focus is on the proposed activation functions and their compatibility with LLMs or transformer architectures. Investigating their effectiveness in other model structures remains an important future research direction.

Dear Reviewer TiU9,

As the review deadline (December 2nd) approaches, we would like to extend our sincere gratitude for the time and effort you have invested in reviewing our work and considering our responses. We are following up to ensure that our replies have fully addressed your concerns and to ask if there is any further clarification or information we can provide.

Thanks for your time and constructive comments. We address your concerns point by point below:

Q1. Does theorem 2 make sense?

A1. We are afraid that you may have misunderstood the meaning of Theorem 2. We restate Theorem 2 for better clarity.

Theorem 2. For a PolyReLU network $g$ of depth $L$, width $K$, and PolyReLU activation of order and Lipschitz constant $\alpha$. Suppose each neuron computes $x \mapsto PolyReLU(a^\top x +b)$ with the parameters pair $(a,b)$ satisfies $\|a\|_1+b\leq 1, PolyReLU:[-1,1]\rightarrow [-1,1]$ (a, b, and PolyReLU are possibly distinct across neurons). 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),$

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

Conversely, there exist some PolyReLU networks that cannot be approximated within error tolerance $\epsilon$ by any ReLU network with a size less than $\Omega \left(KL\ln\left(\frac{1}{\epsilon}\right)\right)$.

Clarifications:

1. The lower bound asserts that there exist some PolyReLU networks that can not be approximated by any ReLU network with a size less than $\Omega \left(KL\ln\left(\frac{1}{\epsilon}\right)\right)$. This does not imply all PolyReLU networks exhibit this behavior, hence the example provided does not contradict our claim.
2. For the upper bound, we mean that for any PolyReLU networks $g$ , there is a ReLU network $f$ of size $O\left(LK\ln^2\left(\frac{L\alpha^L}{\epsilon}\right)\right)$ that can approximate f within error tolerance $\epsilon$. The upper bound provides a constructive guarantee that a ReLU network of the specified size can approximate any PolyReLU network to the desired accuracy. Enlarging the network size unnecessarily weakens the bound’s tightness and is therefore irrelevant to the result.

Q2. Clarifications regarding the proof of Lemma 2.

A2. For the proof of Lemma 2, we are afraid that you may have misunderstood. We respond to your concerns point by point.

1. The size of PolyReLU network, denoted as $g$, can indeed be written as $O(\ln^2(\frac{r}{\epsilon}))$. However, since $r$ is a fixed positive constant, we focus on the approximation error $\epsilon$ and incorporate $r$ into the $O$-notation. Thus, $O(\ln^2(\frac{r}{\epsilon}))$ simplifies to $O(\ln^2(\frac{1}{\epsilon}))$.

2. From Lemma 3.4 in (Telgarsky, 2017), the size of the ReLU network $f$ satisfying the conditions is

   $O(\min(r\ln(r/\epsilon), \ln^2(r/\epsilon))).$

   Using the fact $\min\{a,b\}\leq b$ (based on this, the minimum disappears), we get

   $O(\min(r\ln(r/\epsilon), \ln^2(r/\epsilon))) \leq O( \ln^2(r/\epsilon)).$

   Incorporating $r$ into the $O$-notation, we immediately get the size of the ReLU network $f$ is $O(\ln^2(1/\epsilon))$.

3. Adding the ReLU activation does not introduce additional parameters or modify the layer sizes. Thus, the overall network size remains unchanged.

Q3. Do the PolyCom functions introduce additional learnable parameters?

A3. Yes, PolyCom introduces additional learnable parameters, as explicitly stated in Section 2 (lines 134–136). For instance, a third-order PolyCom incorporates an additional 4L parameters in a transformer with L layers. For a 1B dense model with 24 layers, the supplementary parameter count amounts to 96 (in contrast to the total of 1.3 billion parameters). Similarly, for the MoE model with 16 layers, the additional count is 48 (compared to 6.9 billion total parameters). These increments are quite small and can be considered negligible.

Q4. Do the different models in the comparison have the same number of parameters?

A4. All compared models were designed to maintain an equivalent parameter count. (1.3 billion for the dense model and 6.9 billion for the MoE mode). All compared models were designed to maintain an equivalent parameter count for fair comparison. For instance, SwiGLU’s intermediate size is adjusted to two-thirds that of other activations. More details are provided in Tables 5 and 6 (Appendix D).

Q5. How were the hyperparameters tuned?

A5. Hyperparameters were not tuned specifically for our experiments. We adopted hyperparameters similar to those used in LLaMA 2 [1] for the dense model (note that not all hyperparameters were explicitly reported in their work) and employed the default hyperparameters from OLMoE [2] for the MoE models. The following table summarizes the hyperparameters:

References:

[1] Hugo Touvron, et al. Llama 2: Open foundation and fine-tuned chat models.

[2] Niklas Muennighoff, et al. Olmoe: Open mixture-of-experts language models, 2024

Q5. How do the models compare when trained to convergence?

A5. Achieving complete convergence in pretraining large language models poses a significant challenge due to the immense computational resources required. For example, Meta’s training of the LLaMA-3.2 1B model (https://huggingface.co/meta-llama/Llama-3.2-1B) involved processing 9T tokens and consuming 370,000 H100-GPU hours, which translates to 64 H100 GPUs for 8 months or 64 A100 GPUs for 2 years. For most research aiming at pretraining improvements, the performance gap between models typically stabilizes after reaching a certain threshold of training corpus size. To illustrate this, we conducted additional experiments on dense models by extending the training corpus from 250 billion to 400 billion tokens. These experiments utilized 32 A100 GPUs over more than 10 days. The following tables present the training loss and validation perplexity (PPL) of various models. More detailed analyses can be found in Figure 11 (Appendix F).

Our results demonstrate that models using PolyReLU and PolyNorm consistently outperform those with SwiGLU, even with larger datasets. For large language model pretraining, a training loss difference greater than 0.015 is considered substantial. Moreover, the improvements from PolyReLU and PolyNorm have remained stable beyond the 200 billion token mark. These results indicate that PolyReLU and PolyNorm not only improve training dynamics but also consistently yield better validation performance compared to SwiGLU. This trend suggests that their advantages persist across different scales of training corpus.

Q6. There is no discussion on the runtime overhead of the introduced functions. What does the performance look like as a function of training budget?

A6. We analyze the runtime overhead introduced by the activation functions using a typical feedforward network (FFN) with input tensor $x\in \mathbb{R}^{B \times S \times H}$, where $B$，$S$ and $H$ represent the batch size, sequence length, and hidden size, respectively. The relationship between computational FLOPs and model parameters can generally be regarded as proportional (as discussed in Eleuther AI’s transformer math). Below, we estimate the proportion of the computational cost incurred by activation function calculations within the total computational cost of the FFN matrix computations ($24BSH^2$). The FLOPs ratio is calculated as:

$\text{FLOPs ratio} = \frac{\text{FLOPs for activation}}{24BSH^2}$

The results are summarized in the following table:

Note:

• We assume that the scale of the input tensor is set to [-1, 1]. In this case, the FLOPs for both tanh and exp are approximately 10 each.
• For a fair comparison, the intermediate size for SwiGLU is set to $8/3H$ to ensure a consistent total parameter count across all methods.
• To reduce memory overhead during large language model pretraining, we typically employ gradient checkpointing (refer to PyTorch Docs). Although this approach incurs some additional computational cost, its overall impact on GPU memory and runtime is minimal.

Hence, the overhead is acceptable and there is not much difference in the amount of training budget between them.

Q7. Does this method work also provide benefits when used outside of language modelling?

A7. To evaluate the effectiveness of PolyCom beyond language modeling, we trained ResNet50 on ImageNet following the settings of timm [1]. For comparison, we replaced the ReLU activation in ResNet50 with PolyCom and reported the training loss and top-1/top-5 accuracy on the evaluation set, as shown in the table below. The results demonstrate that PolyCom outperforms ReLU by a significant margin in terms of training loss, top-1 accuracy, and top-5 accuracy.

References:

[1]. Ross Wightman. PyTorch Image Models. https://github.com/rwightman/pytorch-image-models

Thank you for highlighting the typos. We have corrected these errors and hope our explanations address your concerns. Please feel free to let us know if further clarification is needed.

Dear Reviewer aNH6,

As the review deadline (December 2nd) approaches, we sincerely thank you for your time and effort in reviewing our work. We kindly follow up to confirm if our responses have addressed your concerns and whether any further clarification is needed.

We would greatly appreciate it if you could re-evaluate the paper’s rating based on the updated responses.

Thanks for your careful reading and critical review. We address your concerns below and hope you find them satisfactory.

Q1. What is the memory and runtime overhead of switching from, e.g., squared ReLU to PolyCom?

A1: We analyze the runtime overhead introduced by the activation functions using a typical feedforward network (FFN) with input tensor $x\in \mathbb{R}^{B \times S \times H}$, where $B$，$S$ and $H$ represent the batch size, sequence length, and hidden size, respectively. The relationship between computational FLOPs and model parameters can generally be regarded as proportional (as discussed in Eleuther AI’s transformer math). Below, we estimate the proportion of the computational cost incurred by activation function calculations within the total computational cost of the FFN matrix computations ($24BSH^2$). The FLOPs ratio is calculated as:

$\text{FLOPs ratio} = \frac{\text{FLOPs for activation}}{24BSH^2}$

The results are summarized in the following table:

Note:

• We assume that the scale of the input tensor is set to [-1, 1]. In this case, the FLOPs for both tanh and exp are approximately 10 each.
• To reduce memory overhead during large language model pretraining, we typically employ gradient checkpointing (refer to PyTorch Docs). Although this approach incurs some additional computational cost, its overall impact on GPU memory and runtime is minimal.

Q2. Why is a higher effective rank better? Both GeLU and SwiGLU have lower effective ranks than ReLU, but they achieve better performance.

A2. A higher rank in the weight matrices generally indicates a greater capacity for representing complex patterns within the data. However, it is crucial to note that a higher effective rank is a necessary but not sufficient condition for neural networks to achieve better performance. Our experimental results serve as supplementary evidence, demonstrating that PolyCom enables neural networks to capture more intricate patterns.

Regarding the specific case of ReLU as a counterexample, one plausible explanation is that ReLU discards part of the information in activations (values less than 0). This forces the backward gradient flow to encourage the weights to learn more complex patterns, ensuring that the activations after passing through ReLU retain sufficient information. Consequently, the weights in ReLU networks exhibit a larger effective rank.

Thus, while effective rank provides insight into model expressivity, other factors, such as the characteristics of the activation function and its impact on gradient dynamics, also play critical roles in determining performance.

Q3. It's not clear that PolyNorm MoE has lower layer-wise cosine similarity in Figure 7 compared to SwiGLU MoE.

A3. Comparing the rectangles in the lower-left (or upper-right) corner of Figure 7(c) and Figure 7(d), it is evident that the color corresponding to PolyNorm is noticeably redder than that of SwiGLU. This indicates that PolyNorm’s representations for layers 0–8 and 9–15 are less similar compared to SwiGLU’s.

Additionally, examining the lower-right corner of Figure 7(c) and Figure 7(d), we observe that PolyNorm has fewer blue squares. This suggests that the representations across layers 9–15 in PolyNorm are less similar than those in SwiGLU. These observations collectively demonstrate that PolyNorm exhibits lower layer-wise cosine similarity, particularly in the deeper layers, thereby supporting our claims.

Q4. 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?

A4. The table below (also visualized in Figure 13, Appendix H) summarizes the training loss for dense models with SwiGLU, PolyReLU, and PolyNorm activations across a range of model sizes from 110M to 1.3B parameters. It is evident that both PolyReLU and PolyNorm consistently outperform SwiGLU across all model sizes.

Scaling Law Details:

The model configurations for the scaling experiments are detailed below. All models employed the same hyperparameters as those specified for 1B dense models (see Table 7). Models with 110M, 226M, and 502M parameters were trained on a corpus of 200 billion tokens.

This scaling law experiment demonstrates that PolyReLU and PolyNorm provide consistent improvements over SwiGLU as model sizes increase, confirming their benefits across different scales.

Issue regarding the flops and memory overhead

Thank you for the valuable suggestions. We have updated the previous table and will include it in the appendix later, which encompasses the version using gradient checkpointing (this is the configuration we actually used during training). It is important to note that the overhead and proportion often vary for different model sizes, so we have provided the corresponding formulas directly and take $H=1024, B=4, S=4096$, using BF16 precision as an example.

• without gradient checkpointing:

• with gradient checkpointing:

The corresponding code (with gradient checkpointing) is as follows:

• For PolyNorm

import torch
from torch.utils.checkpoint import checkpoint
import torch.nn.functional as F

def _norm(x, eps=1e-6):
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)

def _poly_norm(x, weight, bias, order=3):
    return sum(weight[i] * _norm(x ** (i+1)) for i in range(order)) + bias
    
class PolyNorm(torch.nn.Module):
    def __init__(self):
        super(PolyNorm, self).__init__()
        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
        self.bias = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x, checkpointing=True):
        if checkpointing:
            return checkpoint(_poly_norm, x, self.weight, self.bias, use_reentrant=False)
        return _poly_norm(x, self.weight, self.bias)

• For PolyReLU

def _poly(x, weight, bias, order=3):
    return sum(weight[i] * (x ** (i+1)) for i in range(order)) + bias

class PolyReLU(torch.nn.Module):
    def __init__(self):
        super(PolyReLU, self).__init__()
        self.weight = torch.nn.Parameter(torch.ones(3) / 3)
        self.bias = torch.nn.Parameter(torch.zeros(1))

    def forward(self, x, checkpointing=True):
        x = F.relu(x)
        if checkpointing:
            return checkpoint(_poly, x, self.weight, self.bias, use_reentrant=False)
        return _poly(x, self.weight, self.bias)

We would like to thank the reviewer for your time and constructive comments. We address your concerns below.

Q1. Computational complexity analysis, memory, and overflow issue for higher orders of PolyCom.

A1. We analyze the runtime overhead introduced by the activation functions using a typical feedforward network (FFN) with input tensor $x\in \mathbb{R}^{B \times S \times H}$, where $B$，$S$ and $H$ represent the batch size, sequence length, and hidden size, respectively. The relationship between computational FLOPs and model parameters can generally be regarded as proportional (as discussed in Eleuther AI’s transformer math). Below, we estimate the proportion of the computational cost incurred by activation function calculations within the total computational cost of the FFN matrix computations ($24BSH^2$). The FLOPs ratio is calculated as:

$\text{FLOPs ratio} = \frac{\text{FLOPs for activation}}{24BSH^2}$

The results are summarized in the following table:

Note:

• We assume that the scale of the input tensor is set to [-1, 1]. In this case, the FLOPs for both tanh and exp are approximately 10 each.
• To reduce memory overhead during large language model pretraining, we typically employ gradient checkpointing (refer to PyTorch Docs). Although this approach incurs some additional computational cost, its overall impact on GPU memory and runtime is minimal.

Overflow issues in higher-order polynomials can be particularly problematic in the context of BF16. Due to the limited dynamic range of BF16, operations involving polynomials of higher degrees can lead to intermediate values exceeding this range, causing numerical overflow.

Q2. Empirical results do not seem to be averaged across runs.

A2. This is an excellent point. Historically, in the early stages of deep learning development, smaller-scale models (e.g., ResNet50 with 26M parameters, ELMo with 94M parameters) were more sensitive to random seeds, as they could significantly influence the optimization trajectory. However, with the advent of large-scale models and abundant training resources, the training process has become remarkably stable, provided that parameter initialization follows a consistent distribution.

For instance, our MoE model contains 6.9B parameters and was trained on 200B tokens using 64 A100-SXM-80GB GPUs, with a batch size of 400M tokens over a span of 7 days. This large-scale training ensures a stable gradient optimization direction, effectively nullifying the impact of random seed variability. As a result, most recent advancements in large-scale model architectures do not account for seed variation.

For downstream metrics, we employed a greedy decoding strategy, ensuring stable and reproducible results. Acknowledging the inherent fluctuations during training, we have included metric evolution plots in Figures 1, 4, and 11, which clearly demonstrate the consistent and stable performance improvements achieved by our proposed method. This evidence should alleviate concerns regarding potential variability in our results.

Q3. Other types of polynomial activations.

A3. The resource-intensive nature of experiments (each requires at least 32 A100 GPUs for several days), making it hard to include all activation variants. We have carefully selected the most representative activation functions, focusing on the variants that have demonstrated superior performance. Specifically, we have included PolyReLU due to its simplicity and PolyNorm for its effectiveness.

In Section 5, we have conducted experiments with other activation functions for 100 billion tokens and compared training losses across them. Within the domain of large language models (LLM), the training loss exhibits a strong correlation with downstream task metrics, where a lower loss is indicative of better downstream performance. This relationship is further substantiated by the validation loss and perplexity (PPL) trends, as illustrated in Figures 8 through 10 in Appendix F.

Q4. More discussion of the optimal approximation rate or a formal statement of DeVore et al.’s theorem 4.2.

A4. For clarity, we restate Theorem 4.2 from [1] as follows.

Theorem 4.2 [1]. Let $\mathcal{X}$ be a Banach space $L_q$ on $\mathbb{R}^d$, $1\leq q\leq \infty$. If $F_{n,d}^p=\{f \in \mathcal{X} | \|f\|_{\mathcal{W}^{n,p}} \leq 1\}, 1\leq p\leq q, n \in \mathbb{N}$ , then

$\sup_{f\in F_{n.d}^p} \inf_{\theta \in \mathbb{R}^m} \|f-\mathcal{M}(\theta)\|_q\geq C m^{-\frac{n}{d}},$

where $\mathcal{M}$ be a mapping from $\mathbb{R}^m$ into $\mathcal{X}$ which associate with each $\theta \in \mathbb{R}^m$ the element $\mathcal{M}(\theta) \in \mathcal{X}$, and $C$ is a constant.

Particularly, let $q=p=\infty$ and $\mathcal{X}= L_{\infty}[-1,-1]^d$, the theorem indicates that the approximation error for neural networks with $m$ parameters approximating $F_{n,d}^{\infty}$, i.e., $F_{n,d}$, is bounded below by $C m^{-\frac{n}{d}}$. Given an error tolerance $\epsilon$, we have

Dear Reviewers,

Thank you once again for your valuable and constructive feedback on our submission. We deeply appreciate the time and effort you have dedicated to reviewing our work.

We want to kindly ask if there are any additional questions or aspects of our submission that you would like us to clarify or elaborate on. We are more than happy to provide any further explanations to support your review process.