🤖 AI

Your Pre-trained Diffusion Model Secretly Knows Restoration

本論文は、事前学習済み拡散モデルがテキストエンコーダの出力でプロンプト埋め込みを直接学習することで、微調整や制御モジュールなしに多様な画像・動画復元タスクを高精度に実行できることを示し、学習と推論のダイナミクスを整合させる拡散橋アプローチによりその不安定さを解消した手法を提案しています。

原著者： Sudarshan Rajagopalan, Vishal M. Patel

公開日 2026-04-07

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Sudarshan Rajagopalan, Vishal M. Patel

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

この論文は、**「すでに完成された巨大な AI（拡散モデル）が、実は『画像修復』の天才的な能力を秘めていたが、そのスイッチの場所を誰も知らなかった」**という驚きの発見と、そのスイッチを簡単に見つける方法を紹介するものです。

わかりやすく、3 つのステップで解説します。

1. 発見：「魔法の箱」の奥底に眠る能力

研究者たちは、すでに大量のデータで訓練された「画像生成 AI（Diffusion Model）」を持っています。この AI は、何もないところから美しい絵を描くのが得意です。

これまでの常識：
これまでの研究では、この AI に「画像を修復して」と命令するには、AI の中身（脳の回路）を大幅に書き換えたり（微調整）、複雑な制御装置（ControlNet）を取り付けたりする必要がありました。まるで、新しい車を改造してレースカーにするようなものです。
今回の発見：
しかし、この論文の著者たちは、**「実は AI の中身は変えずに、単に『言葉のスイッチ』を正しい場所に変えるだけで、AI が勝手に画像を修復する能力を発揮する」**ことに気づきました。
- アナロジー：
  巨大な図書館（AI）に「雨の日の写真から雨を消して」と頼むとき、これまで私たちは「司書（AI の中身）を教育し直す」か「特別な指示書（制御装置）を渡す」必要がありました。
  しかし、この研究は**「司書に直接、図書館の奥深くにある『修復マニュアル』のページ番号（埋め込みベクトル）を教えるだけで、司書は瞬時に最高の修復作業をしてくれる」**ことを発見しました。

2. 問題：「言葉」では届かない、そして「道」がズレていた

では、なぜ今まで誰もこのスイッチを見つけられなかったのでしょうか？

言葉の限界：
「雨を消して」「ぼやけを直して」という普通の言葉（テキスト）や、AI が理解しやすいように工夫した言葉（トークン）を使っても、AI は「雨を消す」のではなく、「ノイズを消して元の汚い写真に戻す」ことしかできませんでした。
- メタファー：
  就像是对一个精通修复的工匠说“把这幅画修好”，但他只听到了“把画弄干净”，结果只是把画上的灰尘擦掉，却没能把画本身的裂痕补好。
- 解決策：
  彼らは、言葉そのものではなく、**「言葉が変換された後の『意味の核』（埋め込み空間）」**を直接操作することで、AI が本当に必要な修復作業を行うようにしました。
道（軌道）のズレ：
さらに、単純にこのスイッチを回そうとすると、AI が混乱しました。
- 問題：
  学習するときは「汚れた写真にノイズを足す」練習をしていましたが、実際に使うときは「汚れた写真からきれいな写真へ」進む必要があります。この**「練習の道」と「本番の道」がズレている**ため、AI が迷子になって変な結果を出していました。
- 解決策（ブリッジ）：
  彼らは**「拡散ブリッジ（Diffusion Bridge）」**という新しい道を作りました。
  - アナロジー：
    山頂（きれいな写真）と谷底（汚れた写真）の間に、ただの直線ではなく、「滑り台」のような滑らかな橋を架けました。これにより、AI は学習中も本番中も、同じ滑り台を滑るようにして、確実にきれいな写真へたどり着けるようになりました。

3. 結果：軽くて強力な「魔法の杖」

この方法を使えば、AI の本体（重くて高価なエンジン）を一切変えずに、**「軽量な学習済みプロンプト（魔法の杖）」**という小さな部品を追加するだけで、素晴らしい画像修復が可能になります。

すごい点：
- 汎用性： 霧、雨、雪、暗い写真、ぼやけなど、あらゆる種類の汚れに一つの方法で対応できます（All-in-One）。
- 効率： 巨大な AI を全部書き換える必要がないので、計算コストが安く、新しい機能もすぐに追加できます。
- 性能： 既存の最高峰の修復技術と比べても、劣らず、むしろ自然な質感で勝る結果を出しました。

まとめ

この論文は、**「AI はすでに修復の天才だった。必要なのは、中身を変えることではなく、正しい『意味の鍵』を差し込み、正しい『滑り台』を作ることだった」**という、シンプルで強力なアイデアを提示しています。

まるで、すでに完成された高級車に、新しいナビゲーションシステム（プロンプト）と、最適なルート案内（ブリッジ）を追加するだけで、どんな悪路（劣化）でもスムーズに走り抜けるようになったようなものです。

1. 問題設定と背景

背景:
近年、Stable Diffusion や FLUX、WAN などの大規模な事前学習済み拡散モデルは、自然画像分布に対する強力な事前知識（Priors）を持っており、高品質な画像生成が可能です。これらを画像復元（デハザ、デレイン、低照度補正など）に応用する試みが増えています。

既存手法の課題:
既存の拡散モデルベースの復元手法は、主に以下の 2 つのアプローチに依存しています。

バックボーンの微調整: モデル全体または一部を復元タスク用に再学習させる。
ControlNet 型モジュール: 凍結されたモデルに、タスク固有の条件付けネットワークを追加する。

これらの手法は計算コストが高く、学習データ分布に過剰適合して汎化性能が低下するリスクや、複雑な劣化に対するロバスト性の欠如という問題があります。

核心的な疑問:
「事前学習済みモデルには復元能力が内在しているのか？もしそうなら、従来のテキストプロンプトやトークン最適化ではなぜそれが引き出せないのか？」

2. 提案手法：YPT (Your Pre-trained Diffusion Model)

著者らは、事前学習済みモデルの復元能力は存在するが、従来の条件付け方法ではアクセスできないことを発見し、それを引き出すための新しいフレームワークを提案しました。

2.1 鍵となる発見：埋め込み空間の最適化

テキストプロンプトの限界: 自然言語（例："haze を除去"）や、トークン空間での最適化（Textual Inversion など）では、モデルはノイズを除去するだけで、劣化（ハザや雨など）を除去できず、入力画像の劣化状態に戻ってしまうことが確認されました（Fig. 2）。
埋め込み空間の最適化: 代わりに、テキストエンコーダの出力空間（Context Embedding Space）で直接条件ベクトル $p$ を学習することで、凍結された拡散モデルがノイズ混じりの劣化画像からクリーンな画像へ収束する強力な復元動作を引き出せることを発見しました。これは、復元に関する知識がモデル内部に存在し、トークンというインターフェースを通さずに直接アクセス可能であることを示唆しています。

2.2 課題：訓練と推論の軌道ミスマッチ

単純に埋め込み空間を最適化しても、推論時に不安定になる問題がありました。

原因: 訓練時には、劣化画像 $z_{deg}$ にノイズを加える単純な前方過程（ $z_t = (1-t)z_{deg} + t\epsilon$ ）を使用すると、訓練中の状態は常に $z_{deg}$ にアンカーされます。しかし、推論時には劣化からクリーンな画像 $z_{clean}$ へと遷移する必要があるため、学習されたプロンプトが推論時の中間状態（ $z_{deg}$ 中心ではない状態）に適用され、軌道がズレてアーティファクトが発生します。

2.3 解決策：拡散ブリッジ（Diffusion Bridge）の活用

このミスマッチを解消するため、拡散ブリッジの概念を導入しました。

ブリッジの定義: 劣化状態 $z_{deg}$ とクリーン状態 $z_{clean}$ の間を滑らかに遷移する中間状態を定義します。
EBR (Energy-Oriented Diffusion Bridge) の採用: 既存の DDBM（Denoising Diffusion Bridge Models）は、両端でノイズが小さく、初期段階での軌道修正が不十分になる傾向がありました。そこで、著者らはEBRを採用し、以下の特性を持つブリッジを構築しました。
- 劣化入力にノイズが乗った状態から始まり、クリーンな状態へ向かってノイズレベルが単調に減少する。
- 訓練と推論で同じ中間状態分布を共有し、軌道ミスマッチを解消する。
学習プロセス: このブリッジ上で定義された中間状態 $z_t^{ebr}$ を入力とし、凍結されたバックボーンが $z_{clean}$ を予測するようにプロンプト $p$ を学習します。

2.4 実装詳細

モデル: 画像モデル「FLUX.1」および動画モデル「WAN 1.3B」を使用。
パラメータ効率: バックボーンは完全に凍結。学習するのはプロンプト埋め込みのみ（FLUX で 512 トークン、WAN で 226 トークン）。
残差学習: 直接埋め込みを学習するのではなく、Null-text コンテキストに対する残差プロンプトとして学習し、条件付けインターフェースに注入することで、より効率的かつ高性能を実現。
マルチ劣化対応: 各劣化タイプ（ハザ、雨、雪など）ごとに独立したプロンプトを学習。混合劣化には、サンプリング時に各プロンプトの予測速度を平均化することで対応。

3. 実験結果

広範な実験により、提案手法の有効性が示されました。

性能: 画像・動画の両方で、最先端（SOTA）の AiOR 手法（AutoDIR, PixWizard, FoundIR など）と同等、あるいはそれ以上の性能を達成しました。
指標: PSNR, SSIM などの歪み指標に加え、CLIPIQA や MUSIQ などの知覚的品質指標でも高いスコアを記録。特に動画復元では、事前学習された WAN モデルの時間的事前知識（Temporal Priors）を活かし、DOVER 指標で他手法を凌駕しました。
汎化性 (OOD): 学習データに含まれない劣化（Out-of-Distribution）や、混合劣化、未見の劣化（Under-display camera などの）に対しても、強力な汎化性能を示しました。
アブレーション:
- 単純なプロンプト学習（Naive tuning）や DDBM ベースのブリッジでは性能が劣り、提案する EBR ベースのブリッジが最適であることを確認。
- 埋め込み空間最適化が、トークン最適化より優れていることを再確認。

4. 主な貢献

復元能力の発見と解放: 事前学習済み拡散モデルには復元能力が内在しており、自然言語プロンプトやトークン最適化ではなく、テキストエンコーダ出力空間での埋め込み最適化によってのみ引き出せることを実証。
軌道ミスマッチの解決: 訓練と推論の軌道不一致がプロンプト学習の不安定さを招くことを特定し、**拡散ブリッジ（特に EBR 形式）**を用いた訓練により、一貫したデノイズ経路を確立する手法を提案。
軽量かつ高性能なフレームワーク: 微調整や追加モジュールなしに、FLUX や WAN などの巨大モデルを高性能な復元モデルに変換する軽量なプロンプト学習フレームワークを実装し、SOTA 性能を達成。

5. 意義と結論

この研究は、大規模な生成 AI モデルを「復元タスク」に応用する際のパラダイムシフトを示しています。従来のようにモデル自体を再学習させるか、重い制御モジュールを追加する必要はなく、**「適切な条件付け（埋め込み）を学習するだけで、既存の強力な生成能力を復元に転用できる」**ことを証明しました。

これにより、計算資源を大幅に節約しつつ、多様な劣化に対するロバストで高品質な復元システムを構築できる可能性が開かれました。特に、動画復元において時間的整合性を保ちながら高品質な結果を得られる点は、実用面でも極めて重要です。