AI-Papers- 各デノイズステップの内部で共有Transformerブロックを繰り返し、パラメータを増やさずに計算深度だけを増やす新しいスケーリング軸を提案
- 260MパラメータのLooped-DiT B/16が、6.5倍大きい1.7BのInternVL-Uを6ベンチ平均で2.5ポイント上回り、推論計算量は4.9分の1
- 素朴なループが失敗する原因を中間ループの監督不足と注意の劣化と特定し、deep supervisionと自己変調アテンションで解決
研究の背景
テキストから画像を生成するモデルを強くする方法は、これまでおおむね2つに絞られていました。パラメータを増やしてモデル自体を大きくするか、デノイズ（ノイズを少しずつ取り除く処理）の回数を増やすかです。どちらも効果はありますが、前者は学習と推論の両方のコストを押し上げ、後者は生成の待ち時間を素直に増やします。
この論文が示すのは3つ目の軸です。1回のデノイズステップの内部で、同じTransformerブロックを複数回通すことで、パラメータ数を固定したまま計算の深さだけを伸ばします。言語モデルが推論トークンを並べて考えるのに対し、こちらは明示的なトークンを一切出さずに、潜在空間の中だけで繰り返し推敲を重ねる形になります。
図1が示すように、ループ回数を増やすだけでテキスト指示への忠実度が上がり、しかもその性能は桁違いに大きいモデルより低い推論コストで達成されています。
共有ブロックを繰り返す構造
提案手法のLooped-DiTは、260Mパラメータの拡散Transformerです。17個のTransformerブロックを、前段6ブロック、ループ部5ブロック、後段6ブロックの3つに分けています。学習時はこのうち中央の5ブロックだけをパラメータを共有したまま4回通すため、1デノイズステップあたり実質32ブロック分の計算が走る計算になります。
重みは共有されているので、パラメータ数は増えません。増えるのは1ステップあたりの計算量だけで、パラメータを揃えたベースラインの146 GFLOPsに対しLooped-DiTは267 GFLOPsと1.83倍にとどまります。小さなモデルで大きなモデルの質に迫るという意味では知識蒸留と似た動機ですが、こちらは学習済みの大規模モデルを必要としない点が異なります。
素朴なループが失敗する理由
ただし、単に同じブロックを繰り返すだけでは品質は安定して改善しません。著者らは原因を2つに分解しています。1つは中間のループに対する学習信号が弱いこと、もう1つは繰り返しのたびにアテンション（入力のどこに注目するかを決める仕組み）が過剰に書き込みを行い、各トークンが持っていた位置情報が壊れていくことです。実際、ループ後の特徴から各トークンの位置を復元できるかを調べる検証では、決定係数がループ1回目の0.865から8回目には0.562まで低下していました。
前者に対してはdeep supervisionを導入します。すべてのループが同じノイズ入力と同じタイムステップを共有し、それぞれの出力を同じ正解画像に対して教師付けすることで、途中のループ出力も単独で意味のある予測になるよう訓練します。後者に対しては自己変調アテンション（XSA）を使い、アテンションの更新からそのトークン自身のvalue方向の成分を直交射影で取り除くことで、繰り返しによる過剰な上書きを抑えます。
図9では、XSAを入れた場合にアテンション更新のノルムが残差ストリームに対して適切な大きさに収まり、同じデノイズステップでループが引き起こしていた誤りが回避される様子が示されています。アブレーションでは、ループ単体の寄与が1.0ポイント、deep supervisionの追加でさらに1.3ポイント、XSAで1.6ポイントの改善が得られています。
実験結果
GenEval、DPG-Bench、PRISM、T2I-CoReBench、SpatialGenEval、TIIF-Shortの6つのベンチマークで評価されました。Looped-DiT B/16はGenEvalで87.4、DPG-Benchで87.0を記録し、6ベンチ平均は71.5です。
モデル | パラメータ | 6ベンチ平均 |
|---|---|---|
Looped-DiT B/16 | 0.26B | 71.5 |
InternVL-U | 1.7B | 69.0 |
MiniT2I-L/16 | 0.91B | 67.3 |
CogView4 | 6.4B | 64.0 |
6.5倍のパラメータを持つInternVL-Uを2.5ポイント上回り、推論計算量は4.9分の1です。推論予算を固定した比較では、その予算をデノイズ回数の増加に充てるよりループ深度の増加に充てた方が一貫して高い性能になりました。25ステップから50ステップまでのどの設定でもこの傾向は変わりません。
図5は、ループを重ねる過程で物体の配置関係が徐々に整理され、前のループで生じた矛盾が後のループで修正されていく様子を示しています。この自己修正の挙動こそが、著者らが潜在的な視覚推論と呼ぶものの根拠になっています。
適応的ループと今後の課題
ループ回数を全画像で一律にする必要はありません。各ループの後に続行の価値があるかを判定する小さなゲーティングネットワークを置くことで、簡単なプロンプトは少ないループで打ち切り、難しいものだけ深く回せます。再学習なしで計算量と性能のトレードオフを調整でき、固定ループより平均ループ数を減らしながら高い性能に到達しています。
一方で限界もはっきり書かれています。検証は260Mモデルと512×512解像度に限られ、ブロックの分割方法や学習時のループ回数の網羅的な探索は行われていません。deep supervisionは推論コストこそ増やしませんが学習コストを押し上げます。潜在推論が起きているという主張についても、根拠は振る舞いと表現の観察にとどまり、内部で何が計算されているかの機構的な説明には至っていません。
とはいえ、パラメータでもデノイズ回数でもない第3のスケーリング軸を具体的な設計とともに示した意義は小さくないはずです。ブロック共有とdeep supervision、アテンションの安定化という組み合わせは、拡散モデル全般に移植できる形にまとまっています。
論文情報: "Looped Diffusion Transformer"（Yong Xien Chng et al., 2026） arXiv:2609.40305, CC BY 4.0
