END-TO-END OPTIMIZED IMAGE COMPRESSION 論文閱讀筆記

作者 Johannes Ballé    Google Scientist

論文下載下傳    ​​https://arxiv.org/abs/1611.01704​​

論文首頁（包括源碼、資料集、可視化結果等）    ​​http://www.cns.nyu.edu/~lcv/iclr2017/​​

論文代碼（訓練、測試代碼，其他預訓練模型）    ​​https://github.com/tensorflow/compres​​

論文評論（其他人和作者的評論等） ​​End-to-end Optimized Image Compression | OpenReview​​

作者本人在PCS2108上做的口頭彙報   ​​https://www.youtube.com/watch?v=x_q7cZviXkY​​

​

目錄

​​0 總述​​

​​1 關于GDN​​

​​1.1 為什麼要歸一化​​

​​1.2 常用歸一化方法舉例——BN​​

​​1.2.1 BN層定義​​

​​1.2.2 BN層優點​​

​​1.2.3 BN層缺點​​

​​1.3 GDN的定義和優點​​

​​2 模型優化​​

​​2.1 優化目标函數​​

​​ 2.2 變換子產品​​

​​2.3 量化子產品​​

​​ 2.4 熵編碼子產品​​

​​3 實驗結果​​

0 總述

論文提出一個基于CNN的端到端圖像壓縮架構，它由使用非線性分析變換（編碼器）、均勻量化器（多進制四舍五入）和非線性綜合變換（解碼器）組成。最終RD性能和MS-SSIM都比JPEG和JPEG-2000好。這個架構是後來被大家廣泛使用的超先驗（Hyperprior）架構的基礎。

論文提出的架構如下圖，這是一種通用的非線性變換編碼架構。編碼過程：原圖x，通過參數為Φ的分析變換y=ga(x；Φ)，映射到編碼空間（code space）中量化得離散值的向量q。碼率R是量化後的q直接估計的，R不小于q的熵。

這個架構在Balle2016年ICLR的論文End-to-end optimization of nonlinear transform codes for perceptual quality就提出來了。本文中，作者對此進行了進一步介紹，給出了該架構下的一個實際模型，并給出了充足的理論分析和實驗來驗證該模型的有效性。在此模型中，分析變換（編碼器）共3層，每層用了卷積層+下采樣+GDN的結構優化MSE，如下圖左；而綜合變換（解碼器）采用分析變換的近似逆結構，也是3層不過反着來，每層IGND+上采樣+卷積，如下圖右。量化使用的是均勻的标量量化。

​

 對于RD曲線上的每個點，都直接使用梯度下降優化方法同時聯合優化分析綜合變換。對于均勻标量量化子產品梯度幾乎處處為 000 的問題，作者給出了一種連續的近似均勻标量量化來替代原來的離散均勻量化。最後，作者對量化後的資料使用了一個熵編碼得到實際的比特流，進而計算該模型實際的碼率，以此來驗證該模型的有效性。​​[1]​​

1 關于GDN

其中使用的GDN是作者2015年DENSITY MODELING OF IMAGES USING A GENERALIZED NORMALIZATION TRANSFORMATION提出來的（論文原文看了一眼，翻譯成中文我也看不懂）。GDN可以看做是更适合圖像重建的BN層。GDN在這篇文章被使用後，因為效果不錯，做end-to-end圖像壓縮的也都跟着用。

之前大多數壓縮方法都是建立在正交線性變換基礎上，目的是為了降低資料間的相關性，進而簡化後續的熵編碼。但是線性變換輸出之間的聯合統計特性展現了很強的高階依賴。這個問題可以通過使用局部增益控制操作（最早來源于生物神經元領域）來很好地解決，于是作者使用了他自己之前提出的 GDN 模型來替代線性變換（在那篇文章中作者已經驗證了GDN具有很好的高斯化圖像資料的能力）。​​[1]​​

​​​​​​​1.1 為什麼要歸一化

神經網絡學習過程的本質就是為了學習資料分布，如果我們沒有做歸一化處理，那麼每一批次訓練資料的分布不一樣，從大的方向上看，神經網絡則需要在這多個分布中找到平衡點，從小的方向上看，由于每層網絡輸入資料分布在不斷變化，這也會導緻每層網絡在找平衡點，顯然，神經網絡就很難收斂了。當然，如果我們隻是對輸入的資料進行歸一化處理（比如将輸入的圖像除以255，将其歸到0到1之間），隻能保證輸入層資料分布是一樣的，并不能保證每層網絡輸入資料分布是一樣的，是以也需要在神經網絡的中間層加入歸一化處理。​​[2]​​

歸一化後有什麼好處呢？原因在于神經網絡學習過程本質就是為了學習資料分布，一旦訓練資料與測試資料的分布不同，那麼網絡的泛化能力也大大降低；另外一方面，一旦每批訓練資料的分布各不相同(batch 梯度下降)，那麼網絡就要在每次疊代都去學習适應不同的分布，這樣将會大大降低網絡的訓練速度，這也正是為什麼我們需要對資料都要做一個歸一化預處理的原因。

對于深度網絡的訓練是一個複雜的過程，隻要網絡的前面幾層發生微小的改變，那麼後面幾層就會被累積放大下去。一旦網絡某一層的輸入資料的分布發生改變，那麼這一層網絡就需要去适應學習這個新的資料分布，是以如果訓練過程中，訓練資料的分布一直在發生變化，那麼将會影響網絡的訓練速度。​​[3]​​

1.2 常用歸一化方法舉例——BN

1.2.1 BN層定義

與激活函數層、卷積層、全連接配接層、池化層一樣， BN (Batch Normalization) 層也屬于CNN中網絡的一層。BN本質原理上是将資料歸一化至：均值0、方差為1；BN層是可學習的，參數為γ、β的網絡層​​[4]​​​。BN已經成為許多先進深度學習模型的重要組成部分，特别是在計算機視覺領域。它通過批進行中計算的平均值和方差來規範化層輸入，是以得名​​[5]​​。

BN是針對一個batch裡的資料進行歸一化和尺度操作，且一旦神經網絡訓練完成後，BN的尺度參數也固定了，這就是一個完全變成一個關于歸一化資料的仿射變換，其數學公式描述如下：

​​​​​​​1.2.2 BN層優點

BN層歸一化資料接近于高斯分布，解決了訓練時候中間層資料分布發生改變的問題。而且BN層允許選擇更大的學習率，可以加速訓練​​[6]​​​。此外，BN解決了深度梯度消失和梯度爆炸的問題；BN還可以改善正則化政策：作為正則化的一種形式，輕微減少了對dropout的需求，不用管過拟合中drop out、L2正則項參數的選擇問題​​[7]​​。

​​​​​​​1.2.3 BN層缺點

BN層每次處理的批量的均值和标準差都不同，是以這相當于加入了噪聲，增強了模型的泛化能力。但對于圖像超分辨率重建、圖像生成、圖像去噪和圖像壓縮等生成模型，就不友好了，生成的圖像要求盡可能清晰，不應該引入噪聲​​[8]​​。

當batch size越小，BN的表現效果也越不好（好像最好batch>32），因為計算過程中所得均值和方差不能代表全局​​[9]​​。

注：後針對這個問題，何恺明團隊提出了GN（group normalization），先把通道C分成G組，然後單獨拿出來歸一化，最後把G組歸一化之後的資料合并成CHW，解決了batch很小時BN效果不好的問題​​[10]​​。

​​​​​​​1.3 GDN的定義和優點

注意，GDN本質也是一種歸一化方法，但和經常在神經網絡裡看到的BN不是同一類型的歸一化。BN是針對一個batch裡的資料進行歸一化和尺度操作，且網絡訓練完成後BN的尺度參數也固定了，這就是一個完全關于歸一化資料的仿射變換，BN數學公式描述如下：

 而GDN的歸一化和尺度操作對整幅圖像都是一緻的，相當于一種空間自适應的歸一化操作，而且GDN是非線性的，其公式描述如下

2 模型優化

再看一眼架構圖，整個編碼器是卷積+下采樣+GDN組成的，已經知道GDN是幹啥的了，再看看變換和量化（卷積的參數）是如何選擇的。

2.1 優化目标函數

為了得到參數的最優解，我們想要對損失函數L=λD+R求關于各個子產品的參數的偏導，然後梯度下降求近似最優解

 2.2 變換子產品

 這篇長達27頁的文章并沒有詳細展開說變換怎麼做的，因為作者本人一直在做非線性變換，這裡使用的變換在他2016年發表的文章《End-to-end optimization of nonlinear transform codes for perceptual quality》中有詳細介紹。（其實下面量化子產品使用的方法也是這篇文章中就提出來的，這篇文章還挺硬核充滿數學公式。。。有空該看看的（真的好難看啊））

所采用的變換子產品其實就是GDN作為編碼器的分析變換

，對輸入x進行參數為

的變換後得y，GDN變換

的數學描述為

平方確定在參數值接近0時梯度更小，否則優化可能變得不穩定。為獲得一個明确的映射，我們在每個梯度步驟後将每個

 投影到區間

 上。我們對

 和進行了相同的處理，并另在每一步之後用轉置取

 的平均值，以使其對稱。GND的反變換IGDN參數處理方式相同。

為了消除每個線性變換及其後續非線性(或在綜合變換中為之前的非線性)之間的尺度模糊，我們在每個梯度步驟之後對線性濾波器進行了重新規範化，将每個濾波器除以其平方和系數的平方根。對于分析變換，求和周遊空間和所有輸入通道，對于綜合變換，求和周遊空間和所有輸出通道。其中線性濾波器利用離散餘弦變換(DCT)系數進行參數化。DCT比DFT在加速收斂方面更有效。

2.3 量化子產品

求導的過程中，由于量化是不可導的（如下圖），這篇文章為了使用随機梯度下降優化，訓練加了一個噪聲然後做區間内積分來近似代替多進制四舍五入量化，而測試的時候直接四舍五入量化量化。這種近似代替解決了量化不可微的問題，使得網絡可以訓練，但測試是直接量化和加入均勻噪聲模拟量化的方法出來差得最多好像有0.2dB。

變換後量化前系數通常也不可導，因為它的分布像是個拉普拉斯分布（中間有個奇點），如下圖。為解決這個問題，加入[-0.5, 0.5]之間的随機歸一化噪聲，加了之後得到的看起來像下圖虛線（連續且近似可微）

 ：變換後待量化元素

 ：四舍五入後的變換系數

 ：通過添加噪聲後的變換系數上述是三者的機率密度函數（PDF），離散的

 是機率品質函數（PMF），其中黑點的數值由實線在各自區間内取積分得到的，但是其實不一定等于虛線在區間内的積分，隻能是近似關系。即可以通過這種近似關系，合理等于兩者的微分熵，就是說加均勻噪聲後，不影響碼率估計。這樣積分後相當于每個區間得到一個值，跟量化/采樣是一樣的，從連續區間映射到一個離散值。

實驗結果說明，這種加噪聲後積分的量化方法既不影響碼率，也不影響MSE。左圖為量化帶來的MSE和加均勻噪聲帶來MSE的關系，右圖為離散情況下的bpp以及可微分情況下的bpp的關系。可得，本文提出的量化模型及熵率估計模型與實際中量化引起的誤差和待編碼值得熵率近似一緻。

加了噪聲之後有兩個好處：

首先，y=y+Δy 的密度函數是q的機率品質函數的連續松弛；

其次，獨立的均勻噪聲就其邊際矩而言近似于量化誤差，并且經常被用作量化誤差的模型（Gray and Neuhoff，1998）

針對y的寬松機率模型和熵代碼在code space中假設獨立的邊際，則通過進行參數模組化以減少碼率估計模型與實際模型之間的差距，具體使用精細采樣的分段線性函數，這些函數與一維直方圖的更新類似（請參見附錄）。即通過神經網絡生成編碼點的機率品質函數PMF，通過PMF的積分求得每個特征點對應的可導形式的機率值，通過對機率值求自資訊得到估計的碼率點，在前向傳播中，通過訓練階段生成的神經網絡生成分位點與累積分布函數CDF，将CDF與待編碼點輸出range coder進行熵編碼。

由于編碼數值的機率密度函數被均勻噪聲平滑，這部分會導緻一定的誤差，但是通過減小采樣間隔可以使得模型誤差任意小。

把離散的标量換成連續的近似标量量化後，最終目标損失函數式可寫成如下形式：

 其中向量 ψ(i) 是

 的分段線性逼近的參數。這是連續可微的，适合随機梯度下降優化。

其中各個參數與子產品的對應關系如下圖

 上式的損失用MSE度量，即整體架構圖中的

 是一個固定的變換，直接計算恒等式 

 在之前的《End-to-end optimization of nonlinear transform codes for perceptual quality》中，在感覺域中用的變換

 是作者本人二作的另一篇文章《Perceptual image quality assessment using a normalized Laplacian pyramid》中提出來的NLP (Normalized Laplacian Pyramid)，而不是設定成恒等式直接算MSE。在感覺域中使用NLP做變換可能性能更好（也會增加時間複雜度）。

 2.4 熵編碼子產品

雖然作者提出了一種離散量化的近似連續量化操作，但即用連續的量化操作，仍然不知道量化後資料所服從的分布

 ，我們隻能得到量化後的資料

 ，也無法直接計算出目标損失函數中的碼率項，也就無法優化損失函數。為此，作者又提出了一種近拟計算碼率的方法，即為每個量化區間中加噪聲後積分得到的離散值

 設定一個可學習的參數向量 ψ(i) ，其中的元素看成是從分布

 中采樣得到的，采樣密度為每機關長度采樣 10 個點。然後通過使用 SGD （就是普通随機梯度下降法，stochastic gradient descent）最大化下面的負的期望似然函數來優化參數 ψ(i) :

對于變換和反變換的參數ϕ和θ使用Adam算法來算最優值，使用的Adam算法的初始值

，當R和D都不顯著改變時候，把

縮小10倍，一直到

。Adam和其他SGD變種見【https://zhuanlan.zhihu.com/p/152566066】在計算

 時，将

 表達成分段線性函數，過參數向量 ψ(i) 中的那些點。然後将

 代入即可計算得到

 。整個近似過程都是可微的，故能夠使用梯度下降來優化 ψ(i) 和碼率。

實際熵編碼過程中，作者提出了一種類似CABAC的熵模型。該熵模型主要編碼資訊：

1. 圖像的大小：兩個 16比特的整數（不使用 CABAC 算法）；
2. 圖像類型：灰階圖或 RGB 圖，1比特表示（不使用 CABAC 算法）；
3. 模型參數：使用一個16比特的整數，來表示分析綜合變換的模型參數 λ 的索引（一般訓練好模型參數後，會發給編解碼端）

量化後圖像資料q：将q在空間上光栅掃描并在通道上依次疊代處理形成一串資料流，再使用 CABAC 算法進行壓縮

對量化後的q，用一個決策樹來進行二進制化，對于每個

，我們首先測試編碼的值是否等于分布的模式。如果是這樣，則

 的編碼完成；如果不是，另一個二進制判定決定它是比衆數大還是小。

每次往左決策，即 +0 ；往右決策，即 +1 。最終得到落在區間

 的所有

 值的二進制編碼，再使用 CABAC 進行壓縮。對于不在

 的

 值，使用指數哥倫布碼進行壓縮（也是借鑒于 H.264）。作者在 Kodak 資料集上進行實驗，驗證了這種基于上下文的自适應壓縮算法是優于傳統的确定性壓縮算法的。 

3 實驗結果

使用随機梯度下降法在由6507張圖像組成的ImageNet的一個子集上聯合優化了全套參數φ， θ和所有ψ（做訓練）。每一個λ訓練一次。實驗結果表明，在Kodak資料集上，本文提出的端到端圖像壓縮方案比JPEG和JPEG2000的RD性能要好。