UniDB: A Unified Diffusion Bridge Framework via Stochastic Optimal Control

要約

拡散ブリッジモデルの最近の進歩は、Doobの$ H $ transformを活用して分布間の固定エンドポイントを確立し、画像翻訳と修復タスクで有望な結果を示しています。
ただし、これらのアプローチは頻繁にぼやけたまたは過度に滑らかにされた画像の詳細を生成し、これらの欠点を説明するための包括的な理論的基盤を欠いています。
これらの制限に対処するために、確率的最適制御（SOC）に基づいた拡散ブリッジの統一されたフレームワークであるUNIDBを提案します。
UNIDBは、SOCベースの最適化を介して問題を定式化し、最適なコントローラーの閉じた形式ソリューションを導き出し、それにより既存の拡散ブリッジモデルを統一および一般化します。
Doobの$ H $ -TRANSFORMを使用する既存の拡散橋は、SOCコスト関数の端子ペナルティ係数が無限になる傾向がある場合に出現するフレームワークの特別なケースを構成することを実証します。
調整可能な端子ペナルティ係数を組み込むことにより、UNIDBは制御コストとターミナルペナルティの間の最適なバランスを実現し、詳細な保存と出力品質を大幅に改善します。
特に、UNIDBは既存の拡散ブリッジモデルとシームレスに統合されており、最小限のコード変更のみが必要です。
多様な画像修復タスク全体の広範な実験は、提案されたフレームワークの優位性と適応性を検証します。
私たちのコードは、https：//github.com/unidb-soc/unidb/で入手できます。

要約(オリジナル)

Recent advances in diffusion bridge models leverage Doob’s $h$-transform to establish fixed endpoints between distributions, demonstrating promising results in image translation and restoration tasks. However, these approaches frequently produce blurred or excessively smoothed image details and lack a comprehensive theoretical foundation to explain these shortcomings. To address these limitations, we propose UniDB, a unified framework for diffusion bridges based on Stochastic Optimal Control (SOC). UniDB formulates the problem through an SOC-based optimization and derives a closed-form solution for the optimal controller, thereby unifying and generalizing existing diffusion bridge models. We demonstrate that existing diffusion bridges employing Doob’s $h$-transform constitute a special case of our framework, emerging when the terminal penalty coefficient in the SOC cost function tends to infinity. By incorporating a tunable terminal penalty coefficient, UniDB achieves an optimal balance between control costs and terminal penalties, substantially improving detail preservation and output quality. Notably, UniDB seamlessly integrates with existing diffusion bridge models, requiring only minimal code modifications. Extensive experiments across diverse image restoration tasks validate the superiority and adaptability of the proposed framework. Our code is available at https://github.com/UniDB-SOC/UniDB/.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google