🤖 machine learning

Looped Diffusion Language Models

本論文は、標準的なマスクド拡散モデルと比較して優れた訓練効率と推論性能を達成するとともに、推論時の柔軟な計算スケーリングを可能にするために、トランスフォーマーの初期から中間層を選択的にループさせる新しい手法であるLoopMDMを導入する。

原著者： Sanghyun Lee, Chunsan Hong, Seungryong Kim, Jonghyun Lee, Jongho Park, Dongmin Park

公開日 2026-05-26

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

原著者： Sanghyun Lee, Chunsan Hong, Seungryong Kim, Jonghyun Lee, Jongho Park, Dongmin Park

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

複雑なパズル、例えば数独や難解な数学の問題を解こうとしていると想像してください。あなたには、答えを導き出そうとする専門家チーム（コンピュータモデル）がいます。

AI 言語モデルの世界では、これらのチームが機能する主な 2 つの方法があります：

「ワンパス」チーム（自己回帰型）： 彼らはパズルを左から右へ読み進め、単語を一つずつ推測します。もし初期段階で誤りを犯すと、後で修正することが困難になります。
「スクラッチパッド」チーム（マスク拡散型）： 彼らは、すべての答えが隠された（マスクされた）パズルから始めます。すべての隠された箇所を一度に推測し、その作業を確認してから推測を洗練させます。パズルが解決するまで、この「推測と確認」のサイクルを繰り返します。

この論文は、「スクラッチパッド」チーム向けの新しいトリック、LoopMDM を紹介します。

問題：作業過多で頭脳不足

通常、「スクラッチパッド」チームをより賢くするには、より多くの専門家（ニューラルネットワークの層を追加）を雇う必要があります。しかし、人を増やすことは高価で時間がかかります。研究者たちは問いかけました：より多くの人を雇わずに、既存のチームをより賢くすることはできるか？

解決策：「ループ」戦略

著者らは、パズル解決プロセスの最中に、チームが新しい部屋（ネットワークの新しい層）へ移動する必要はないことに気づきました。代わりに、彼らは中間の部屋にとどまり、行き来しながらアイデアを何度も何度も洗練させることができます。

犯罪を解決する探偵グループのように考えてみてください：

標準モデル： 探偵 A が証拠を見て、ファイルを探偵 B に渡し、探偵 B が探偵 C に渡します。一度手元を離れると、彼らは考えを変えることができません。
LoopMDM： 探偵 A が証拠を見て、ループして再度それらを確認し、論理を二重チェックするために 3 回目にループしてから、ファイルを探偵 B に渡します。

彼らはこれを**「選択的ループ」**と呼びます。すべての探偵にループさせるのではなく、プロセスの中間にいる探偵だけにそれをさせます。これは、チームが真剣に考えなければならないが、最終的な答えをまだ決定していないという絶妙なポイントです。

彼らが発見したもの（結果）

この論文は、この単純なトリックが驚くほど強力であると主張しています：

トレーニングコストの削減： チームが新しい部屋を建設する代わりに同じ「中間の部屋」を再利用するため、標準モデルと同じレベルのスキルに達するために必要な計算資源（FLOPs）を3.3 倍削減してモデルをトレーニングできます。新しい車を買うのではなく、既存の車をチューニングするだけでフェラーリのエンジンが手に入るようなものです。
数学能力の向上： 数学の問題（GSM8K ベンチマークなど）において、この方法は同じサイズの標準モデルと比較して8.5 ポイントの精度向上をもたらしました。これは、このループトリックを使用していない物理的に大きい（深い）モデルさえも凌駕しました。
柔軟な速度： モデルがループする回数を変更することで、モデルの「賢さ」を制御できます。
- 素早い答えが必要ですか？1 回ループさせます。
- 難しい数学問題に対する完璧な答えが必要ですか？12 回ループさせます。
- モデルは、その中間ループで「考える」ことをより多く許容するほど、性能が向上します。

なぜ機能するのか（「なぜ」）

研究者らは、これがなぜ機能するかを証明するために「数独」実験を行いました。

彼らがモデルに本を読むように厳格な順序で数独を解かせたとき、標準モデルは惨めに失敗しました。
しかし、LoopMDM はそれを完璧に解くことができました。なぜなら、「ループ」によりモデルは隠された箇所（空のマス）を共有ワークスペースとして扱えるからです。
モデルがループするにつれ、隠された箇所同士が「会話」します。一つの推測が間違っていれば、ループによりチーム全体がそれを認識し、最終的な答えを確定する前に一緒に修正することができます。答えを列伝のように叫ぶのではなく、チームが一緒に消しゴムで消して書き直すことができるホワイトボードを持っているようなものです。

「適応的」ボーナス

この論文は、これを使うための賢い方法として**「適応的ループ」**も提案しています。
すべての単語に対してモデルを正確に 12 回ループさせる代わりに、モデルは自身の信頼性を確認できます。

答えが明白な場合、1 回ループさせて次に進みます。
答えが難解な場合、12 回ループさせます。
これにより、品質を維持しながらエネルギーを節約できます。

まとめ

LoopMDM は、処理の途中で「円を描くように考える」ことを可能にすることで、AI モデルをより賢く、より速くする手法です。より大きく高価な脳を構築する代わりに、既存の脳に作業を二重チェックするために数秒余分に費やさせるだけです。その結果、より速く学習し、難しい数学の問題をより良く解決し、より少ない電力を使用するモデルが生まれます。

技術的サマリー：ループド拡散言語モデル（LoopMDM）

問題定義

マスクド拡散モデル（MDM）は、訓練目的、ノイズスケジューリング、アーキテクチャの進歩を通じて性能格差を縮小し、言語モデリングにおける自己回帰モデル（ARM）の有力な代替手段として登場しました。しかし、MDM 用に特に設計されたトランスフォーマーアーキテクチャの有効な設計は、まだ十分に探求されていません。ARM は、パラメータ効率と推論時の計算スケーリングのために、共有ブロックを繰り返し適用して深度をループに変換する「ループドトランスフォーマー」を成功裡に利用してきましたが、このアプローチが MDM において有用かどうかは不明です。MDM はすでにノイズ除去ステップ全体で反復計算を実行しており、内部ループを追加することが追加の利益をもたらすのか、それとも既存のメカニズムを単に重複させるだけなのかは不確かです。さらに、ループに関する先行研究はパラメータ効率や推論時の精緻化に焦点を当てており、MDM フレームワーク内で固定パラメータ数における「訓練計算効率」と性能をどのように向上させるかという点に関する理解にはギャップが残されています。

手法：LoopMDM

著者らは、MDM のノイズ除去ネットワークに選択的ループを組み込んだアーキテクチャである「LoopMDM（ループドマスクド拡散モデル）」を提案します。この設計は、以下の 2 つの原則に基づいています：

選択的適用：ループ回数は、ネットワーク全体に均一に適用するのではなく、訓練 FLOPs 単位あたりの利益が最も高い特定の層に選択的に適用されます。
スケーラビリティ：このモデルはループ数（ $S$ ）に応じた効果的なスケーリングをサポートし、推論時に $S$ が増加するにつれて単調な性能向上をもたらします。

アーキテクチャ

ノイズ除去トランスフォーマーは、3 つの連続したグループに分割されます：

ヘッド： $n_h$ 層。1 フォワードパスあたり 1 回適用（入力埋め込み）。
ミッドブロック： $n_m$ 層。これらはループドスタックを形成します。これらの層は重みを共有し、 $S$ 回連続して適用されます。
テール： $n_o$ 層。1 回適用（出力投影と LM ヘッド）。

フォワードパスは、実効深度 $n_h + S \cdot n_m + n_o$ 層に展開されますが、パラメータ数は $n_h + n_m + n_o$ 層を持つ非ループドネットワークと同じままです。

訓練戦略

確率的ループ数：訓練中、ループ数 $S$ は各ステップで $\{1, \dots, S_{max}\}$ から一様にサンプリングされます。これにより、共有ミッドブロックはさまざまな実効深度にさらされ、反復回数全体にわたって安定した表現を確保し、 $S_{max}$ 以上の値を含む任意の $S$ での展開を可能にします。
計算マッチング：公平な比較を確保するため、LoopMDM はベースラインよりも比例的に少ないステップ数で訓練され、総訓練 FLOPsを一致させます。これにより、アーキテクチャの違いを生の計算増加から分離します。

適応的推論

本論文では、ループ数が固定されず、サンプリング中に動的に決定される適応戦略を導入しています。プロセスは、連続するループ反復間の相対的隠れ状態の変化が閾値 $\epsilon$ 未満になった場合、または最大予算に達した場合に停止します。これにより、更新が大きいタイムステップ（通常は中間ノイズレベル）にはより多くの計算を割り当て、安定したタイムステップには少ない計算を割り当てます。

主要な貢献と結果

1. 訓練計算効率

3 つの事前学習コーパス（FineWeb-Edu、OpenWebText、LM1B）において、LoopMDM は、最大 3.3 倍少ない訓練 FLOPsを使用して、同サイズの非ループド MDM と同等かそれ以上のテスト負対数尤度（NLL）を達成します。

総 FLOPs を一致させて訓練した場合、LoopMDM はゼロショットパープレキシティ、生成パープレキシティ、および下流タスクの精度において、ベースラインを一貫して上回ります。
アブレーション調査：利益は、どこを、どの程度ループするかに強く依存します。少数の早期から中間の層（具体的にはゼロベースインデックスでの層 1〜2）を選択的にループすることで、利益の大部分を捉えることができます。ネットワーク全体をループしたり、入力/出力層の近くにループを配置したりすると、性能が低下するか、利益が得られません。

2. 深度スケーリングを超えた性能

GSM8K数学推論ベンチマークにおいて、LoopMDM は同サイズの非ループドベースラインを**+8.5 ポイント**上回ります。

重要なのは、LoopMDM がステップあたりの訓練 FLOPsを一致させた場合、より深い非ループド MDM（例えば 21 層モデル）を上回る点です。これは、利益が単にモデル深度を増加させることではなく、共有計算の反復適用（反復的精緻化）に起因することを示しています。
総訓練トークン数（FLOPs 一致のための訓練ステップ数の減少による）が少ないにもかかわらず、LoopMDM は下流の知識転送を維持または向上させており、ループに費やされるステップごとの計算が、トークン露出の減少を補っていることを示唆しています。

3. 機械的洞察

マスク位置をワークスペースとして：制限された（左から右への）生成順序を用いた数独実験を通じて、著者らは LoopMDM が浅いベースラインでは解決できないグローバルに制約されたタスクを解決できることを示しました。モデルは、マスク位置を並列ワークスペースとして使用し、初期に局所的には妥当だがグローバルに矛盾する予測を修正します。
アテンションダイナミクス：分析により、マスク間アテンションがループ数とともに増加することが明らかになりました。特に中間タイムステップにおいて顕著です。これは、ループがマスク位置間の相互作用を促進し、モデルがトークン予測をコミットする前に隠れ状態を精緻化することを可能にしていることを示唆しています。
適応的効率：適応的ループ戦略は、平均ループ数を 12 から約 5 に削減しながら性能を維持し、すべてのタイムステップが同じ計算予算を必要としないことを実証しています。

意義と主張

本論文は、選択的ループがマスクド拡散モデルを改善するためのシンプルかつ強力なアーキテクチャ選択であると主張しています。その意義は、以下の点にあります：

効率性：MDM はパラメータ数を増加させることなく、大幅な訓練効率の向上（最大 3.3 倍の FLOP 削減）を達成できます。
品質：ループの利益は単に深度増加の代理指標ではなく、拡散モデルに固有の反復計算とマスクワークスペースとの特定の相互作用に由来します。
スケーラビリティ：このアプローチは推論時に柔軟な計算スケーリングを可能にし、適応戦略によりさらなる効率化を提供します。

著者らは、この研究を MDM のより良いアーキテクチャ設計への一歩として位置づけ、選択的ループが固定パラメータおよび計算予算の下で訓練効率と汎化を改善するための標準的なアプローチとなり得ると示唆しています。また、現在の設計はループ位置の手動指定に依存しているという限界を認め、将来的にはこれらの戦略を動的に学習することや、このアプローチを数十億規模のモデルに拡張することを検討できる可能性に言及しています。