🤖 machine learning

Learned Relay Representations for Forward-Thinking Discrete Diffusion Models

本論文は、マスクド拡散モデルが可微分なトークンごとのチャネルを介してノイズ除去ステップ間で潜在情報を伝播可能にする「学習型リレー表現（Relay）」という手法を導入し、これにより計画およびコーディングタスクにおける性能を向上させつつ推論遅延を大幅に削減するものである。

原著者： Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

公開日 2026-05-25

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

CC BY 4.0

原著者： Benjamin Rozonoyer, Jacopo Minniti, Dhruvesh Patel, Neil Band, Avishek Joey Bose, Tim G. J. Rudner, Andrew McCallum

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

以下は、論文「Learned Relay Representations for Forward-Thinking Discrete Diffusion Models」を平易な言葉と日常的な比喩を用いて解説したものです。

問題点：「ハードリセット」による記憶喪失

複雑なパズル、例えば数独やコードの作成に取り組んでいると想像してください。あなたは非常に賢いアシスタント（AI モデル）に助けを求めています。

現在、このアシスタントは、単一の文を終える瞬間にすべてのことを忘れる「記憶喪失」の人のように機能しています。

アシスタントは、疑問符でいっぱいの空白のページを見つめます。
必死に考え、複雑な頭の中での計算を行い、1 つまたは数文字を埋めると決めます。
ハードリセット： その文字を書き込むとすぐに、ページ全体の残りの部分について抱いていたすべての複雑な思考、計算、そして「感覚」を捨て去ってしまいます。
次のステップでは、今書き込んだ文字だけを見てゼロからやり直し、先ほど行った豊かな内面的な作業を完全に忘れ去ります。

これを**「ハードリセット」問題**と呼びます。この論文は、これが非効率であると主張しています。まるで料理人がスープを味見して「塩が必要だ」と書き留め、塩を加える前にスープの味を即座に忘れてしまうようなものです。彼らは毎回、食材を加えるたびに、スープ全体を最初から再度味見しなければなりません。

解決策：「リレー」のバトン

著者たちは「Relay（リレー）」と呼ばれる新しい手法を提案しています。

リレー競技を想像してください。ランナーが立ち止まって競技を忘れ、最初からやり直すのではなく、次のランナーにバトンを渡します。このバトンには、前のランナーの勢い、戦略、そして疲労が乗っており、次のランナーはちょうど自分が中断した場所から引き継ぐことができます。

AI の世界において、この「バトン」とは、モデルが先へ運ぶ**情報の連続した流れ（隠れ状態）**です。

リレー以前： モデルは計算し、トークンを記述し、その計算を削除します。
リレーあり： モデルは計算し、トークンを記述し、さらに 次のステップに向けて自分自身に「メモ」を渡します。このメモには、「ねえ、私はこの特定のパターンについて考えていたし、次の部分は 80% 確信しているよ」と書かれています。

仕組み：バトンパスを学ぶ

この仕組みを機能させるために、論文はいくつかの重要な工夫を紹介しています。

「先を見据えた」トレーニング：
通常、AI モデルは「現在の」答えを正しく出すことだけを訓練されます。著者たちはモデルを「先を見据えた」存在にするよう訓練しました。彼らはモデルに教えました。「今日の正しい文字を書くだけでなく、明日の正しい文字を書くのを助けるメモ（リレー）を書きなさい」と。
切り捨てられた逆伝播（BPTT）：
これは「先を見据えて学ぶ」ための高度な数学用語です。ダンスの振り付けを練習していると想像してください。1 つの動きだけを練習するのではなく、3 つの動きの連続を練習します。3 つ目の動きでつまずいた場合、「あ、3 つ目を準備しなかったせいで、1 つ目の動きをミスしてしまった」と気づきます。
論文では、モデルがトレーニング中に数ステップ先を見て学習できるようにする切り捨てられた BPTTという手法を使用しています。これにより、モデルは将来のステップをより容易にするために、情報を前方へ伝える方法を学習します。
「ソフト」なチャネル：
前方へ伝えられる情報は、単語や文字だけではありません。「ソフト」な数値（連続値）です。オンオフのスイッチではなく、調光スイッチのようなものです。これにより、モデルは最終的な決定にまだ至っていない微妙な手がかりや確率を運ぶことができます。

結果：より速く、より賢く

著者たちはこれを 2 つの分野でテストしました。

数独パズル： 彼らは数独を計画タスクとして扱いました。「リレー」モデルは、標準的なモデルよりも少ない試行回数（少ない「前方パス」）でパズルを解き、間違いも少なくなりました。グリッドを埋めながら「全体像」を保持する能力に優れていました。
コーディング（ソフトウェア作成）： 最先端のコーディング AI（Fast-dLLM v2）を採用し、そこにリレーシステムを追加しました。
- 精度の向上： より良いコードを書きました。
- 速度の向上： 同じ結果を得るために、「思考エンジン」を実行する回数が32% 減少しました。

なぜこれが重要なのか

この論文は、「ハードリセット」を停止させ、モデルがリレーのバトンのように自らの「思考」を前方へ運ぶことを可能にすることで、AI モデルを以下のようにできると主張しています。

より賢く： 複雑なコードや論理パズルのような長い系列にわたって、より良い推論が可能になります。
より速く： 毎回すべてをゼロから再計算する必要がなくなり、時間とエネルギーを節約できます。

要約すると、この論文は AI モデルに、自らの作業を忘れ去るのをやめ、人間が難しい問題を解決する際のように、一歩一歩それを積み重ねることを教えています。

技術的サマリー：先見的離散拡散モデルのための学習型リレー表現

1. 問題定義：マスク拡散モデルにおける「ハードリセット」

マスク拡散モデル（MDM）は、反復的なデノイジングを通じて離散系列を生成する。完全にマスクされた系列から開始し、モデルは反復的に位置の割合をマスク解除する。しかし、標準的な MDM 推論には重大な非効率性が存在する。各フォワードパスの後、モデルはマスクされた位置に対して生成された豊富な内部連続計算（隠れ状態）を破棄してしまう。次のステップは、新たにコミットされた離散トークンのみから開始される。

著者はこれをハードリセット問題と命名する。MDM は生成ごとに複数のフォワードパスを実行するが、中間の連続計算を再利用できないため、モデルはステップ間での推論を効果的に償却できない。これは再帰的計算と対照的である。固定パラメータモデルをステップにわたってアンロールすることは、モデルが近似できる関数クラスを拡大し、複雑な推論タスクにおける性能向上に寄与することが示されている。ここで扱われる中心的な問いは、MDM が中間潜在状態を破棄することなく、連続的なマスク解除ステップにわたってより豊富な情報を運ぶ再帰的計算をサポートできるようにするにはどうすればよいか、という点である。

2. 手法：学習型リレー表現（Relay）

提案される解決策Relayは、デコードステップ間で前方へ伝達される連続的な微分可能な潜在状態を導入し、モデルを「先見的」にする。

中核メカニズム

拡張状態軌道: 推論ステップ $k$ において、モデルは拡張状態 $s_k = (x_{t_k}, h_k)$ を維持する。ここで、 $x_{t_k}$ は現在の部分的にマスクされた系列、 $h_k$ は学習されたリレー状態（連続的な潜在ベクトル）である。
アーキテクチャ: モデルはバックボーン $f_\theta$ $f_{θ}$ 、トークン埋め込み $Emb_\theta$ $E m b_{θ}$ 、アン埋め込みヘッド $UnEmb_\theta$ $U n E m b_{θ}$ 、および特定のリレーモジュール $R_\theta$ $R_{θ}$ で構成される。
- 次の隠れ状態は以下のように計算される： $h_{k+1} = f_\theta(Emb_\theta(x_{t_k}) + R_\theta(h_k))$ 。
- リレー状態 $h_k$ は投影され、バックボーンに入る前に埋め込まれたトークンに加算される。
- 出力ロジット $\ell_k$ は $UnEmb_\theta$ を介して $h_{k+1}$ から導出される。
切断された BPTT による学習: 独立したステップで学習される標準的な MDM と異なり、Relay は**時間方向の切断された誤差逆伝播（Truncated Backpropagation Through Time, BPTT）**を用いて学習される。
- モデルは完全にマスクされた系列から開始し、 $K$ ステップ（例： $K=2$ ）のロールアウトを実行する。
- ロールアウト中、位置はポリシー（例：信頼度閾値）に基づいてマスク解除され、コミットされたトークンは誤差蓄積を防ぐためにロールアウト中に真値に対して「教師強制」される。
- 損失は $K$ ステップにわたる交差エントロピー損失の合計である。
- 勾配はリレー経路（ $h_k \to h_{k+1}$ ）を通じて逆伝播され、モデルは現在のステップを予測するだけでなく、将来のステップに最大限有用となるように $h_k$ に情報を符号化することを学習する。

推論

推論時、プロセスはマスク解除スケジュールやサンプリングポリシーの点では大きく変化しない。唯一の追加は、モデルがコミットされたトークン alongside にリレー状態 $h_k$ を前方へ運ぶ点である。リレー状態は、まだ離散トークンにコミットされていない情報に対する微分可能なメモリチャネルとして機能する。

3. 主要な貢献

Relay フレームワーク: 切断された BPTT を介して、学習された潜在リレーをデコードステップ間で前方へ伝達するようにモデルを学習させることで、MDM に再帰的計算を組み込むための汎用的な手法。既存の MDM 加速技術（ブロック拡散や KV キャッシングなど）と互換性があり、アーキテクチャに依存しない。
スケーラビリティと適応: Fast-dLLM v2 (1.5B) の全パラメータ適応を通じて、大規模言語モデル（LLM）スケールでの Relay の検証。この手法は、最小限の追加学習 FLOPs で最先端の離散言語モデル（DLM）を成功裡に適応させる。
性能とレイテンシの改善: Relay が性能 - レイテンシのパレートフロンティアを前進させ、コーディングタスクにおいて標準的な教師あり微調整（SFT）を上回る性能を発揮しつつ、推論レイテンシを大幅に削減することを実証。
設計空間の分析: ロールアウト手順、リレーチャネル、BPTT 学習目的の寄与を分離する、難易度の高い数独ベースの計画タスクに関する広範なアブレーション研究。

4. 実験結果

数独計画タスク

Sudoku-Extreme データセットを用いて、4 つの学習目的を比較した。

MLM: 標準的な一様マスク拡散（単一ステップ、リレーなし）。
Rollout: 教師強制によるマルチステップロールアウトだが、リレーチャネルなし（ $R_\theta \equiv 0$ ）。
Relay (sg): リレー経路を含むが、ストップグラディエントを使用（リレーを通じた BPTT なし）。
Relay: $K=2$ の BPTT をリレー通じて行う完全な手法。

知見:

精度-NFE フロンティア: Relay は、精度とフォワード評価回数（NFE）の間のトレードオフにおいて最良の結果を達成した。
構成要素の影響: MLM から Rollout へ移行することで性能が向上した。リレーチャネルを追加（Relay sg）することで大幅な飛躍が見られた。最後に、BPTT を有効化（Relay）することで最良の結果が得られ、モデルがより積極的なマスク解除下でも部分的な盤面の自己整合性を保つことを学習していることが示された。
効率性: 固定の信頼度閾値において、Relay は各フォワードパスでより多くのセルをコミットし、合法性（行/列/ボックス違反なし）を維持しながら、ベースラインと比較して同じ精度に到達するための総フォワードパス回数を削減した。

事前学習モデルの適応（Fast-dLLM v2）

著者は、コードと数学のデータセットの混合に対して Relay を用いて Fast-dLLM v2 (1.5B) を適応させた。

コーディング性能: HumanEval および MBPP ベンチマークにおいて、Relay は標準的な SFT を上回った。
- HumanEval: Relay は 42.1% の精度（SFT は 38.4%）を達成し、NFE は 88.3（SFT は 130.7）であった。
- MBPP: Relay は 46.6% の精度（SFT は 43.9%）を達成し、NFE は 78.8（SFT は 84.8）であった。
レイテンシ: Relay は精度を向上させながら、標準的な SFT と比較して推論レイテンシを最大 32% 削減した。
メモリオーバーヘッド: A100 GPU でのプロファイリングにより、Relay は 2 番目のフォワードパス中にライブ活性化メモリを増加させるが、ボトルネックがリレー状態の保存ではなく語彙投影の逆伝播であるため、ピークメモリ使用量はバニラ SFT とほぼ同一（約 20.1 GiB 対 21.2 GiB）であることが示された。

5. 意義と主張

本論文は、Relayが中間連続計算を蓄積することを可能にすることで、反復的非自己回帰生成の性質を根本的に変えることを主張する。

先見性: この手法は、モデルが各ステップを新たな予測問題として扱うのではなく、将来のマスク解除決定のために中間表現を明示的に最適化することを可能にする。
パレートフロンティアの前進: 目標精度に到達するために必要なフォワードパス回数を削減することで、Relay は性能 - レイテンシフロンティアを前進させ、複雑な推論や生成タスクに対する離散拡散モデルの効率性を高める。
互換性: このアプローチは、既存の最先端 DLM や加速技術（ブロック拡散、KV キャッシングなど）とアーキテクチャの大規模な変更なしに動作する、ドロップイン型の強化として提示される。

著者は、Relay が非貪欲で先見的な生成に向けた意味のある一歩であり、解釈可能性や他のモダリティ（画像、分子など）への応用においてさらに探求可能な、潜在推論のための連続的な基盤を提供すると結論づけている。