FreeU: Free Lunch in Diffusion U-Net

Introduction

• 既存の拡散モデルに対して、no costs, no training, no additional parameter なしで生成の質を向上させることができる "FreeU" という手法を提案した論文についてまとめる
• 要は skip connections 側にハイパスフィルターを入れることで画像の質を向上させたという内容

拡散モデルについて

• 通常拡散モデルでは、difussion process と denoising process からなる
  • diffusion process では入力データに対してガウシアンノイズを段階的に付与していき、最終的にはノイズデータになる
  • denoising process ではノイズデータから元の入力データを復元するように、ノイズ除去が行われている
    • この過程には U-Net がよく用いられている

逆拡散過程（denoising process）

• まず、U-Net を用いた denoising process の特性理解を行った
  • 上段は通常の画像生成の iteration
  • 中段、下段はそれぞれ、フーリエ変換による low-frequency と high-frequency の画像を示している
• フーリエ変換による画像の特性について（一般的な性質）
  • 低周波領域の情報は画像全体の特徴（レイアウト、色）を表現している
  • 高周波領域の情報は画像のエッジ、形状を表現している

• low-frequency
  • 各ステップでの変化は緩やかなものだった
  • low-frequency 部分が大きく変化すると画像の特性が変わってしまい、うまく denoising できなくなってしまう
• high-frequency
  • 各ステップでの変化は速かった
  • 細やかな特徴はノイズの影響を受けやすく、そのためノイズが入ったときには high-frequency 部分に現れることになる

U-Net

• Fig.2 での考察から、low-freq. と high-freq. の情報には拡散過程での振る舞いが異なっていることが分かった
• さらに考察を進め、U-Net の構造による影響を見た
  • U-Net は backbone と skip features からなり、これらを繋げるような構造を取っている
  • backbone, skip ともに $b$ と $s$ という係数を導入して、これらの影響と出力画像との関係性を見た（Fig.5）

The backbone

• backbone features とは、入力情報をいったん backbone のネットワークですべて畳み込んでから up sampling しているもの
• Fig.5 上段をみると、backbone features の影響を大きくしてくと（$b$ を大きくすると）、生成される画像の質が向上していることが分かる
  • backbone は denoising の働きを持っていることが示唆された

The skip connections

• Fig.5 下段を見ると、skip features の影響は生成される画像の質には反映されていないことが分かる
• skip features は high-frequency な情報を持っていることが分かった

Free lunch in diffusion U-Net

• ここまでの実験を元に、backbone と skip features を適宜調整することで高品質な画像を生成できることが分かった
  • そこでそれぞれに係数を導入する
• $l$-th ブロックにおいて、backbone feature を $x_l$、skip feature を $h_l$ をする
  • backbone feature map には $b_l$、skip feature map には $s$ を導入する
  • $b_l$ は backbone の増幅を、$s_l$ は skip の減衰を目的としている

The backbone

• backbone のチャンネルすべてを増幅すると生成画像に oversmoothed texture（ぼやけるということ？）が生じることが分かった
  • そこで半分のチャンネルのみの増幅を行うこととした

$$x_{l,i}^\prime = \begin{cases} {b_l \cdot x_{l,i}, ~~~~ (\textrm{if}~~i < C/2)}\\ {x_{l,i}, ~~~~ \textrm{otherwise}} \end{cases}$$

The skip

• skip feature に対して low-frequency な情報を削減するために以下の計算を行う
  • FFT、IFFT はフーリエ変換、逆変換を示す

$$\begin{align} \mathcal{F}(h_{l,i}) &= \textrm{FFT}(h_{l,i}) \\ \mathcal{F}^\prime(h_{l,i}) &= \textrm{FFT}(h_{l,i}) \odot \alpha_{l,i} \\ h_{l,i}^\prime &= \textrm{IFFT}(F^\prime(h_{l,i})) \end{align}$$

• $\alpha_{l,i}$ は以下で定義されている

$$\alpha_{l,i}(r) = \begin{cases} {s_l ~~~~ (\textrm{if}~~ r< r_{\text{thresh}})}\\ {1 ~~~~ \textrm{otherwise}} \end{cases}$$

ここまでのまとめ

• 一度ここで論文の趣旨をまとめようと思う

低周波・高周波領域の振る舞いについて

• まず Fig.2 から、低周波・高周波領域の情報の denoising に関する振る舞いの違いが見られることが分かった
  • 低周波領域は画像全体の特徴を持っていて、高周波領域はエッジなどの形状の情報を持っている（これは一般的な2次元フーリエ変換に関する知見）
  • 低周波領域の情報は denoising process で大きな変化はない $\to$ ノイズ耐性がある
  • 高周波領域の情報は denoising process で大きく変化する $\to$ ノイズ耐性がない

backbone、skip feature の違いについて

• Fig.5 から、backbone と skip feature の出力画像への影響を調査した
  • backbone を強めると、きれいな画像が（よくノイズが除去された画像が）出力される
    • ただし backbone だけ強めていいかと言われると、$b=1.4$ を例に分かるように、backbone の影響が強くなると画像全体が oversmoothing （のっぺりとした）風潮の画像になることが分かった
  • skip を強めても出力画像の質にはあまり影響がなかった
• Fig.6 からbackbone を強めていくと、高周波領域の情報抑制されていく（弱まっていく）ことが分かった
  • backbone features は low-frequency な情報を持っていることが分かる
• Fig.7 から skip features は high-frequency な情報を持っていることも分かる
  • skip features は U-Net デコーダーに対して高周波の情報を付与する働きをしている

FreeU の構造について

• 以上の考察から、backbone featues の情報を強めて、skip features の情報を弱める（係数の導入だけ）構造を提案した
  • 特別な追加での学習も必要なく、Free lunch と呼んでいるゆえん
  • low-freq, high-freq の話はある程度こじつけ感はあった (...?)
• FreeU とはまとめると、backbone と skip maps のre-qweighting に他ならない

• skip connection のほうでは FFT $\to$ IFFT をしている
  • 途中で $\alpha$ を掛けていて、ハイパスフィルターの役割を作っている^[1]
  • "we introduce skip feature scaling factors, aiming to reduce low-frequency information and alleviate the problem of texture oversmoothing." と言っている

まとめ

• backbone, skip features を re-weight するだけで生成画像の質を向上させることができていて、すごい！
  • backbone を強めて画像の質そのものを高めている
  • skip features から low-frequency を取り除いて high-frequency にすることでよりくっきりとした画像を生成する
• ただ主張がいったりきたりしていて、何がいいたいのか謎すぎた
  • とくに low-frequency と high-frequency の件とハイパスフィルターの件がうまく結びつくようなロジックになっていなかった（お察しください感が強かった気が...）

出典

• FreeU: Free Lunch in Diffusion U-Net

脚注

1. これがハイパスなのかローパスなのかは論文中に分かりやすく明言されていなかった（読めば書いている程度）。この主張が重要なのではと思った次第。