더욱 표현력이 뛰어난 피드포워드 레이어: Part I. 토큰 적응형 활성화 혼합

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

피드포워드 네트워크(FFN) 레이어는 트랜스포머 기반의 대규모 언어 모델(LLM)에서 상당한 비율의 파라미터와 비선형 표현력을 담당합니다. ReLU 및 GELU에서 SwiGLU와 같은 게이티드 방식으로 발전해 왔음에도 불구하고, 대부분의 FFN 설계는 여전히 단일 고정 활성화 함수를 사용하며 모든 토큰에 동일한 비선형 변환을 적용합니다. 본 연구에서는 입력 의존적인 게이트를 사용하여 사전 정의된 활성화 함수의 집합을 혼합하는 토큰 적응형 FFN 설계인 Mixture of Activations (MoA)를 제안합니다. 또한, 입력과 독립적인 방식으로 ReLU 유형 및 SwiGLU 유형의 FFN 모두에서 활성화 함수들의 선형 조합을 형성하는 학습 가능한 활성화(LA) 방식을 소개합니다. 이론적으로, 고정 활성화 FFN, LA, MoA 간의 엄격한 표현력 차이를 분석했습니다. LA는 고정 활성화 FFN을 포함하며, MoA는 LA를 포함하며, 입력 의존적인 비선형 혼합을 통해 추가적인 표현력을 제공합니다. 실험적으로, 0.12B에서 2B 파라미터 범위의 밀집 및 MoE 언어 모델에 대해 다양한 토큰 할당량, 옵티마이저 및 학습률 스케줄 하에서 광범위한 사전 훈련 실험을 통해 MoA를 평가했습니다. MoA는 일관되게 더 낮은 최종 손실 값을 달성하고 잘 조정된 기준 모델보다 더 나은 확장성을 보였으며, 파라미터 및 계산 오버헤드는 최소화되었습니다. 이러한 결과는 토큰 적응형 활성화 혼합이 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.