[論文レビュー] Falcon Mamba: The First Competitive Attention-free 7B Language Model

Jingwei Zuo, Maksim Velikanov+5 more|arXiv (Cornell University)|Oct 7, 2024

DNA and Biological ComputingBiochemistry, Genetics and Molecular Biology被引用数 3

ひとこと要約

Falcon Mamba 7B は、5.8兆トークンで学習された、Mambaアーキテクチャに基づく最初の競争力のある完全なアテンションフリーな7Bパラメータ言語モデルである。これは、線形な推論コストと長文における低メモリ使用量を維持しながら、Mistral 7B や Llama3.1 8B、Falcon2 11B といった先端的なオープンソースモデルを上回る性能を発揮し、純粋なMambaモデルがスケールアップにおいてTransformerベースの最先端モデルと同等またはそれを上回ることを実証した。

ABSTRACT

In this technical report, we present Falcon Mamba 7B, a new base large language model based on the novel Mamba architecture. Falcon Mamba 7B is trained on 5.8 trillion tokens with carefully selected data mixtures. As a pure Mamba-based model, Falcon Mamba 7B surpasses leading open-weight models based on Transformers, such as Mistral 7B, Llama3.1 8B, and Falcon2 11B. It is on par with Gemma 7B and outperforms models with different architecture designs, such as RecurrentGemma 9B and RWKV-v6 Finch 7B/14B. Currently, Falcon Mamba 7B is the best-performing Mamba model in the literature at this scale, surpassing both existing Mamba and hybrid Mamba-Transformer models, according to the Open LLM Leaderboard. Due to its architecture, Falcon Mamba 7B is significantly faster at inference and requires substantially less memory for long sequence generation. Despite recent studies suggesting that hybrid Mamba-Transformer models outperform pure architecture designs, we demonstrate that even the pure Mamba design can achieve similar, or even superior results compared to the Transformer and hybrid designs. We make the weights of our implementation of Falcon Mamba 7B publicly available on https://huggingface.co/tiiuae/falcon-mamba-7b, under a permissive license.

研究の動機と目的

• 純粋なMambaベースの言語モデルが、スケールアップにおいて最良のTransformerモデルと同等またはそれを上回る性能を達成できることを示すこと。
• アテンションフリーなアーキテクチャが、一般の言語理解および生成においてハイブリッドMamba-Transformerモデルと同等またはそれを上回ることを検証すること。
• 定常メモリと高い推論効率を備えた純粋なMambaモデルが、長文生成において実現可能で利点があることを検証すること。
• 効率的でスケーラブルなシーケンスモデリングの分野における研究を前進させるために、完全にオープンソースで、許可が広いライセンスに基づく7Bモデルを提供すること。

提案手法

• 64層、4096の隠れ層次元、16次元の状態表現を持つ状態空間モデル（SSM）を備えた、完全なMambaアーキテクチャを採用。
• 7Bスケールでの柔軟性とパフォーマンス向上を図るため、入力と出力の埋め込みを非結合（untied）に設定。
• アテンション機構を一切使用せず、慎重に選別されたデータミックスを用いて5.8兆トークンで学習。
• 生成中の干渉を防ぐために、因果的畳み込みの前後でパディングトークンに対して隠れ状態をゼロにリセットし、左側パディングを適用。
• 訓練の安定性を高めるために、慎重な重み初期化と学習率への感受性の低減を実施し、Mamba訓練で一般的に見られる損失の急上昇を是正。
• Hugging Faceおよびllama.cpp経由でのデプロイを実施し、GGUF量子化とCPU推論をサポートすることで、広範なアクセシビリティを実現。

Figure 1: Data mixtures across training stages

実験結果

リサーチクエスチョン

• RQ1 純粋なMambaベースの言語モデルは、7Bパラメータで、最先端のオープンソースTransformerモデルと同等の性能を達成できるか？
• RQ2 完全にアテンションフリーなアーキテクチャが、ハイブリッドアテンションコンponentsを含まない状態でも、長文長さにわたって線形な推論コストとメモリ効率を維持できるか？
• RQ3 純粋なMambaモデルの性能は、ハイブリッドMamba-TransformerモデルやRWKVやGriffinといった他のSSMベースのアーキテクチャと比べてどうか？
• RQ4 MambaモデルがTransformerと比較して、インライン学習の限界がある場合に、データ品質とスケーリングがどれほどそれを補うことができるか？
• RQ5 純粋なMambaモデルで安定した訓練と高いパフォーマンスを実現するための、主な訓練およびアーキテクチャ的実践は何か？

主な発見

• Falcon Mamba 7B は、Open LLM LeaderboardでMistral 7B、Llama3.1 8B、Falcon2 11B を上回り、アテンション機構を一切使用しないにもかかわらず、競争力のある性能を示した。
• RecurrentGemma 9B や RWKV-v6 Finch 7B/14B と同等またはそれ以上の性能を発揮し、多様なアーキテクチャ設計において強力な結果を確認した。
• シーケンス長に関係なく、定常的なメモリ使用量とスループットを維持しており、効率的な長文生成を可能にした。
• 中程度のコンテキスト長（8kトークン）で学習されたにもかかわらず、強力な一般化能力と推論能力を示しており、長文最適化の可能性を示唆した。
• 特に長文に対して、同等のTransformerモデルと比較して、推論がはるかに速く、メモリ使用量も顕著に少なかった。
• 責任ある使用ポリシーを含む、許可が広いApache 2.0ベースのライセンスで公開されており、Hugging Faceおよびllama.cpp経由で広範な研究および展開へのアクセスを可能にした。

Figure 2: We vary the context length of the prompt to determine the maximum sequence length that could be processed without encountering an out-of-memory (OOM) error. To ensure a fair comparison, all models were configured with a rescaled vocabulary size.