AI-Papers
- 拡散モデルのタイムラインを「局所的な複雑性が均等」になるよう自動分割し、複数のサブネットワークに処理を振り分ける「複雑性均衡拡散分割（CBS）」を提案
- SiT-XL/2でFIDが6.24から4.03へ約35%改善し、論文の計測によると1ステップあたりの計算量は基準線と同じ118.64GFLOPsを維持
- de Boorの等分配原理を理論的根拠に持ち、SiT・JiT・UNetの複数アーキテクチャで効果を確認済みで手動チューニングが不要
研究の背景と課題
拡散モデルは、ノイズから徐々にデータを復元していく生成AIの代表的な手法です。最近の研究では、生成タイムラインを複数の区間に分割し、区間ごとに異なるサブネットワークを割り当てる「時間的分割」アプローチが注目されています。タイムステップによって処理の複雑さが異なるという観察に基づいた考え方です。
しかし従来手法では、分割点を手動で決める必要があり、理論的な根拠が乏しいまま試行錯誤が続いてきました。単純な等間隔分割では各区間の計算負荷がアンバランスになる問題もあります。本研究はこの課題に対して、数値解析の理論を活用した自動分割の仕組みを提案しています。
CBSの提案手法
本研究が提案する複雑性均衡拡散分割（CBS: Complexity-Balanced Diffusion Splitting）は、数値解析の理論を拡散モデルに応用し、分割点を自動的に求める手法です。中心となるのは、数学者Carl de Boorが提唱した「等分配原理」です。
この原理は、関数近似の誤差が最小になる分割の条件を数式で与えます。各区間での「複雑性の積分値」が等しくなるよう分割すると、全体の近似誤差が最小化されるという結論です。CBSはこの原理を、拡散モデルのODE（常微分方程式）ソルバーの誤差解析に適用しました。
2つの複雑性指標
分割点の決定には、各タイムステップの「局所的な複雑性」を測るモニタ関数が必要です。CBSは2種類の指標を提案しています。
- Dirichletエネルギー：フロー場の空間的な変動の激しさを測る指標。細かい構造を多く含む区間ほど高い値になる
- 軌跡加速度：サンプリング経路の2階微分で曲率を捉える指標。経路が大きく曲がる区間ほど高い値になる
どちらの指標も、本番の大型モデルより計算コストが低い補助モデルを使って事前に計算します。実験では軌跡加速度がわずかに優れたFIDスコアを達成しており、論文では主要指標として位置付けられています。計算した複雑性をもとにde Boorの等分配条件に従って分割点を数値的に求めると、理論的に最適な区間分割が自動で得られます。
実験結果
CBSの効果はImageNet-256でのSiT-XL/2を主要な評価対象として検証されました。classifier-free guidance（CFG）適用時にFIDが6.24から4.03へと改善し、ナイーブな等間隔分割との比較で約35%の向上を達成しています。
論文の計測によると、この改善はSiT-XL/2の1ステップあたりの計算量を基準線と同じ118.64GFLOPsに保ったまま得られています。つまり推論コストを増やすことなく生成品質が向上しています。複数アーキテクチャへの適用結果は以下のとおりです。
- SiT-B/2（中型モデル、ImageNet-256）：FID 16.51 → 10.72（約35%改善）
- JiT（ImageNet-64）：FID 16.41 → 13.93（約15%改善）
- UNet（CIFAR-10）：FID 3.55 → 2.72（約23%改善）
アーキテクチャが異なっても一貫した改善が見られており、手法の汎用性が裏付けられています。
まとめと今後の展望
CBSは、数学的に根拠のある等分配原理から分割点を自動決定することで、拡散モデルの時間的分割を実用レベルに引き上げた研究です。推論コストを変えずに生成品質を向上させるという実用的な価値は、AAD-1のような蒸留ベースの手法とは異なる方向性を示しています。蒸留がステップ数そのものを削減するのに対し、CBSは各ステップの利用効率を最大化するアプローチです。
現状の限界として、補助モデルによる複雑性計算には事前準備のコストが発生します。また、デフォルト設定のN=3（3分割）の妥当性はアーキテクチャによって変わる可能性があります。コードはプロジェクトページで公開されており、他の拡散モデルへの応用が今後の研究課題として期待されます。
