💬 NLP

Triplet-Block Diffusion RWKV

本論文は、トリプレットブロック配置を通じて因果 RWKV モデルの $O(L)$ 推論効率と並列双方向離散拡散を統合する新たなアーキテクチャ $B^3D-RWKV$ を導入し、既存モデルと同等の精度を達成しつつ、デコードスループットで 1.6 倍の高速化を実現する。

原著者： Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

公開日 2026-05-26

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

CC BY 4.0

原著者： Ke Lin, Yiyang Luo, Zhaolong Su, Yunya Song, Anyi Rao

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

以下は、論文「Triplet-Block Diffusion RWKV」（B3D-RWKV）の解説です。単純な概念と創造的なアナロジーを用いて分解して説明します。

大きな問題：「一方通行の道」対「グループプロジェクト」

物語を書く二つの異なる方法を想像してみてください。

厳密な因果モデル（「一方通行の道」）：
従来の AI を、すでにタイプした単語しか見られない作家だと考えてください。彼らは左から右へと単語を一つずつ書き進めます。次の内容を見ることはできません。
- 良い点： 新しい単語を追加するたびにすべてを再読する必要がないため、長い文書の読み込みが非常に速いです。
- 悪い点： 一度に二つの単語を書くことができないため、テキストの生成は遅くなります。単語#2 を始める前に単語#1 を完了させなければなりません。
拡散モデル（「グループプロジェクト」）：
一方、別の AI は、「ガベージ」や「マスク」（[MASK] のようなもの）で満たされた空白のページから始めます。単語を一つずつ書くのではなく、ページ全体を一度に見て、欠落している単語が何であるか推測し、いくつかを修正し、物語が意味をなすまでこれを繰り返します。
- 良い点： 多くの単語を同時に修正できる（並列処理）ため、はるかに高速になる可能性があります。
- 悪い点： これを行うためには、前後の文脈（前の内容と後の内容）を同時に見る必要があります。これには通常、非常に高価で遅いコンピュータ・アーキテクチャが必要です。

対立点： これら二つを簡単に組み合わせることはできません。「一方通行の道」の作家は先を見ることができませんが、「グループプロジェクト」はすべてを一度に見る必要があるからです。

解決策：「トリプレット・ブロック」のトリック

著者の Ke Lin と同僚たちは、「トリプレット・ブロック・レイアウト」と呼ばれる巧妙なトレーニング・トリックを考え出しました。彼らは、作家の脳（アーキテクチャ）を変更することなく、「一方通行の道」の作家を「グループプロジェクト」のチームのように振る舞わせる方法を発見しました。

このトリックがどのように機能するかを、リハーサルのアナロジーを用いて説明します。

あなたは左から右にしか台本を読めない俳優（AI）だと想像してください。欠落したセリフを推測しなければならないシーンを学ぶ必要がありますが、正しく推測するには、現在の場所の後に来るセリフを知る必要があります。

著者たちは、すべてのシーンに対して3 部構成のリハーサルを設定しました。

パート 1（マスク付きコピー）： シーンを読み進めますが、半分は黒いテープ（[MASK]）で覆われています。これを左から右に読みます。
パート 2（損失対象のマスク付きコピー）： 同じ黒いテープで覆われた、全く同じシーンをもう一度読みます。ここでテストが行われます。欠落したセリフを推測しなければなりません。
- 魔法： パート 1 を読み終えたばかりなので、脳（AI の内部メモリ）はパート 1 からすべての可視セリフをすでに吸収しています。パート 2 を厳密に左から右に読んでいても、脳はパート 1 に保存されていたため、シーンの右側からの文脈をすでに「知っている」ことになります。
- 結果： 左から右に読み進めながら、「疑似双方向」の知識（過去と未来の両方を知っている状態）で欠落したセリフを推測することができます。
パート 3（クリーンコピー）： 完全な完璧なシーンを最後に一度読みます。これで脳がリセットされ、次のシーンに備えることができます。

このトリプレットパターン（マスク付き → マスク付き/テスト → クリーン）を繰り返すことで、AI は「未来」からの情報（実際にはトレーニングシーケンスの前のブロックだったもの）を使用して欠落した単語を予測することを学びます。これにより、元の「一方通行の道」の速度を維持したまま学習できます。

結果：高速かつ高精度

チームはこの手法を用いてB3D-RWKV-7.2Bというモデルを構築しました。彼らが発見した点は以下の通りです。

速度： 「一方通行の道」のアーキテクチャ（RWKV）を維持したため、デコードが驚くほど高速です。彼らは、同じモデルの標準バージョンよりも1.6 倍高速であると主張しています。
賢さ： 一般的なタスク（質問への回答や物語の作成など）において、他のトップクラスのモデルと同等のパフォーマンスを発揮します。
トレードオフ： 論文は、完璧なステップバイステップの論理を必要とする非常に複雑な数学の問題については、拡散の「推測」性質が時として小さな誤差を生む可能性があることに言及しています。しかし、ほとんどのタスクにおいては、その能力を発揮します。

要約

この論文は、**「いかにして高速な左から右への作家を、賢く全知の編集者のように振る舞わせるか」**というパラドックスを解決します。

彼らは、作家にシーンを連続して 3 回リハーサルさせることでこれを実現しました。最初のリハーサルで作家のメモリに文脈を埋め込み、2 番目でそのメモリを使って欠落部分を推測する練習を行い、3 番目で次のシーンに備えて盤面をクリアします。これにより、線形作家の速度を維持しつつ、拡散モデルの並列処理能力を獲得することが可能になりました。

彼らが主張していないこと：

医療診断や臨床用途にこれが機能すると主張していません。
すべての AI 問題に対する魔法の弾丸であると主張していません。複雑な数学的構造に対してはわずかに苦労すると認めています。
モデルの安全性機能を変更していないことを明確に述べており、元のモデルが持っていたバイアスをそのまま継承すると明記しています。

技術概要：トリプレットブロック拡散 RWKV（B3D-RWKV）

問題定義

因果的トランスフォーマー言語モデルは、主に 2 つの構造的限界に直面している。すなわち、厳密な逐次デコードによる並列化の不可能性と、ステップごとのアテンションコストの二次関数的増加であり、これらは長文脈推論を計算集約的にする。線形時間因果モデル（例：RWKV）と離散拡散モデルはそれぞれこれらの弱点を独立して解決するが、それらの統合は本質的に矛盾していた。拡散モデルはトークンブロックを並列にノイズ除去するために双方向アテンションを必要とするのに対し、因果モデルは厳密に単方向である。既存の統合試みは、線形時間バックボーンの効率性を損なうアーキテクチャ変更を必要としたり、ゼロから再学習を要したりすることが多かった。

手法

著者らは、 $O(L)$ の推論効率と、基盤となる因果的バックボーンを変更することなく、並列かつ双方向の離散拡散を統合する拡散スタイルの RWKV 変種であるB3D-RWKVを提案する。中核的な革新は、厳密な因果的アーキテクチャ内で双方向アクセスをシミュレートするように設計されたトレーニングおよび推論戦略であるトリプレットブロックレイアウトである。

1. トリプレットブロックレイアウト（トレーニング）

因果モデル上で拡散を可能にするため、サイズ $B$ の論理的生成ブロックは、トレーニングサンプル内で 3 つの連続する物理ブロックに拡張される。

$b_1$ （マスク済みコピー）： 論理ブロックのマスク済みバージョン。
$b_2$ （損失計算用マスク済みコピー）： ノイズ除去損失が計算される、同一のマスク済みコピー。
$b_3$ （クリーンな正解コピー）： 後続のブロックに対して再帰状態を更新する、マスクされていないバージョン。

疑似双方向アクセス：
因果的バックボーンがトークンを厳密に左から右へ処理するため、 $b_2$ 内の任意のマスク済み位置に到達する隠れ状態は、すでに $b_1$ 全体を処理済みである。 $b_1$ と $b_2$ は同じマスクパターンを共有するため、 $b_1$ 内のマスクされていないトークンは、モデルが $b_2$ のマスク済みトークンの予測を試みる前に、隠れ状態に実質的に「吸収」される。これにより、モデルは論理ブロック内の自身のマスクされていないコンテキストに対して疑似双方向アクセスを得ることとなり、RNN/SSM バックボーンの単方向性を維持しつつ、拡散の双方向要件を満たす。

トレーニング目的：
モデルはマスク予測目的を用いてトレーニングされる。損失は以下の構成からなる。

交差エントロピー損失（ $L_{CE}$ ）： トークンがマスクされている監督位置で計算される。
信頼度感知並列損失（ $L_{CAP}$ ）： LLaDA-2.0 に続き、この項は既に正しく予測された位置における確率分布を鋭くし、推論サンプリャがステップごとに多くのトークンをコミットすることを促す。

2. 推論：ブロック単位反復ノイズ除去

推論中、モデルは 1 つの論理ブロックずつ生成する。このプロセスには以下が含まれる。

全てのトークンがマスクされた状態でブロックを初期化する。
反復的なノイズ除去ステップを実行し、モデルがマスク済み位置のトップ 1 確率を予測する。
トップ 1 確率が閾値 $\tau$ を超えるトークンをコミットする。
閾値を超えるトーク数が $k_{min}$ 未満の場合、進捗を確保するためのフォールバック機構により、最も信頼度の高いトップ $k_{min}$ 個のトークンをコミットする。
全ての位置がコミットされると、クリーンなブロックがプレフィックスに追加され、次の論理ブロックが開始される。

主要な貢献

B3D-RWKV-7.2B： マスク予測目的でトレーニングされた、7B スケールでトレーニングされた初の拡散スタイルの線形時間 RNN 言語モデル。元のパラメータを変更することなく、RWKV-7 バックボーンに並列トークン選択を統合した。
トリプレットブロックフレームワーク： 厳密な因果的バックボーン（厳密な因果性と前方伝播状態特性を満たす）をブロック拡散モデルに変換する汎用的なトレーニングレシピ。このアプローチは、バックボーンを双方向ブロックに変更するアーキテクチャ的複雑さを回避する。
効率性と性能： 標準的な因果デコードに対してスループットを大幅に改善しつつ、競争力のある推論能力を維持する。

実験結果

モデルは、一般知識（MMLU、ARC、PIQA、RACE）および推論/数学タスク（GSM8K、MATH、GPQA）を含む 8 タスクスイートで評価された。

精度： B3D-RWKV-7.2B は、RWKV-7 ベースラインおよび同規模の他の拡散モデル（例：LLaDA-8B、Dream-7B）と同等の精度を達成する。特に、疑似双方向知覚が推論を強化することに起因し、ARC-Challenge や RACE などの推論タスクにおいてベースラインを顕著に上回る。
スループット： 1K から 512K トークンまでの文脈長にわたり、RWKV-7 ベースラインと比較してデコードスループットで平均 1.6 倍の高速化を示す。サンプリングパラメータを調整することで、わずかな品質のトレードオフを伴いながら 2.02 倍の高速化が可能である。
数学推論： 複雑な数学タスク（例：MATH）における性能は、因果的ベースラインと比較してわずかに低下する。著者らはこれを、部分的な評価が与えられない記号数学問題に必要な高度に構造化されたステップバイステップの検証に、並列デコードが苦戦する性質に起因すると指摘する。

意義と主張

本論文は、アーキテクチャ変更なしに事前学習済み因果言語モデルを拡散言語モデルに変換する効率的な経路を提供すると主張する。トリプレットブロックレイアウトを活用することで、B3D-RWKV は、線形時間モデルの $O(L)$ 推論効率と、拡散の並列デコード能力を成功裏に統合した。

著者らは、この研究を因果的および拡散的パラダイム間の非互換性に対する構造的解決策として位置づけている。単一のバックボーン（RWKV-7）での現在の検証、トレーニング中の 3 倍の物理的シーケンス長さコスト、および RL 整合を伴わない比較的小規模な SFT データセットの使用（これが特定タスクでの精度低下に寄与する可能性がある）といった限界を認めている。しかし、中核的な貢献は、アーキテクチャの再設計ではなくデータレイアウトを通じて、厳密な因果的かつ線形時間のバックボーン上で高スループットな並列拡散生成が実現可能であることを示した点にある。