💬 NLP

DREAMSTATE: Diffusing States and Parameters for Recurrent Large Language Models

本論文は、条件付きDiffusion Transformerを用いてRWKVの再帰状態を編集可能な知識表現としてモデル化するフレームワークであるDREAMSTATEを紹介し、グローバルなコンテキストに基づいて再帰パラメータを動的に調整することで適応性を高める、並列DiTを備えた新しいハイブリッドアーキテクチャを提案する。

原著者： Liu Xiao

公開日 2026-01-28

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Liu Xiao

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

非常に賢く、高速で、物語を一つ一つの単語ごとに読み進めるロボットを想像してみてください。このロボットは特別で、物語全体を一度に覚える必要はありません。その代わりに、新しい単語が現れるたびに更新される、極めて圧縮された小さな「メモ」を頭の中に保持します。これが、現代のリカレントニューラルネットワーク（前述のRWKVモデルのようなもの）が機能する仕組みです。

この論文は、このロボットをさらに賢く、柔軟にするための2つの主要な手法を備えた、DREAMSTATEと呼ばれる新しいシステムを紹介しています。以下に、シンプルな比喩を用いてその内訳を説明します。

1. 問題点：ロボットの「メモ」が硬直的すぎる

ロボットの内部的な「メモ」（状態）を、それまでに読んだすべての内容をまとめた「心のスナップショット」だと考えてください。

問題: 現在、ロボットがこのメモを更新するために使用するルールは固定されています。それは、材料を混ぜ合わせるための、たった一つの変更不可能なレシピを持つシェフのようなものです。シェフがスパイシーなカレーを作っていようが、甘いケーキを作っていようが、混ぜるスピードや順番は全く同じです。
結果: この「静的な」レシピは、まあまあ機能しますが、完璧ではありません。論文ではこれを**「構造的ノイズ」**と呼んでいます。それは、眼鏡が少しぼやけているような状態です。ロボットは見えてはいますが、内部のルールが文脈に合わせて変化しないため、世界を完璧な明瞭さで捉えることはできていません。

2. 第一部：「状態」を「夢見る」ことを学ぶ

研究者たちはまず、次のように問いかけました：コンピュータに、これらの内部的なメモがどのような姿をしているかを理解させることはできるだろうか？

比喩: ロボットの内部的なメモを、異なる「気分」や「性格」のようなものだと想像してください。もしロボットがプログラミングについて読んでいれば、そのメモは「プログラマー」のようになります。もし詩について読んでいれば、そのメモは「詩人」のようになります。
実験: チームは、特殊なAIツール（拡散トランスフォーマー、またはDiTと呼ばれます）を使用して、何千ものこれらのメモを調査しました。彼らは、メモがランダムではなく、公園の特定の場所に集まる趣味の似た人々のように、明確なクラスター（集まり）を形成していることを見出しました。
画期的な発見: 彼らは、AIにゼロからこれらのメモを生成することを教えました。ロボットが物語を読んで「プログラマー」のメモを構築するのを待つのではなく、今や「プログラマー」のメモを即座に作成し、それをロボットに手渡すことができるのです。
結果: ロボットに一般的なメモの代わりに「ストーリーテラー」のメモを与えると、ロボットは即座により優れた、よりクリエイティブな物語を書き始めました。それは、ロボットに特定の「帽子」を被せることで、考え方を瞬時に変えさせるようなものです。

3. 第二部：ルールを動的にする

メモを生成できることが証明された後、彼らはこう問いかけました：「混ぜ合わせるレシピ」も、その場で変更できるだろうか？

比喩: たった一つの変更不可能なレシピを持つシェフを思い出してください。研究者たちは、新しいキッチンのセットアップを提案しました。
- メインシェフ (RWKV): 引き続き、単語ごとに料理を作ります（高速かつ効率的です）。
- スーシェフ（副料理長） (拡散モデル): この新しい助っ人は、調理が始まる前に、メニューの全体（グローバルな文脈）を見渡します。
革新: スーシェフは物語全体を観察し、その物語に特化した「新しいレシピ」を書き上げます。そしてメインシェフに、「このスパイシーなカレーには、もっと速く混ぜて熱を加えてください」とか、「この甘いケーキには、優しく混ぜてください」と指示を出します。
仕組み: システムは、並列AI（このDiT）を使用して、大きな全体像を把握し、その特定のタスクに必要な具体的な「混ぜ合わせのルール」（パラメータ）を生成します。そして、これらの新しいルールを古い安定したルールと融合させます。
結果: ロボットは、もはや一つの硬直した思考スタイルに縛られません。自らの内部的な「法則」を状況に合わせて適応させることができ、前述の「ぼやけた眼鏡」（構造的ノейズ）を効果的に打ち消すことができます。

4. 効果はあったのか？

研究者たちは、この新しいハイブリッドシステムをテストしました。

視覚的な証明: 彼らは、ロボットが生成する「メモ」が、実際に（公園のクラスターのように）整理され、意味のあるものであることを示しました。
学習の証明: 彼らはシステム全体を一緒に学習させ、システムはクラッシュすることなく学習に成功しました。ロボットは、独自のカスタムルールを作成することを学びながら、同時に次の単語を予測する能力も向上させました。

まとめ

要するに、DREAMSTATEは、ロボットの脳を固定されたルールを持つ静的な機械として扱うのをやめるための方法です。

ロボットの内部メモリを、芸術作品のように生成・編集できるものとして扱います。
ヘルパーAIが全体の文脈を監視し、その特定の瞬間にメインロボットが使用するための完璧な「動作ルール」を設計するという、動的なシステムを作り出します。

これにより、ロボットはより柔軟になり、「モード」（プログラマーや詩人など）を瞬時に切り替え、目の前のタスクに合わせて思考スタイルを適応させることができるようになるのです。

技術要約: DREAMSTATE

問題提起

現代的な再帰型ニューラルネットワーク（RNN）、特にRWKVのようなアーキテクチャは、標準的なTransformerと比較して推論効率や固定サイズの状態表現において大きな利点を提供します。しかし、本論文は現在のRNN設計における2つの決定的な限界を指摘しています。

未活用の状態表現ポテンシャル: RNNの状態は短距離のモデリングにおいては強力ですが、編集可能な知識表現として扱われることは稀です。内部状態の分布をモデリングすることで、直接的な生成や操作を可能にする研究は不足しています。
静的な再帰における構造的ノイズ: Transformerにおける「アテンション・ノイズ」（モデルが関連のないコンテキストに注意を過剰に割り当ててしまう現象）と同様に、RNNも「構造的ノイズ」に苦しんでいます。RWKVにおけるコアな再帰メカニズムは、学習後に固定される静的なWKVパラメータ（ $\theta_{WKV}$ ）に依存しています。この静的な性質は重み付けのバイアスとして作用し、多様なコンテキストに対して最適とは言えず、モデルが顕著な情報に焦点を当てるために「動作規則」を動的に調整することを妨げています。

著者らは、もし内部状態がコンテキストの構造化された表現であるならば、それを直接モデル化し操作することが可能であると主張しています。さらに、構造的ノイズを軽減するために、再帰パラメータの静的な仮定に挑戦すべきであると考えています。

手法

本論文は、これらの問題に対処するために、2段階のアプローチを通じて生成モデリング（具体的には拡散モデル）に基づいたフレームワークであるDREAMSTATEを提案しています。

1. 生成的状態合成 (DREAMSTATE フレームワーク)

著者らはまず、RWKVの内部状態を、学習可能な多様体（ $\mathcal{M}$ ）からサンプリングされたデータポイントとして扱います。

モデル: 条件付きDiffusion Transformer (DiT) が、状態 $S$ とテキスト埋め込み $c$ の条件付き確率分布 $p(S|c)$ をモデル化するように学習されます。
プロセス: 高次元の状態テンソル（フラット化されたマルチヘッド状態）は、一連のパッチとして扱われます。DiTは、入力コンテキストに条件付けられた状態に加えられたノイズを予測することを学習します。
目的: 状態の多様体が適切に定義されていることを検証し、特定の動作モード（例：特定のペルソナ）のための制御された状態生成を可能にすることです。

2. 動的パラメータ合成 (ハイブリッド・アーキテクチャ)

構造的ノイズに対処するため、著者らは、静的な再帰パラメータが動的に生成される新しいハイブリッド・アーキテクチャを提案しています。

メカニズム: 「パラメータDiT」は、可変長の入力コンテキスト（初期実装では最初のトークンの埋め込み $x[:,0,:]$ ）からグローバルな特徴を追跡し、コンテキスト固有のパラメータベクトル $\theta_{gen}$ を生成します。
対象パラメータ: DiTは、WKVパラメータを構成するフラット化された線形射影行列（ $W_r, W_k, W_v$ ）を生成します。
融合: 生成されたパラメータは、補間によって元の静的なベースパラメータ（ $\theta_{static}$ ）と融合されます：
$\theta_{WKV-final} = \alpha \cdot \theta_{static} + (1 - \alpha) \cdot \theta_{gen}$
ここで $\alpha$ は学習可能なハイパーパラメータです。これにより、モデルは安定した知識を保持しながら、コンテキスト固有の変調を導入することができます。
学習: システムは、以下のマルチオブジェクティブ損失関数を用いてエンドツーエンドで学習されます：
$L_{total} = \lambda_1 L_{LLM}(\theta_{WKV-final}) + \lambda_2 L_{param\_diff}(\psi)$
これは、次トークン予測（標準的なクロスエントロピー）と、パラメータ生成のための拡散損失とのバランスを取ります。

主な貢献

DREAMSTATE フレームワーク: RWKVモデルの動的な状態を生成的にモデリングする手法。著者らは、内部状態が特定の専門家ペルソナ（例：「プログラマー」対「詩人」）に対応する明確なクラスターを形成する、構造化された多様体であることを検証しました。これにより、状態操作による制御された推論が可能になります。
動的パラメータ・アーキテクチャ: 再帰パラメータがグローバル・コンテキスト・トラッカーに基づいて生成される、新しいハイブリッド設計。これは、再帰の動作規則を静的なものから動的なものに変えることで、構造的ノイズに対抗するための原則的な方法として提案されています。
マルチオブジェクティブ学習パラダイム: パラメータ生成と言語モデリングを共同で最適化する学習戦略を定義し、その成功を経験的に検証しました。これは、このハイブリッドなアプローチの実現可能性を証明しています。
コンセプト単位の自己回帰: 従来のトークン単位の自己回帰と、「コンセプト単位の自己回帰」（状態/パラメータ生成を介したもの）を組み合わせることが、将来のモデル設計における実行可能な方向性であることを示す概念実証を提供しました。

実験結果

状態多様体の可視化: t-SNEを用いて、実際のRWKV-7の状態が特定の専門家ペルソナ（例：Linuxターミナル、Ethereumデベロッパー、ストーリーテラー）に対応する明確なクラスターを形成していることを示しました。これにより、状態が構造化された、明確に定義されたオブジェクトであることが確認されました。
制御された推論: DREAMSTATE（または補間された状態）によって生成された状態でRWKV-7モデルを初期化することにより、モデルはベースラインと比較して、著しくテーマ的一貫性と創造性の高いテキストを生成しました。例えば、「忍耐」というプロンプトと「砂漠の気候変動」というプロンプトの間を補間することで、両方の概念を比喩的に融合させた物語が生成されました。
学習の安定性: ハイブリッド・アーキテクチャは、The Pile のサブセットを用いて正常に学習されました。学習損失曲線は、言語モデリング損失とパラメータ拡散損失の両方において、安定かつ一貫した減少を示しており、共同最適化プロセスの安定性を裏付けています。

意義と主張

本論文は、RNNの状態表現を、単なる一時的な変数から、操作可能な確率的オブジェクトへと進化させる新しい研究の方向性を切り開くと主張しています。

制御可能性: 内部状態を直接生成・操作することで、モデルを特定の動作モードに向けてプライミング（準備）できることを示しており、推論に対する前例のない制御を実現します。
適応性: 再帰パラメータをグローバル・コンテキスト・トラッカーによって生成される動的な変数として扱うことで、より柔軟で表現力豊かな世界モデルの構築を目指しています。
アーキテクチャのリファレンス: 本論文は、拡散ベースのパラメータ合成と再帰的推論のハイブリッドを、将来の大型言語モデル（LLM）のための具体的なアーキテクチャのリファレンスとして位置づけており、状態のモデリングの「方法」こそがLLMの未来の鍵であることを示唆しています。

著者らは、自らの研究が概念実証であることを強調しています。それは、構造的ノイズを軽減するためにパラメータを動的に合成する設計の実現可能性を検証するものであり、既存のすべてのモデルに対してすべてのベンチマークですぐに優位性を持つと主張するものではなく、提案されたパラダイムの生存可能性を確立することに主眼を置いています。