XFT: Unlocking the Power of Code Instruction Tuning by Simply Merging Upcycled Mixture-of-Experts

Overview

• Introduced a simple yet powerful training scheme called XFT
• XFT merges upcycled Mixture-of-Experts (MoE) to improve the performance of instruction-tuned code Large Language Models (LLMs)
• Vanilla sparse upcycling fails to improve instruction tuning, but XFT introduces a shared expert mechanism and a novel routing weight normalization strategy to significantly boost instruction tuning
• XFT also includes a learnable model merging mechanism to compile the upcycled MoE model back to a dense model, achieving upcycled MoE-level performance with only dense-model compute

Plain English Explanation

XFT is a new training method that aims to enhance the performance of large language models (LLMs) when they are used for code-related tasks, such as generating or understanding code. The key idea behind XFT is to combine the power of Mixture-of-Experts (MoE) models with a process called "upcycling."

Upcycling is a technique that takes a pre-trained model and transforms it into a more specialized and powerful model. However, the researchers found that simply applying upcycling to LLMs did not improve their performance on code-related tasks. To address this, XFT introduces a new mechanism where the experts in the MoE model share information and a novel way of normalizing the weights used to route inputs to the experts.

After this specialized training, XFT then uses a learnable merging process to combine the upcycled MoE model back into a dense model. This allows the model to maintain the performance benefits of the upcycled MoE, but with the efficiency of a dense model. The researchers were able to create a new state-of-the-art "tiny" code LLM (less than 3 billion parameters) using this approach, as well as improve the performance of existing models on a variety of code-related benchmarks.

Technical Explanation

The key innovation in XFT is the way it combines upcycling and MoE models to improve the performance of instruction-tuned code LLMs. Vanilla sparse upcycling, where a pre-trained model is fine-tuned on a specialized task, failed to significantly improve the models' performance on code-related tasks.

To address this, XFT introduces a "shared expert" mechanism, where the experts in the MoE model are allowed to share information with each other. This is combined with a novel "routing weight normalization" strategy, which helps the model learn better how to route inputs to the appropriate experts.

After this upcycling process, XFT then uses a learnable model merging mechanism to compile the upcycled MoE model back into a dense model. This allows the model to maintain the performance benefits of the upcycled MoE, but with the efficiency of a dense model.

The researchers applied XFT to a 1.3 billion parameter model, creating a new state-of-the-art "tiny" code LLM (less than 3 billion parameters) with strong performance on code-related benchmarks like HumanEval and HumanEval+. They also showed that XFT can consistently improve the performance of supervised fine-tuning on other code-related datasets, such as MBPP+, MultiPL-E, and DS-1000.

Critical Analysis

The researchers note that XFT is fully orthogonal to existing techniques like Evol-Instruct and OSS-Instruct, meaning it can be combined with these approaches to further improve instruction-tuned code LLMs.

However, the paper does not provide a detailed analysis of the computational cost and training time required for the XFT approach, which could be an important consideration for real-world deployment. Additionally, the researchers only tested XFT on a relatively small model (1.3 billion parameters), and it's unclear how well the approach would scale to larger, more powerful LLMs.

Further research could also explore the generalizability of XFT beyond code-related tasks, as well as investigate potential biases or limitations introduced by the shared expert mechanism and routing weight normalization strategy.

Conclusion

XFT introduces a novel training scheme that combines upcycling and Mixture-of-Experts models to significantly improve the performance of instruction-tuned code LLMs. By introducing a shared expert mechanism and a learnable model merging process, XFT is able to create a new state-of-the-art "tiny" code LLM and consistently boost the performance of existing models on a variety of code-related benchmarks.

This work opens up new possibilities for enhancing the capabilities of large language models in specialized domains, and the researchers suggest that XFT can be combined with other techniques to further advance the field of instruction-tuned LLMs. While the paper raises some questions about the practical considerations of the approach, the core ideas behind XFT demonstrate the potential for innovative training schemes to unlock the full potential of these powerful language models.