🤖 machine learning

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models

本論文は、マスクされた拡散言語モデルにおけるトークンのアンマスク順序を最適化する原理的なフレームワークであるSelf-Aware Scheduling（SAS）を導入するものであり、これはデコーディングの不一致に関する扱いやすい上界を導出することで軽量なポリシーを学習させ、ヒューリスティックなスケジュールの比較において、数独や数学的推論といったタスクでの生成品質を大幅に向上させるものである。

原著者： Jiawei Xu, Minghui Liu, Aakriti Agrawal, Yifan Chen, Furong Huang

公開日 2026-06-23

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

CC BY 4.0

原著者： Jiawei Xu, Minghui Liu, Aakriti Agrawal, Yifan Chen, Furong Huang

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

大きなアイデア：ロボットの「考え方」が重要である

巨大なジグソーパズルに挑戦しているところを想像してみてください。ただし、箱に描かれた完成図は見えません。あなたは、すべてのピースがどこにハマるかを推測しなければなりません。

ほとんどのAIモデル（エッセイやコードを書くものなど）は、本を読む人のように動作します。つまり、左から右へと、単語を一つ、次へと、また一つと、厳格な順序で書いていきます。これは**自己回帰的（Autoregressive）**な生成と呼ばれます。

しかし、**拡散モデル（Diffusion Model）**と呼ばれる新しいタイプのAIは、働き方が異なります。すべてのピースが黒い正方形（「マスク」）で覆われているパズルを想像してください。AIの仕事は、パズル全体が明らかになるまで、これらのピースを一つずつ剥がしていくことです。

問題点： AIには選択肢があります。どのような順番でピースを剥がすかは自由です。

コーナーのピースから始める（簡単）。
真ん中のピースから始める（難しい）。
ただランダムにピースを選ぶ。

以前は、科学者たちはAIに対し、「今、最も確信が持てるピースを常に選ぶ」といった単純なルールに基づいてピースを選ぶよう指示していました。論文では、これらのことを**ヒューリスティック・スケジュール（Heuristic Schedules）**と呼んでいます。著者たちは、これらのルールは「近視眼的（myopic）」であると主張しています。つまり、簡単なピースを先に選んでしまいますが、それが後のパズルを解く作業をより困難にしてしまう可能性があるのです。

解決策：「自己認識型スケジューリング（SAS）」

著者らは、**自己認識型スケジューリング（Self-Aware Scheduling: SAS）**と呼ばれる新しい手法を開発しました。硬直したルールに従うのではなく、AIが自分自身の「思考の順序」を学習する方法です。

これは、テストを受けている学生に例えられます：

従来の方法（ヒューリスティック）： 学生はテストを見て、自分が100%確信を持っている問題にすぐ答え、難しい問題は最後に回します。もし簡単な問題が難しい問題を解くためのヒントを与えてくれなかった場合、彼らは行き詰まってしまいます。
SASによる方法： 学生はテスト全体を見渡し、「もし問題5に先に答えたら、それは問題10を解くためのヒントを与えてくれるだろうか？」と自問します。学生は、たとえその問題自体が今すぐ解くのが簡単ではなくても、残りのテストを解く上で最も有利になるような順番を選ぶことを学習します。

その仕組み（「秘伝のタレ」）

この論文では高度な数学を使用していますが、概念はシンプルです：

「自己認識」報酬： AIにはすでに訓練された「脳（モデル）」があります。新しい部分は、順番を決定する「マネージャー（ポリシー）」です。
テスト： マネージャーは、パズルを剥がしていくさまざまな順番を試します。
スコア： パズルが解けたかどうかを最後まで待つ（これは時間がかかり、稀なケースです）代わりに、マネージャーはこうチェックします。「今選んだ順番において、私の『脳』はどれくらい正確に答えを説明できているか？」
- もしその順番によって「脳」が自信を持ち、論理的になれるのであれば、マネージャーには高いスコアが与えられます。
- もしその順番が「脳」を混乱させるのであれば、スコアは低くなります。
学習： マネージャーはこのスコアを使用して、どの順番が最適かを学習します。これは、コーチが選手に対して「ただ速く走るのではなく、チーム全体が勝つのに役立つ方向に走りなさい」と教えるようなものです。

研究結果

研究者たちは、3つの異なる種類のパズルでこれをテストしました：

数独（論理パズル）：
- 10億パラメータのAIを使用。
- 従来の方法： 最良の「ルールベース」の手法でも、パズルの約**82%**を解きました。
- SASによる方法： AIは独自の順序を学習し、パズルの**91.8%**を解きました。
- 驚きの事実： 人間のエキスパートによる「論理的な」順序（最も簡単な数字から解く方法）でさえ、AIが学習した方法よりも成績が悪かったのです！AIは、自身の「脳」にとってより効果的な、異なる経路を見つけ出しました。
数学問題（GSM8K）：
- より大きな80億パラメータのAIを使用。
- 従来の方法： 問題の**64%**を正解。
- SASによる方法： **76%**を正解。
コーディング（MBPP）：
- 従来の方法： 問題の**39.5%**を解決。
- SASによる方法： **41%**を解決。

「第2段階」のボーナス

この論文は、AIが一度「最適な思考順序」を学習すると、その経路に特化した追加の訓練をAIに与えることができることも発見しました。

これは、ミュージシャンが曲の練習における最善の方法を学び、その後、それをマスターするまでその方法「だけ」を練習するようなものです。
数独においては、この追加ステップによって、精度が**91.8%から97.5%**へと向上しました。

なぜこれが重要なのか

この論文は、AIが「何を」知っているかと同じくらい、「どのように」考えるか（情報の開示順序）が重要であると主張しています。AIに硬直したルールを強制するのではなく、自らの思考経路を計画させることで、数学、論理、コーディングといった複雑な問題を解く能力を大幅に向上させることができるのです。

要約すると： この論文は、AIに対し、単に「次の単語を予測する」のではなく、「答えへの最善の経路を計画する」ことを教えているのです。

技術要約：スケジューリングに関する考察 (SAS)

1. 問題提起

マスク型拡散モデル (MDM) は、トークンを反復的にアンマスク（復元）することで離散的なシーケンスを生成する。このプロセスは、自己回帰モデルの固定された左から右への因数分解に代わる柔軟な選択肢を提供する。しかし、この柔軟性は決定的な自由度をもたらす。それが「思考の順序」（アンマスク・スケジュール）である。

現在のアプローチは、ヒューリスティックなスケジュール（例：局所的な確信度、マージン、またはエントロピーに基づくトークンの選択）に依存している。著者らは、これらのヒューリスティックは「近視眼的」であると主張する。なぜなら、それらは次の一歩を最適化するだけであり、グローバルなデコーディングの軌跡を最適化するものではないからである。その結果、局所的には確実であっても、将来のステップに向けて制約を効果的に伝播させることに失敗する「容易な」トークンを選択してしまうことが多い。さらに、専門家が設計した論理的な順序（例：人間の数独の解法戦略）であっても、特定の学習済み拡散モデルにとっては最適ではない可能性がある。なぜなら、モデルの条件付き予測は、人間の論理とは異なる経路において最も信頼性が高くなる場合があるからである。

核心となる問題は、ベースとなる拡散モデルを再学習することなく、生成品質を最大化する最適なアンマスク順序を学習することであり、これは $n!$ 通りの可能な順序を探索するという困難な課題に対処するものである。

2. 手法：自己認識型スケジューリング (SAS)

本論文では、ベースとなる拡散モデル ( $p_\theta$ ) を凍結したまま、軽量な順序ポリシー ( $v_\phi$ ) を学習するプラグアンドプレイのフレームワークである自己認識型スケジューリング (Self-Aware Scheduling: SAS) を提案している。

理論的基礎

著者らは、拡散デコーディングを潜在的な順序生成プロセスとして定式化している。彼らの目的は、真のデータ分布 $\pi(x)$ と、デコーディング手順によって誘導される分布 $P^{seq}_{\theta,\phi}(x)$ との間のカルバック・ライブラー (KL) ダイバージェンスを最小化することである。

計算不可能性: 全ての $n!$ 通りのデコーディング順序を平均化する必要があるため、この KL ダイバージェンスを直接最小化することは困難である。
計算可能な上界: 著者らは、このミスマッチに対する計算可能な上界を導出した。彼らは、KL ダイバージェンスを最小化することが、データ・ポリシーの結合分布とモデル・ポリシーの結合分布の間の結合 KL ミスマッチを最小化することと等価である（定数の差を除いて）ことを示している。
自己認識報酬: この導出により、凍結されたデノイザーにおけるターゲットシーケンスのパスワイズ対数尤度 (pathwise log-likelihood) として定義される、高密度でモデル自己認識型の報酬関数が得られる。
$R_\theta(x, \sigma) = \log p_\theta(x | \sigma) = \sum_{t=1}^n \log p_{i_t}^\theta(x_{i_t} | \tilde{x}^{(t-1)}(\sigma))$
この報酬は、特定の順序でトークンを公開するように強制された際に、モデルがデータをどれだけうまく説明できるかに基づいて、各公開決定にクレジットを割り当てる。

アルゴリズム

ポリシーのパラメータ化: 順序ポリシー $v_\phi$ は、現在マスクされている位置に対するカテゴリカル分布を出力する軽量なニューラルネットワーク（MLPやTransformerなど）である。
最適化: ポリシーは Group Relative Policy Optimization (GRPO) を用いて訓練される。プロセスは有限ホライゾンのマルコフ決定過程 (MDP) として定式化される：
- 状態 (State): 部分的に公開されたシーケンス。
- 行動 (Action): 次にアンマスクするインデックスの選択。
- 報酬 (Reward): 現在のコンテキストにおける正解トークンの対数尤度（教師強制/Teacher-forced）。
二段階フレームワーク:
1. 順序学習: $\theta$ を凍結して $\phi$ を最適化し、既存のモデルにとって最適な軌跡を見つける。
2. 第二段階のファインチューニング: 学習された順序ポリシーを固定し、学習された軌跡に沿って拡散モデル $\theta$ をファインチューニングすることで、エラーをさらに減少させる（オプション）。

3. 主な貢献

理論: 本論文は、デコーディングのミスマッチ（KL ダイバージェンス）をデノイザーの累積パスワイズ対数尤度に結びつける、計算可能な上界を導出した。これにより、モデル自身の尤度を、順序選択のための高密度な報酬として使用することの理論的正当性が確立された。
アルゴリズム (SAS): GRPO を用いてアンマスク順序ポリシーを学習する、プラグアンドプレイのフレームメントを導入した。これは、あらゆる順序の逐次デコーディング (any-order sequential decoding) と 準自己回帰（ブロック）デコーディング (semi-autoregressive/block decoding) の両方に適用可能である。
実証的検証: 様々な推論タスク（数独、数学、コード）およびモデルスケール（1B から 8B パラメータ）において、SAS がヒューリスティックや専門家のスケジュールを大幅かつ一貫して上回る改善を実現することを実証した。

4. 実験結果

数独 (1B モデル)

性能: SAS は、パズルの正解率を 82.0%（最良のヒューリスティック：Confidence）から 91.8% へと向上させた。
比較: 学習されたポリシーは、人間の専門家の順序（76.6%）を大幅に上回った。これは、不完全な拡散モデルにとって、最適な経路は必ずしも人間の論理的な経路ではないことを示唆している。
分析: 可視化により、ポリシーが「易しいものから難しいものへ」というカリキュラムを採用していることが明らかになった。つまり、制約を伝播させるために、確信度の高いトークン（裸のシングルなど）を最初に解決するが、有益な場合には人間の論理から逸脱する。
第二段階: 学習された軌跡に沿ってモデルをファインチューニングすることで、正解率はさらに 97.5% まで上昇した。

数学的推論 (LLaDA-8B)

GSM8K: Pass@1 が 64% から 76% に向上。
MBPP (コード): Pass@1 が 39.5% から 41% に向上。
汎用性: 学習されたポリシーは、様々な生成長（128 から 512）やブロックサイズ（準自己回帰デコーディング）に対しても堅牢であり、左から右へのヒューリスティックを含むベースラインを一貫して上回った。

ドメイン間転移

数学タスクで訓練されたポリシーは、コードタスクへ効果的に転移し（およびその逆も）、テストされたすべての設定において最強のヒューリスティック・ベースラインを上回った。これは、ポリシーがタスク固有のパターンを記憶しているのではなく、情報密度に関する転移可能なスケジューリング原理を学習していることを示している。

報酬の比較

希薄な 0/1 の終端報酬（正解/不正解）と比較して、高密度な自己認識報酬はより安定し、しばしば優れたパフォーマンスを提供した。これは、希薄な終端フィードバックよりも、モデルと一貫した高密度な信号の方が有益であることを強調している。

5. 重要性と主張

本論文は、デコーディング順序を制御可能で最適化可能な「推論のプリミティブ（基本要素）」として位置づけている。著者らは、何を生成するかを学習するだけでなく、各情報に対していつコミットするかを学習することが可能であると主張している。

原理的な最適化: 局所的な不確実性を最適化するヒューリスティックなスケジュールとは異なり、SAS はグローバルな軌跡の尤度を最適化し、デコーディングパスをモデルの予測能力に適合させる。
効率性: このアプローチは計算効率が高く、フルサイズの拡散バックボーンに対する高価な強化学習の更新を必要とせず、軽量なポリシーの訓練のみを必要とする。
自己改善: このフレームワークは、反復的なスケジューリングとモデルの洗練を通じて、自己改善する拡散型 LLM への実用的な道筋を示唆している。ここでは、モデルが自身にとって最も信頼できる推論経路を学習する。

著者らは、教師強制による訓練と自由走行（free-running）推論との間の潜在的なギャップや、現在のポリシー学習におけるタスクごとの訓練への依存といった限界についても認めている。ただし、ドメイン間の転移は有望であることも述べており、まだ完全には特徴付けられていないとしている。