JetMoE: Reaching Llama2 Performance with 0.1M Dollars
JetMoE-8B achieves Llama2 performance with under $0.1M using a sparsely-gated mixture of experts architecture.
Key Findings
Methodology
JetMoE-8B employs a Sparsely-gated Mixture-of-Experts (SMoE) architecture, applying sparse activation to both attention and feedforward layers. Trained on 1.25T open-source data using 30,000 H100 GPU hours, it activates only 2B of its 8B parameters per input, reducing inference computation by 70%.
Key Results
- JetMoE-8B scored 80.5 on Hellaswag, surpassing Llama2-7B's 78.6, demonstrating significant performance gains.
- On the MMLU test, JetMoE-8B achieved 49.2, outperforming Llama2's 46.9, showcasing its multi-task learning advantage.
- JetMoE-8B scored 34.2 on MBPP (Pass@1), indicating strong code generation capabilities.
Significance
JetMoE-8B demonstrates the feasibility of training high-performance LLMs on a limited budget, advancing both academic and industrial efforts in cost-effective model development. Its open-source datasets and code enhance research transparency and collaboration.
Technical Contribution
JetMoE-8B introduces sparse activation in both attention and feedforward layers, significantly reducing computational costs. Compared to dense models, it offers new engineering possibilities, especially in resource-constrained environments.
Novelty
JetMoE-8B is the first to apply sparse activation in both attention and feedforward layers, significantly reducing computational demands compared to traditional MoE models. Its open-source and cost-effective nature is innovative in the LLM field.
Limitations
- Due to budget constraints, no ablation studies on the model architecture were conducted, which may affect understanding of component contributions.
- The data mixture strategy is empirically chosen and may not be optimal.
- The model is relatively weak in coding and extraction tasks.
Future Work
Future work could include more in-depth ablation studies on the model architecture, optimizing data mixture strategies, and exploring performance improvements across various tasks.
AI Executive Summary
JetMoE-8B is a novel large language model designed to achieve high performance at a cost of under $0.1 million. Existing large language models typically require extensive computational resources, limiting their widespread application in academia and industry. JetMoE-8B significantly reduces computational costs while maintaining excellent performance through its sparsely-gated mixture of experts architecture.
The core technical innovation of JetMoE-8B lies in its sparse activation mechanism, applied to both attention and feedforward layers, allowing the model to activate only a portion of its parameters during inference, thus reducing computational load by 70%. Trained on 1.25T open-source datasets, JetMoE-8B surpasses existing models like Llama2 in multiple benchmark tests.
Despite JetMoE-8B's significant advancements in performance and cost efficiency, there is room for improvement in coding and extraction tasks. Future research can further optimize data mixture strategies and explore the model's potential across more tasks. JetMoE-8B's open-source nature provides a solid foundation for researchers, fostering collaboration and innovation in cost-effective, efficient model development.
Deep Analysis
Background
Large language models (LLMs) have made significant strides in recent years, particularly in natural language processing tasks. However, these models typically require substantial computational resources and data, limiting their application in academia and industry. The Mixture-of-Experts (MoE) architecture, which reduces computational costs through sparse activation, has emerged as a popular solution.
Core Problem
Existing LLMs demand excessive computational resources, making them costly for widespread application. The core challenge is to maintain high performance while reducing computational costs, a critical issue in current research.
Innovation
JetMoE-8B's core innovation is its sparsely-gated mixture of experts architecture, which applies sparse activation to both attention and feedforward layers for the first time. This innovation significantly reduces computational costs, making high-performance LLM training feasible on a limited budget.
Methodology
- �� Utilizes a sparsely-gated mixture of experts architecture to reduce computation.
- �� Applies sparse activation in both attention and feedforward layers.
- �� Trained using open-source datasets to ensure model openness and reproducibility.
- �� Employs a two-phase training strategy to optimize data mixture.
Experiments
Experiments used 1.25T open-source datasets, trained with 30,000 hours of H100 GPU. Benchmarks include Hellaswag, MMLU, and MBPP, evaluating the model's performance in multi-task learning and code generation tasks.
Results
JetMoE-8B outperformed Llama2 in benchmarks like Hellaswag and MMLU, demonstrating its multi-task learning advantage. In the MBPP test, JetMoE-8B also showed strong code generation capabilities.
Applications
JetMoE-8B is suitable for resource-constrained environments, such as academic research and small to medium enterprises. Its cost-effective and efficient nature makes it widely applicable in natural language processing and code generation tasks.
Limitations & Outlook
Due to budget constraints, no ablation studies were conducted, which may affect understanding of component contributions. The data mixture strategy is empirically chosen and may not be optimal. The model is relatively weak in coding and extraction tasks.
Plain Language Accessible to non-experts
Imagine you're in a factory with many machines, each with different functions. Traditionally, all machines are turned on at once during production, consuming a lot of power and resources. JetMoE-8B is like a smart factory that only activates certain machines as needed, saving resources. This approach not only reduces costs but also increases efficiency, allowing the factory to produce high-quality products even on a limited budget.
ELI14 Explained like you're 14
Imagine you're playing a game with lots of characters, each with different skills. Usually, you'd use all the characters' skills at once, but that drains your energy. JetMoE-8B is like a smart player who knows when to use which character's skills, saving energy while scoring high! That's the cleverness of JetMoE-8B, doing the best with limited resources!
Glossary
Sparsely-gated Mixture-of-Experts (SMoE)
An architecture that reduces computational costs by selectively activating parts of the model.
Used in JetMoE-8B to optimize computation in attention and feedforward layers.
Attention Mechanism
A technique for focusing on important information when processing sequence data.
JetMoE-8B optimizes the attention mechanism through sparse activation.
Feedforward Layer
A layer in neural networks for processing input data and generating output.
In JetMoE-8B, computation is reduced by sparse activation in feedforward layers.
Open-source Dataset
Publicly available datasets that anyone can access and use.
JetMoE-8B is trained on open-source datasets to ensure model openness.
Benchmark Test
Standard test sets used to evaluate model performance.
JetMoE-8B surpasses existing models in multiple benchmark tests.
Open Questions Unanswered questions from this research
- 1 How to further improve model performance without increasing computational costs?
- 2 How to optimize data mixture strategies to enhance model generalization?
- 3 What is the potential of sparse activation in other tasks?
Applications
Immediate Applications
Academic Research
JetMoE-8B's cost-effective nature makes it suitable for academic research, especially in resource-limited settings.
Long-term Vision
Industrial Applications
JetMoE-8B's efficient computation can be widely applied in industry, particularly in natural language processing and code generation tasks.
Abstract
Large Language Models (LLMs) have achieved remarkable results, but their increasing resource demand has become a major obstacle to the development of powerful and accessible super-human intelligence. This report introduces JetMoE-8B, a new LLM trained with less than $0.1 million, using 1.25T tokens from carefully mixed open-source corpora and 30,000 H100 GPU hours. Despite its low cost, the JetMoE-8B demonstrates impressive performance, with JetMoE-8B outperforming the Llama2-7B model and JetMoE-8B-Chat surpassing the Llama2-13B-Chat model. These results suggest that LLM training can be much more cost-effective than generally thought. JetMoE-8B is based on an efficient Sparsely-gated Mixture-of-Experts (SMoE) architecture, composed of attention and feedforward experts. Both layers are sparsely activated, allowing JetMoE-8B to have 8B parameters while only activating 2B for each input token, reducing inference computation by about 70% compared to Llama2-7B. Moreover, JetMoE-8B is highly open and academia-friendly, using only public datasets and training code. All training parameters and data mixtures have been detailed in this report to facilitate future efforts in the development of open foundation models. This transparency aims to encourage collaboration and further advancements in the field of accessible and efficient LLMs. The model weights are publicly available at https://github.com/myshell-ai/JetMoE.