More Expressive Feedforward Layers: Part I. Token-Adaptive Mixing of Activations

📄 arXiv: 2605.26647v1 📥 PDF

作者: Mingze Wang, Jinbo Wang, Yikuan Xia, Kai Shen, Shu Zhong

分类: cs.LG, cs.AI, stat.ML

发布日期: 2026-05-26

备注: 31 pages

💡 一句话要点

提出Token自适应激活混合(MoA)方法，提升Transformer FFN层表达能力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer 前馈网络 激活函数 自适应 语言模型 MoA 可学习激活

📋 核心要点

  1. 现有Transformer模型中的FFN层通常使用固定的激活函数，限制了其表达能力，无法根据不同token进行自适应调整。
  2. 论文提出MoA方法，通过轻量级的输入相关门控机制，混合多个激活函数，实现token级别的自适应非线性变换。
  3. 实验表明，MoA在多种规模的语言模型上均能取得更低的损失，并展现出更好的模型扩展性，同时计算开销很小。

📝 摘要（中文）

前馈网络(FFN)层在基于Transformer的大型语言模型(LLM)中占据了大部分参数和非线性表达能力。尽管从ReLU和GELU发展到SwiGLU等门控变体，但大多数FFN设计仍然使用单一的固定激活函数，对所有token应用相同的非线性变换。本文提出了激活混合(MoA)，一种token自适应的FFN设计，它使用轻量级的输入相关门控来混合激活函数字典，同时共享相同的线性投影。作为输入无关的对应方法，我们还引入了可学习激活(LA)，它为ReLU型和SwiGLU型FFN形成激活函数的线性组合。理论上，我们建立了固定激活FFN、LA和MoA之间严格的有限宽度表达分离：LA严格包含固定激活FFN，而MoA严格包含LA，额外的表达能力来自输入相关的非线性混合。在经验上，我们通过在0.12B到2B参数的密集和MoE语言模型上进行广泛的预训练实验来评估MoA，这些实验在不同的token预算、优化器和学习率调度下进行。MoA始终实现更低的终端损失，并表现出比良好调整的基线更有利的缩放行为，且参数和计算开销最小。这些结果表明，token自适应激活混合是提高LLM中FFN表达能力的一种简单有效的机制。

🔬 方法详解

问题定义：现有Transformer模型中的FFN层，通常采用ReLU、GELU或SwiGLU等固定的激活函数。这种固定激活方式无法根据不同token的特性进行自适应调整，限制了FFN层的表达能力，阻碍了模型性能的进一步提升。

核心思路：论文的核心思路是引入token自适应的激活函数混合机制。通过学习输入相关的门控权重，动态地混合多个预定义的激活函数，使得FFN层能够根据不同token的特征，选择最合适的非线性变换，从而提升模型的表达能力。

技术框架：MoA方法的核心在于激活函数的混合。对于每个token，首先通过一个轻量级的线性层计算出混合权重，然后使用这些权重对预定义的激活函数字典中的激活函数进行加权求和，得到最终的激活函数。这个过程可以表示为：Activation = sum(gate_i * activation_function_i)，其中gate_i是门控权重，activation_function_i是激活函数字典中的第i个激活函数。

关键创新：MoA的关键创新在于引入了token自适应的激活函数混合机制。与传统的固定激活函数相比，MoA能够根据不同token的特征动态地调整激活函数，从而提升模型的表达能力。此外，MoA使用轻量级的门控机制，避免了引入过多的参数和计算开销。

关键设计：MoA的关键设计包括：1) 激活函数字典的选择，可以选择常用的ReLU、GELU、Swish等激活函数；2) 门控机制的设计，可以使用简单的线性层或者更复杂的神经网络；3) 混合权重的归一化方式，可以使用softmax或者sigmoid等函数进行归一化，确保权重的和为1。此外，论文还提出了输入无关的可学习激活(LA)作为对比，LA通过学习激活函数的线性组合，但混合权重不依赖于输入。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，MoA在多种规模的语言模型上均能取得更低的终端损失，并展现出比良好调整的基线更有利的缩放行为。例如，在2B参数的语言模型上，MoA相比于基线模型，在相同的训练步数下，能够取得更低的损失。此外，MoA的参数和计算开销很小，使得其易于集成到现有的Transformer模型中。

🎯 应用场景

MoA方法可以广泛应用于各种基于Transformer的自然语言处理任务，例如文本分类、机器翻译、文本生成等。通过提升FFN层的表达能力，MoA有望提高这些任务的性能，并降低模型对数据量的依赖。此外，MoA的token自适应特性使其在处理长文本和复杂语义关系时具有优势。

📄 摘要（原文）

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.