Mixture of Hidden-Dimensions: Not All Hidden-States' Dimensions are Needed in Transformer

Transformer models encounter inefficiency when scaling hidden dimensions due to the uniform expansion of parameters. When delving into the sparsity of hidden dimensions, we observe that only a small subset of dimensions are highly activated, where some dimensions are commonly activated across tokens, and some others uniquely activated for individual tokens. To leverage this, we propose MOHD (Mixture of Hidden Dimensions), a sparse architecture that combines shared sub-dimensions for common features and dynamically routes specialized sub-dimensions per token. To address the potential information loss from sparsity, we introduce activation scaling and group fusion mechanisms. MOHD efficiently expands hidden dimensions with minimal computational increases, outperforming vanilla Transformers in both parameter efficiency and task performance across 10 NLP tasks. MOHD achieves 1.7% higher performance with 50% fewer activated parameters and 3.7% higher performance with 3× total parameters expansion at constant activated parameters cost. MOHD offers a new perspective for scaling the model, showcasing the potential of hidden dimension sparsity. … MoHD Multi-Head Attention RMSNorm MoHD Feed Forward RMSNorm Layer L MoHD Blocks