Reproduce “ModuleFormer: Modularity Emerges from Mixture-of-Experts” on OpenResearch

About this paper

Yikang Shen, Zheyu Zhang, Tianyou Cao, Shawn Tan, Zhenfang Chen, Chuang Gan

ModuleFormer, developed by researchers from MIT-IBM Watson AI Lab, Tsinghua University, and Mila, presents a modular Transformer architecture that leverages Sparse Mixture of Experts to improve the efficiency and adaptability of large language models. The model achieves comparable performance to dense LLMs with significantly fewer active parameters and demonstrates enhanced continual learning capabilities and effective module-level pruning.

Key insights

• Modularity can emerge naturally during pre-training on general, uncurated data by employing a Mutual Information loss that balances expert utilization and encourages specialization.
• A sparse, modular architecture inherently promotes better resistance to catastrophic forgetting during continual learning, outperforming dense models and existing parameter-efficient fine-tuning methods like LoRA.
• Specific module manipulation techniques, such as load concentration loss for fine-tuning and routing regularization for new module insertion, enable effective pruning for specialization and seamless knowledge expansion.

Summary by alphaXiv.