More Expressive Feedforward Layers: Part I. Token-Adaptive Mixing of Activations
Abstract
Feedforward network (FFN) layers account for a large fraction of parameters and nonlinear expressivity in Transformer-based large language models (LLMs). Despite the evolution from ReLU and GELU to gated variants such as SwiGLU, most FFN designs still use a single fixed activation function, applying the same nonlinear transformation to all tokens. In this work, we propose Mixture of Activations (MoA), a token-adaptive FFN design that mixes a dictionary of activation functions using lightweight input-dependent gates while sharing the same linear projections. As an input-independent counterpart, we also introduce learnable activations (LA), which form linear combinations of activation functions for both ReLU-type and SwiGLU-type FFNs. Theoretically, we establish strict finite-width expressive separations among fixed-activation FFNs, LA, and MoA: LA strictly contains fixed-activation FFNs, while MoA strictly contains LA, with the additional expressivity arising from input-dependent nonlinear hybridization. Empirically, we evaluate MoA through extensive pre-training experiments on dense and MoE language models ranging from 0.12B to 2B parameters under different token budgets, optimizers, and learning rate schedules. MoA consistently achieves lower terminal loss and exhibits more favorable scaling behavior than well-tuned baselines, with minimal parameter and computational overhead. These results suggest that token-adaptive activation mixing is a simple and effective mechanism for improving FFN expressivity in LLMs.
Related
- Sharp and Adaptive Cluster Recovery in Slowly Mixing Gaussian Hidden Markov Models
- Radon Measure Representations for Infinite-Width Neural Networks with Singular Activations
- What Do Sparse Autoencoder Feature Discovery Scores Identify? From Activations to Interventions
- Fully Deterministic Shallow Neural Networks with Localized Activations. Part II: Condition Number Estimates
- Periodically kicked feedforward chains of simple excitable FitzHugh–Nagumo neurons
BibTeX
@misc{wang2026more,
title = {{More Expressive Feedforward Layers: Part I. Token-Adaptive Mixing of Activations}},
author = {Mingze Wang and Jinbo Wang and Yikuan Xia and Kai Shen and Shu Zhong},
howpublished = {arXiv},
year = {2026},
month = aug,
doi = {10.48550/arxiv.2605.26647},
eprint = {2605.26647},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2605.26647},
}
RIS
TY - UNPB
TI - More Expressive Feedforward Layers: Part I. Token-Adaptive Mixing of Activations
AU - Mingze Wang
AU - Jinbo Wang
AU - Yikuan Xia
AU - Kai Shen
AU - Shu Zhong
PY - 2026
DA - 2026/08/28
PB - arXiv
DO - 10.48550/arxiv.2605.26647
UR - https://arxiv.org/abs/2605.26647
AB - Feedforward network (FFN) layers account for a large fraction of parameters and nonlinear expressivity in Transformer-based large language models (LLMs). Despite the evolution from ReLU and GELU to gated variants such as SwiGLU, most FFN designs still use a single fixed activation function, applying the same nonlinear transformation to all tokens. In this work, we propose Mixture of Activations (MoA), a token-adaptive FFN design that mixes a dictionary of activation functions using lightweight input-dependent gates while sharing the same linear projections. As an input-independent counterpart, we also introduce learnable activations (LA), which form linear combinations of activation functions for both ReLU-type and SwiGLU-type FFNs. Theoretically, we establish strict finite-width expressive separations among fixed-activation FFNs, LA, and MoA: LA strictly contains fixed-activation FFNs, while MoA strictly contains LA, with the additional expressivity arising from input-dependent nonlinear hybridization. Empirically, we evaluate MoA through extensive pre-training experiments on dense and MoE language models ranging from 0.12B to 2B parameters under different token budgets, optimizers, and learning rate schedules. MoA consistently achieves lower terminal loss and exhibits more favorable scaling behavior than well-tuned baselines, with minimal parameter and computational overhead. These results suggest that token-adaptive activation mixing is a simple and effective mechanism for improving FFN expressivity in LLMs.
ER -