ModuleFormer: Modularity Emerges from Mixture-of-Experts (2306.04640v2)

Abstract: LLMs have achieved remarkable results. However, existing models are expensive to train and deploy, and it is also difficult to expand their knowledge beyond pre-training data without forgetting previous knowledge. This paper proposes a new neural network architecture, ModuleFormer, that leverages modularity to improve the efficiency and flexibility of LLMs. ModuleFormer is based on the Sparse Mixture of Experts (SMoE). Unlike the previous SMoE-based modular LLM, which requires domain-labeled data to learn domain-specific experts, ModuleFormer can induce modularity from uncurated data with its new load balancing and concentration losses. ModuleFormer is a modular architecture that includes two different types of modules: new stick-breaking attention heads and feedforward experts. Different modules are sparsely activated conditions on the input token during training and inference. In our experiment, we found that the modular architecture enables three important abilities for large pre-trained LLMs: 1) Efficiency, since ModuleFormer only activates a subset of its modules for each input token, thus it could achieve the same performance as dense LLMs with more than two times throughput; 2) Extendability, ModuleFormer is more immune to catastrophic forgetting than dense LLMs and can be easily extended with new modules to learn new knowledge that is not included in the training data; 3) Specialisation, finetuning ModuleFormer could specialize a subset of modules to the finetuning task and the task-unrelated modules could be easily pruned for a lightweight deployment.

• The paper introduces a novel architecture that leverages customized load balancing and concentration losses to induce modularity from uncurated data.
• The paper demonstrates that selective module activation via stick-breaking attention heads and feedforward experts doubles throughput while maintaining competitive performance.
• The paper shows improved resistance to catastrophic forgetting and supports continual learning by enabling the addition of new, task-specific modules.

An Overview of "ModuleFormer: Modularity Emerges from Mixture-of-Experts"

The paper introduces ModuleFormer, a novel architecture aimed at enhancing the efficiency and flexibility of LLMs by incorporating principles of modularity. This approach builds on the Sparse Mixture of Experts (SMoE) framework, adapting it to promote modular learning from uncurated data. The significant contributions of this research lie in the architectural design and training mechanisms that allow for efficient computation and extendability in LLMs, addressing common challenges such as high computational costs and catastrophic forgetting.

Architectural Innovations

ModuleFormer diverges from previous SMoE-based models by eliminating the dependency on domain-labeled data. It achieves this through two new loss functions: load balancing and load concentration. These allow the model to automatically induce and leverage modularity from uncurated data. The architecture uses two distinct module types: stick-breaking attention heads and feedforward experts, each selectively activated based on input tokens. This ensures computation is only dedicated to the most relevant modules, significantly increasing throughput while maintaining performance.

Key Abilities Enabled by Modularity

1. Efficiency: ModuleFormer activates only a subset of its modules per input token, achieving comparable performance to dense LLMs while doubling processing throughput. This selective activation reduces latency and memory usage.
2. Extendability: The architecture exhibits increased resistance to catastrophic forgetting and facilitates the addition of new modules for learning additional knowledge domains. This feature is crucial for continuous learning in dynamic environments where knowledge quickly evolves.
3. Specialisation and Pruning: The modular design allows for specific modules to be fine-tuned for particular tasks, with unrelated modules pruned to allow for a lighter, more efficient deployment. This adaptability is beneficial for users constrained by computational resources.

Experimental Validation

The experimental results, spanning tasks such as LLMing and code generation, demonstrate that ModuleFormer achieves parity with dense models while offering efficiency benefits. The model's structure supports the integration of new knowledge domains without sacrificing previously learned information, showcasing its potential in continual learning scenarios.

Implications and Future Directions

ModuleFormer's introduction of modularity into LLMs presents significant advantages in terms of computational efficiency and lifelong learning capabilities. By circumventing the issues associated with dense model finetuning and domain-specific expert labeling, ModuleFormer offers a framework that could redefine how LLMs are deployed and maintained.

Future work in this area could explore the optimization of gating strategies to improve the effectiveness of module selection further. Additionally, examining the application of these principles across various domains beyond textual data could elucidate the broader applicability of modular architectures in AI.

In summary, ModuleFormer provides a robust proof-of-concept for integrating modular design into LLMs, setting a foundation for more efficient and versatile AI models. This research elucidates a path toward scalable, flexible, and continually learning models, advancing the fields of natural language processing and machine learning.