Timestep Rescheduling in Diffusion Inversion

技術要約：拡散逆転におけるタイムステップ・リスケジューリング（TRDI）

問題提起
拡散逆転（Diffusion Inversion）とは、実画像 $z_0$ を拡散モデルのガウス潜在空間（$z_T$）へと写像するプロセスであり、画像の再構成や編集といったタスクを可能にするものである。Denoising Diffusion Implicit Models (DDIM) は決定論的かつ効率的な定式化を提供しているが、この逆転プロセスは本質的に近似的である。標準的な手法では、デノイジング・ステップを解析的に逆転させることでこれを行うが、これは偏差を導入し、時間の経過とともに蓄積されるため、顕著な再構成誤差や編集の忠実度の低下を招く。

既存の手法は、主に固定点問題を解くことで逆転方程式を洗練させ、これらの誤差を軽減しようと試みている。しかし、これらのアプローチは、タイムステップ・スケジューリング（拡散タイムステップの選択と間隔）が逆転の忠実度に与える影響という極めて重要な要因を、ほとんど見落としている。ほとんどのパイプラインは一様に間隔の開いたタイムステップを利用しているが、著者らはこれが最適ではないと主張している。本論文は、タイムステップの分布がグローバルな逆転誤差の蓄積に大きく影響すると断じているが、拡散サンプリングや学習における役割と比較して、この関係性は未だ十分に探求されていない。

手法：TRDI (Timestep Rescheduling in Diffusion Inversion)
著者らは、推論段階での追加パラメータや計算オーバーヘッドを必要とせずに、全体の逆転誤差を最小化するように設計された非一様タイムステップ・スケジューラーであるTRDIを提案する。この手法は、逆転誤差の理論的分析に基づき、「粗から密へ（coarse-to-fine）」の2段階戦略を通じて実装されている。

1. 理論的洞察:
   著者らは、タイムステップのジャンプ $\Delta t$ によって誘発される逆転誤差が、スケーリングされた固定点問題として再定式化できることを導出した。誤差の大きさ $\delta$ は、スケーリング係数 $c_\alpha(t, \Delta t)$ と、連続するモデル予測の差分（$\Delta \epsilon_\theta$）の積に比例する。

   • 重要な観察: このスケーリング係数は、タイムステップ指数に対して放物線状の傾向を示す。誤差は、小さいタイムステップ（初期の拡散ステップ）と大きいタイムステップ（後半のステップ）の両方に集中しており、中間部で減少する。さらに、ステップサイズ（$\Delta t$）が大きいほど、本質的に大きな誤差を誘発する。

2. グローバル・リスケーリング（粗い調整）:
   この放物線状の誤差分布に対処するため、著者らはグローバルなリスケーリング機構を導入する。彼らは、ハイパーパラメータ $\gamma$ によって制御されるべき乗則変換を用いて、元の一様なタイムステップ列 $t_k$ を修正する：
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • $\gamma > 1$ の場合、タイムステップは初期のステップに向かって引き伸ばされ（高誤差の初期領域により多くの解像度を割り当てる）、
   • $\gamma < 1$ の場合、タイムステップは終盤に向けて密度が高くなる。
     このステップは、ステップの割り当てをグローバルに調整することで、微細な最適化のためのスケジュールを準備するものである。

3. 動的計画法によるローカル・リスケジューリング（細かい調整）:
   蓄積される誤差をさらに最小化するために、著者らはウィンドウ化された動的計画法（DP）を採用している。

   • コストマップ: コストマップ $E[k, t]$ は、$k$ 番目の逆転ステップにおいてステップインデックス $t$ に到達するための最小累積誤差を追跡するように初期化される。
   • 再帰: 現在のステップの周囲にある幅 $2d+1$ のスライディングウィンドウ内で、アルゴリズムは、直前の累積誤差と現在のステップの誤差係数の和を最小化するような、隣接するタイムステップ間の局所的に最適な遷移を再帰的に探索する。
   • バックトラッキング: 最適なタイムステップ列は、最終ステップから最初のステップへとバックトラッキングすることによって取得される。

このリスケジューリングプロセス全体は、既知のノイズスケジュールの基づいてオフライン（または前処理ステップとして）実行されるため、実際の画像逆転や編集の実行時における計算コストは無視できるほど小さい。

主な貢献

• 理論的定式化: タイムステップのジャンプと逆転誤差との理論的なつながりを導出し、誤差をノイズスケジュールの依存するスケーリングされた固定点問題として枠付けた。
• 適応型スケジューラー: 低誤差領域には大きなステップを、高誤差領域には細かいステップを適応的に割り当てることで、効率と精度のバランスを効果的に取る、粗から密へのリスケジューリング・フレームワーク（TRDI）を提案した。
• プラグアンドプレイの強化: 追加のパラメータ、学習、または大幅な計算オーバーヘッドを必要としない「オフザシェルフ」の拡張機能として設計されており、既存の逆転パイプラインとの互換性がある。

実験結果
著者らは、様々な拡散モデル（Stable Diffusion v1.5, SDXL, SDXL Turbo）およびベースラインの逆転手法（DDIM, ReNoise, NPI, GNRI）を用い、画像再構成（MSCOCOデータセット）および画像編集（PIE-Bench）に関する広範な実験を通じてTRDIを検証した。

• 再構成: SD v1.5において、TRDIを統合することで、すべてのベースラインにおいて一貫して指標が向上した。例えば、DDIMベースラインにおいて、PSNRは20.07から20.21へと上昇し、LPIPSは（知覚的品質の向上を示すように）193.97から187.85へと減少した。定性的な結果では、ベースラインが正確に再構成できなかった電車の窓や果物の形状といった構造的な詳細の修正が示された。
• 編集: PIE-Benchにおいて、TRDIは高いCLIP類似性を維持しながら、構造保持と背景保持を大幅に改善した。特に、背景の劣化が深刻になりやすいSDXL Turbo（4ステップ）のような少ステップモデルにおいて、TRDIは背景保持において顕著な改善（例：DDIMにおいて構造距離を17.43%削減）を達成した。
• アブレーション研究: スケーリング係数 $\gamma \approx 1.05$ およびウィンドウ幅 $d=8$ が、構造的忠実度と意味的な編集可能性の間の最良のトレードオフを提供することを研究で確認した。

意義と主張
本論文は、タイムステップ・スケジューリングが、拡散逆転の忠実度における重要かつこれまで見過ごされてきた要因であることを主張している。タイムステップのサイズとインデックス、および誤差蓄積の間の放物線状の関係を明らかにすることで、著者らは、計算ステップを戦略的に非一様に割り当てることが、グローバルな逆転誤差を大幅に減少させられることを示している。

著者らは、TRDIを、より強力な逆転定式化（固定点反復や確率的逆転など）の代替物としてではなく、補完的で軽量なスケジューラーレベルのレトロフィットとして位置づけている。これにより、既存の決定論的な逆転パイプラインが、計算コストを増大させたりモデルの再学習を必要としたりすることなく、優れた再構成および編集結果を実現できるようになる。本研究は、ステップを「どのように」取るかと同じくらい、「いつ」取るかを最適化することが重要であることを強調している。