Google對(duì)基于循環(huán)網(wǎng)絡(luò)的模型的改進(jìn)

這一兩年來，基于神經(jīng)網(wǎng)絡(luò)的圖像壓縮進(jìn)展十分迅速。2016年，基于神經(jīng)網(wǎng)絡(luò)的模型首次取得了和JPEG相當(dāng)?shù)谋憩F(xiàn)。而到了2017年，神經(jīng)網(wǎng)絡(luò)在圖像壓縮方面的表現(xiàn)已經(jīng)超過了現(xiàn)代工業(yè)標(biāo)準(zhǔn)（WebP、BPG）。本屆CVPR 2018上，就有三篇基于神經(jīng)網(wǎng)絡(luò)進(jìn)行圖像壓縮的論文，分別來自Google、蘇黎世聯(lián)邦理工學(xué)院、港理工和哈工大。下面我們將介紹這三篇論文的主要內(nèi)容。

Google對(duì)基于循環(huán)網(wǎng)絡(luò)的模型的改進(jìn)

2016年，Google的研究人員使用循環(huán)神經(jīng)網(wǎng)絡(luò)（一個(gè)混合GRU和ResNet的變體）在圖像壓縮上取得了和JPEG相當(dāng)?shù)谋憩F(xiàn)。（arXiv:1608.05148）。

而本次在CVPR 2018上提交的論文（arXiv:1703.10114），Google對(duì)之前的模型進(jìn)行了三大改進(jìn)，將其表現(xiàn)提升至超越WebP的水平。

Kodak數(shù)據(jù)集上相同MS-SSIM下的壓縮率比較，藍(lán)色為Google新提出的模型

對(duì)比JPEG2000、WebP、BPG 420

網(wǎng)絡(luò)架構(gòu)

整個(gè)網(wǎng)絡(luò)的架構(gòu)和Google之前的模型類似，如下圖所示：

上圖中，Ei為編碼器，Dj為解碼器，中間的國際象棋棋盤圖案表示binarizer（二值化輸入為1、-1）。每個(gè)網(wǎng)絡(luò)層對(duì)應(yīng)的分辨率標(biāo)識(shí)于下方，深度標(biāo)識(shí)于上方。其中的數(shù)字（I/H）分別表示輸入（I）和隱藏狀態(tài)（H）的卷積核大小。如3/1表示3x3輸入卷積，1x1隱藏卷積，1/0表示1x1（前饋）卷積。另外，前饋卷積單元（H = 0）使用tanh激活，其他層為卷積GRU層。

以上示意的是單次迭代過程。在每次迭代中，上圖中的循環(huán)自動(dòng)編碼器編碼之前的重建圖像和原始圖像之間的殘差。在每一步，網(wǎng)絡(luò)從當(dāng)前的殘差中提取新信息，接著合并到循環(huán)層的隱藏狀態(tài)中存儲(chǔ)的上下文。每次迭代后，在量化瓶頸模塊保存相應(yīng)的比特，從而生成輸入圖像的漸進(jìn)式編碼。

了解了模型的基本架構(gòu)后，下面我們來看看Google的研究人員所做的三大改進(jìn)。

隱藏狀態(tài)引火

初次迭代時(shí)，每個(gè)GRU層的隱藏狀態(tài)初始化為零。在試驗(yàn)中，Google的研究人員發(fā)現(xiàn)，起初的幾個(gè)迭代過程中，圖像質(zhì)量的提升十分明顯。因此，研究人員假設(shè)，缺乏一個(gè)良好的隱藏狀態(tài)初始化導(dǎo)致模型在早期碼率上表現(xiàn)不佳。由于編碼器和解碼器都堆疊了許多GRU網(wǎng)絡(luò)層序列，編碼器的binarizer和解碼器的重建都需要好幾次迭代才能觀察到首層GRU的隱藏狀態(tài)改進(jìn)。因此，研究人員使用了隱藏狀態(tài)引火（hidden-state priming）技術(shù)為每個(gè)GRU層生成了更好的初始隱藏狀態(tài)。

所謂隱藏狀態(tài)退火，或者叫“k-退火”（k-priming），單獨(dú)增加了編碼器和解碼器網(wǎng)絡(luò)的首次迭代的循環(huán)深度（額外增加了k步）。為了避免占用額外的帶寬，這些步驟是單獨(dú)運(yùn)行的，編碼器產(chǎn)生的額外位元并不會(huì)加入實(shí)際的碼流。對(duì)編碼器而言，這意味著多次處理原始圖像，拋棄生成的位元，不過保存編碼器循環(huán)單元隱藏狀態(tài)的變動(dòng)。對(duì)解碼器而言，這意味多次生成解碼圖像，不過僅僅保留最終的圖像重建（當(dāng)然，同時(shí)保留解碼器隱藏狀態(tài)的變動(dòng)）。

下圖為網(wǎng)絡(luò)在0.125 bpp時(shí)重建的圖像對(duì)比，可以看到，引火后的重建效果明顯好于未引火時(shí)。

左：原圖；中：未引火；右：引火

此外，我們還可以在中間的迭代過程中進(jìn)行引火，研究人員稱其為發(fā)散（diffusion）。

同樣，我們也能從視覺上直接看出發(fā)散給圖像重建帶來的質(zhì)量提升。

從左往右，依次為0-3次發(fā)散

空間自適應(yīng)碼率

如前所述，在不同的迭代中，模型生成不同碼率的圖像表示。然而，每張圖像應(yīng)用的碼率是常量，而沒有考慮圖像的內(nèi)容。實(shí)際上，圖像的不同局部的復(fù)雜程度是不一樣的，比如，圖像上方可能是晴朗的天空，而圖像下方則可能是繁復(fù)的花海。

因此，Google的研究人員引入了空間自適應(yīng)碼率（Spatially Adaptive Bit Rates，SABR），根據(jù)目標(biāo)重建質(zhì)量動(dòng)態(tài)地調(diào)整局部的碼率。

消融測(cè)試印證了SABR的效果：

左為Kodak數(shù)據(jù)集，右為Tecnick數(shù)據(jù)集

SSIM加權(quán)損失

訓(xùn)練有損壓縮圖像網(wǎng)絡(luò)可以說是左右為難。最直接的方式是直接比較和參考圖像（即原圖）像素間的差異，比如，使用經(jīng)典的L1或L2損失。然而，直接比較像素差異沒有考慮到主觀感知因素，畢竟人眼對(duì)不同類型的像素差異（失真）的敏感程度不同。而基于感知的指標(biāo)卻不可微或者條件梯度不良。

為了兼顧兩方面的需求，Google的研究人員提出了一種加權(quán)L1參數(shù)：

其中，x為參考圖像（原圖），y為fθ(x)的解壓縮圖像（θ為壓縮模型的參數(shù)）。S(x, y)為衡量圖像不相似性的感知指標(biāo)，ˉS為基線。具體來說，ˉS為S(x, y)的移動(dòng)平均。移動(dòng)平均不是常量，但是在短暫的訓(xùn)練窗口中基本可以視作常量。在Google研究人員的試驗(yàn)中，移動(dòng)平均的衰減為0.99. 然后，將w(x, y)視為固定值，這樣就可以更新梯度了。

具體而言，Google研究人員使用的S(x, y)基于結(jié)構(gòu)相似性指標(biāo)（SSIM）。研究人員首先將圖像切分為8 x 8的小塊。然后在每個(gè)小塊上使用以下方法計(jì)算局部權(quán)重：

整個(gè)圖像的損失為所有局部加權(quán)損失之和。

內(nèi)容加權(quán)圖像壓縮

之前我們提到過，Google研究人員的三大改進(jìn)之一是空間自適應(yīng)碼率（SABR），其背后的直覺是圖像的不同局部復(fù)雜程度不同，應(yīng)該分配不同的碼率。無獨(dú)有偶，香港理工大學(xué)和哈爾濱工業(yè)大學(xué)的研究人員Li、Zuo等在CVPR 2018上提交的論文“內(nèi)容加權(quán)圖像壓縮”，同樣是基于圖像局部的復(fù)雜性采用不同的碼率（arXiv:1703.10553）。

老鷹應(yīng)該比天空占用更多的碼率

SABR根據(jù)圖像的重建質(zhì)量調(diào)整碼率，使用的是啟發(fā)式的算法。而港理工和哈工大的研究人員則使用一個(gè)三層卷積網(wǎng)絡(luò)學(xué)習(xí)圖像的重要性映射（importance map），然后通過量化生成重要性掩碼（importance mask），并應(yīng)用于之后的編碼過程。

另外，模型生成的重要性映射可以適應(yīng)不同的bpp。如下圖所示，壓縮得很厲害時(shí)，重要性映射僅僅在明顯的邊緣分配更多的碼率。而隨著bpp的升高，重要性映射給紋理分配了更多碼率。這和人眼的感知是一致的。

基于上下文模型并行學(xué)習(xí)

循環(huán)網(wǎng)絡(luò)之外，基于自動(dòng)編碼器（auto encoder）的圖像壓縮模型也是一個(gè)很有希望的方向。

2017年3月，Twitter的Theis等提出了基于自動(dòng)編碼器的模型（arXiv:1703.00395），表現(xiàn)與JPEG 2000相當(dāng)。

Theis等提出的壓縮自動(dòng)編碼器架構(gòu)

2017年4月，ETHZ（蘇黎世聯(lián)邦理工學(xué)院）的Agustsson等，提出了soft-to-hard熵最小化訓(xùn)練方法（arXiv:1704.00648），改進(jìn)了上述自動(dòng)編碼器模型。2017年5月，WaveOne的Rippel和Bourdev提出的自動(dòng)編碼器架構(gòu)使用了金字塔分解（pyramidal decomposition）編碼器、自適應(yīng)算術(shù)編碼（adaptive arithmetic coding）、自適應(yīng)碼長正則化（adaptive codelength regularization），此外還使用了對(duì)抗訓(xùn)練（arXiv:1705.05823）。該模型的表現(xiàn)超越了現(xiàn)代圖像壓縮的工業(yè)標(biāo)準(zhǔn)（WebP、BPG）。

在CVPR 2018上，ETHZ的Mentzer、Agustsson等提交的論文（1801.04260），則在自動(dòng)編碼器的訓(xùn)練過程中使用一個(gè)輕量上下文模型（三維CNN）來估計(jì)潛圖像表示的熵，基于熵?fù)p失并行學(xué)習(xí)，從而更好地控制重建誤差（失真）和熵（壓縮率）之間的折衷。

在Kodak數(shù)據(jù)集上，模型的表現(xiàn)超越了現(xiàn)代工業(yè)標(biāo)準(zhǔn)，與前述WaveOne提出的模型相當(dāng)

量化

ETHZ研究人員的這項(xiàng)工作使用了之前提到的soft-to-hard熵最小化中的量化方法，不過對(duì)其進(jìn)行了簡化。研究人員使用最近鄰分配計(jì)算：

為了在反向傳播階段計(jì)算梯度，研究人員使用以下可微逼近：

注意，以上可微逼近只在反向傳播時(shí)應(yīng)用，以免還要選擇退火策略硬化逼近（軟量化）。

模型架構(gòu)

讓我們看下整個(gè)模型架構(gòu)的示意圖。

示意圖上部為編碼器，下部為解碼器。深灰色塊表示殘差單元。編碼器中，k5 n64-2表示核大小5、輸出頻道64、步長2的卷積層，其他卷積層同理。相應(yīng)地，在解碼器中，它表示反卷積層。所有卷積層使用batch norm和SAME補(bǔ)齊。Normalize表示將輸入歸一化至[0, 1]，歸一化基于訓(xùn)練集的一個(gè)子集的均值和方差。Denormalize為其逆操作。Masked quantization（掩碼量化）采用了之前提到過的重要性映射，不過，ETHZ簡化了重要性映射的生成方法，沒有使用一個(gè)單獨(dú)的網(wǎng)絡(luò)，相反，直接在編碼器的最后一層增加了一個(gè)額外的單頻道輸出y作為重要性映射，之后將其轉(zhuǎn)換為掩碼：

其中，yi,j表示空間位置(i,j)處y的值。k值的選取需滿足掩碼轉(zhuǎn)換在0到1之間平滑過渡。

下圖可視化了自動(dòng)編碼器的潛表示，可以看到重要性映射的效果：

M：加入重要性映射；M'：未加入重要性映射

整個(gè)訓(xùn)練過程如下：

從編碼器E得到壓縮（潛）表示z和重要性映射y：(z,y) = E(x)

使用剛剛提到的公式轉(zhuǎn)換重要性映射y至掩碼m

使用分素相乘掩碼z。

量化（Q）。

使用四層的三維CNN網(wǎng)絡(luò)計(jì)算上下文（P）。

解碼（D）。

訓(xùn)練過程中為自動(dòng)編碼器（E、D）和量化器（Q）計(jì)算如下?lián)p失函數(shù)（折衷碼率和失真）：

相應(yīng)的上下文模型P的損失函數(shù)為：

其中，C為潛圖像表示的編碼代價(jià)：

類似地，MC為掩碼編碼代價(jià)：

各個(gè)模型均可在GPU上并行訓(xùn)練。

非現(xiàn)實(shí)圖像

ETHZ研究人員順便測(cè)試了模型在非現(xiàn)實(shí)圖像上的表現(xiàn)。

Manga109數(shù)據(jù)集樣本

如上圖所示，BPG壓縮的黑白漫畫，文字更銳利，而ETHZ研究人員新提出的模型則保留了更多臉部的細(xì)微紋理。

CLIC挑戰(zhàn)

另外，CVPR 2018還舉辦了一場(chǎng)學(xué)習(xí)圖像壓縮挑戰(zhàn)（CLIC），以鼓勵(lì)這一領(lǐng)域的進(jìn)展。

有三個(gè)團(tuán)隊(duì)在CLIC取得優(yōu)勝，其中來自圖鴨科技的TucodecTNGcnn4p在MOS和MS-SSIM兩項(xiàng)指標(biāo)上均獲第一。

TucodecTNGcnn4p基于端到端的深度學(xué)習(xí)算法，其中使用了層次特征融合的網(wǎng)絡(luò)結(jié)構(gòu)，以及新的量化方式、碼字估計(jì)技術(shù)。