Correction Training：拡散モデルの予測を反復修正し、生成品質を改善する学習法

画像を作る拡散モデルは、ノイズを含む状態から画像を予測し、その予測を使って次の状態へ進む。この処理を繰り返して生成する。学習では実画像と独立したノイズから入力を作れるが、生成中の入力は、それまでのモデル自身の不完全な予測に左右される。著者らが問題にするのは、この学習時と生成時の入力のずれだ。

例えば、商品紹介に使う画像を生成する場面を考える。途中の予測に形の不自然さが生じると、その影響が後の段階にも残る可能性がある。ノイズ除去を進める処理は学んでいても、自分が作った誤りを含む予測を直すことは、通常の学習では明示的な課題になっていない。ただし、入力のずれが品質低下をどこまで引き起こすかは、この論文でも理論的に確定していない。

既存の品質改善策として、本論文はCFGと比較する。CFGは、条件付き予測と無条件予測の差を増幅する手法だ。条件に沿った画像を得る助けになる一方、多様性の低下や過度な彩度とのトレードオフがある。商品画像の例なら、指示への誘導を強めることと、途中で生じた形の誤りを修正することは、異なる働きとして考えられる。

Correction Trainingは、基礎モデルの予測そのものを受け取り、修正する写像を学ぶ。過去の予測を追加入力として使う「自己条件付け」とも、この点で異なる。何を手掛かりに次を予測するかに加えて、出てきた予測をどう直すかを学習課題にする。もっとも、論文が評価するのは画像生成の指標であり、商品画像としての業務上の適否を直接確かめた研究ではない。

主な評価指標のFIDは、生成画像と実画像の分布の隔たりを測り、低いほどよい。ただし、FIDが下がったことだけで、生成の多様性が保たれたとは言えない。そこで著者らはprecisionとrecallも報告する。前者は実画像らしい範囲への収まり、後者は実画像の多様な範囲を覆う度合いを見る指標だ。

もう一つ区別したいのが、生成品質と復元精度である。復元精度は、予測が元の正解画像にどれだけ近いかを測る。本論文には、見た目が改善しても、正解画像からの誤差は増える結果がある。「もっともらしい画像を作る」と「元の画像を正確に復元する」を、同じ成果として読むことはできない。

学習の出発点は通常のノイズ除去だ。実画像にノイズを混ぜ、基礎モデルに正解画像を予測させる。次に、その予測を修正器、つまりcorrectorへ渡し、同じ正解画像へ近づけるよう学習する。正解は共通だが、基礎モデルはノイズを含む入力から、修正器は基礎モデルの予測から復元する役割を持つ。

標準設定では、修正器を事前学習済みの基礎モデルの独立したコピーから始める。修正器には元のノイズ状態を追加入力として渡さず、元の時刻と条件を保つ。ここで時刻とは、ノイズ除去のどの段階かを示す情報である。修正器の学習によって基礎モデルへ直接修正を迫るのではなく、基礎モデル自身も通常のノイズ除去の目的で更新を続ける。

生成時は、まず基礎モデルが画像を予測する。その予測を同じ修正器へ繰り返し通し、修正後の予測を使って次のノイズ除去段階へ進む。修正回数をKとすると、学習で教師ありにするのは通常K＝1だが、生成時はそれを超えて反復する。この学習深度を超える反復が、大きな改善の鍵だったと著者らは述べる。商品画像の例なら、次の生成段階へ進む前に、現在の画像予測を何度か見直す動きに相当する。ただし、毎回必ず正解へ近づく保証はない。

主実験は256×256の画像生成で、ImageNet上のDiT、SiT、SiT-REPA、JiTを扱う。DiT・SiT・JiTには10万回、SiT-REPAには5万回の追加更新を行った。比較ではNFE、すなわちネットワークを評価する回数を約250に揃える。修正とCFGの両分岐も数えるため、修正を加える側はノイズ除去の段階数を減らしている。

例えばDiTとSiTは、各段階で2回修正し、83段階で249 NFEとする。SiT-REPAとJiTでは、修正する区間と回数を別に選ぶ。したがって主比較は、同じ段階数へ無制限に処理を足した結果ではなく、ほぼ同じ評価回数を別の処理へ配分した結果として読める。ただし、NFEは処理時間そのものではなく、モデルごとに設定も異なる。

テキストからの画像生成では、BLIP-3o mixtureのJourneyDB、Short-Caption、Long-Caption部分集合を使い、同じ5万回更新のチェックポイントで比較する。評価はGenEvalとDPG-Benchで、文章の指示と生成画像の一致を見る。これらの数値と、人が感じる品質や多様性の評価は分けて扱う必要がある。

著者らは、反復修正の利点が単なるモデルの深さで説明できるかも調べた。ImageNetで8万回、初めから学習する実験では、Transformerの層を評価する総回数を6,720に揃えて比較する。パラメータ数を整合させた実験では、28層のノイズ除去器のFIDが8.93だったのに対し、14層の基礎モデルと14層の修正器は1回修正で8.74、3回修正で5.66だった。この値は8千枚でのFIDである。

別の実効的な深さを揃えた実験では、双方240段階とし、5万枚でのFIDが深いノイズ除去器の4.88に対して反復修正は4.21だった。二つの実験群は評価枚数やパラメータ条件が異なるため、数値を横断して順位を付けるべきではない。それでも、検証した条件では、深いモデルへ計算を使う以外の配分に利点があった。

反復修正には追加の処理が必要になる。そこで著者らは、修正した予測を目標に学生モデルを学習させる蒸留も試した。SiTの固定した教師が2回修正した予測を作り、学生を追加で10万回更新する。この段階では正解画像への損失を使わず、教師の予測を学習する。学生は生成の各段階で1回の順伝播、つまりネットワークを一度通す処理だけを行う。

ImageNetのクラス条件付き生成で、教師のFIDは249 NFEで1.92、学生は250 NFEで2.04、83 NFEで2.26だった。同じ表の基礎モデルは250 NFEで8.62である。追加修正パスを省いても改善の多くを残せた結果だ。一方、基礎モデルと修正器のパラメータを共有する実験では、別ネットワーク版のFID 1.92に対し共有版は2.14だった。共有はモデルの大きさを抑えるが、修正ごとの追加処理は残る。

著者報告では、ImageNetのガイダンスなしクラス条件付き生成で、約250 NFEに揃えた4モデルすべてのFIDが改善した。基準モデルから修正ありへの変化は、DiTが9.58から2.07、SiTが8.55から1.92、SiT-REPAが5.53から1.31、JiTが14.86から2.96だった。CFGの参照値は順に2.51、2.11、1.35、2.35で、JiTではCFGの方が低い。修正がCFGを一律に上回った結果ではない。

同時にrecallは、順に0.676から0.603、0.669から0.608、0.739から0.703、0.672から0.601へ低下した。画像分布としての改善は大きいものの、多様性が完全に保たれたとは言えない。

クラスを指定しない無条件生成でも、SiT-REPA-XL/2の250 NFE比較で、FIDは基準の24.82から4回修正の5.37へ改善し、別の誘導手法autoguidanceの9.37も下回った。ただしrecallは基準の0.699から0.595へ低下する。この比較は修正回数とautoguidanceの強さをそれぞれ探索し、最低FIDの設定を選んだ結果であり、任意の設定で同じ優位性が出るという意味ではない。

テキスト画像生成の主比較では、CFGなしのGenEvalが28.30から50.09、DPG-Benchが70.71から76.82へ改善した。基準は250 NFE、修正はK＝6で245 NFEである。CFGの強さをw＝4とした比較では、修正してからCFGを適用し、GenEvalは68.08から70.48、DPG-Benchは83.38から83.69へ上昇した。こちらは250 NFEに対して246 NFEで、追加改善の幅は小さい。

人手評価でも、掲載された設定では修正を組み合わせた側が品質で60.75%、多様性で60.25%の票を得た。対象はGenEvalの単一物体80プロンプト、各指標800回答。品質は単画像、多様性は4画像の群を比べた。ただし基準のCFGはw＝6、修正側はw＝4で、CFGの後に8回修正する。双方50段階だがNFEは整合しておらず、定量評価とも処理順が異なる。

さらに人手評価の主掲載値は複数設定を探索した中の一つで、再評定による票率の変動があり、提示資料に信頼区間はない。強いCFGとの併用も万能ではない。50段階固定、w＝6の追加評価では、6回修正するとGenEvalは71.70から69.53、DPG-Benchは84.44から83.42へ低下した。

反復回数には設定ごとの適量がある。修正しすぎるとFIDやrecallが悪化し、過度に規則的な形や不自然な模様が現れることがある。しかも、実画像とノイズから作った学習分布の入力では、3回以上の修正で正解画像への平均二乗誤差、MSEが基礎モデルより増える結果が示された。見た目の改善を、復元誤差が単調に減った結果とは説明できない。

入力のずれについては、SiTの途中状態10万件を使い、学習由来か生成由来かを分類器に見分けさせた。修正後は概して識別が難しくなる方向へ変化し、両者の分布差が縮まることを示唆する。ただし、これは経験的な検査であり、その縮小が生成改善を引き起こしたという因果関係は確立していない。

このため、厳密な画像復元を目的とする用途では、FIDだけで採用を判断すべきではない。元画像への忠実性が必要なら、その目的に合った評価が別途必要だ。また、選ばれた画像例だけから、生成分布全体の多様性や、その一部に生成が偏る現象の有無を結論することもできない。

この研究は、追加学習できる画像拡散モデルで、ノイズ除去と予測修正への計算配分を比較する根拠になる。共有や蒸留も、その選択肢を広げる。ただし、蒸留の提示結果はImageNet上のSiTであり、他のモデルや用途へ同じ成果を約束するものではない。表現空間や時刻条件の比較にも限定設定があり、普遍的な最適構成は決まっていない。

実務判断に必要な空白は、実測レイテンシ、スループット、金銭コストである。提示資料はNFE、パラメータ数、学習更新数を比較しているが、それらの実測は報告していない。蒸留で追加修正パスを省いても、教師と学生の追加学習は必要になる。次の証拠として見るべきなのは、対象用途でも品質とrecallの両立が再現するか、そして計算配分の変更が実際の待ち時間にどう表れるかだ。