更具表达力的前馈层：第一部分. 激活函数的Token自适应混合
文章背景与核心概要
在前馈网络（FFN）的设计中，传统的Transformer大语言模型长期依赖于在所有Token上统一应用单一固定的激活函数（如ReLU、GELU或SwiGLU）。尽管这些激活函数在实践中取得了巨大成功，但它们在捕捉不同Token的复杂动态特征时仍显单一。本文指出了传统固定激活函数在表达能力上的局限性，并创新性地提出了两种全新的FFN设计：通过轻量级输入相关门控机制混合激活函数字典的“激活函数混合（MoA）”，以及作为其输入无关对标方案的“可学习激活函数（LA）”。
从理论层面来看，作者严格证明了固定激活函数、LA和MoA在有限宽度下的表达能力层级：LA严格包含了固定激活函数，而MoA又严格包含了LA，其核心增量来自于输入相关的非线性混合机制。而在实证方面，通过在0.12B至2B参数规模的稠密型和混合专家（MoE）语言模型上进行的大规模预训练实验，MoA在带来极小参数和计算开销的同时，展现出了更低的终端损失和更优越的缩放行为（scaling behavior）。
📌 摘要 (Summary)
前馈网络（FFN）层占用了基于Transformer的大语言模型（LLMs）中绝大部分的参数量和非线性表达能力。尽管激活函数经历了从ReLU、GELU到诸如SwiGLU等门控变体的演进，但大多数FFN设计仍在使用单一的固定激活函数，对所有Token应用相同的非线性变换。在这项工作中，我们提出了激活函数混合（MoA，Mixture of Activations），这是一种Token自适应的FFN设计，它通过轻量级的输入相关门控来混合激活函数字典，同时共享相同的线性投影。作为其输入无关的对应方案，我们还引入了可学习激活函数（LA，Learnable Activations），它为ReLU型和SwiGLU型的FFN构建了激活函数的线性组合。
在理论上，我们确立了固定激活函数、LA和MoA之间严格的有限宽度表达能力分离：LA严格包含固定激活函数，而MoA严格包含MoA，其额外的表达能力源于输入相关的非线性混合。在实验上，我们在0.12B到2B参数规模的稠密型和MoE语言模型上，通过在不同Token预算、优化器和学习率调度下的广泛预训练实验对MoA进行了评估。在极小的参数和计算开销下，MoA始终能够实现比精心调优的基线更低的终端损失，并表现出更有利的缩放行为。这些结果表明，Token自适应的激活函数混合是提升LLM中FFN表达能力的一种简单而有效的机制。
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.
🔗 快速链接 (Quick Links)
📄 摘要原文 (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.