M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

本論文は、Transformer の計算複雑性の限界を克服し、行列値状態を持つ非線形 RNN である M²RNN を提案することで、大規模言語モデルにおいて効率的な長期依存関係の追跡と高い性能を実現することを示しています。

要約

この論文は、人工知能（AI）の「記憶力」と「思考力」を劇的に向上させる新しい仕組み**「M2RNN（マトリクス・ツー・マトリクス・RNN）」**というものを提案するものです。

専門用語を抜きにして、日常の例え話を使って解説します。

1. 今の AI の「悩み」：記憶が薄れる、複雑な計算が苦手

現在の AI の主流は「トランスフォーマー」という仕組みです。これは**「全員の話を一度に聞いて、整理する」**という非常に効率的な方法ですが、2 つの大きな弱点があります。

1. 記憶の限界（長文の理解）： 本を 1 冊読んだ後、最初のページの細かい内容を思い出そうとすると、後半の情報が邪魔をして忘れっぽくなります。
2. 複雑な作業が苦手： 「チェスの駒を動かす」「コードを書く」といった、一連のルールに従って状態を正確に追跡する作業が、理論的に苦手です。

一方で、昔からある「RNN（リカレント・ニューラルネットワーク）」という仕組みは、**「一歩ずつ前に進みながら記憶を更新する」という方法です。これは複雑な作業が得意ですが、「一度に大量の情報を処理するのが遅い」**という弱点がありました。

2. M2RNN のアイデア：「メモ帳」を「巨大なホワイトボード」に変える

この論文の著者たちは、「RNN の弱点は、記憶する場所（状態）が小さすぎるからだ」と気づきました。

• 従来の RNN： 記憶する場所が**「小さなメモ帳（ベクトル）」**のようなもの。ここに書き込める情報量には限りがあり、新しい情報が入ると古い情報が消えてしまいます。
• M2RNN の革新： 記憶する場所を**「巨大なホワイトボード（行列）」**に変えました。

具体的な仕組み（3 つのポイント）

1. 巨大なホワイトボード（行列状態）：
   従来の「メモ帳」ではなく、横に縦に広がる「ホワイトボード」を使います。これにより、一度に大量の情報を整理して保存できるようになります。

   • 例え： 1 人の秘書（メモ帳）がメモするより、100 人の秘書が巨大なホワイトボードに書き込む方が、複雑なプロジェクトの進捗管理ができるのと同じです。

2. 「忘れん坊」防止のスイッチ（フォゲット・ゲート）：
   ホワイトボードに新しいことを書くと、古いことが消えてしまうのを防ぎます。M2RNN は「これは重要だから残そう」「これはもう要らないから消そう」と判断する**「忘れん坊防止スイッチ」**を持っています。

   • 例え： 冷蔵庫の整理。新しい食材を入れる時、古い野菜を捨てるか、新しいものを追加するかの判断を AI が自動で行い、必要な情報だけを新鮮な状態で保ちます。

3. 効率的な書き込み（外積）：
   ホワイトボードに書く際、従来の方法だと「1 行ずつ」書く必要がありましたが、M2RNN は**「1 回で全体をアップデートする」**ような効率的な書き込み方をします。これにより、AI の計算速度（ハードウェアの性能）を最大限に活かせるようになりました。

3. 実験結果：なぜすごいのか？

この新しい仕組みを試したところ、驚くべき結果が出ました。

• 長文の理解力が向上： 訓練時に見たことのない長い文章でも、重要な情報（「藁の中の一本の針」）を正確に見つけ出すことができました。
• 複雑なタスクが得意： 文字の並び替えや、コードの実行など、これまでの AI が苦手だった「状態を追跡する」タスクで、完璧に近い成績を収めました。
• ハイブリッドな最強チーム：
  最も面白いのは、「トランスフォーマー（全体を見る）」と「M2RNN（一歩ずつ深く考える）」を混ぜたモデルを作ったことです。
  • 例え： 料理で言えば、「トランスフォーマー」は「全体の味見をするシェフ」、「M2RNN」は「一つ一つの食材を丁寧に調理するシェフ」です。この 2 人をチームにすると、**「全体も完璧に、細部も完璧」**な料理が作れます。
  • 実際の実験では、既存の高性能モデル（Mamba-2 や Gated DeltaNet）と混ぜるだけで、わずかな計算コストの増加で、精度が劇的に向上しました。

4. まとめ：AI の未来への一歩

この論文が伝えたいことはシンプルです。

「AI に『小さなメモ帳』ではなく『巨大なホワイトボード』を持たせ、賢く整理させることで、より長く、より複雑なことを考えられるようになった」

これにより、AI は単なる「言葉の予測」だけでなく、**「長い物語の理解」「複雑なプログラミングの作成」「論理的な推論」**など、より高度なタスクをこなせるようになります。また、この仕組みは既存の AI にも簡単に組み込めるため、すぐに実用化される可能性が高いと期待されています。

つまり、M2RNN は、AI が「賢く、記憶力抜群のパートナー」に進化するための重要な鍵となる技術なのです。

技術概要

M2RNN: 行列値状態を持つ非線形 RNN によるスケーラブルな言語モデリング

技術的サマリー（日本語）

本論文は、Transformer の並列性は高いものの、計算複雑性クラス（TC0）の制約により、エンティティ追跡やコード実行などの複雑な状態追跡タスクに本質的に限界があるという問題に焦点を当てています。これに対処するため、著者らは**M2RNN（Matrix-to-Matrix RNN）**という、行列値の隠れ状態と表現力豊かな非線形状態遷移を持つ新しい RNN アーキテクチャを提案しました。

以下に、問題定義、手法、主要な貢献、実験結果、および意義について詳細にまとめます。

---

1. 背景と問題定義

既存のモデルには以下の課題がありました。

• Transformer の限界: 並列計算に優れていますが、状態追跡タスク（例：置換群の合成、コード評価）において、理論的に必要な表現力（NC1 以上）を欠いています。
• 線形 RNN/SSM（Mamba-2, Gated DeltaNet など）の課題:
  • 状態追跡能力の不足: 線形な状態遷移のため、非線形な状態追跡タスク（S5 置換群タスクなど）を解くことが証明的に困難です。
  • コンテキスト内検索（In-Context Retrieval）の性能低下: 状態容量が限られているため、長いコンテキストからの特定情報の検索精度が低下します。
• 従来の非線形 RNN（LSTM, GRU）の課題:
  • 言語モデリング性能の低さ: 状態サイズ（ベクトル値）が小さすぎるため、大規模な言語モデリングタスクで Transformer や線形 RNN に劣ります。
  • ハードウェア効率の悪さ: シーケンシャルな計算のため並列化が難しく、バッチサイズに応じたパディングにより Tensor Core の利用率が低下し、不要な FLOPs が発生します。

2. 提案手法：M2RNN

M2RNN は、非線形 RNN の表現力と、線形 RNN の大規模な状態容量を融合させたアーキテクチャです。

2.1 核心的な設計

• 行列値の状態（Matrix-Valued States）:
  • 従来のベクトル状態 $h_t \in \mathbb{R}^d$ ではなく、行列状態 $H_t \in \mathbb{R}^{K \times V}$ を採用します。
  • 外積による状態拡張: 線形アテンションや SSM で用いられる外積 $k_t v_t^\top$ を状態更新に組み込むことで、パラメータ数を大幅に増やさずに状態容量を劇的に拡大します。
• 非線形遷移とフォゲットゲート:
  • 状態更新式：$Z_t = \tanh(H_{t-1}W + k_t v_t^\top)$
  • 入力独立のフォorgetゲート: LSTM/GRU と異なり、ゲート値 $f_t$ が前状態 $H_{t-1}$ に依存せず、入力 $x_t$ のみに依存して計算されます。これにより、ゲート計算の並列化が可能になります。
  • 更新式：$H_t = f_t H_{t-1} + (1 - f_t)Z_t$
• ハイブリッド構成:
  • 計算コストの高い M2RNN レヤーを、アテンション層や線形 RNN（Mamba-2, Gated DeltaNet）と組み合わせた「ハイブリッドモデル」を構築します。

2.2 システム最適化

• ハードウェア効率: 行列値の状態を用いることで、バッチサイズに依存せず Tensor Core（NVIDIA GPU）を効率的に利用できます。従来の FlashRNN などのパディングによる FLOPs の無駄を排除しました。
• 分散学習: Triton を用いたカーネル実装と、2 種類のテンソル並列（TP）戦略（トポロジー依存型と独立型）を提案し、大規模モデルの学習を可能にしています。

3. 主要な貢献

1. 完全な状態追跡能力の証明:
   • M2RNN は非線形ベクトル RNN が表現できるすべてのタスク（正規言語）を表現可能であることを理論的に示しました。
   • 実験では、訓練時に見ていない長さのシーケンスに対しても、S3 置換群タスクなどで 99.5% 以上の精度を達成し、完全な一般化を実現しました。
2. 状態サイズ拡大の重要性の立証:
   • 非線形 RNN の性能不足は「非線形性」そのものではなく、「状態サイズが小さすぎる」ことに起因することを示しました。M2RNN は行列状態によりこのボトルネックを解消し、言語モデリング性能を飛躍的に向上させました。
3. ハイブリッドモデルによる SOTA 性能:
   • 既存のハイブリッドモデル（Mamba-2 や Gated DeltaNet とアテンションの組み合わせ）に、M2RNN レヤーを 1 層だけ追加するだけで、言語モデリング、コンテキスト内検索、長文コンテキスト一般化において顕著な性能向上が見られました。

4. 実験結果

4.1 億パラメータの Dense モデルと、70 億パラメータ（アクティブ 11 億）の MoE モデルで評価を行いました。

• 言語モデリング（Perplexity）:
  • 単体モデルでは、M2RNN は Mamba-2 や Gated DeltaNet と同等かそれ以上の性能を示しました。
  • ハイブリッドモデル: Hybrid M2RNN は、同等の Hybrid Gated DeltaNet よりも 0.4〜0.5 ポイントの Perplexity 改善を達成しました。
• コンテキスト内検索（In-Context Retrieval）:
  • RULER ベンチマークや実世界データ（SQuAD, NQ など）において、Hybrid M2RNN は最善のハイブリッドモデルとなりました。
  • 特に、Hybrid Gated DeltaNet に M2RNN を 1 層追加するだけで、検索精度が大幅に向上しました。
• 長文コンテキスト一般化（LongBench）:
  • 長文要約やコードタスクにおいて、Hybrid Gated DeltaNet + M2RNN は、SOTA のハイブリッド線形アテンションアーキテクチャを最大 8 ポイント上回る性能を示しました。
• スループット:
  • M2RNN 単体は計算コストが高いですが、ハイブリッド構成で 1 層のみ使用した場合、訓練スループットはベースラインの 94% 以上（6% 以内の低下）を維持しつつ、精度を大幅に向上させることができました。

5. 意義と結論

M2RNN は、非線形 RNN が持つ「高い表現力（状態追跡能力）」と、線形 RNN が持つ「大規模な状態容量・効率的な推論」を両立させた画期的なアーキテクチャです。

• 理論的意義: 非線形 RNN が TC0 を超える計算能力を持ち、長文コンテキストでの状態追跡に有効であることを実証しました。
• 実用的意義: 既存の Transformer や線形 RNN ベースのモデルに、M2RNN レヤーを少量（1 層など）組み込むだけで、長文コンテキスト処理や複雑な推論タスクの性能を大幅に向上させる「高効率なビルディングブロック」として機能します。
• 将来展望: 計算コストのさらなる削減や、より大規模なスケールでの検証を通じて、次世代の効率的かつスケーラブルな言語モデルの基盤となる可能性を秘めています。

本論文は、非線形 RNN レヤーが、単なる代替案ではなく、効率的でスケーラブルな言語モデル構築における不可欠な構成要素となり得ることを示しました。