長くて複雑なシーケンスを処理する需要がますます高まるにつれて、大規模言語モデル (LLM) の限界が押し広げられています。従来の Transformer ベースのアーキテクチャは、強力ではあるものの、シーケンス長に関して二次的な複雑さを持つため、重大なスケーリングの問題に直面します。この制限は、拡張されたコンテキスト入力を扱う場合に特に顕著になり、シーケンスの遠い部分からの情報を効果的にキャプチャして利用する能力が妨げられます。この課題に対応して、長いシーケンスの処理において線形複雑性を達成することを目的とした、革新的なアプローチの波が生まれています。
これらの方法には、Linear Attention モデル、State Space Models (Mamba など)、Linear RNNs (DeltaNet など)、および RWKV が含まれます。これらのアーキテクチャはそれぞれ、二次複雑性問題に対する独自の方法を提供し、長いシーケンスのより効率的な処理を可能にします。ただし、これらの線形アーキテクチャは、長文脈情報を完全に理解して活用するのに苦労することがよくあります。
たとえば、RWKV-7 (2.9B パラメータ モデル) は、最大 28K トークンまでのパスキー検索タスクで高い精度を示します。ただし、そのパフォーマンスはこのしきい値を超えると急速に低下します。128K 長のデータを使用した継続的な事前トレーニングを行っても、長文脈の制限は依然として残ります。この問題は RWKV に特有のものではなく、Mamba などの他のアーキテクチャにも当てはまり、このクラスのモデルにとって根本的な課題となっています。拡張されたコンテキストでパフォーマンスを維持するための苦労は、線形複雑性言語モデルの改善のための重要な領域を強調しています。
線形複雑性言語モデルの現状
線形複雑性言語モデルは、Transformer ベースのアーキテクチャの魅力的な代替手段として浮上しており、長いシーケンスの処理に固有の二次的な計算負荷を回避しています。このドメインで際立っている RWKV モデル ファミリは、トレーニング中の Transformer の並列処理能力と RNN のような再帰的状態表現を巧みに組み合わせています。
RWKV の進化は、基礎となる RWKV-4 から始まり、RWKV-5、RWKV-6 と進み、RWKV-7 で最高潮に達するまで、いくつかの反復に及びます。各反復は、モデルの機能を強化し、制限に対処する、改良と改善をもたらしました。さらに、Jamba、Zamba、MiniMax などのハイブリッド言語モデルは、独自のハイブリッド設計を導入することでその足跡を残し、線形複雑性モデルの状況をさらに豊かにしました。
効率的な長文脈処理の追求は、革新的な注意機構の開発にもつながりました。たとえば、Native Sparse Attention は、トークンを時間ブロックに編成し、3 つの異なる注意パスを採用します。グローバル コンテキスト用に圧縮された粗粒度のトークン、ローカルの詳細用に選択的に保持された細粒度のトークン、およびローカル コンテキスト情報をキャプチャするためのスライディング ウィンドウです。その他の注目すべき注意機構には、SeerAttention と Block Attention (MoBA) があり、それぞれが長いシーケンス内の関連情報に注意を払うための独自の方法を提供しています。
RWKV-X: 拡張された長距離コンテキスト モデリングのためのハイブリッド アーキテクチャ
広東省人工知能およびデジタル経済研究所 (SZ)、深セン、河海大学、南京、深セン大学、青海大学、西寧の研究者は、RWKV-X と呼ばれる新しいハイブリッド アーキテクチャを導入しました。このアーキテクチャは、RWKV の短距離依存関係のモデリングにおける効率と、長距離コンテキストをキャプチャするために特別に設計されたスパースな注意機構を巧妙に組み合わせています。
以前のハイブリッド アプローチとは異なり、RWKV-X は、トレーニング中に線形時間複雑性を、推論デコード中に定数時間複雑性を実現します。これにより、長いシーケンスの処理に非常に効率的になります。このモデルは、64K トークン シーケンスで継続的に事前トレーニングされた場合、64K パスキー検索ベンチマークでほぼ完璧な精度を示します。短文脈タスクで強力なパフォーマンスを維持しながら、長文脈ベンチマークで以前の RWKV-7 モデルを一貫して上回ります。
RWKV-X の革新は、長文脈言語モデリングの課題に対処する上で大きな進歩を表しています。再帰モデルとスパースな注意機構の長所を組み合わせることで、RWKV-X は効率と精度のバランスを取り、拡張されたシーケンスのより効果的な処理への道を開きます。
RWKV-X: アーキテクチャとトレーニング
RWKV-X は、RWKV-7 ブロックとスパースな注意ブロックを統合して、両方のアプローチの強みを活用するハイブリッド アーキテクチャを具体化しています。ゼロからトレーニングする代わりに、RWKV-X は、LLaMA Pro に触発されたインターリーブ ブロック拡張アプローチとゼロ初期化メカニズムを使用して、既存のモデルを基に構築します。
トレーニング プロセスは 2 つの段階で構成されており、モデルの短期および長期コンテキストでのパフォーマンスを最適化するように慎重に設計されています。
- 短文脈事前トレーニング: 最初に、モデルは MiniPile データセットから抽出された短い 1024 トークンのコンテキストでトレーニングされます。この段階では、新しく追加されたブロックのパラメータを除くすべてのパラメータが凍結され、ベース RWKV-7 モデルからの事前トレーニングされた知識が確実に保持されます。これにより、新しく追加されたブロックは、事前トレーニングされた表現を中断することなく、既存のアーキテクチャに適応できます。
- 長文脈継続事前トレーニング: 第 2 段階では、ProLong-64K データセットと 64K トークンのコンテキスト長を使用して、長文脈継続事前トレーニングを行います。合計で約 10 億トークンを処理します。このフェーズでは、すべてのパラメータが凍結解除され、共同で最適化されるため、モデルは表現を微調整し、長距離依存関係を学習できます。トレーニングでは、トークンの重要度に基づいて動的にトークンに重み付けする長文脈クロスエントロピー (LongCE) 損失を使用します。この損失関数は、モデルがシーケンスの最も関連性の高い部分に焦点を当て、長距離関係をキャプチャする能力を向上させるのに役立ちます。
2 段階のトレーニング プロセスにより、RWKV-X は、短距離モデリングに対する RWKV-7 の効率と、スパースな注意機構の長距離コンテキスト認識を効果的に組み合わせることができます。最初に短いコンテキストで事前トレーニングし、次に長いコンテキストで微調整することにより、モデルはシーケンスのさまざまな部分からの情報を効果的に統合することを学習します。
RWKV-X: 評価とパフォーマンス
短文脈評価では、RWKV-X が標準ベンチマーク全体で競争力のあるパフォーマンスを維持し、短いシーケンスを効果的に処理する能力を示していることが明らかになりました。小型の RWKV-X (0.22B) は、RWKV-7 の 51.8 と同等の平均スコア 51.0 を達成しています。大規模なスケールでは、RWKV-X (3.6B) は 71.9 に達し、RWKV-7 (2.9B, 72.8) および Qwen2.5-3B (71.4) とほぼ一致し、LLaMA3.2-3B (69.7) を上回っています。これらの結果は、RWKV-X が短いコンテキストでのパフォーマンスを犠牲にすることなく、汎用 LLM バックボーンとして効果的であることを確認しています。
さらに、効率分析では、RWKV-X の長いシーケンスに対する優れたスケーリング特性が示されています。128K トークンでは、RWKV-X は Flash-Attention v3 よりも 1.37 倍の高速化を実現しており、この利点はコンテキスト長が増加するにつれて拡大します。これは、シーケンス長が長くなるにつれて、RWKV-X が他の注意機構と比較してますます効率的になることを示しています。
短文脈と長文脈の両方での RWKV-X の優れたパフォーマンスは、言語モデルとしての汎用性と効率性を強調しています。短いシーケンスで競争力のあるパフォーマンスを維持しながら、長いシーケンスで大幅な高速化を実現する能力は、幅広いアプリケーションにとって有望なアーキテクチャとなっています。
RWKV-X: 制限と今後の方向性
RWKV-X は、短距離依存関係のモデリングに対する RWKV の効率と、長距離コンテキスト モデリング専用に設計された新しいスパースな注意機構を効果的に組み合わせたハイブリッド言語モデルとして登場しました。RWKV-X は、長文脈言語モデリングにおいて強力なパフォーマンスと効率を示していますが、いくつかの制限が残っています。
第一に、上位 k 個のチャンク選択に依存するスパースな注意機構は、意味的に関連性のある依存関係を見落とす可能性のあるヒューリスティックなアプローチを採用しています。上位 k 個の選択戦略は、シーケンス内の最も重要な情報を常にキャプチャできるとは限らず、パフォーマンスが最適化されない可能性があります。
第二に、現在の実装では、スパースな注意デコードがバニラ RWKV よりも遅く実行されていることが示されており、パフォーマンスを最適化するには、さらなるエンジニアリング作業が必要であることを示しています。RWKV-X は、長いシーケンスで他の注意機構と比較して大幅な高速化を実現していますが、そのスパースな注意デコードは、依然としてバニラ RWKV よりも遅く、その実装には改善の余地があることを示唆しています。
今後の研究では、より高度なスパースな注意機構の探索、スパースな注意デコードの実装の最適化、および代替トレーニング戦略の調査を通じて、これらの制限に対処することに焦点を当てる可能性があります。これらの課題を克服することで、RWKV-X は、長文脈アプリケーションにとってさらに強力で効率的な言語モデルになる可能性を秘めています。