この記事からさらに予備知識を減らしたはじめての深層学習画像圧縮という記事を書きました。こちらの記事は実際にコードを動かしたり、論文を細かいところまで理解したい方向けです。

はじめに

この記事では、深層学習を用いた画像圧縮 (Learned Image Compression, LIC）について、基本的な原理を紹介します。

LICは、畳み込みニューラルネットワークやTransformerを用いて画像をより小さい潜在表現に変換し、各要素の出現確率をモデル化することで、高効率の画像圧縮を実現する技術です。

よく使われるLICのライブラリには、PytorchベースのCompressAIと、tensorflowベースのtensorflow/compressionが存在します。今回はCompressAIのコードを一通り理解することを目標に、LICの基本的論文、End-to-end Optimized Image Compression (Ballé et al. 2016. https://arxiv.org/abs/1611.01704) と Variational Image Compression with Scale Hyperprior (Ballé et al. 2018. https://arxiv.org/abs/1802.01436) の発想やアーキテクチャについて、できるだけ詳しく説明します。

前提知識① 変分オートエンコーダ(VAE)

変分オートエンコーダ(Variational Autoencoder; VAE) とは、Kingma & Welling 2013 (https://arxiv.org/abs/1312.6114) によって提唱されたニューラルネットワークの構造です。詳しい説明は他の記事などに譲るので、ここではざっくりとした説明と要点にとどめます。オートエンコーダ(AE)は潜在変数モデルを用いたアーキテクチャで、入力画像（今回は画像に限るものとします）をCNNに通すことで、その裏側にある何らかの潜在変数(latent variable)を求めることが目標です。
VAEでは、この潜在変数はさらに確率分布としてモデル化されています。具体的には各潜在変数が$\mathcal{N}(\mu, \sigma^2)$という正規分布としてモデル化され、エンコーダはこの$\mu$と$\sigma$の値を直接出力します。

VAEの構造。(EuginioTLの作品。CC BY-SA 4.0）

大事なのは、VAEを学習させることで、画像からある潜在変数（または分布）を求めることができて、デコーダはこの変数からもとの画像を復元できるという点です。
ならば、できるだけ内側の潜在変数のサイズを小さくしたら、画像を圧縮できるのではないか？という発想に至ります。

前提知識②　エントロピー符号化

何らかのシンボルの列 $\boldsymbol{y} = \lbrace y_1, y_2, \cdots y_n\rbrace$ があって、さらに各シンボルの出現確率$p(y_1), p(y_2), \cdots p(y_n)$ があるとする。これをできるだけ短い10110...というビット列に変換して、逆に元に戻せるようにしたい。

この操作をエントロピー符号化と呼びます。シャノンの情報源符号化定理より、このときに達成できるビット長の下限はシンボル列のエントロピーに等しいです。実際に算術符号化やハフマン符号化を用いればこの限界ビット長に近い値まで圧縮ができます。

量子化と符号化

従来の画像圧縮では、フーリエ変換（または離散コサイン変換やウェーブレット変換）とその逆変換を用いて、画像を周波数領域に移す→量子化→符号化、という流れで画像を圧縮します。VAEなどのオートエンコーダは、従来のネットワークの各種変換部分にあたります。

ということで、ニューラルネットワークで変換をした後には、量子化と符号化をしないといけません。画像圧縮のロスの多くはここに由来します。
ここで二つの問題があります。

量子化をする（例えばtorch.round)と、その時点で勾配が消失する
符号化のために、量子化された潜在変数ベクトル $\hat{\boldsymbol{y}}$の各要素$\hat{y}_i$の出現分布$p(\hat{y}_i)$を求めないといけない
しかも、この分布は送信側（符号化側）と受信側（復号化側）で共有していないといけない

なので、(1)勾配が途切れないような量子化をする、(2)ネットワークで各シンボルの出現確率をモデル化する、という問題が現れます。Balle 2016 はこの二つの問題を解決しました。

Balle 2016のアイデア……ノイズ付与による量子化 + Factorized Entropy Model

1. ノイズ付与による量子化

Balle 2016では、量子化の部分に訓練時のみノイズを付加するというアイデアが提案されました。

\begin{align}
\boldsymbol{u} &\sim \mathcal{U}\left( -\frac{1}{2}, \frac{1}{2}\right) \\
\tilde{\boldsymbol{y}} &= \boldsymbol{y} + \boldsymbol{u}
\end{align}

ただし、 $\mathcal{U}$ は一様分布です。これは、マイナス0.5からプラス0.5までのノイズを加えることを意味します。

量子化を行うと、最大で0.5だけ値が変化します。この代わりに、0.5だけ値が変化するような一様なノイズを加えることで、値の変化に対しても出力が一定であるようにします。訓練時は量子化をせずにノイズ付与だけを行うので、勾配を消失させずに訓練ができます。

2. Factorized Entropy Model

次に、yの生起確率を考えます。これは最終的に符号化機に渡すので、例えばyが1,2,3の値を取るなら、$p(y=1) =0.5, p(y=2) = 0.2, p(y=3) = 0.3$のようになっていないといけません。ここで、Balle 2016 は、CNNを使って生起確率をモデル化するような構造を考えました。すなわち、あるネットワーク$f(\boldsymbol{w})$を用いて、$p(\tilde{y}) = f(\boldsymbol{w})$ とします。生起確率はyの値には依存しないように設計されています。なので、どのような入力に対しても同じ確率分布を返します。このため、入力と出力で確率分布を共有することができます。

生起確率はもちろん、実際の出現確率にできるだけ近い必要があります。実際の値と離れた確率でも符号化はできますが、圧縮率が落ちます。ここをうまく近似できるように、ネットワークを使って訓練させます。

Balle 2016 が用いた確率モデルはCompressAIに実装がなく、おそらくあまり使われていないので、かわりにそれを改良したBalle 2018 のモデルを紹介します。これはEntropy Bottleneck layerという名前でCompressAIに実装されています。（compressai.entropy_models.EntropyBottleneck)

確率モデルなので本来は$y$のpdf（確率密度関数）を求めるところですが、代わりにcdf（累積分布関数）を使います。この理由はあとで説明します。

さて、累積分布関数 $c(\cdot)$ は、次の条件を満たさないといけません。

c(-\infty) = 0, \space c(\infty) = 1, \space p(x) = \frac{\partial c(x)}{\partial x}

まずは、$c(\cdot)$をいくつかの関数の合成として定義します。

c(x) = f_1\circ f_2\circ ... \circ f_K\ (x)

その上で、各$f_k$を次のようなアフィン変換型の関数とします。

\begin{align*}
f_k(\boldsymbol{x}) =g_k(\boldsymbol{H}^{(k)}\boldsymbol{x} + \boldsymbol{b}^{(k)}) &&  1\leq k< K
\\
f_K(\boldsymbol{x}) =\text{sigmoid}(\boldsymbol{H}^{(K)}\boldsymbol{x} + \boldsymbol{b}^{(K)})
\\
g_k(\boldsymbol{x}) = \boldsymbol{x} + \boldsymbol{a}^{(k)}\odot\tanh(\boldsymbol{x})
\end{align*}

出力は[0,1]の間でないといけないので、最終層にだけシグモイドを使います。ここで現れたパラメータ、$\boldsymbol{H}, \boldsymbol{b}, \boldsymbol{a}$ はこの層が持つ学習可能なパラメータです。¹

さて、これによって各シンボルの出現確率をモデル化できました。送信側と受信側に同じ確率モデルを持たせておけば、問題なく符号化を行うことができます。

ところで、なぜcdfをモデル化するのでしょうか。これを考えるために、まず$y_i \sim p_{y_i}(y_i)$とします。$u \sim \mathcal{U}(-0.5,0.5)$より、$\tilde{y} = y + u$はふたつの独立な確率分布に従う確率変数の足し合わせになります。このとき、$\tilde{y}$が従う分布は二つの分布の畳み込み演算の結果です。すなわち、$y$の従う分布を$p(y)$としたときに、

\begin{align*}
p(\tilde{y}) &= \int_\mathbb{R} p(y)\cdot \mathcal{U}(\tilde{y} - y | -0.5, 0.5)dy \\
&= \int_{\tilde{y}-0.5}^{\tilde{y} + 0.5} p(y)dy \\
&= \left[c(y)\right]^{\tilde{y} + 0.5}_{\tilde{y}-0.5}\\
&= c(\tilde{y}_i+0.5)- c(\tilde{y}_i - 0.5)
\end{align*}

最終的に、$\tilde{y}$の従う分布の確率密度関数（すなわち$\tilde{y}$の生起確率）は、$c(\tilde{y}_i+0.5)- c(\tilde{y}_i - 0.5)$ ただし$c$は$p(y)$の累積分布関数(cdf)となります。先ほどpdfではなく直接cdfをモデル化したのはこのためです。

これによって、各要素${\tilde{y}_i}$ の出現確率がわかりました。すべての積を取ることで、全体としての同時密度関数を求めることができます。これを用いてエントロピーを計算することができます。

p_{\boldsymbol{\tilde{y}}}(\boldsymbol{\tilde{y}}) = \prod_i p_{\tilde{y_i}}(\tilde{y_i})

さて、これで訓練の様子については説明し終わりました。各$\tilde{y_i}$に対して、その確率密度関数$p(\tilde{y_i})$を求めて、そこから全体のエントロピーを計算します。実際に圧縮をしたらエントロピーに近づくはずなので、このエントロピーが小さくなるように、なおかつ画像の歪みも小さくなるように訓練をすればOKです。

実際にはエントロピーは$E[-\log p_{\boldsymbol{\tilde{y}}}(\boldsymbol{\tilde{y}})]$ で表現されます。各要素に対して出現確率のlogを取り、最後に足し合わせることでエントロピーが求められます。これを評価関数の片方として使います。

圧縮

さて、それでは圧縮と解凍の話に移りましょう。訓練と圧縮／解凍の違いは、符号化を行うか否かです。エントロピー符号化は可逆変換なので、訓練時はこれをスキップしてしまいます。したがって各要素の出現確率を計算する必要はなく、全体のエントロピーだけ分かれば問題ありません。

ところが、今回は実際に符号化を行います。よって、符号化機(ハフマン符号化、算術符号化、レンジコーダー、rANSなど、何でも構いません)に各シンボルの出現確率を渡す必要があります。

ここで思い出して欲しいのは、最終的に$p_{\boldsymbol{\tilde{y}}}$に入る値は$\tilde{y}$ ではなく $\hat{y}$、つまり量子化された整数値だということです。なので、今まで連続型確率変数としてpdfとcdfの議論をしてましたが、これをpmf（確率質量関数）にしても問題ないことになります。

ここで、pdfをpmfに変換する場合、考慮しないといけないこと一つあります。それは最大値と最小値です。pdfは全区間の積分値が1になればよいですが、pmfはある区間[min, max] で出現確率の和が1になっていないといけません。

CompressAIのEntropy Bottleneckの実装では、分位数(Quantile)というパラメータを設定しておいて、チャンネル毎に符号化の際の最大値と最小値を求めます。初期値では[-10,10] に設定されています。この区間からはみ出た値はおそらく消えてしまいますが、yが数マス消えたぐらいならデコーダのCNNが空気を読んで復元してくれるので、厳密にすべての値が符号化されるように設定しなくてもよいのだと思います。

ここまででEntropy Bottleneck の構造を説明し終わりました。この層は学習可能なパラメータとしてH,b,fという値を持ちます。
compressai.entropy_models.EntropyBottleneck には以下のような関数が定義されています。

likelihood関数では、H, b, fによってyの値からその確率モデルを推定し、その尤度を返します。
quantize関数は訓練中は一様分布に従うノイズを加え、圧縮時は量子化します。(圧縮の時は勾配を更新しないので、勾配が切れても良いです)
訓練時に使うforward 関数では、quantizeとlikelihoodの2つの関数を呼んで、H, b, aの値を誤差逆伝播によって更新します。
update関数は、各値の出現確率を計算し、pmfを生成して保存しておきます
compress関数はpmfとyの値を符号化機に渡し、decompress関数は符号化後のビット列とpmfを使ってyを復元します。

Entropy Bottleneckから Hyperprior + Gaussian Conditionalへ

Balle(2016)のモデルはすでにJPEGを超える性能ですが、まだ改善の余地がありました。それは各要素$y_i$が従う分布を、入力に関わらず事前に推定して、なおかつ同じチャンネルの全要素に同じ分布を仮定していることです。

どんな要素がどれだけの割合で出現するか、というのは画像によって違うはずです。なおかつ、画像の隣接要素には相関があり、それをCNNに通したあとの潜在表現の各隣接要素にも相関があることがわかっています。これを使えば、もっと正確な出現確率のモデル化ができそうです。

ただし、これによって今までの方法は使えなくなります。各発生確率に相関があると、先程用いた方法が使えません。1次元の正規分布なら計算できますが、多次元正規分布など、高次元の分布を用いると、そのcdfを計算するのは容易ではありません。²

ここで、Balle 2018 が提唱したのがhyperprior + gaussian conditonalモデルです。今まで用いていたFactorized Priorモデルでは、$y_i$の従う分布をH, b, a といった変数でモデル化しました。これに対して、gaussian conditionalモデルは、各要素$y_i$ の分布を正規分布$\mathcal{N}(\mu_i, \sigma_i^2)$ としてモデル化します。

しかし、これによって先程の方法は使えなくなります。送信側はこの分布にしたがって圧縮ができても、受信側はこの分布の情報を持っていません。なので、どうにかして受信側にパラメータの情報を送ってあげる必要があります。

それではパラメータ$\Phi_i = (\mu_i, \sigma_i)$はどう予測すればよいでしょうか。この部分の予測に使われるのがHyperpriorです。

Hyperpriorは、「yの値から正規分布のパラメータを出力する」ネットワーク構造です。次の図のように、今までのネットワークの右に一本パスを増やす形です。

ここで、hyperpriorの情報$z$は副情報(sub-information) として、yと一緒に受信側に送られます。yと比べてサイズが小さいので、hyperpriorの分のビットレートの増加よりも、効果的な圧縮による減少効果のほうが勝ちます。

Entropy Bottleneckでは、$y$の出現確率のpmfをモデル内に保存しておけますが、Hyperpriorは各$y_i$に異なる分布を仮定するので、これができません。分布の情報がないと符号化機は復号ができないので、先に$z$を復号したのちに、その情報を用いて$\Phi_i = (\mu_i, \sigma_i)$を推定します。

また、実際の実装では、zが送信側でも一度圧縮・解凍されています。これは、$\Phi_i = (\mu_i, \sigma_i)$が送信側と受信側で完全に同じである必要がある（この値が違うと最終的に生成されるpmfが違ってしまい、正しく復号ができない）ためです。

以上をまとめると、次のようになります。
まず、Balle 2016と同じように、入力画像$x$をエンコーダ $g_a(x)$ に通して、潜在表現$y = g_a(x)$ を得ます。ここで各$y_i$の従う分布は正規分布$\mathcal{N}(\mu_i, \sigma_i^2)$です。このパラメータを推定するために、Hyperprior Network である$h_a(y)$と $h_s(z)$ を用いて、$ (\mu_i, \sigma_i) = h_s(Q(h_a(y_i + u)))$とします。ただし$Q$は量子化を指します。

これによって$\mu_i, \sigma_i$の値がわかるので、この情報を用いて$\tilde{y} = y + u $の発生確率を計算します。$p(\tilde{y}_i) = c(\tilde{y}_i+ 0.5) - c(\tilde{y}_i - 0.5)$となって、訓練時はこれを用いてエントロピー$\mathcal{R} = E[-\log p(\boldsymbol{y})]$と、誤差項$\mathcal{D} = d(\boldsymbol{x}, g_s(Q(g_a(\boldsymbol{x})))$($d$はMSEなどの関数)を計算し、評価関数$\mathcal{R} + \lambda \mathcal{D}$ を最適化します。$\lambda$はハイパーパラメータで、これが大きいほど歪みが増える代わりに圧縮率が下がり、これが小さいと圧縮率が上がる代わりに歪みが増加します。

CompressAIを用いた実験

この二つのモデルはCompressAIというパッケージに実装されています。https://github.com/InterDigitalInc/CompressAI/tree/master/compressai

訓練済みモデルや、訓練と評価用のコード、あとはデモなども公開してくれています。試しに今の二つのモデルを走らせて性能を比較してみましょう。まずはcompressAIをインストールします。

pip install compressai

もろもろのパッケージをimportして、画像を読み込みます。ネットワークに入力するために、torchvision.transformを用いてtorch.tensorに変換します。

import PIL
import torch
from torchvision import transforms
from IPython.display import display
import numpy as np
import math
import matplotlib.pyplot as plt
img = PIL.Image.open("/path/to/your/image/kodak/kodim05.png").convert("RGB")
x = transforms.ToTensor()(img).unsqueeze(0)

モデルはcompressai.zooから読み込むことができます。

Factorized_model = compressai.zoo.bmshj2018_factorized(quality = 4, pretrained = True)
Hyperprior_model = compressai.zoo.bmshj2018_hyperprior(quality = 3, pretrained = True)
f_out = Factorized_model(x)
h_out = Hyperprior_model(x)

いくつか評価指標を計算するための関数を用意します。PSNRはMSE（平均2乗誤差)をもとにした指標で、高いほど元の画像との誤差が小さいことを表します。

def compute_psnr(a, b):
    mse = torch.mean((a - b)**2).item()
    return -10 * math.log10(mse)

def compute_bpp(out_net):
    size = out_net['x_hat'].size()
    num_pixels = size[0] * size[2] * size[3]
    return sum(torch.log(likelihoods).sum() / (-math.log(2) * num_pixels)
              for likelihoods in out_net['likelihoods'].values()).item()

%matplotlib inline
fix, axes = plt.subplots(1, 3, figsize=(16, 12))
for ax in axes:
    ax.axis('off')
    
axes[0].imshow(img)
axes[0].title.set_text('Original')

axes[1].imshow(transforms.ToPILImage()(f_out["x_hat"].clamp_(0, 1).squeeze().cpu()))
axes[1].title.set_text(f'Factorized Network Result, \n PSNR = {compute_psnr(x, f_out["x_hat"]):3f}, \
    bpp = {compute_bpp(f_out):3f}')

axes[2].imshow(transforms.ToPILImage()(h_out["x_hat"].clamp_(0, 1).squeeze().cpu()))
axes[2].title.set_text(f'Hyperprior Network Result, \n PSNR = {compute_psnr(x, h_out["x_hat"]):3f}, \
    bpp = {compute_bpp(h_out):3f}')

plt.show()

結果として、Hyperpriorを用いたモデルの方がPSNRが向上していることが確かめられました。

TODO

図版をもっと増やす
Gaussian Conditional について詳しく説明する

これらはすべて微分可能な関数なので、逆伝播も定義することができます。また、実際は他の条件を満たすために$a$にはtanh関数を、$H$にはsoftplus関数をかけています。元論文を参照してください。 ↩
のちに発表されたモデルではベクトル量子化と多次元正規分布を用いたものもあります。Zhu et al. Unified Multivariate Gaussian Mixture for Efficient Neural Image Compression. CVPR 2022. ↩