Falcon Mamba 7B リリースノート

Falcon Mamba 7B は、Technology Innovation Institute (TII) が提供する新しいオープンアクセスモデルで、State Space Language Models (SSLMs) が、アテンション機構に伴う計算とメモリコストの線形増加なしに、最先端 (SoTA) トランスフォーマーに匹敵する性能を実現できることを示しています。

純粋な Mamba アーキテクチャとシーケンススケーリング

Falcon Mamba は、元の Mamba アーキテクチャに基づき、選択的な状態空間を利用してトランスフォーマーのシーケンススケーリングの制限を克服しています。大規模での安定した学習を確保するため、TII はベースの Mamba 設計に追加の RMS 正規化層を加えました。

このアーキテクチャは、アテンションベースのモデルに対して 2 つの主要な技術的利点を提供します:

  • 任意のシーケンス長: モデルはメモリストレージの増加に伴わず任意の長さのシーケンスを処理でき、単一の A10 24GB GPU に収まります。
  • 一定の生成時間: 新しいトークンを生成するのに必要な時間は、コンテキストのサイズに関係なく一定です。

モデルのトレーニングとデータ

Falcon Mamba は約 5,500GT のデータでトレーニングされました。データセットは主に RefinedWeb データで構成され、公共のソースからの高品質な技術データとコードが補完されています。トレーニングプロセスは、ほとんどの期間で一定の学習率を使用し、その後、最終的な性能向上のために高品質なキュレーションデータの一部を組み込んだ短い学習率減衰ステージを行いました。

パフォーマンスベンチマーク

lm-evaluation-harness(Hugging Face スコア正規化付き)および lighteval による評価により、Falcon Mamba-7B が現在までで最も強力な純粋な SSM モデルであり、ハイブリッド SSM-アテンションモデルや複数のトランスフォーマーベースのモデルと競合またはそれを上回る性能を示すことが分かりました。

新しいリーダーボードバージョンの結果

モデル名 IFEval BBH MATH LvL5 GPQA MUSR MMLU-PRO 平均
Falcon Mamba-7B 33.36 19.88 3.63 8.05 10.86 14.47 15.04
recurrentgemma-9b 30.76 14.80 4.83 4.70 6.60 17.88 13.20
Falcon2-11B 32.61 21.94 2.34 2.80 7.53 15.44 13.78
Meta-Llama-3.1-8B 12.70 25.29 4.61 6.15 8.98 24.95 13.78
gemma-7B 26.59 21.12 6.42 4.92 10.98 21.64 15.28

LLM リーダーボード(第一版)結果

モデル名 ARC HellaSwag MMLU Winogrande TruthfulQA GSM8K 平均
Falcon Mamba-7B 62.03 80.82 62.11 73.64 53.42 52.54 64.09
recurrentgemma-9b 52.00 80.40 60.50 73.60 38.60 42.60 57.95
Falcon2-11B 59.73 82.91 58.37 78.30 52.56 53.83 64.28
Meta-Llama-3-8B 60.24 82.23 66.70 78.45 42.93 45.19 62.62

メモリとスループットの分析

Falcon Mamba は、シーケンス処理のプレフィル段階とデコード段階の両方で、トランスフォーマーモデルに比べて大幅な効率向上を示します。

プレフィル効率

プロンプト(プレフィル)を処理する主な方法は 2 つあります:

  1. 並列プレフィル: GPU の利用率を最大化するためにプロンプト全体を並列に処理します。このモードでは、隠れ状態を保存する必要があるため、メモリ要件はプロンプト長に比例して増加します。それでも、Falcon Mamba はトランスフォーマーよりも長いシーケンスを処理できます。
  2. シーケンシャルプレフィル: プロンプトをトークン単位またはチャンク単位で処理します。この方法により、SSM はトランスフォーマーで見られるメモリスケーリングの問題なく、任意の長さのプロンプトを処理できます。

生成スループット

プロンプト長 1、最大 130k トークン生成の H100 GPU を使用したテストでは、Falcon Mamba は一定のスループットと安定した CUDA ピークメモリを維持しました。対照的に、トランスフォーマーモデルはシーケンス長が増加するにつれてピークメモリ使用量が増大し、生成速度が低下しました。

実装と利用可能性

Falcon Mamba は Hugging Face の transformers ライブラリ(バージョン >4.45.0)に統合されており、AutoModelForCausalLMpipeline API と互換性があります。

モデルバリエーション

  • ベースモデル: 標準の事前学習モデル。
  • インストラクトモデル: 指示タスクの性能向上のために、50 億トークンの教師ありファインチューニング(SFT)データで微調整されたバージョン。
  • 量子化バージョン: ベースモデルとインストラクトモデルの両方の 4 ビット変換版が、bitsandbytes 対応 GPU を持つユーザー向けに提供されています。

最適化

ユーザーはモデルをロードした後に torch.compile(model) を適用することで、推論を高速化できます。

Sources