[論文レビュー] RecurrentGemma: Moving Past Transformers for Efficient Open Language Models

Aleksandar Botev, Soham De+60 more|arXiv (Cornell University)|Apr 11, 2024

Topic Modeling被引用数 4

ひとこと要約

RecurrentGemma-2B は、グリフィンアーキテクチャに基づくオープンソース言語モデルを紹介する。このアーキテクチャは、グローバルアテンションの代わりに線形再帰と局所アテンションを採用し、固定サイズの状態表現を実現する。これにより、長文シーケンスにおいて顕著に高速な推論が可能となり、メモリ使用量も低減される。Gemma-2B と同等のパフォーマンスを達成しているが、学習に使用したトークン数は33%少ない。

ABSTRACT

We introduce RecurrentGemma, a family of open language models which uses Google's novel Griffin architecture. Griffin combines linear recurrences with local attention to achieve excellent performance on language. It has a fixed-sized state, which reduces memory use and enables efficient inference on long sequences. We provide two sizes of models, containing 2B and 9B parameters, and provide pre-trained and instruction tuned variants for both. Our models achieve comparable performance to similarly-sized Gemma baselines despite being trained on fewer tokens.

研究の動機と目的

• 長文シーケンスにおいて推論速度とメモリ効率の両面でトランスフォーマーを凌駕する、高効率なオープンソース言語モデルの開発。
• 線形再帰に基づくモデルが、Gemma-2B などの最先端トランスフォーマーと同等のパフォーマンスを達成できることの実証。
• KVキャッシュの線形増加を固定サイズの状態に置き換えることで、長文生成を可能にする。
• 研究およびリソース制約のある環境での利用を想定し、事前学習済みおよびインstructチューニング済みバージョンの両方をリリースする。
• Gemma の責任あるAI実践に準拠し、厳密なベンチマークと人間評価を通じて、モデルの安全性と整合性を検証する。

提案手法

• グリフィンアーキテクチャを採用し、線形再帰（RG-LRU）と局所アテンション（ウィンドウサイズ2048）を組み合わせることで、グローバルアテンションを排除したシーケンスモデリングを実現する。
• 入力シーケンスを圧縮する固定サイズの状態ベクトルを用いることで、自己回帰的生成における成長するKVキャッシュの必要性を排除する。
• 学習の安定化を図るため、Gemma の設計に準拠し、入力埋め込みに可学習スケーリング係数（sqrt(モデル幅)）を適用する。
• TPU上で効率的な推論を実現するための特別なパラスカーネルを採用し、同時に参照用のPyTorch実装も提供する。
• Gemma-2B と同じデータセットから2Tトークンを用いて学習し、二段階の事前学習プロセス（一般混合データから高品質データへ）を実施する。
• インストラクションチューニングと独自のRLHFアルゴリズムを用いて、制御トークンを含む定義済み対話フォーマットに準拠したインストラクション対応の微調整を実施する。

実験結果

リサーチクエスチョン

• RQ1 線形再帰と局所アテンションを組み合わせた再帰的アーキテクチャが、標準的なNLPベンチマークでGemma-2B などのトランスフォーマーと同等のパフォーマンスを達成できるか？
• RQ2 KVキャッシュを固定サイズの状態に置き換えることで、従来のトランスフォーマーと比較して、長文シーケンスにおける推論が著しく高速かつメモリ効率的になるか？
• RQ3 2Tトークン（Gemma-2B の3Tに比べ）で学習したモデルが、大規模な事前学習スケールを維持しながら、同等のパフォーマンスを達成できるか、その程度はいかほどか？
• RQ4 人間評価において、Mistral 7B v0.2 Instruct と比較して、RecurrentGemma-2B-IT はどの程度優れた整合性と使いやすさを示すか？
• RQ5 グリフィンベースのアーキテクチャは、メモリ制約のない任意長の生成をサポートできるか。また、スループットへのスケーリング特性はいかがなっているか？

主な発見

• 学術ベンチマークスイートにおいて、RecurrentGemma-2B は平均44.6%のスコアを達成した。Gemma-2B の45.0%とほぼ同等であり、33%少ないトークンで学習したにもかかわらず、競争力のあるパフォーマンスを示した。
• 人間評価において、RecurrentGemma-2B-IT は1,000件のインストラクション対応プロンプトに対して、Mistral 7B v0.2 Instruct と比較して43.7%の勝率を記録し、優れた整合性と使いやすさを示した。
• TPUv5eデバイス上での推論スループットは、全シーケンス長にわたり、RecurrentGemma が Gemma を常に上回り、シーケンス長の増加に伴う性能劣化が見られなかった。
• 自己回帰的サンプリング中、RecurrentGemma は6kトークン/秒の高いスループットを維持しているが、Gemma は成長するKVキャッシュの影響により、特に長文シーケンスでスループットが著しく低下した。
• プロンプト処理速度は両モデルでほぼ同等（約40kトークン/秒）であり、性能優位性が自己回帰的生成に特化していることを確認した。
• 固定サイズの状態により、モデルは計算力とコンテキストウィンドウの制限を除き、任意長の生成が可能である。これに対して、トランスフォーマーはメモリ制限によるKVキャッシュの成長によって制限を受ける。