Looped-DiT: 260Mモデルが6.5倍大きいモデルを上回る

この方法では、共有ブロック群がデノイズステップごとにN回ループし、自己変調アテンションがループを調整し、深層教師あり学習が各パスの後の状態をトレーニングします。

単純なループが失敗する理由

同じブロックを2回実行しても、拡散transformerが確実に改善するわけではありません。論文はその原因を2つの問題に求めています。中間ループ間の教師信号の弱さと、ループ回数が増えるにつれて局所的な情報を徐々に侵食するアテンション更新です。Looped-DiTはこれを解決するために2つのコンポーネントを追加しています。

• 深層教師あり学習（Deep Supervision） は、各ループ後の状態をループ後のブロックを通してデコードし、それらの予測のそれぞれを同じクリーン画像のターゲットに対してトレーニングします。
• 自己変調アテンション（Self-Modulating Attention） は、ループ内のアテンション更新を調整し、排他的自己アテンション（XSA）またはヘッド単位のアテンションゲートを使用します。

バックボーンはMiniT2Iのピクセル空間MMDiTで、凍結されたFLAN-T5-Largeで条件付けされています。リポジトリにはB/32、B/16、L/16のトレーニング、任意のループ深さでの推論、6つのベンチマークでの評価、各トレーニングセット用のデータセット準備スクリプトが含まれています。

結果

スコアはEMA重み、Euler 100ステップ、ガイダンス6.0、ループ深さ4を使用しています。

重要な主張は論文のアブストラクトにあります。パラメータ数と計算量を揃えた設定では、ループ設計が非ループのベースラインを一貫して上回り、固定の推論予算のもとでは、深いループは追加のデノイズステップよりも大きな効果をもたらします。著者らはまた、深いループが以前のループで生じた誤りを段階的に修正することも報告しており、これを潜在的な推論を示唆する挙動と表現しています。

実行方法

両方のチェックポイントはPyTorchファイルとして公開されており、sensenova/Looped-DiT-B16 と sensenova/Looped-DiT-B32 です。推論は単一のモジュール呼び出しです。--loops でループ深さを設定し、複数の深さを渡すとそれぞれ1行ずつ出力されるため、単一のプロンプトで深さを比較できます。

hf download sensenova/Looped-DiT-B16 looped-dit-b16.pt --local-dir checkpoints

python -m looped_dit.sample --checkpoint checkpoints/looped-dit-b16.pt \
    --prompt "a red cube on top of a blue sphere" --loops 1 2 3 4 --out loops.png

論文の主要な結果では、bfloat16のEMA重みを使い、Euler 100ステップ、ガイダンス6.0、ループ深さ4、512x512 を使用しています。深さ4はトレーニング済みの深さですが、著者らは再トレーニングなしで他の深さも機能すると述べているため、ループ回数はチェックポイントの固定された性質ではなく推論時の調整ノブとなります。

入手状況

これは研究向けリリースであり、ComfyUIとの統合ではありません。現時点で ComfyUIノード、Comfy-Orgによる再パッケージ、diffusersパイプラインは存在しない ため、実行するにはリポジトリ独自のPython環境が必要です。PyTorch 2.1以降のCUDAビルド、requirements.txt、そして評価用スタック（mmdet、vLLM、modelscope）は、ベンチマークデータとMask2Formerの重みが同梱されていないため個別にインストールする必要があります。リリース日と、これまでに集めた51個のGitHubスターを考えると、コミュニティによる移植が登場しても不思議ではありませんが、執筆時点では何も存在しません。