Mixture of Hidden-Dimensions Transformer (2412.05644v3)

Abstract: Transformer models encounter challenges in scaling hidden dimensions efficiently, as uniformly increasing them inflates computational and memory costs while failing to emphasize the most relevant features for each token. For further understanding, we study hidden dimension sparsity and observe that trained Transformers utilize only a small fraction of token dimensions, revealing an "activation flow" pattern. Notably, there are shared sub-dimensions with sustained activation across multiple consecutive tokens and specialized sub-dimensions uniquely activated for each token. To better model token-relevant sub-dimensions, we propose MoHD (Mixture of Hidden Dimensions), a sparse conditional activation architecture. Particularly, MoHD employs shared sub-dimensions for common token features and a routing mechanism to dynamically activate specialized sub-dimensions. To mitigate potential information loss from sparsity, we design activation scaling and group fusion mechanisms to preserve activation flow. In this way, MoHD expands hidden dimensions with negligible increases in computation or parameters, efficient training and inference while maintaining performance. Evaluations across 10 NLP tasks show that MoHD surpasses Vanilla Transformers in parameter efficiency and task performance. It achieves 1.7% higher performance with 50% fewer activation parameters and 3.7% higher performance with a 3x parameter expansion at constant activation cost. MOHD offers a new perspective for scaling the model, showcasing the potential of hidden dimension sparsity to boost efficiency

• The paper introduces a novel architecture that leverages hidden dimension sparsity to optimize computational efficiency in Transformer models.
• It employs a dynamic routing mechanism with shared and specialized sub-dimensions to maintain performance with fewer activation parameters.
• Empirical validation across 10 NLP tasks shows up to a 3.7% performance gain while significantly reducing computational overhead.

Mixture of Hidden-Dimensions Transformer: A Technical Overview

The paper "Mixture of Hidden-Dimensions Transformer" proposes a novel architecture, the Mixture of Hidden-Dimensions (MOHD), designed to enhance the efficiency of Transformer models. By addressing the challenges of hidden dimension sparsity, MOHD offers an innovative approach to model scaling that reduces the computational and memory overhead associated with Transformers.

Architecture and Methodology

MOHD introduces a sparse conditional activation framework that leverages hidden dimension sparsity observed in large-scale LLMs. The central concept of this architecture is the differentiation between shared and specialized sub-dimensions. Shared sub-dimensions are consistently activated across multiple tokens, capturing common features, while specialized sub-dimensions are selectively activated for token-specific characteristics.

The architecture utilizes a dynamic routing mechanism, which enables efficient activation of relevant sub-dimensions based on token input, thereby maintaining parameter performance without a proportional increase in parameter count. This mechanism is complemented by activation scaling and grouped fusion techniques to preserve activation flow and mitigate information loss due to sparseness.

Empirical Findings

The efficacy of MOHD is empirically validated across ten diverse NLP tasks. The results demonstrate that MOHD outperforms traditional Vanilla Transformers, indicating higher parameter efficiency and task performance. Specifically, MOHD achieves a 1.7% performance improvement with only 50% of the activation parameters, and a 3.7% improvement with a threefold increase in parameter expansion, where activation costs remain constant.

The paper highlights significant hidden dimension sparsity, where 50% of dimensions account for over 92% of activation magnitude. Insights from these observations drive the MOHD design, showing that traditional Transformers underutilize many dimensional aspects, thereby presenting opportunities for efficiency enhancement.

Theoretical and Practical Implications

The MOHD architecture exemplifies a strategic advancement in leveraging hidden dimension sparsity to optimize the scalability and efficiency of Transformer models. Theoretical implications suggest that modeling complexity and capacity in Transformers critically depend on informed activation strategies that prioritize meaningful sub-dimensions, rather than uniformly scaling hidden dimensions.

Practically, the MOHD architecture paves the way for more resource-efficient AI systems capable of robust performance with reduced computation. This is highly beneficial for real-world applications where computational cost and speed are crucial, such as in large-scale deployment of NLP models.

Future Directions

Looking forward, the possibilities for refining and extending the MOHD architecture are vast. One potential avenue is the exploration of more nuanced activation strategies that dynamically adjust shared and specialized sub-dimension ratios based on contextual and domain-specific requirements. Additionally, integrating MOHD with emerging model pruning techniques could further enhance model efficiency without compromising performance.

In conclusion, the Mixture of Hidden-Dimensions Transformer provides a promising new perspective in efficiently scaling the hidden dimensions of Transformer models, aligning computational demand with the asymmetries in activation sparsity. Such advancements underscore the evolving landscape of model architecture design, emphasizing efficiency without detracting from efficacy.