Approximating Two-Layer Feedforward Networks for Efficient Transformers (2310.10837v3)

Abstract: How to reduce compute and memory requirements of neural networks (NNs) without sacrificing performance? Many recent works use sparse Mixtures of Experts (MoEs) to build resource-efficient LLMs (LMs). Here we introduce several novel perspectives on MoEs, presenting a general framework that unifies various methods to approximate two-layer NNs (e.g., feedforward blocks of Transformers), including product-key memories (PKMs). Leveraging insights from this framework, we propose methods to improve both MoEs and PKMs. Unlike prior work that compares MoEs with dense baselines under the compute-equal condition, our evaluation condition is parameter-equal, which is crucial to properly evaluate LMs. We show that our MoEs are competitive with the dense Transformer-XL on both the WikiText-103 and enwiki8 datasets at two different scales, while being much more resource efficient. This demonstrates that MoEs are relevant not only to extremely large LMs but also to any-scale resource-efficient LMs. Our code is public.

A Critical Evaluation of the Sigma-MoE Framework for Efficient LLMing

The paper presents an empirical investigation into the efficacy of Mixture of Experts (MoEs) architectures, specifically introducing a novel variant called the sigma-MoE. The paper aims to challenge the prevailing belief that MoEs underperform compared to dense models under parameter-equivalent conditions, such as Transformer-XL. The central thesis advocates that the sigma-MoE framework can achieve competitive performance while maintaining computational efficiency.

Contribution of Sigma-MoE Framework

The sigma-MoE model diverges from traditional MoEs by introducing several innovative architectural components. These components include:

1. Non-competitive Selection Function: The model employs a sigmoid activation function in place of the conventional softmax. This choice aligns sigma-MoE with the non-competitive dynamics seen in standard feedforward networks, effectively approximating top-k selections in feedforward layers.
2. Global Entropy Regularization: Leveraging a regularization scheme that prioritizes balanced expert utilization while avoiding complex and arbitrary constraints. This methodology regulates the entropy of selection scores computed within batches, aiming for each expert to contribute to processing.
3. Expert Dropout: To prevent the collapsing phenomenon, where only select experts are engaged predominantly, sigma-MoE incorporates expert dropout, ensuring the distribution of tasks across experts is more uniform.
4. Normalized Initialization of Selection Mechanism: By normalizing the row vector lengths of the projection matrix, the authors mitigate biases in expert selection that could stem from unequal norms. This strategy intends to maintain uniform selection conditions across experts.

Empirical Evaluation

The paper conducts comprehensive experiments on diverse datasets, including C4 and peS2o, highlighting the robustness of the sigma-MoE framework. The results demonstrate that sigma-MoE can achieve comparable or superior perplexity scores to dense and other baseline models—an achievement underscored particularly when considering parameter-equivalence.

• For instance, in experiments using the C4 dataset with a model size of 1024 and 262 million parameters, sigma-MoE achieved perplexity scores that were slightly better or on par with dense model counterparts.
• Moreover, sigma-MoE's capability to retain robust performance with reduced computational resource expenditure, achieving a 75-87.5% reduction in FLOPs, underscores the framework's efficiency.

Discussion on Computational Efficiency and Trade-offs

The sigma-MoE framework is evaluated against standard metrics such as execution time and memory usage benchmarks, demonstrating similar computational footprints across MoE variants due to shared selection mechanism dynamics. The paper also introduces new tables detailing FLOPs and memory reduction, addressing reviewer concerns and ensuring clarity in computational efficiency comparisons.

The theoretical and practical implications of sigma-MoE are significant, as they provide a scalable solution for reducing the bottleneck associated with two-layer feedforward blocks in transformers—posing an essential step toward resource-efficient LLMs.

Future Directions

While sigma-MoE demonstrates promise, further exploration could involve extending its applications to downstream tasks to evaluate model transferability. Additionally, examining the impacts of various hyperparameters in different deployment scenarios may unveil new insights into optimizing MoE architectures for diverse LLMing tasks.

In summary, this paper enriches the discourse on efficient computational models for language processing by challenging entrenched assumptions about MoEs. Its empirical results, theoretical insights, and methodical evaluations mark a significant stride toward making transformer architectures more adaptable and efficient without compromising performance. As such, the sigma-MoE framework holds potential as a practical tool in advancing the development of scalable AI systems.