基于非對(duì)稱(chēng)離散高斯似然的深度圖像壓縮方法*

羅 春，何小海，卿粼波，任 超，熊淑華

（四川大學(xué) 電子信息學(xué)院，四川 成都 610065）

0 引言

近些年來(lái)，深度學(xué)習(xí)技術(shù)在圖像壓縮領(lǐng)域大放光彩。不同于傳統(tǒng)圖像壓縮算法，深度學(xué)習(xí)方法學(xué)習(xí)一個(gè)端到端的壓縮框架，涉及變換、量化和熵編碼過(guò)程。相較于傳統(tǒng)壓縮編碼技術(shù)，如JPEG 和BPG，深度學(xué)習(xí)方法具有更優(yōu)的率失真性能和視覺(jué)感受。

隨著深度學(xué)習(xí)技術(shù)在圖像壓縮領(lǐng)域的應(yīng)用，許多先進(jìn)的壓縮算法被提出。Ballé 等人[1]先提出了端到端的圖像壓縮算法用于壓縮自然圖像，后續(xù)又提出了超先驗(yàn)?zāi)Ｐ蚚2]，對(duì)隱變量用高斯分布建模。Minnen 等人[3]引入了自回歸先驗(yàn)預(yù)測(cè)分布參數(shù)，但使得模型只能串行運(yùn)行，限制了模型的計(jì)算速度。Cheng 等人[4]在此基礎(chǔ)上采用了混合高斯分布模型對(duì)隱變量進(jìn)行更細(xì)粒度的建模，搭配非局部注意力機(jī)制和自回歸模型用于熵編碼。由于率失真和卷積的限制，純卷積網(wǎng)絡(luò)無(wú)法真正去除隱變量間的相關(guān)性，導(dǎo)致率失真性能的瓶頸。為此，研究者們將transformer 引入圖像壓縮領(lǐng)域，突破率失真性能的瓶頸，實(shí)現(xiàn)高效的圖像壓縮編碼技術(shù)。Zou 等人[5]提出了即插即用的窗口注意力構(gòu)建編解碼器，結(jié)合局部和全局相關(guān)性的同時(shí)依舊保持較高的計(jì)算效率，使得壓縮的圖像滿足預(yù)期。Liu 等人[6]結(jié)合transformer 的非局部特性和卷積網(wǎng)絡(luò)的局部特性，構(gòu)造編碼器和熵模型，相比于純卷積神經(jīng)網(wǎng)絡(luò)或純transformer 具有更大的感受野，在Kodak 數(shù)據(jù)集上達(dá)到了最優(yōu)水平。

許多工作致力于研究如何高效地將圖像映射至隱空間，由于卷積和transformer 天然存在的局部和全局相關(guān)性，隱變量間普遍地存在統(tǒng)計(jì)依賴(lài)。因此，在隱空間中如何選取隱變量分布是一個(gè)需要解決的問(wèn)題。此前的工作是選取高斯分布建模隱變量，在給定均值和方差的情況下高斯分布具有最大的熵，這對(duì)建模未知分布是有益的，但是也伴隨熵增加帶來(lái)的額外碼流。

針對(duì)上述問(wèn)題，本文提出一種結(jié)合語(yǔ)義信息的圖像壓縮算法。該算法引入Spike and Slab[7]先驗(yàn)，并結(jié)合語(yǔ)義信息建模隱變量，可以有效地降低圖像碼率，同時(shí)保持高質(zhì)量的解碼圖像。

1 深度圖像壓縮技術(shù)

深度圖像壓縮的經(jīng)典流程為先經(jīng)過(guò)編碼器E得到隱變量的表示y，隨后經(jīng)過(guò)量化操作Q得到離散化的隱變量表示，最后經(jīng)過(guò)解碼器D得到解碼后的圖像，整個(gè)過(guò)程可以表示為：

式中：φ和θ分別為編碼器和解碼器的參數(shù)。在此基礎(chǔ)上，引入超先驗(yàn)信息建模隱變量，其過(guò)程可以表示為：

式中：φh和θh分別為超先驗(yàn)編碼器Eh和超先驗(yàn)解碼器Dh的參數(shù)。超先驗(yàn)?zāi)Ｐ途唧w的原理如圖1（a）所示。圖1（b）展示了一種改進(jìn)后的深度圖像壓縮編碼框架——上下文超先驗(yàn)?zāi)Ｐ停涫褂米曰貧w模型，可以有效地學(xué)習(xí)到隱變量參數(shù)，缺陷是限制了模型的并行化部署。

圖1 模型原理

后續(xù)的深度圖像壓縮技術(shù)大多沿用了超先驗(yàn)的框架，主要研究編解碼網(wǎng)絡(luò)框架和熵模型的構(gòu)建。比如，文獻(xiàn)[8]、文獻(xiàn)[9]和文獻(xiàn)[10]使用殘差網(wǎng)絡(luò)高效地壓縮隱變量，從而獲得更高質(zhì)量的解碼圖像。文獻(xiàn)[11]引入注意力機(jī)制來(lái)重構(gòu)編解碼器，例如，使用空間注意力機(jī)制學(xué)習(xí)特征通道間的交互，有效地提取空間依賴(lài)，以及利用自注意力克服卷積的歸納偏置的缺陷，建立全局的相關(guān)依賴(lài)。熵模型是一種收發(fā)雙方共享的模型，也是影響深度圖像壓縮性能的重要環(huán)節(jié)。文獻(xiàn)[12]使用通道自回歸模型，相較于空間自回歸模型有優(yōu)勢(shì)。文獻(xiàn)[13]和文獻(xiàn)[14]結(jié)合空間通道自回歸模型，在不影響運(yùn)行速度的前提下實(shí)現(xiàn)了性能的提升。

隨著研究的推進(jìn)，編解碼器和熵模型的效果已逐漸達(dá)到瓶頸，依靠增加計(jì)算量實(shí)現(xiàn)性能提升的方法逐漸變得低效。因此，如圖1（c）所示，從隱變量分布出發(fā)，引入非對(duì)稱(chēng)高斯分布，即Spike and Slab 分布，以實(shí)現(xiàn)隱變量分布的熵減過(guò)程。

2 非對(duì)稱(chēng)深度圖像壓縮方法

2.1 問(wèn)題方程

提出的非對(duì)稱(chēng)深度圖像壓縮結(jié)構(gòu)如圖2 所示。如其他壓縮模型一樣，輸入圖像經(jīng)過(guò)編碼器后得到隱空間表示，但并不馬上進(jìn)行量化，而是經(jīng)過(guò)一個(gè)“熵減”過(guò)程，具體地，可以分為語(yǔ)義信息提取和稀疏化過(guò)程。語(yǔ)義信息提取學(xué)習(xí)二進(jìn)制掩碼s，經(jīng)過(guò)稀疏化后的隱變量具有高度語(yǔ)義相關(guān)的稀疏性，將限制學(xué)習(xí)的高斯分布參數(shù)，趨向熵減少的方向。具體地，經(jīng)過(guò)熵減階段后的隱變量分布可以表示為：

圖2 非對(duì)稱(chēng)離散高斯似然的深度圖像壓縮方法網(wǎng)絡(luò)結(jié)構(gòu)

式中：σ1＞＞σ0。對(duì)于那些語(yǔ)義冗余的隱變量（s=0），學(xué)習(xí)的高斯分布趨向于一個(gè)沖激函數(shù)，離散化后的隱變量具有較小的熵。

為了訓(xùn)練端到端的圖像壓縮框架，使用下列的損失函數(shù)更新網(wǎng)絡(luò)參數(shù)：

式中：λ為一個(gè)超參數(shù)，用于比特率和失真率間的權(quán)衡；d(,x)為衡量重建圖像和原始圖像差異的函數(shù)，通常使用均方差（Mean Square Error）函數(shù)；R()和R()為傳輸離散隱變量和超先驗(yàn)所需要的碼率。

2.2 語(yǔ)義信息提取

借助圖像語(yǔ)義信息，可以獲得文本相關(guān)的稀疏結(jié)構(gòu)，有助于去除變量間的相關(guān)性。具體地，如圖3 所示構(gòu)造了卷積—通道注意力級(jí)聯(lián)的掩碼提取器G_mask，對(duì)輸入圖像的各個(gè)通道進(jìn)行不同的非線性映射，自適應(yīng)地調(diào)整各個(gè)通道間的注意力權(quán)重，強(qiáng)化重要特征，削弱無(wú)關(guān)特征，最終得到聯(lián)系緊密的特征圖。經(jīng)過(guò)掩碼頭Mask_head 和閾值Threshold模塊的二值化處理，將得到語(yǔ)義掩碼S∈{0,1}N。具體的流程可表示為：

圖3 語(yǔ)義信息提取過(guò)程

獲得語(yǔ)義掩碼后，和隱變量進(jìn)行交織操作，拉近隱變量和對(duì)應(yīng)掩碼的空間距離。具體地，將對(duì)應(yīng)通道平面堆疊形成新的張量ym，用于后續(xù)的稀疏化處理。

2.3 含參稀疏化過(guò)程

結(jié)合卷積的局部特性與transformer 的全局特性[15]，本文使用卷積構(gòu)造含參稀疏化過(guò)程，如圖4所示。具體地，將輸入的張量ym經(jīng)過(guò)卷積網(wǎng)絡(luò)Q，K和V映射到特征空間得到Q,K,V∈?B×C×H×W，經(jīng)過(guò)矩陣轉(zhuǎn)換操作后變成適合操作的矢量形式Q,K,V∈?B×L×C，最終得到捕獲相關(guān)性的注意力特征為：

圖4 含參稀疏化過(guò)程

式中：γ和dk分別為縮放因子和矢量的維度。稀疏過(guò)程可以降低計(jì)算的復(fù)雜度，提高計(jì)算效率。相較于無(wú)參稀疏化過(guò)程，含參稀疏化過(guò)程可以自適應(yīng)地調(diào)節(jié)稀疏程度，面對(duì)復(fù)雜環(huán)境具有較強(qiáng)的適應(yīng)性，適合端到端的學(xué)習(xí)方式的部署。

3 實(shí)驗(yàn)分析

3.1 實(shí)驗(yàn)設(shè)置

模型訓(xùn)練中使用Vimeo90K 數(shù)據(jù)集作為訓(xùn)練集，它包含153 939 張裁剪空間尺寸為256×256 的圖像，通用數(shù)據(jù)集Kodak 作為測(cè)試集。在訓(xùn)練的初期，使用Adam 優(yōu)化器對(duì)壓縮網(wǎng)絡(luò)的參數(shù)進(jìn)行優(yōu)化，batchsize 設(shè)置為24。初期的學(xué)習(xí)率設(shè)置為0.000 1，在訓(xùn)練的后期將學(xué)習(xí)率設(shè)置為0.000 01。

模型使用式（4）優(yōu)化率失真曲線。具體地，使用均方差（Mean Square Error，MSE）作為失真項(xiàng)，超參數(shù)λ選擇為{0.001 8,0.003 5,0.013}，并采用pytorch 框架，在GTX2080Ti 上進(jìn)行訓(xùn)練。

3.2 實(shí)驗(yàn)結(jié)果分析

本文比較了當(dāng)前的先進(jìn)的圖像壓縮算法，包括Cheng[4]、BPG、hyperprior[2]和多功能視頻編碼（Versatile Video Coding，VVC），從客觀指標(biāo)（率失真曲線）和主觀感受兩個(gè)方面評(píng)價(jià)參與比較的壓縮算法的優(yōu)劣。

圖5 展示了不同壓縮算法在通用數(shù)據(jù)集Kodak上的率失真性能。從圖中可以看出，當(dāng)bpp 位于0.2～0.3 的區(qū)間時(shí)，本文所提方法和VVC、Cheng 具有相似的性能，但在bpp 小于0.2 大于0.3 的區(qū)間內(nèi)時(shí)，本文所提算法相較于VVC 和Cheng 具有明顯優(yōu)勢(shì)。

圖5 不同壓縮算法在Kodak 數(shù)據(jù)集上的率失真曲線

圖6 展示了不同壓縮方法在Kodak 數(shù)據(jù)集圖像kodim23 上的解碼圖像。相較于其他方法，本文提出的方法在保持紋理細(xì)節(jié)方面具有優(yōu)勢(shì)。

圖6 圖像kodim23 的可視化解碼結(jié)果（bpp/PSNR/SSIM）

3.3 隱變量分析

圖7 展示了隱變量的可視化結(jié)果。可以看出，隨著隱變量的通道層數(shù)的增加，通道特征逐漸變得稀疏，意味著分配的碼流越來(lái)越少。在紋理細(xì)膩的區(qū)域，如帽子、圍巾區(qū)域，一直都在分配碼流；而臉部等高頻成分較少的區(qū)域，只在前幾個(gè)通道分配碼流。這些實(shí)驗(yàn)結(jié)果驗(yàn)證了本文提出的方法可以有選擇地改善碼流分配。

圖7 掩碼后的隱變量

4 結(jié)語(yǔ)

本文主要研究了深度圖像壓縮技術(shù)中的率失真優(yōu)化問(wèn)題，主要思想是在隱空間借助通道注意力提取語(yǔ)義信息，并基于自注意力機(jī)制構(gòu)建含參稀疏化過(guò)程，實(shí)現(xiàn)單高斯分布向非對(duì)稱(chēng)高斯分布的遷移，以節(jié)約碼流。實(shí)驗(yàn)結(jié)果表明，本文提出的方法在節(jié)省碼率、優(yōu)化率失真性能上具有良好表現(xiàn)。