NVIDIAがHymba 1.5Bを発表: 効率的NLPモデルへのハイブリッドアプローチ

原文リンク(2025-01-03)

NVIDIAの研究者は、トランスフォーマーと状態空間モデル（SSM）アーキテクチャを組み合わせて前例のない効率とパフォーマンスを実現するオープンソース言語モデル、Hymba 1.5Bを発表した。NVIDIAの最適化されたトレーニングパイプラインで設計されたHymbaは、従来のトランスフォーマーの計算およびメモリの制約に対処しつつ、SSMの記憶能力を向上させている。

従来のトランスフォーマーベースの言語モデルは長期記憶と並列化に優れているが、二次計算量の複雑さと大きなメモリ要求という大きな課題に直面しています。一方、MambaやMamba-2のようなSSMは一定の計算量の複雑さとハードウェアの最適化を提供するが、記憶想起タスクでは性能が劣る。Hymbaは両方のアーキテクチャの強みを組み合わせることでこれらのトレードオフを解決している。

Hymbaのハイブリッド・ヘッドモジュールは高解像度記憶想起のためのアテンションヘッドと効率的なコンテキスト要約のためのSSMヘッドを融合し、両方のコンポーネントが順次ではなく並列で動作できるようにしている。この設計によりパフォーマンスを犠牲にすることなく、計算能力とメモリ要求を削減している。

Hymba 1.5Bアーキテクチャはいくつかの革新的なメカニズムを導入することで、精度を維持しながら効率を向上させることにフォーカスしている：

• アテンション・オーバーヘッドの削減：アテンション計算の50％以上をSSM処理に置き換え、タスク精度を維持しながらコストを削減している。
• ローカルアテンションの優位性：ローカルアテンションとSSMの組み合わせでグローバル情報を十分要約できるため、グローバルアテンションが最小限に抑えられている。
• KVキャッシュ最適化：Hymbaは、層間キャッシュの冗長性を削減し、同等のトランスフォーマーモデルと比較して最大10倍のメモリ使用量削減を実現する層間KVキャッシュ共有を導入している。
• メタ・トークン：プロンプトの前に付加される128個の学習可能な埋め込みが、メモリの初期化子として機能する。

出典NVIDIAブログ

学習可能なメタ・トークンの役割が議論されている。Superlinked社の機械学習スペシャリスト Daniel Svonava氏が質問を投げかけた：

学習可能なメタ・トークンが従来の方法と比較してどのようにアテンションメカニズムの焦点を改善するのか説明できますか?

データサイエンティスト Marek Barák氏はこのように説明した：

アテンションには、文中の最初のトークンに過度にフォーカスするという問題があります。最初のトークンは実際には多くの情報を持っていないため、これはほとんど意味的な理由がありません。メタ・トークンを使用すると、トークン全体に対してよりバランスの取れたソフトマックス分布が得られます。

Hymba 1.5Bは、Llama 3.2 1B、OpenELM 1B、Qwen 2.5 1.5Bなど、20億パラメータ未満の主要モデルとの直接比較においてトップパフォーマーであることを証明した。MMLU、ARC-C、Hellaswag、SQuAD-Cなどのベンチマークにおいて、Hymbaは競合を上回る性能を示した。

出典：https://arxiv.org/pdf/2411.13676

NVIDIAはタスクのパフォーマンスと効率のバランスをとるために、Hymbaのトレーニングパイプラインを最適化している。事前学習戦略は2段階のプロセスを含んでいる：初期段階では多様でフィルタリングされていないデータセットで学習し、その後高品質データでファインチューニングを行った。教師ありファインチューニング（SFT）や直接選好最適化（DPO）を用いた強化学習などの段階を経て、指示ファインチューニングはモデルの能力を向上させている。

Hymba 1.5BはオープンソースリリースとしてHugging Face とGitHub で公開されており、研究者や開発者が実世界のアプリケーションでその能力をテストすることができるようになっている。