BASE Layers: Simplifying Training of Large, Sparse Models (2103.16716v1)

Abstract: We introduce a new balanced assignment of experts (BASE) layer for LLMs that greatly simplifies existing high capacity sparse layers. Sparse layers can dramatically improve the efficiency of training and inference by routing each token to specialized expert modules that contain only a small fraction of the model parameters. However, it can be difficult to learn balanced routing functions that make full use of the available experts; existing approaches typically use routing heuristics or auxiliary expert-balancing loss functions. In contrast, we formulate token-to-expert allocation as a linear assignment problem, allowing an optimal assignment in which each expert receives an equal number of tokens. This optimal assignment scheme improves efficiency by guaranteeing balanced compute loads, and also simplifies training by not requiring any new hyperparameters or auxiliary losses. Code is publicly released at https://github.com/pytorch/fairseq/

• The paper presents BASE layers, which reformulate token-to-expert routing as a balanced linear assignment problem to ensure equal expert workload.
• It employs a soft mixing mechanism through modified residual connections, eliminating the need for complex routing heuristics and additional hyperparameter tuning.
• Experimental evaluations on models up to 110B parameters reveal improved compute efficiency and performance compared to traditional sparse expert methods.

BASE Layers: Simplifying Training of Large, Sparse Models

The paper introduces an innovative approach for training LLMs using what the authors term as the BASE (Balanced Assignment of Sparse Experts) layer. The primary focus is on the simplification and efficiency enhancement of sparse model training without the need for complex routing heuristics or additional hyperparameter tuning.

Core Contributions

The BASE layer offers a structured solution to the challenge of balanced token-to-expert routing in large, sparse models. Unlike previous methodologies that employed auxiliary loss functions or heuristic-based routing, the BASE layer formulates the routing of tokens to experts as a linear assignment problem. This ensures each expert receives an equal number of tokens, optimizing compute load distribution and simplifying the training process. By employing a straightforward yet effective soft mixing mechanism through modified residual connections, the BASE layer facilitates expert specialization seamlessly.

Methodology

1. Architecture Design: BASE layers consist of a set of experts, each parameterized by position-wise functions and expert embeddings. A linear assignment algorithm is utilized for training to balance the token distribution across experts, which mitigates the complexity associated with custom loss functions or auxiliary parameters.
2. Training and Inference: During training, a balanced assignment is computed to ensure optimal expert workload. At inference, a greedy approach assigns tokens to the highest scoring experts to maintain efficiency without data leakage from future contexts.
3. Implementation Efficiency: The linear assignment problem is efficiently solved using the auction algorithm, enabling effective load balance without overwhelming computational overhead. This contributes to enhanced performance metrics and compute efficiency, allowing the nesting of expert layers within a standard transformer architecture.

Experimental Evaluation

Extensive experiments demonstrate the effectiveness of BASE layers in models with up to 110B parameters, outperforming traditional dense and model-parallel strategies across various compute environments (8, 32, and 128 GPUs). These sparse models not only achieve higher performance but also enhance compute efficiency compared to current sparse expert models such as Sparsely Gated Mixtures of Experts (MoE) and Switch transformers.

BASE layers demonstrate resilience across different model configurations, showing robustness in balancing experts even without explicit balancing loss terms. The capability of BASE layers to outperform in large computational regimes highlights their applicability in scaling neural architectures without proportional cost surges.

Implications and Future Directions

The achievements presented underscore the BASE layer's potential in advancing the training of sparse models, evidenced by substantial performance improvements and reduced communication overhead. On a broader scale, this work paves the way for further exploration into sparse training methodologies that can be easily integrated into existing architectures, reducing the reliance on dense parameter utilization.

Future avenues of research might focus on:

• Exploring finer token assignment granularity and its impact on expert specialization.
• Developing more computationally efficient algorithms for dynamic linear assignments.
• Investigating the application of BASE layers beyond NLP, such as in computer vision or reinforcement learning contexts.

In conclusion, the BASE layer represents a significant step forward in simplifying and improving the efficiency of training large-scale sparse models. Its deployment can lead to more resource-efficient models without compromising performance, a critical need in modern AI development.