Local Epistemic Uncertainty Guided Active Sampling for Plug-and-play Diffusive Image Restoration

あなたは、友人のぼやけた、傷だらけの写真を手直ししようとしているところだと想像してください。あなたには、人間の顔が本来どうあるべきかを知っている超スマートなAIアシスタントがいますが、そのアシスタントは、この特定の写真におけるあなたの友人の具体的な姿までは知りません。これが、画像修復（image restoration）の世界です。これはコンピュータビジョンの分野の一つであり、科学者たちは、損傷した写真の中に欠落したディテールを「幻覚（hallucinate）」として描き戻す方法を機械に教えています。長い間、これらのAIアシスタントは、硬直した組み立てラインのように機能してきました。彼らは一歩進み、写真を確認し、また同じステップを繰り返し、写真が鮮明になるまでこれを何百回も繰り返しました。彼らは、空のような単純な部分であっても、複雑で乱れた顔であっても、全く同じ量の「推測パワー」を使用していました。

問題は、現実の世界は一様ではないということです。写真の一部（青い空など）は修正するのが簡単ですが、他の部分（動きによるブレで隠れた顔など）は悪夢のようなものです。従来の方法は、簡単な部分も難しい部分も全く同じように扱っていたため、空に対して時間を浪費したり、顔の処理を急ぎすぎたりして、結果として奇妙な歪みが生じたり細部が失われたりすることがよくありました。この論文は、このプロセスに対する新しい考え方を提示しており、AIが、その瞬間に自分がどれほど「混乱」しているかに基づいて、いつ速度を落とし、いつ加速すべきかを知る、好奇心旺盛な探偵のように振る舞うべきであることを示唆しています。

---

探偵のジレンマ：いつ推測し、いつ確認するか

LEADer（Local Epistemic Uncertainty Guided Active Sampling）を紹介しましょう。AIが画像を修復しようとするプロセスを、ミステリーを解く探偵だと考えてみてください。探偵には「事前知識（prior）」、つまり顔がどのようなものかという一般的な概念がありますが、証拠（損傷した写真）は乱れています。

昔の探偵たちは、厳格で退屈なスケジュールに従っていました。彼らは、たとえその手がかりが単純な指紋であろうと、複雑で判読困難な筆跡であろうと、あらゆる手がかりを正確に10秒間ずつ調べました。手がかりが簡単であれば時間を無駄にし、難しければ時間が足りず、事件を台無しにしてしまいました。これが、論文で「固定されたデータ制約と一様なステップサイズ」と呼ばれているものです。これは硬直しており、非効率的で、しばしばぼやけた顔や奇妙なアーティファクト（不自然な跡）を引き起こします。

LEADerは、ステップごとに次のような単純な問いを投げかけることで、ゲームのルールを変えます。「自分はこの画像の部分について、どの程度確信を持っているだろうか？」

著者らはこれを**「局所的なエピステミック不確実性（Local Epistemic Uncertainty）」**と呼んでいます。平易な言葉で言えば、それはAI自身の「混乱度」を測定する指標です。

• 低不確実性（高信頼度）： AIは空のパッチを見て、「これは何であるか正確に分かっている。青色だ」と言います。
• 高不確実性（低信頼度）： AIはぼやけた目を見て、「このブレの下に何があるのか全く分からない。もっと深く考える必要がある」と言います。

LEADerの2つのスーパーパワー

LEADerはこの「混乱メーター」を使用して、空間（画像自体）と時間（修復のためのステップ）に関する2つの巧妙な仕組みを実行します。

1. スマート・アジャスター（空間ドメイン）
あなたが絵を描いているところを想像してください。もしあなたが澄んだ青空を描いているなら、筆を強く押し付けすぎて滑らかな色を台無しにしたくないはずです。しかし、複雑で乱れた木を描いているなら、細部を正しく表現するために強く押す必要があります。
従来の方法は、あらゆる場所で同じ力で筆を押していました。しかし、LEADerは自身の「混乱メーター」を確認します。

• もしAIが特定のピクセルについて混乱している（高不確実性）場合、LEADerはAIにこう指示します。「自分の推測に頼りすぎず、与えられた元のぼやけた写真により忠実であれ」。これにより、AIがそこに存在しない偽のディテールを作り出すことを防ぎます。
• もしAIが自信を持っている（低不確実性）場合、LEADerはこう言います。「さあ、想像力を使ってエッジを鋭くさせよう！」 これにより、AIが「これは現実である」と知っている微細なディテールを保持します。
  これは**不確実性校正型事前分布変調（Uncertainty-Calibrated Prior Modulation: UCPM）**と呼ばれます。これは、AIの想像力と写真の現実とのバランスを取り、自然でシャープな結果を出しつつ、偽の構造を作り出さないようにします。

2. タイムトラベラー（時間ドメイン）
次に、探偵が犯罪現場を歩いているところを想像してください。部屋によってはすべてが明白なので、速く歩けます。しかし、暗くて混乱している部屋では、隅々までチェックしながらゆっくり歩く必要があります。
従来の方法は一定の速度で歩き、作業を終えるために100の小さなステップを踏んでいました。LEADerは**状態認識型軌道プルーニング（State-Aware Trajectory Pruning: SATP）**と呼ばれる手法を使用します。

• AIが自信を持っている（低不確実性）とき、AIは「すべてのステップを確認する必要はない」と判断します。そして、大きなジャンプをして先へとスキップします。
• AIが混乱している（高不確験性）とき、AIは速度を落とし、慎重に小さなステップを踏みます。
  論文では、このスキップを行ってもミスが蓄積しないことが数学的に証明されています。それは、映画の退屈な部分を飛ばして、エキサイティングなシーンだけを高画質で観るようなものです。結果はどうでしょうか？ AIは品質を損なうことなく、より速く仕事を完了します。

研究結果

研究者たちは、LEADerを既存のいくつかのAI画像修復ツールに組み込んでテストを行いました。彼らは単に推測したのではなく、数値を算出しました。

• より良い画像： LEADerを他の手法に加えると、修復された画像はより鮮明になりました。標準的なテストセットであるCelebA-HQ 1Kにおいて、鮮明度（PSNR）は約**0.88%から2.38%**向上し、よりリアルに見える（LPIPSスコアの低下）ようになりました。
• より速いスピード： LEADerは不要なステップをスキップするため、作業をより速く完了させます。平均して、**3.04%から8.42%**の時間を節約できました。DiffPIRのようなケースでは、7秒以上かかっていた時間を4秒未満に短縮しました。
• 追加メモリなし： 最も素晴らしい部分の一つは、LEADerが「プラグアンドプレイ」であることです。大規模な新しいコンピュータや追加のメモリを必要としません。論文によれば、LEADerを追加してもメモリ使用量の増加は極めてわずか（0.3%未満）であり、既存のツールと共に容易に使用できることが示されています。

結論

この論文は、画像の修復における従来の方法――写真のあらゆる部分とプロセスのあらゆる瞬間を同じように扱うこと――は間違いであると主張しています。AIに自身の「混乱」に耳を傾けさせることで、LEADerはよりスマートで、速く、注意深い画像修復プロセスを実現します。それは単に推測するのではなく、いつ推測し、いつ確認すべきかを知っているのです。著者らは、このアプローチがさまざまな種類の損傷（ブレ、ノイズ、欠損）に対して有効であり、現代的なほとんどの画像修復AIに、AIをゼロから再学習させることなく、より高性能かつ高速にするために追加できることを示しています。それは、AIの思考における小さな変化ですが、最終的な画像の仕上がりには大きな違いをもたらします。