DBRX (language model)

From Systems Analysis Wiki
Jump to navigation Jump to search

DBRX is an open large language model (LLM) developed by Databricks (through its Mosaic AI research team) and released on March 27, 2024. The model is built on a fine-grained Mixture-of-Experts (MoE) architecture and, at release, set a new standard for quality and efficiency among openly available models, outperforming established open models such as LLaMA 2, Mixtral, and Grok-1 on standard benchmarks and matching or exceeding the closed GPT-3.5 on many of them[1].

History of Creation

The development of DBRX is directly linked to Databricks' acquisition of the startup MosaicML in June 2023 in a deal valued at roughly $1.3 billion; MosaicML's team and technology became the foundation of the company's Mosaic AI research division[2]. The project was led by the Mosaic team, including Databricks' Chief Scientist (Neural Networks), Jonathan Frankle. According to the developers, the final training of the model took roughly two to three months and cost about $10 million[3][4].

The release of DBRX was aimed at demonstrating that a company could build a world-class MoE model from scratch using its own tools and, in doing so, provide the community and enterprises with a powerful open alternative to proprietary APIs[1].

Architecture and Technical Specifications

DBRX is a transformer-based, decoder-only model trained using next-token prediction. Its principal feature is the fine-grained MoE architecture[1]:

  • Parameters: The model has 132 billion total parameters, of which only 36 billion are active for any given input, which significantly increases inference speed.
  • Experts: Unlike other open MoE models such as Mixtral and Grok-1, which use 8 experts and select 2, DBRX uses 16 experts and selects 4 for each input. This fine-grained approach provides roughly 65 times more possible combinations of experts, which improves model quality.
  • Additional techniques: The architecture uses Rotary Position Encodings (RoPE), Gated Linear Units (GLU), and Grouped Query Attention (GQA); the model also uses the GPT-4 tokenizer (tiktoken).
  • Context window: The model supports a context length of up to 32,768 tokens.

Training

DBRX was pre-trained on a carefully curated dataset of 12 trillion tokens of text and code. The developers note that this data was at least twice as effective, token-for-token, as the data used to train their earlier MPT models[1]. Training was performed on a cluster of 3,072 NVIDIA H100 GPUs connected via 3.2 Tbps InfiniBand[4]. During training, curriculum learning was applied, changing the composition of the data mix during the process to improve model quality.

Model Versions

Two main versions of the model were released:

  • DBRX Base: The base pre-trained model, intended for further fine-tuning.
  • DBRX Instruct: A version fine-tuned to follow instructions and for use in dialogue and question-answering tasks[1].

Performance and Benchmarks

DBRX Instruct demonstrated leading results among open models and strong competitiveness with closed models at the time of its release.

Comparison of DBRX Instruct with open models
Benchmark DBRX Instruct Next Best (Mixtral / Grok-1)
MMLU (general knowledge) 73.7% 71.5% (Mixtral Instruct)
HellaSwag (commonsense reasoning) 89.0% 87.6% (Mixtral Instruct)
HumanEval (programming) 70.1% 63.2% (Grok-1)
GSM8k (mathematics) 66.9% 62.9% (Grok-1)
Open LLM Leaderboard (avg.) 74.5% 72.7% (Mixtral Instruct)

In comparison with closed models, DBRX Instruct surpasses GPT-3.5 on most benchmarks — for example, on general knowledge as measured by MMLU (73.7% vs. 70.0%) and on programming as measured by HumanEval (70.1% vs. 48.1%)[1]. In terms of overall response quality, the model is competitive with Gemini 1.0 Pro and Mistral Medium — scoring higher than Gemini 1.0 Pro on the Inflection-corrected MTBench, MMLU, HellaSwag, and HumanEval, while Gemini 1.0 Pro is stronger on GSM8k. Databricks did not claim that DBRX reaches the level of frontier models such as GPT-4 or Claude 3 Opus, which it trails; the company noted only that, on a narrow internal SQL use case, early versions of DBRX were competitive with GPT-4 Turbo[1][5].

License and Availability

DBRX is distributed under the Databricks Open Model License, which permits both commercial and research use. However, the license contains a restriction similar to Meta's Llama 2 license: companies whose products had more than 700 million monthly active users (MAU) at the time of the model's release must obtain a separate license from Databricks[5].

The model is available for download on Hugging Face (DBRX Base and DBRX Instruct) and on GitHub. It can also be accessed through the Databricks platform and major cloud providers, including AWS, Google Cloud, and Microsoft Azure (via Azure Databricks)[1].

Reception and Legacy

DBRX was well received as a demonstration that an enterprise could train a competitive MoE model from scratch, and as a validation of Databricks' training stack (built on open-source projects such as MegaBlocks, Composer, and LLM Foundry) that customers could use to build their own models[1]. At the same time, the rapidly moving field soon narrowed its advantage: Meta released Llama 3 several weeks later, in April 2024, and other open models followed. Databricks did not release a successor frontier base model, instead focusing its research on efficient model training, fine-tuning, and, more recently, reinforcement learning and agent systems, while continuing to make DBRX available to the community[6].

Literature

  • Vaswani, A. et al. (2017). Attention Is All You Need. arXiv:1706.03762.
  • Shazeer, N. et al. (2017). Outrageously Large Neural Networks: The Sparsely‑Gated Mixture‑of‑Experts Layer. arXiv:1701.06538.
  • Fedus, W.; Zoph, B.; Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
  • Gale, T. et al. (2022). MegaBlocks: Efficient Sparse Training with Mixture‑of‑Experts. arXiv:2211.15841.
  • Su, J. et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
  • Shazeer, N. (2020). GLU Variants Improve Transformer. arXiv:2002.05202.
  • Ainslie, J. et al. (2023). GQA: Training Generalized Multi‑Query Transformer Models from Multi‑Head Checkpoints. arXiv:2305.13245.
  • Jiang, A. et al. (2024). Mixtral of Experts. arXiv:2401.04088.

References

  1. 1.0 1.1 1.2 1.3 1.4 1.5 1.6 1.7 1.8 "Introducing DBRX: A New State-of-the-Art Open LLM". Databricks Blog. [1]
  2. "MosaicML Acquisition Generative AI Platform Developed by Kempner's Jonathan Frankle Acquired by Databricks in $1.3 Billion Deal". Kempner Institute, Harvard University. [2]
  3. Wiggers, Kyle (March 27, 2024). "Databricks spent $10M on new DBRX generative AI model, but it can't beat GPT-4". TechCrunch. [3]
  4. 4.0 4.1 "DBRX: A new open LLM by Databricks". SuperAnnotate. [4]
  5. 5.0 5.1 "Databricks Claims a Performance Lead with a New Open LLM". Synthedia. [5]
  6. "Jonathan Frankle". Personal website. [6]