DBRX (language model) — DBRX

DBRXは、Databricks社内のMosaic AI研究チームによって開発されたオープンな大規模言語モデル（LLM）です。このモデルは2024年3月27日に正式にリリースされ、企業向けの高性能ソリューションとして位置づけられています。

DBRXは、きめ細かい専門家混合（Mixture of Experts, MoE）アーキテクチャに基づいて構築されており、高いパフォーマンスと効率的なトレーニングおよび推論を両立させています。リリース時点で、DBRXは主要なベンチマークにおいて、LLaMA 2、Mixtral、Grok-1などのモデルを上回り、すべてのオープンモデルの中で最高の結果を示し、GPT-3.5 Turboレベルのクローズドモデルとの競争力も証明しました。

開発経緯

DBRXの登場は、オープンな生成モデルを開発するというDatabricksの戦略の延長線上にあります。2023年6月、Databricksは大規模モデルのトレーニングを専門とするスタートアップMosaicMLを買収し、その基盤の上にMosaic AI部門を設立しました。

ニューラルネットワークの主任アーキテクトであるジョナサン・フランクルが率いるMosaic AIチームは、最高のプロプライエタリシステムに匹敵する品質をオープンなフォーマットで達成することを目標に、新しい大規模LLMの開発に着手しました。このプロジェクトはDBRXと名付けられました。モデルの開発と事前学習には約2.5ヶ月を要し、推定コストは約1000万ドルでした。

アーキテクチャ

DBRXはデコーダーのみ（decoder-only）のトランスフォーマーモデルであり、きめ細かい（fine-grained）専門家混合（MoE）アーキテクチャを実装しています。

アーキテクチャの主な特徴：

• 総パラメータ数: 1320億。
• エキスパート: モデルは16の小規模な専門サブモデル（「エキスパート」）で構成されています。
• 活性化メカニズム: 各入力トークンに対して、16のエキスパートのうち4つだけが活性化されます。これにより、推論時にアクティブなパラメータは360億に抑えられ、高速性と効率性が確保されます。この方式は、Mixtralモデル（8エキスパート、2つ活性化）と比較して65倍多くのエキスパートの組み合わせを可能にします。
• コンポーネント: 回転位置埋め込み（RoPE）、ゲート付き線形ユニット（gated linear units, GLU）、グループ化クエリアテンション（grouped query attention, GQA）などの最新のアーキテクチャソリューションが使用されています。
• コンテキスト長: 32,768トークン。

このようなアーキテクチャにより、モデルは膨大なパラメータ数（知識を保存するため）の利点と、小規模モデルの効率性（推論速度のため）を両立させることができます。

トレーニング

DBRXの事前学習は、テキストとコードで構成される12兆トークン規模の、慎重にキュレーションされたデータセットで行われました。データ品質が最優先事項とされ、開発者はDatabricksのクラウドプラットフォーム（Apache Spark、Databricks Notebooks、Unity Catalog）を使用して、データのクリーニング、準備、監査を行いました。

トレーニングでは、さまざまな段階でデータタイプの比率を変更するカリキュラム学習（curriculum learning）の手法が採用されました。例えば、トレーニングの最終段階では、複雑なタスクを段階的に導入することに重点が置かれ、開発者によると、これにより品質が著しく向上しました。トレーニングは3072基のNvidia H100 GPUからなるクラスターで実施されました。

事前学習後、ベースモデルはさらなるファインチューニング（instruction tuning）を経て、ユーザーの指示を実行するために最適化された対話型バージョンDBRX Instructが作成されました。

パフォーマンス

リリース時点で、DBRXは幅広いベンチマークにおいてオープンLLMの品質の新たな基準を打ち立てました。

オープンモデルとの比較

主要なベンチマークにおけるDBRX Instructの結果

ベンチマーク	タスク	DBRX Instruct	次点の最良モデル (Mixtral/Grok-1)
Hugging Face Open LLM Leaderboard (AVG)	一般知識	74.5%	72.7% (Mixtral Instruct)
HumanEval	プログラミング	70.1%	63.2% (Grok-1)
GSM8K	数学的推論	66.9%	62.9% (Grok-1)
MMLU	一般知識	73.7%	71.5% (Mixtral Instruct)

DBRXは、総合ランキングHugging Face Open LLM Leaderboardと、包括的なテストDatabricks LLM Gauntletの両方で1位を獲得し、先行モデルに対して大きな差をつけました。

クローズドモデルとの比較

DBRX Instructは、MMLU（73.7%対70.0%）やHumanEval（70.1%対48.1%）を含むいくつかの主要な指標でGPT-3.5 Turboを上回っています。一部のベンチマーク（例：MTBench）における回答の質では、Gemini 1.0 Proや初期バージョンのGPT-4のレベルに迫っています。

トレーニングと推論の効率性

• トレーニング効率: MoEアーキテクチャの採用により、同等の品質を持つ密なモデルと比較して、FLOPSでのコストを2～4倍削減することができました。
• 推論効率: 360億のパラメータのみを活性化することで、DBRXは同等サイズの密なモデル（例：LLaMA2-70B）と比較して2～3倍高いスループット（推論速度）を実現します。

ライセンスと利用可能性

DBRXは、特別に策定されたDatabricks Open Model Licenseの下で配布されています。このライセンスは、商用利用を含む自由な使用と改変を許可していますが、いくつかの制限が含まれています。特に、LLaMA 2のライセンスと同様に、DBRXベースのサービスが月間アクティブユーザー7億人を超えるオーディエンスに使用される場合、Databricksからの別途許可が必要となります。

モデルの事前学習済み重み（ベース版とInstruct版）は、Hugging Faceのリポジトリからダウンロードできます。

参考文献

• Mosaic Research Team. (2024). Introducing DBRX: A New State‑of‑the‑Art Open LLM. Databricks Blog.
• Databricks. (2024). Databricks Open Model License (DBRX). Online specification.
• Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
• Lepikhin, D.; et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arXiv:2006.16668.
• Ainslie, J.; et al. (2023). Grouped‑Query Attention: Efficient Training of Generalized Multi‑Query Transformers. arXiv:2305.13245.
• Su, J.; et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
• Dao, T. (2023). FlashAttention‑2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Cai, W.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
• Chen, Y.; et al. (2024). Scaling Laws for Fine‑Grained Mixture of Experts. arXiv:2402.07871.
• Kundu, A.; et al. (2024). Strategic Data Ordering: Enhancing Large Language Model Training via Curriculum Learning. arXiv:2405.07490.