Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models

Authors: Zhijian Zhuo, Ya Wang, Yutao Zeng, Xiaoqing Li, Xun Zhou, Jinwen Ma
Primary area: foundation or frontier models, including LLMs
Keywords: activation function, transformer, pre-training, large language models

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

Decision: accept (poster)
Reviewer scores: 6, 8, 8, 6

Reviews and discussion:
<thread>
creation date:
1730609472620

thread_type:
official_review

<thread_object>
creator: reviewer_tiu9

<content>
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.

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.

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

questions:
Check the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680960997

thread_type:
official_review

<thread_object>
creator: reviewer_nvau

<content>
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.

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

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

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730705461937

thread_type:
official_review

<thread_object>
creator: reviewer_nayg

<content>
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.

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.

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.

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730717503478

thread_type:
official_review

<thread_object>
creator: reviewer_anh6

<content>
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.

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

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?

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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732086310681

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
A Summary of Paper Updates

comment:
We sincerely thank all reviewers for their constructive feedback. Based on your valuable suggestions, we have refined the paper to enhance clarity and rigor. The key updates include:

- Section 3.2: Restated the lower bound portion of Theorem 2 for improved clarity.
- Appendix B: Expanded discussion on the optimal approximation rate.
- Appendix D2: Added hyperparameters for the 1B dense model and MoE-1B-7B in Table 7.
- Appendix E: Introduced computational complexity analysis.
- Appendix H: Presented scaling curves for different models.

We hope these adjustments meet your expectations and contribute to the improved readability and comprehension of our paper. If there are any further questions or suggestions, please feel free to let us know. We are looking forward to providing additional clarification and discussion.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732558809108

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response to the Computational Overhead and Memory Footprint

comment:
We would like to thank again to all the reviewers for their insightful comments!

One common question raised by the reviewers is: what are **the computational overhead and memory footprint** incurred by those activation functions?


To address this, we conducted a thorough analysis using a typical feedforward network 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 transformer model parameters can generally be regarded as proportional (as discussed in [Eleuther AI’s transformer math](https://blog.eleuther.ai/transformer-math/)). Below, we estimate the proportion of the computational cost incurred by activation functions within the total computational cost of the FFN matrix computations ($24BSH^2$) (actually, when considering all transformer modules, the ratio would be even smaller). 

The FLOPs ratio is:

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

It is important to note that the overhead and proportion often vary for different model sizes, so we provide the corresponding formulas directly and take $H=1024, B=4 (\texttt{each\ device}), S=4096$, using BF16 precision as an example.

- without gradient checkpointing:

| Method               |     ReLU      |    GeLU     |    SwiGLU     |   ReLU$^2$    | 3rd-order PolyNorm | 3rd-order PolyReLU |
| -------------------- | :-----------: | :---------: | :-----------: | :-----------: | :----------------: | :----------------: |
| Intermediate Size    |      4H       |     4H      |     8/3H      |      4H       |         4H         |         4H         |
| FLOPs for activation |     4BSH      |    72BSH    |   112/3BSH    |     8BSH      |       72BSH        |       40BSH        |
| FLOPs ratio          | 1/(6H)=0.016% |  3/H=0.29%  | 14/(9H)=0.15% | 1/(3H)=0.032% |     3/H=0.29%      |    5/(3H)=0.16%    |
| Memory Overhead      |  4BSH=128MB   | 10BSH=320MB |  8BSH=256MB   |  8BSH=256MB   |    12BSH=384MB     |     8BSH=256MB     |

- with gradient checkpointing:

| Method               |     ReLU      |   GeLU    |    SwiGLU     |   ReLU$^2$    | 3rd-order PolyNorm | 3rd-order PolyReLU |
| -------------------- | :-----------: | :-------: | :-----------: | :-----------: | :----------------: | :----------------: |
| Intermediate Size    |      4H       |    4H     |     8/3H      |      4H       |         4H         |         4H         |
| FLOPs for activation |     8BSH      |  144BSH   |   224/3BSH    |     16BSH     |       144BSH       |       80BSH        |
| FLOPs ratio (H=1024) | 1/(3H)=0.033% | 6/H=0.59% | 28/(9H)=0.30% | 2/(3H)=0.065% |     6/H=0.59%      |   10/(3H)=0.33%    |
| Memory Overhead      |       0       |     0     |       0       |       0       |         0          |         0          |

Note:
- We assume the scale of the input 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.
- To reduce memory overhead during pretraining, we typically employ gradient checkpointing (refer to [PyTorch Docs](https://pytorch.org/docs/stable/checkpoint.html)). Although this approach incurs some additional computational cost, its overall impact on GPU memory and runtime is minimal.

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

Additionally, we have provided the code with gradient checkpointing as follows:

- For PolyNorm
```python 
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
```python 
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)
```
</content>
</thread_object>

</thread>
<thread>
creation date:
1733309731187

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers and Area Chairs,

We are truly grateful for the thorough insights and constructive feedback offered by the reviewers during the review process, as well as the valuable guidance from the Area Chairs. As we reach the end of the rebuttal phase, we would like to highlight the key contributions of our manuscript, "Polynomial Composition Activations: Unleashing the Dynamics of Large Language Models," and share our reflections on the discussions surrounding the feedback we received.

**Main Contributions**

We propose a novel activation function, PolyCom, which is a composition of polynomial and other types of functions. Furthermore, we introduce two instances of PolyCom: PolyReLU and PolyNorm, and detail their seamless integration into the transformer architecture. Our activation functions highlight the following key advantages:
1. **Comparable Computational Overhead and Memory Footprint**: the overhead and memory footprint are acceptable, and there is negligible difference in the training budget required compared to the widely used SwiGLU.
2. **Optimal Theoretical Approximation Rates**: Theoretically, we derive bounds on the number of trainable parameters required for PolyReLU networks to approximate ReLU networks, and vice versa. Moreover, we demonstrate that a PolyReLU network of size $O(\epsilon^{-d/n})$ can approximate any function in Sobolev spaces with error tolerance $\epsilon$, thereby achieving the optimal approximation rates.
3. **Better Convergence Speed**: Empirically, we validate the efficacy of this novel activation function on LLMs with 1B dense models and MoE models with 1B active and 7B total parameters. The results of both models reveal that PolyCom can attain a remarkable 1.5x convergence speedup when contrasted with SwiGLU.

These contributions were positively acknowledged by the reviewers:
- "...exhibit non-trivial performance gains for training language models with >1B parameters, even compared to strong baselines such as SwiGLU and squared ReLU." by Reviewer NayG.
- "...along with improved convergence rates in the learning curves with fixed parameter size model..." by Reviewer "TiU9"
- "...the empirical performance improvements appear promising..." by Reviewer aNH6
- "...has strong theoretical guarantees, ...and has an optimal approximation rate for general smooth functions in Sobolev spaces." by Reviewer NayG
- "...without adding more trainable parameters to the model..." by Reviewer NvaU


**Summary of Revisions**

In response to the constructive feedback from the reviewers, we have made several significant adjustments and additions:

**Addressing common responses:**
1. **Expanded Theoretical Clarifications:** Based on reviewers' comments, we clarified the theoretical contributions, particularly regarding Theorem 2 and Lemma 2, ensuring the distinction between supremum and upper bounds was accurately presented.
2. **Added Computational Complexity Analysis:** We provided detailed analyses of the runtime and memory overhead for the proposed activation functions, including FLOPs ratios and memory consumption, which were included in the revised appendix.
3. **Clarification on Training Stability:** We expanded the discussion on training stability for PolyReLU and PolyNorm, attributing stability to normalization operators and demonstrating consistent results across transformer-based and ResNet50 architectures. Furthermore, we incorporate some practical usage recommendations.
4. **Included Additional Experiments:** In response to requests for broader evaluations, we conducted experiments on ResNet50 for non-transformer settings and extended scaling law evaluations to compare performance across model sizes.

**Addressing individual responses:**
1. **Reviewer aNH6:** We addressed the misunderstandings about Theorem 2, explained impracticality of full convergence in LLM pretraining, and clarified minimal parameter increments for PolyReLU and PolyNorm.
2. **Reviewer NayG:** We provided FLOPs and memory analyses, discussed effective rank's role in expressivity, and added scaling law experiments showing consistent performance improvements.
3. **Reviewer NvaU:** We compared computational overhead and generalizability beyond transformers with new ResNet50 experiments; clarified Theorem 4.2 for optimal approximation rates.
4. **Reviewer TiU9:** We explained how PolyNorm mitigates instability in FP16/BF16 training and confirmed no exploding gradient issues due to normalization and clipping strategies.

We believe that we have addressed the reviewers' key concerns effectively through thoughtful revisions and clarifications. These improvements have strengthened the clarity, rigor, and overall quality of our paper.

Once more, we extend our heartfelt thanks to the reviewers and area chairs for their constructive feedback and valuable guidance. Your time and effort have been pivotal in refining our work, and we are truly grateful for your support.

Best regards,

Authors of Submission #10038
</content>
</thread_object>

</thread>


Metareview:
metareview:
**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.