Mixture of Modular Experts: Distilling Knowledge from a Multilingual Teacher into Specialized Modular Language Models (2407.19610v1)

Abstract: This research combines Knowledge Distillation (KD) and Mixture of Experts (MoE) to develop modular, efficient multilingual LLMs. Key objectives include evaluating adaptive versus fixed alpha methods in KD and comparing modular MoE architectures for handling multi-domain inputs and preventing catastrophic forgetting. KD compresses LLMs into smaller, efficient models, while MoE enhances modularity with specialized tasks. Experiments showed similar performance for both KD methods, with marginal improvements from adaptive alpha. A combined loss approach provided more stable learning. The router, trained to classify input sequences into English, French, German, or Python, achieved 99.95% precision, recall, and F1 score, with Logistic Regression being the most effective classifier. Evaluations of modular MoE architectures revealed that Pre-trained Language Experts (PLE) and Joint Expert Embedding Training (JEET) performed similarly, while the MoE with Common Expert (MoE-CE) setup showed slightly lower performance. Including a common expert in MoE-CE improved its performance. Studies on catastrophic forgetting indicated that sequential training led to significant forgetting, while single-session training with balanced batches and the MoE approach mitigated this issue. The MoE architecture preserved knowledge across multiple languages effectively. The research contributes open-sourced resources including the dataset (https://zenodo.org/doi/10.5281/zenodo.12677631), a balanced dataset creation tool (https://github.com/padas-lab-de/multi-language-dataset-creator), and the research codebase (https://github.com/ModMaamari/mixture-modular-experts).

• The paper introduces an innovative framework combining knowledge distillation and mixture of experts to distill a large multilingual teacher into efficient modular models.
• The research compares adaptive and fixed alpha methods for KD and employs a high-precision router to allocate resources effectively across language experts.
• The findings demonstrate that modular MoE architectures mitigate catastrophic forgetting and enable scalable, multi-domain language processing without full retraining.

Overview of "Mixture of Modular Experts: Distilling Knowledge from a Multilingual Teacher into Specialized Modular LLMs"

This paper examines the integration of Knowledge Distillation (KD) and Mixture of Experts (MoE) frameworks to create multilingual LLMs that are both modular and efficient. It explores various strategies to enhance the model performance while maintaining its ability to handle multi-domain inputs. The paper focuses on several distinct areas: evaluating different KD approaches, training an effective router for language classification, and assessing the performance of different MoE architectures in preventing catastrophic forgetting.

Methodology

The research commences by utilizing KD to compress a large-scale GPT-2 Medium teacher model, which possesses 340 million parameters, into smaller student models. This process is streamlined using both adaptive and fixed alpha methods, with an aim to examine their relative advantages in KD performance. The adaptive alpha method yielded a small yet measurable improvement over its fixed counterpart, although both demonstrated comparable results.

The core contributions involve exploring multiple MoE architectures, specifically the Pre-trained Language Experts (PLE), Joint Expert Embedding Training (JEET), and MoE with Common Expert (MoE-CE). These architectures leverage specialized models—or 'experts'—that can dynamically process inputs based on language classification, facilitated by a highly accurate router component which achieves a precision and recall of 99.95% using Logistic Regression for classification. The router's capability ensures efficient resource allocation during inference.

Results and Evaluation

Empirical evaluations reflect nuanced differences among the MoE architectures. The PLE and JEET architectures perform comparably across various languages, with PLE slightly outperforming JEET in English and German, while JEET excels in French and Python. Notably, the MoE-CE setup does not surpass the performance of PLE and JEET unless a common expert is included, at which point its effectiveness closely aligns with the other architectures across languages.

The paper meticulously addresses the issue of catastrophic forgetting, a critical challenge in continual learning domains like multilingual NLP. Sequential training is confirmed to exacerbate forgetting, while both a balanced batching strategy in single-session training and the MoE system effectively mitigate this issue, maintaining stability and preserving previously learned knowledge.

Implications and Future Directions

The findings imply significant potential for the development of modular LLMs capable of efficient multi-language scenario handling. The modular MoE architectures exhibit flexibility conducive to expanding their capabilities by integrating additional specialists without necessitating a comprehensive retraining—thereby conserving computational resources.

The implications of this research are far-reaching, suggesting the maturation of AI models that can adapt to evolving language datasets while maintaining stability and performance. Future research should seek to scale the current approach to larger datasets encompassing more diverse languages, potentially broadening the models' applicability and robustness. Moreover, further refinement of adaptive loss methods and robust evaluation of alternate MoE strategies may yield significant advancements in versatility and efficiency.

This paper’s systematic effort in integrating KD and MoE not only improves the modularity of LLMs but also paves the way for more adaptable and resilient AI systems, consolidating the foundations for future exploration in LLM specialization and knowledge preservation.