はじめに

この記事では、深層学習を用いた画像の圧縮について紹介します。
以前も似た記事を書いたのですが、今回はそれをさらに簡単にしたものです。できるだけ数式を少なく、事前知識を少なく理解できるように努めました。

ニューラルネットワークの画像処理

まずは一般に、ニューラルネットワークを用いた画像処理について話しておきます。2010年代に AlexNet が画像分類タスクで優れた性能を残したことで、画像処理には 畳み込みニューラルネットワーク(CNN) が広く用いられるようになりました。

画像は縦横とRGBの3つの次元を持ったデータです。RGBの次元を一般化して チャンネル と呼びます。CNNは、画像のような形の[チャンネル、縦、横] のデータを受け取って、 フィルタ や カーネル という窓を使って画像を走査する処理を行います。

(https://neurohackademy.github.io/convolutional-neural-networks/03-convolutions/)

例えば分類タスクの場合、この畳み込みの処理を何層も重ねて、最終的には画像より小さいサイズの 特徴ベクトル を出力します。これに線形層やSoftMaxによる後処理をすると、各クラスの確率が現れます。

どうやって圧縮をするのか？

画像圧縮は、分類に比べるとかなり面倒なタスクです。というのも、入力は画像ですが、出力は 0000111100101.... というビット列だからです。

画像圧縮も分類と同じで、最初は画像を特徴量ベクトルに変換するところから始めます。特徴量ベクトルは連続値の数値データをもったベクトルです。
数値データをビット列に変換するには、 符号化 という処理を行います。ランレングス符号化やハフマン符号化が有名です。これらの符号化手法は、離散値のデータを受け取り、その出現頻度に合わせて符号化を行います。なので、 量子化 という四捨五入のような処理をして、連続値の特徴ベクトルを離散値のデータに変換します。

ここまでの流れをまとめます。
圧縮のためには、CNNを用いて画像を特徴量ベクトルに変換し、その後に量子化を行います。これで離散値のデータを手に入れたので、圧縮ができます。ランレングス符号化などの手法なら、このまま圧縮することが可能です。

（なお、量子化をすると勾配がゼロになるので、実はもう一段階工夫が必要です。この記事ではスキップします）

圧縮したデータを解凍した後は、特徴量ベクトルから画像に戻すような変換をCNNで作り、学習させればよいです。これでシンプルな画像圧縮機を作ることができました！

pytorch風の擬似コードを書くと次のようになります。

class SimpleCompressor():
    def __init__(self):
        self.cnn_a = CNN() #[3,H,W] -> [c,h,w] に変換する、何らかのニューラルネット
        self.cnn_s = CNN() #[c,h,w] -> [3,H,W] に変換する、何らかのニューラルネット
    def forward(self, x):
        y = self.cnn_a(x) #特徴量ベクトル
        y_hat = quantize(y) #量子化
        x_hat = self.cnn_b(y_hat) #画像を復元
        return x_hat
    def compress(self, x)
        y = self.cnn_a(x) #特徴量ベクトル
        y_hat = quantize(y) #量子化
        binary = ranlength_encoding(y_hat) #符号化
        return binary
    def decompress(self, binary)
        y_hat = runlength_decoding(binary) #逆符号化
        x_hat = cnn_s(y_hat) #画像を復元
        return x_hat

すこし記法の説明をします。入力画像は $x$, 特徴量ベクトルは $y$ で表します。y_hat は $\hat{y}$ のことで、「量子化されたもの」を意味しています。
訓練時は forward 関数を用いて、圧縮や解凍をするときは compress decompressを用います。

エントロピー符号化を用いる

ただし、これだとあまり性能がよくありません。
すこし情報理論の話をします。実は、データを圧縮するときに、どれだけの圧縮率が達成できるかは エントロピー というものを用いて計算することができます。 シャノンの情報原符号化定理 と呼ばれます。

理想的な符号化は、このエントロピーの下限に近い値で圧縮をするものです。そのためには、特徴量ベクトルの各シンボルの出現確率がわかっていないといけません。どうすればよいでしょうか？

手法1: 頻度を数える

特徴量ベクトルのシンボルを、1が10個、0が20個……と数えていけばよいと思うかもしれません。しかし、そうやって頻度情報を用いて圧縮をしても、復号側はその情報を持っていません。また、この手法はニューラルネットとして学習するにはあまり適していません。

手法2: 分布を固定する

次に思いつくのは、適当な分布を仮定する方法です。例えば、[c,h,w] サイズの特徴量ベクトルの各ピクセルが、標準正規分布 $\mathcal{N}(0,1)$ にi.i.d. に従っているとしたらどうでしょう？
こうすると、符号化側と復号側で分布を共有できる上、ニューラルネットとして学習することもできます。実際の各シンボルと正規分布を用いてエントロピーを計算し、これが小さくなるように学習できるためです。

手法3: 分布を学習してから固定する

上で言った「ピクセルごとに分布を決める」のは良いアイデアですが、分布を固定するのは今ひとつです。実際にデータが標準正規分布に従っている保証は特にありません。

どちらかというと、たくさんのデータを食べさせてみて、これが良さそうだという分布をひとつ固定した上で、これを使って圧縮・解凍を行うのがよさそうです。

ということで、ニューラルネットワークを用いて分布を作ることを考えます。とある都合で、今回は確率密度関数(pdf) ではなく累積分布関数(cdf) を作ることにします。
入力は特徴量ベクトル $y$ で、出力は同じサイズの累積分布関数です。

エントロピー符号化を用いた圧縮モデル

では、実際に疑似コードを書いて作ってみます。

class EntropyModelCompressor():
    def __init__(self):
        self.cnn_a = CNN() #[3,H,W] -> [c,h,w] に変換する、何らかのニューラルネット
        self.cnn_s = CNN() #[c,h,w] -> [3,H,W] に変換する、何らかのニューラルネット
        self.entropy_model = EntropyModel() #CDFを生成するニューラルネット
    def forward(self, x):
        y = self.cnn_a(x) #特徴量ベクトル
        y_hat = quantize(y) #量子化
        
        cdf = entropy_model.generate() #cdfを生成 (NEW!)
        likelihoods = calculate_entropy(cdf, y) #エントロピーを計算 (NEW!)
        
        x_hat = self.cnn_b(y_hat) #画像を復元
        return x_hat, likelihoods
    def compress(self, x)
        y = self.cnn_a(x) #特徴量ベクトル
        y_hat = quantize(y) #量子化
        cdf = entropy_model.generate()
        binary = arithmetic_encoding(y_hat, cdf) #符号化
        return binary
    def decompress(self, binary)
        cdf = entropy_model.generate()
        y_hat = arithmetic_decoding(binary, cdf) #逆符号化
        x_hat = cnn_s(y_hat) #画像を復元
        return x_hat

先ほどとは違って、ランレングス符号化ではなくエントロピー符号化を用いています。そのため、 arithmetic_encoding に渡す引数には cdf が増えています。
さて、今私たちが作った EntropyModel には、良さそうな分布を学習してもらう必要があります。そのためにはロス関数が必要です。どんなロス関数に最適化させればよいでしょうか？

答えは エントロピー です。シャノンの情報源符号化定理より、圧縮されて出てくるビットの下限はエントロピーです。そして、算術符号などのエントロピー符号化は、おおむねこのエントロピーに近い値で圧縮をしてくれます。

ということで、エントロピーができるだけ小さくなるようにすればよいです。詳細には触れませんが、これはちゃんと微分可能なかたちで計算することが可能です。私たちは再構成画像が正確になるように、同時に分布のエントロピーが低くなるようにモデルを訓練すればよいです。

おわりに

実はこれで、LICの最初の論文の本質的なアイデアは説明し終わったことになります。 End-to-end Optimized Image Compressinon という2017年の論文です。
このシンプルな手法の性能は普段使っているJPEGより上で、さらにその進化版のJPEG2000も超えることがわかっています。

現在に至るまでより高性能なLIC手法が開発され続け、いまでは従来の（非深層学習の）画像圧縮手法のうち最高性能のもの、VVC-intraと呼ばれる手法も上回るようになっています。

加えて、最近になってJPEG-AIという規格の標準化も進められています。画像を保存する時に、JPEG, webp, HEICなどに並んでAIベースの拡張子が現れることもそう遠くないかもしれません。

この記事で省いた箇所や、より最新の手法に興味のある向きはもう一つの記事もチェックしてみて下さい。