Revisiting L1 Loss in Super-Resolution: A Probabilistic View and Beyond を読んで

はじめに

こんにちは。2022年に誕生したAI Labというチームで、主に図面解析をしている中村遵介です。

趣味が料理と画像を4倍に拡大することなので、今日は最近読んだ「Revisiting $l_1$ Loss in Super-Resolution: A Probabilistic View and Beyond[1]」という、画像の拡大で利用される損失関数に関する論文を紹介したいと思います。 趣味以外の理由として、CADDiでは図面画像の解析を行なっておりノイズ除去や画像拡大などの分野に注目しているという点もあります。

畳み込みニューラルネットに関する知識は必要ですが、画像の拡大に関する知識は必要としないように書いたつもりです。

論文の概要

いったん細かい話を置いておいて、論文の概要をざっくりご説明します。 この論文が取り組んだ課題は以下の点になるかと思います。

  • 入力された画像を拡大する超解像分野において、低画質画像と高画質画像を1対1で学習させる既存手法を拡大し、低画質画像から「対応する可能性のある高画質画像の分布」を学習する手法を提案した
  • 提案した損失関数が既存手法で頻繁に使用される $l_1$ loss を下限に持ち、同時に不可能決定問題としてのランダム性も捉えられることを明らかにした
  • 分布を学習することの副作用として、生成画像に対する画素ごとの不確実性も同時に予測できることを明らかにした

それでは詳細についてみていきましょう。

前置き

Single Image Super-Resolution について

論文のタイトルは「Super-Resolution」とだけありますが、報告されている内容は Single Image Super-Resolution(SISR) と呼ばれる分野の話になります。 SISRは、日本語で「単一画像超解像」と呼ばれ、1枚の画像を入力とし、対応する1枚の拡大された画像を出力するタスクを指します。

SISR.png

SISRではバイリニアフィルタによるシンプルな手法もあれば、畳み込みニューラルネット(Convolutional Neural Network: CNN)を用いた手法もあり、画像処理分野では今も盛んに議論が行われる領域です。 これの一つの理由が、SISRが不可能決定問題であるというところにあると思います。つまり、ある低画質画像 $x$ が与えられた際、縮小すると $x$ になる画像、すなわち求めたい綺麗な高画質画像 $y$ は複数存在します。そのため、SISRという問題設計においては正解の手法が存在せず、今後も既存手法より良い結果を得られる手法が登場する可能性が残り続けます。

ill-posed.png

SISRでよく用いられる損失関数と評価指標

最近では、CNNを使用した教師付き学習を利用した数多くのCNNが提案されています。学習方法やネットワーク構造など様々な点が議論されていますが、学習時に利用する損失関数に着目すると、大きく4つに分けられるかと思います。

  • 生成画像と教師画像の $l_2$ loss
  • 生成画像と教師画像の $l_1$ loss
  • 生成画像と教師画像の perceptual loss
  • 生成画像と教師画像の discriminator loss

このうち$l_2$ loss と $l_1$ loss は非常に似ており、生成画像 $y$ と 教師画像 $\hat{y}$ に対して画素ごとに二乗誤差もしくは絶対誤差を計算して、その合計値を最小化することでより教師画像に近い画像を生成しようと試みます。 3番目のperceptual loss は特徴空間での比較を行います。具体的には、生成画像と教師画像をなんらかの学習済みCNNに通し、途中で得られる深層特徴について、その二乗誤差もしくは絶対誤差を最小化しようと試みます。 最後の discriminator loss はGenerative Adversarial Network(GAN)ベースの手法で利用される損失関数で、生成画像と教師画像を見分けるネットワークを同時に訓練し、生成ネットワークはなるべく見分けがつかないような画像を出力しようと学習していきます。 これらの損失関数は単独で使用されることもあれば、組み合わせて使用されることもあります。

一方で評価指標は大きく分けて3つ存在します。

  • Peak Signal-to-Noise Ratio(PSNR
    • 二乗誤差ベースで生成画像と教師画像の近さを比較します。値が大きいほど2つの画像が近いことを意味します
  • Structural Similarity Index Measure(SSIM)
    • パッチの統計情報ベースで生成画像と教師画像の近さを比較します。値が1に近いほど2つの画像が近いことを意味します
  • MOS(Mean Opinion Score)
    • 複数の評価者による5段階評価で画像の綺麗さ・自然さを評価します。5が理想的な高画質画像であることを意味します

L1 loss の限界

SISRに対するCNNベースの教師付き学習では、$l_2$ loss もしくは $l_1$ lossが使用されることが多いです。特に近年では $l_1$ lossがPSNR/SSIM/見た目の点で $l_2$ lossより優れているというのが実験的に言われ、$l_1$ lossが使用される傾向が強まっています。

やや横道に逸れて歴史の話をすると、最初にCNNがSISRに適用されたのがSRCNN[2]という手法でした。これは $l_2$ lossを損失関数に利用した3層のネットワークで、当時の手法から大きくPSNRを向上させました。その後に登場するESPCN[3]やVDSR[4]・DRCN[5]では高速化やネットワークの層数の増加が行われましたが、やはり $l_2$ lossを使用する点は変わりませんでした。 PSNRが二乗誤差を元に計算されるため、この傾向は非常に自然な流れであったと思われます。

流れが変わったのが SRGAN[6] と呼ばれるGANベースの手法の登場です。SRGANは discriminator loss を採用することで、不可能決定問題に対してGANで「低画質画像 $x$ に対応し、かつ高画質画像として自然な分布から生成されたかのような綺麗な画像 $y$」を生成するという一つの解を提案しました。その中では学習の安定化のために、perceptual loss、そしてやはり $l_2$ loss が使用されていました。このSRGANの生成部分のネットワークはSRResNetと呼ばれ残差構造を積極的に取り入れた手法だったのですが、このネットワーク構造がSISRに向いていることがわかり、後の多くの手法がこのSRResNetを参考にしています。

その中でも注目を集めたのがEDSR[7]という手法で、EDSRはSRResNetを巨大化させたネットワークです。徒に巨大化させたわけではなく、例えばそれまで使用されてきたバッチ正規化を「実際に学習させてみた結果、ほぼ正規化能力が失われていた」ということで削除したり、より残差構造を取り入れたりといった工夫をしています。そして、はっきりと $l_1$ lossによる学習がPSNR/SSIMの両者において良い結果を示したことを報告しました。

We train our networks using L1 loss instead of L2. Minimizing L2 is generally preferred since it maximizes the PSNR. However, based on a series of experiments we empirically found that L1 loss provides better convergence than L2. EDSR[7]より引用

以降は多くの場合で $l_2$ lossの代わりに $l_1$ lossを損失関数に採用することが増えています。安定的に高パフォーマンスを出せるRCAN[8]も $l_1$ lossで学習を行なっています。

話を戻します。

$l_1$ loss が $l_2$ lossより実験的に優れていること自体は良いのですが、その本質はあまり変わりがありません。縮小したらある低画質画像 $x$ になる高画質画像の集合を 仮に $y$ とすると、$l_2$ loss での学習は、これら$y$の平均値を学習することと等しく、$l_1$ lossは中央値を学習することに等しいです。

つまり、$l_2$ loss も $l_1$ loss も解空間のある一点だけを捉えようとする損失関数であり、解空間そのものを捉えるには不十分であると言えます。

そこで、論文ではパラメータ $W$ と入力 $x$ が与えられた時の高画質画像 $y$ の事後分布を明示的にモデリングすることで、解空間を推定することに注力します。

確率的モデリング

多くの既存手法では以下の尤度関数の最大化を目的とします。

\max_{W}{L\left(W\mid\hat{y}\right)} = P\left(\hat{y}\mid x;W\right).

$W$がパラメータで、$\left( x, \hat{y} \right)$ が学習データとして用意された低画質画像と高画質画像のペアです。論文の根幹は事後分布 $P(y\mid x;W)$ の明示的な推定です。 ここでいう $y$ は高画質画像として可能性のある複数の自然画像を意味しています。このとき、 $\hat{y}$ は $P(y\mid x;W)$ に従って高い確率で生成された画像のうちの1枚とみなすことができます。 つまり、最終的には以下の確率を最大化したいことになります。

P\left(\hat{y}, y\mid x;W\right) = P\left(\hat{y}\mid y\right)P\left(y\mid x;W\right).

日本語で言えば「入力$x$とパラメータ$W$が与えられた条件での、高画質画像の確率分布に対し観測された$\hat{y}$がその分布から生成される確率」を最大にしたい、という感じでしょうか。 ここでは$y$が無数に存在するので扱いやすくするため期待値を考えることにしています。

\mathbb{E}_{y\sim P\left(y\mid x;W\right)}\left[P\left(\hat{y}\mid y\right)\right]

あとはこれをいくつかの仮定をおきながら分解して考えるだけです。まず第一に「入力$x$とパラメータ$W$が与えられた条件での、高画質画像の確率分布($P\left(y\mid x;W\right)$)」を噛み砕いていみます。 論文では、ある画素は似た領域の加重平均で表せるという仮定と中心極限定理から、この分布は多変量正規分布に従うという仮定を置いています。

P\left(\hat{y}, y\mid x;W\right) \sim N\left(\mu_{\left(x;w\right)};\Sigma_{\left(x;w\right)}\right)

ここで出てきた $\mu$ と $\Sigma$ はパラメータ$W$を用いてCNNで実際に推論するものになります。また、$P\left(\hat{y}, y\mid x;W\right)$ が多変量正規分布に従うという仮定を置きましたが、実際にこれを利用して学習する場合、分布からのサンプリングを行う必要があります。これは微分可能な操作ではないという問題がありますが、これに関してはすでに Variational Auto-Encoder[9] が Reparameterization Trick と呼ばれる方法での回避を提案しているので、そのまま利用することにしています。

すなわち、標準正規分布から生成される $z$ を用いて

\mathbb{E}_{y\sim P\left(y\mid x;W\right)}\left[P\left(\hat{y}\mid y\right)\right] = \mathbb{E}_{z\sim N\left(0, 1\right)}\left[P\left(\hat{y}\mid \mu + \sigma \ast z\right)\right].

のように$y$を捉え直すことで微分可能にします。

$P\left(y\mid x;W\right)$ については考え終わったので、残っている $P\left(\hat{y}\mid y\right)$ をみていきます。 ここでは、[10]に従ってこの分布をボルツマン分布に従うと仮定します。

P\left(\hat{y}\mid y\right) \propto \prod_{i}^{H\times W} \exp \left( - \frac{\|\hat{y}_i - y_i\|_1}{kT} \right)

これを先程の式に代入することで最終的な目的関数を得ることができます。

\mathbb{E}_{z\sim N\left(0, 1\right)}\left[\prod_{i}^{H\times W} \exp \left( - \frac{\|\hat{y}_i - \left(\mu_i + \sigma_i \ast z\right)\|_1}{kT} \right)\right]

あとは扱いやすいように負の対数尤度関数にして、最大化問題を最小化問題へと変更しています。

\min\mathbb{E}_{z}\left[\frac{1}{kT} \sum_{i}^{H\times W}{\|\hat{y}_i - \left(\mu_i + \sigma_i \ast z\right)\|_1}\right]

$kT$ は定数なので無視することにすると、結局のところは上の式を最小化する $\mu$ と $\sigma$ の2つををCNNで推定してあげれば良いことがわかります。その際に $z$ がランダム性を与えてくれるイメージです。

ちなみに $z$ は標準正規分布なのでその期待値は0になります。つまり、上の式は(Jensenの不等式にそのまま代入すると)下限が$\hat{y}$と$\mu$の$l_1$ lossになります。

こうして振り返ってみると、$l_1$ lossが分布の1点だけしか捉えられていなかったことがよくわかります。

さて上式の最小化ですが1点だけ困った点があり、それはネットワークとしては $\sigma$を0としてしまう(ランダム性を排除して$l_1$ lossと同等にする)のが最も最小化できる戦略になる、ということです。 論文ではこれを実験的にも確かめ、$\sigma$が0に落ちる様子をプロットしてくれています。

sigma-to-zero.png

そこで、$\sigma$を0にせず(ランダム性を失わず)良い感じに勾配降下法で常識を最小化するための設計を提案しています。

分散の実践的な設計

論文では、実際に学習させるための$\sigma$の設計を2パターン挙げており、それぞれを「入力データに依存しない$\sigma$」と「入力データに依存する$\sigma$」と呼んでいます。 最終的に選択しているのは後者の方ですが、せっかくですので両方見ていきます。

データに依存しない分散

データに依存しない分散$\sigma$の設計は実に簡単で、「学習によって$\sigma$が0に収束してしまうなら、$\sigma$を学習しなければいい」というものです。 すなわち、あらかじめ$\sigma$を非常に小さな定数$k$として設定する、という方法です。この場合、最適化対象は以下のように読み替えられます(余計な定数は除外しました)。

\mathbb{E}_{z}\left[\sum_{i}^{H\times W}{\|\hat{y}_i - \left(\mu_i + \sigma_i \ast z\right)\|_1}\right] = \mathbb{E}_{z}\left[\sum_{i}^{H\times W}{\|\left(\hat{y}_i + k \cdot z\right) - \mu_i\|_1}\right].

これは、すなわち教師データである$\hat{y}$ に対して、平均0、分散 $k$ のガウシアンノイズをかけてから $l_1$ loss で学習させていることと同義です。 この考え方自体は Noise2Noise[11] と似ています。Noise2Noiseは、(ノイズ画像、綺麗な画像)というペアを用いでデノイジングを学習するのではなく、(ノイズ画像、ノイズ画像)のペアを用いても同等の性能が得られる、ということを示した論文です。ただし、事前条件としてガウシアンノイズであれば平均が0の分布に従う必要があります。また、理論的には$l_2$ lossでの精度を保証したものですが($l_2$ lossが平均値を推定するため)、実験では$l_1$ lossでも良い性能を示すことを報告してくれています。今回は乗せるノイズは平均0のガウシアンノイズなので問題なしです。

しかし、この方法だとノイズによっては$\mu$が$\hat{y}$に対して遠ざかるような方向の勾配を産む可能性があります。これは先程の式を$\mu$で偏微分すれば明らかで、余計な定数項を除くと

\frac{\partial \| \hat{y} + k\cdot z - \mu \|_1}{\partial \mu}

となりますが、これは$\hat{y}$と$\mu - k \cdot z$との大小関係で符号が入れ替わります。 一方で、一般的な$l_1$ lossに関しては

\frac{\partial \| \hat{y} - \mu \|_1}{\partial \mu}

なので$\hat{y}$と$\mu$との大小関係で符号が入れ替わります。このため、$\mu - k \cdot z$と$\mu$の間に$\hat{y}$が存在するような$z$では、$l_1$ lossによる学習と、提案する学習では勾配の方向が逆転します。 この勾配の逆転による学習の不安定化を防ぐためには、十分な数の$z$のサンプリングが必要になりますが、これは学習コストが肥大化を意味し、論文では使用されませんでした。

データに依存する分散

やや天下り的なのですが、$\sigma$ を定数$k$ではなく $|\hat{y} - \mu|$ に設定してみます。 これによって、ネットワークが求めるべき$\mu$は以下であれば良いということになります。

\min_{\mu}\mathbb{E}_{z}\left[\frac{1}{kT} \sum_{i}^{H\times W}{\|\hat{y}_i - \left(\mu_i + |\hat{y}_i - \mu_i| \ast z\right)\|_1}\right]

これを$\mu$で偏微分すると(余分な定数項は除外)

\frac{\partial \mathbb{E}_{z}\left[\sum_{i}^{H\times W}{\|\hat{y}_i - \left(\mu_i + |\hat{y}_i - \mu_i| \ast z\right)\|_1}\right]}{\partial \mu_i}

となります。これは$\hat{y}$と$\mu$との大小関係でのみ符号が入れ替わるので、勾配の符号だけ見れば$l_1$ lossで学習するのと同義です。また、$\sigma$が0に収束するということは、$\mu$と$\hat{y}$が完全に一致することを意味しますが、それはSISRの不可能決定性に反するので、過学習しない限りは起きないことになります。 論文ではこれを $l_\mathbb{E}$ loss として呼んでいます。

ここで$l_\mathbb{E}$ lossの意味合いを振り返ると、考え方としては$y$の分布の平均$\mu$を予測しつつ、同時に$\hat{y}$と$\mu$の誤差も分散として使用することを意味します。

そこで、単に上式を最適化するだけでなく、ネットワーク上では1つのブランチで$y$の分布の平均$\mu$を予測しつつ、同時に別のブランチで$\hat{y}$と$\mu$の誤差を予測し、それを$\sigma$として扱うマルチタスクラーニングとしても良さそうです。$\sigma$ブランチでは

\min_{\sigma} \| |\hat{y} - \mu| - \sigma \|_1

を求めれば良いことになります。論文ではこれをauxillary lossとして命名しています。 この$\sigma$を推論しておくメリットとして、$\sigma$の値が予測の自信度として使用できる、という点を挙げています。

実験

理論が出来上がったので後は実験して評価をします。 論文では一般的なSISRの実験と同じ設定を用いていました。 学習にはDIV2Kデータセットのtraining setを使用し、評価データセットではSet5 / Set14 / BSD100 / Urban100 / Manga109データセットを使用します。Manga109は比較的新しいデータセットなので、3,4年より前のSISR論文では評価に使用していない場合もあります。 また、Set5 / Set14は慣習的に使用しますが、枚数がそれぞれ5枚 / 14枚なので評価値としてはブレが大きい印象があります。 拡大倍率については、2 / 3 / 4倍について実験を行なっています。最近だと8倍についても実験する論文を見かけることがあります。 評価指標は輝度に対するPSNRとSSIMを使用しています。

学習の詳細についてはあまり触れませんが、入力$x$が48px x 48pxになるようにクロップしてリサイズしています。これは、SISRでは広範囲のコンテキストを必要としないという前提に基づくことが多く、他の多くの論文でもあまり大きな画像を入力に用いることはありません。 そのため、RCANのような大きいモデルであってもV100を1枚で十分学習が可能です。

最終的な推論結果としては$\mu$を使用しています。

提案手法は損失関数を調整し、また最後に$\sigma$推論用のブランチを追加するだけで実現できるため、基本的にどのようなCNNにも適用できます。そこで、論文では、VDSR / SRResNet / EDSR / RCAN といった有名どころのCNNに適用し、$l_1$ loss との比較を行なっています。

軽量なモデルに対して提案手法を適用した結果の比較(EDSR-baselineはEDSRの中で提案されている比較的軽量なモデルです) table-lightweight 重いモデルに対して提案手法を適用した結果の比較 table-complex.png

PSNR/SSIMともに上がり幅としては僅かに見えますが、CNNベースのSISRの精度は年々上がり幅が小さくなっているので、それを踏まえて見ると悪くない結果のように思えます。

quantitive.png

比較実験

提案手法の有効性を確認するために、以下の実験を行なっています。

  • 損失を $l_1$ lossだけにして学習
  • 損失を $l_\mathbb{E}$ loss だけにして学習(論文ではEq. (8)として登場)
  • 損失を auxillary loss だけにして学習(論文ではEq. (10)として登場)
  • 損失を $l_1$ loss と auxillary loss で学習
  • 損失を $l_\mathbb{E}$ loss と auxillary loss で学習

ablation-study.png

それぞれについて、ベースネットワークとしてEDSRを使用し、Set14でPSNR/SSIMを評価しています。学習時間はやはり伸びますが、PSNR/SSIMとしては提案手法が最も良い結果を示していました。

予測の自信度の可視化

推測した$\sigma$を使用してモデルの自信度を画素ごとに表示させています。 当然と言えば当然なのですが、画像のエッジ部分(高周波部分)でモデルが自信をなくしているのが分かります。

uncertainity.png

感想

$l_1$ lossは近年のCNNベースのSISRで積極的に利用されてきましたが、きちんと焦点を当てている論文は中々限られているように思っていたので新鮮でした。 一方で、最終的な推論結果を$\mu$のみから導き出していた点がやや疑問点として残りました。結果として分布の平均値をピックアップせざるを得ないのかな、という印象です。$\sigma$の結果を利用して、自信がない領域だけ生成的な手法で推測してあげる、みたいなことが出来たら面白いのかもしれないなぁ...というような未来を感じさせる論文でした。

引用

  • [1] He, Xiangyu, and Jian Cheng. "Revisiting L1 Loss in Super-Resolution: A Probabilistic View and Beyond." arXiv preprint arXiv:2201.10084 (2022).
  • [2] Dong, Chao, et al. "Image super-resolution using deep convolutional networks." TPAMI 2015.
  • [3] Shi, Wenzhe, et al. "Real-time single image and video super-resolution using an efficient sub-pixel convolutional neural network." CVPR 2016.
  • [4] Kim, Jiwon, Jung Kwon Lee, and Kyoung Mu Lee. "Accurate image super-resolution using very deep convolutional networks." CVPR 2016.
  • [5] Kim, Jiwon, Jung Kwon Lee, and Kyoung Mu Lee. "Deeply-recursive convolutional network for image super-resolution." CVPR 2016.
  • [6] Ledig, Christian, et al. "Photo-realistic single image super-resolution using a generative adversarial network." CVPR 2017.
  • [7] Lim, Bee, et al. "Enhanced deep residual networks for single image super-resolution." CVPR workshops 2017.
  • [8] Zhang, Yulun, et al. "Image super-resolution using very deep residual channel attention networks." ECCV 2018.
  • [9] Kingma, Diederik P., and Max Welling. "Auto-encoding variational bayes." ICLR 2014.
  • [10] Bruna, Joan, Pablo Sprechmann, and Yann LeCun. "Super-resolution with deep convolutional sufficient statistics." ICLR 2016.
  • [11] Lehtinen, Jaakko, et al. "Noise2Noise: Learning image restoration without clean data." ICML2018