[論文メモ] Simple Baselines for Image Restoration
画像修復タスクで非線形の活性化関数なしのシンプルなSOTAモデルを提案。
既存のSOTA手法は複雑だが、このモデルの複雑さをinter-blockとintra-blockに分けて考えてみる。
inter-blocはアーキテクチャ自体についてで、図2のようにマルチスケールやマルチステージのモデルが提案されている。
intra-blockはアーキテクチャを構成するブロックについてで、Swin Transformer BlockやHINBlockなどが提案されている。
これらのinter・intra blockの中でも単純なblockでSOTAを達成できないか？というお気持ち。
手法
「必要な構造」だけの非常に単純な形のベースラインモデルを作成する。
「必要な構造」はタスク(SIDDデータセットのdenoising、GoProデータセットのdeblurring)の評価で決定。
inter-blockの複雑さが小さいアーキテクチャとしてシングルステージのUNet(図2(c))を採用。
各アーキテクチャはintra-blockをスタックして全体を構成するがそのBlockをどう設計するか。
intra-blockに関しては3つを考える
PlainNet's Block
まずはシンプルにConv(Depth-wise Convも)とReLUとショートカットのBlockとする(論文では図3(a)となっているが多分図3(b)の間違い)。
transformer blockを使わないのは
- transformer blockのパフォーマンスは良いとされているがSOTAモデルのために必須では無いという主張がある
- Depth-wise Conv(DWConv)のほうがself-attentionより単純な構造
- ConNetとTransformerを比較するのが目的ではない
といった理由から。
このBlockをPlain Blockと予備、これを使ったアーキテクチャをPlainNetとする。
Baseline's Block
ベースラインモデルを作成するために先ほどPlain BlockにNormalizationとAttentionの追加とActivationの変更を行う。
Normalization
BatchNormはバッチサイズが小さいと不安定で、代わりに提案されたInstanceNormは必ずしもプラス効果があるわけではなくチューニングが必要。
そこでLayerNormを採用した。
これにより10倍大きい学習率でも学習がスムーズになり、SIDDとGoProデータセットでもパフォーマンスが向上した。
Activation
ReLUはよく使われている活性化関数だが、近年のSOTAでは主流がReLUからGELUに置き換わってきている。
そこでPlain BlockのReLUをGELUに変更したら一部のパフォーマンスが向上したのでGeLUを採用する。
Attention
ViTのようなself-attentionは計算コストが高い。Swin Transformerのwindow-basedなattentionはグローバルな情報が欠落し、そもそもローカルな情報はPlainBlockのDWConvで十分。
しかし、Channel Attentionについては低コストかつグローバルな情報を扱えるので採用した。パフォーマンスも向上。
これがBaseline Block(図3(c))で、Baseline Blockを使ったアーキテクチャをBaselineとする。
Baselineモデルのパフォーマンスは表1参照。
Nonlinear Activation Free Network
このBaselineモデルは結果的に低コストでSIDDとGoProデータセットでSOTAを超えたが、現状のようにシンプルな構造のままよりパフォーマンスを向上させられないか？もしくはパフォーマンスを維持したままよりシンプルにできないか？
BaselineモデルのGELUとCAに注目する。
GELU
いくつかのSOTAモデルを見るとGated Liner Units(GLU)が採用されていて、これがキーになる可能性が高い。
GLUは特徴マップを、線形変換を
、活性化関数
として以下の式(入力を分割したり線形変換が1つとか種類はあるが)。
はアダマール積。活性化関数はSigmoidが多い。
それに対してBaselineで採用したGeLUは以下の式。
以下のように実装されることが多い。
を恒等写像、
を
とするとGELUはGLUの特殊なケースと考えられる。
なのでGELUをGLUに置き換えられそう。また、GLUは非線形な計算を含んでいて非線形活性化関数は必須では無いのでこれも取り除く。
そうするとといった非線形なものが含まれない形になる。
これらから、入力をチャンネル方向で2分割しチャンネル次元でアダマール積を取るSimpleGateでGLUを置き換える。
は二分割され同じ形の特徴マップ。
SimpleGateのイメージを図4(c)に示す。
Channel Attention(CA)
CAとしてSqueeze-and-excitation networks(SENet)のものを採用したが(図4(a))、CAは以下の式に書ける。
はチャンネル毎の積。
がpoolingやconvをひっくるめたもので、これは式(1)GLUに非常に似ていて、これもGLUの特徴ケースとみなせる。poolingによるグローバル情報の集約は必須なので、それを追加してSimple Channel Attention(SCA)を提案。
は線形変換。図4(b)がSCAのイメージ。
BaselineモデルのGELUとCAをそれぞれSimpleGateとSCAに置き換えることでパフォーマンスを落とさず、よりシンプルなモデルにできた。
しかも非線形活性化関数を完全に取り除けた。
なのでこれを Nonlinear Activation Free Network(NAFNet)とする。
