字節(jié)提出非對(duì)稱圖像重采樣模型,JPEG、WebP上抗壓縮性能領(lǐng)先SOTA
圖像重采樣 (Image Rescaling,LR) 任務(wù)聯(lián)合優(yōu)化圖像下采樣和上采樣操作,通過(guò)對(duì)圖像分辨率的下降和還原,可以用于節(jié)省存儲(chǔ)空間或傳輸帶寬。在實(shí)際應(yīng)用中,例如圖集服務(wù)的多檔位分發(fā),下采樣得到的低分辨率圖像往往會(huì)進(jìn)行有損壓縮,而有損壓縮往往導(dǎo)致現(xiàn)有算法的性能大幅下降。
近期,字節(jié)跳動(dòng) - 火山引擎多媒體實(shí)驗(yàn)室首次嘗試了有損壓縮下的圖像重采樣性能優(yōu)化,設(shè)計(jì)了一種非對(duì)稱可逆重采樣框架,基于該框架下的兩點(diǎn)觀察,進(jìn)一步提出了抗壓縮圖像重采樣模型 SAIN。該研究將一組可逆網(wǎng)絡(luò)模塊解耦成重采樣和壓縮模擬兩個(gè)部分,使用混合高斯分布建模分辨率下降和壓縮失真造成的聯(lián)合信息損失,結(jié)合可微的 JPEG 算子進(jìn)行端到端訓(xùn)練,大幅提升了對(duì)于常見(jiàn)壓縮算法的魯棒性。
目前對(duì)于圖像重采樣的研究,SOTA 方法是基于可逆網(wǎng)絡(luò) (Invertible Network) 構(gòu)建一個(gè)雙射函數(shù) (bijective function),其正運(yùn)算將高分辨率 (HR) 圖像轉(zhuǎn)換為低分辨率 (LR) 圖像和一系列服從標(biāo)準(zhǔn)正態(tài)分布的隱變量,逆運(yùn)算則隨機(jī)采樣隱變量,結(jié)合 LR 圖像進(jìn)行上采樣還原。
由于可逆網(wǎng)絡(luò)的特性,下采樣和上采樣算子保持高度的對(duì)稱性 (symmetry),這就導(dǎo)致受壓縮的 LR 圖像難以通過(guò)原本學(xué)得的上采樣算子進(jìn)行還原。為了增強(qiáng)對(duì)于有損壓縮的魯棒性,該研究提出了一個(gè)基于非對(duì)稱可逆框架的抗壓縮圖像重采樣模型 SAIN (Self-Asymmetric Invertible Network)。
SAIN 模型的核心創(chuàng)新點(diǎn)如下:
- 提出了非對(duì)稱 (asymmetric) 式的可逆圖像重采樣框架,解決了先前方法中由于嚴(yán)格對(duì)稱性導(dǎo)致的性能下降問(wèn)題;提出了增強(qiáng)可逆模塊 (E-InvBlock),在共享大量參數(shù)和運(yùn)算的前提下,增強(qiáng)模型擬合能力,同時(shí)建模壓縮前和壓縮后的兩組 LR 圖像,使模型能夠通過(guò)逆運(yùn)算進(jìn)行壓縮恢復(fù)和上采樣。
- 構(gòu)建可學(xué)習(xí)的混合高斯分布,建模分辨率下降和有損壓縮造成的聯(lián)合信息損失,通過(guò)重參數(shù)化技巧直接優(yōu)化分布參數(shù),更加符合隱變量的實(shí)際分布。
SAIN 模型進(jìn)行了 JPEG 和 WebP 壓縮下的性能驗(yàn)證,在多個(gè)公開(kāi)數(shù)據(jù)集上的表現(xiàn)大幅領(lǐng)先 SOTA 模型,相關(guān)研究已經(jīng)入選 AAAI 2023 Oral。
- 論文地址:https://arxiv.org/abs/2303.02353
- 代碼鏈接:https://github.com/yang-jin-hai/SAIN
非對(duì)稱重采樣框架
圖 1 Dual-IRN 模型圖。
為了提升抗壓縮性能,該研究首先設(shè)計(jì)了一種非對(duì)稱式的可逆圖像重采樣框架,提出了 baseline 方案 Dual-IRN 模型,深入分析了該方案的不足之處后,進(jìn)一步優(yōu)化提出了 SAIN 模型。如上圖所示,Dual-IRN 模型包含兩個(gè)分支,其中 D-IRN 和 U-IRN 為兩組可逆網(wǎng)絡(luò),分別學(xué)習(xí) HR 圖像與壓縮前 / 壓縮后 LR 圖像之間的雙射。
在訓(xùn)練階段,Dual-IRN 模型通過(guò)可微的 JPEG 算子傳遞兩個(gè)分支之間的梯度。而在測(cè)試階段,模型通過(guò) D-IRN 進(jìn)行下采樣獲得高質(zhì)量的 LR 圖像,經(jīng)過(guò)現(xiàn)實(shí)環(huán)境中的真實(shí)壓縮后,再通過(guò)具有壓縮感知 (compression-aware) 的 U-IRN 完成壓縮恢復(fù)和上采樣。
這樣的非對(duì)稱式框架,使得上采樣和下采樣的算子避免了嚴(yán)格的可逆關(guān)系,從根源上解決了壓縮算法破壞上下采樣過(guò)程對(duì)稱性帶來(lái)的問(wèn)題,相對(duì)于 SOTA 的對(duì)稱式方案大幅提升了抗壓縮性能。
隨后,研究人員對(duì) Dual-IRN 模型展開(kāi)進(jìn)一步的分析,觀察到以下兩個(gè)現(xiàn)象:
- 第一,衡量 D-IRN 和 U-IRN 兩個(gè)分支中間層特征的 CKA 相似性。如上圖 (b) 所示,D-IRN 最后一層的輸出特征(即網(wǎng)絡(luò)生成的高質(zhì)量 LR 圖像)與 U-IRN 淺層的輸出特征具有高度的相似性,表明 U-IRN 的淺層行為更貼近采樣損失的模擬,而深層行為更貼近壓縮損失的模擬。
- 第二,統(tǒng)計(jì) D-IRN 和 U-IRN 兩個(gè)分支中間層隱變量的真實(shí)分布。如上圖 (c)(d) 所示,不具壓縮感知的 D-IRN 的隱變量整體上滿足單峰的正態(tài)分布假設(shè),而具有壓縮感知的 U-IRN 的隱變量則呈現(xiàn)了多峰的形態(tài),表明有損壓縮帶來(lái)的信息損失形式更加復(fù)雜。
基于以上分析,研究人員從多個(gè)方面對(duì)模型進(jìn)行優(yōu)化,最終得到的 SAIN 模型不僅網(wǎng)絡(luò)參數(shù)量下降了將近一半,并且還取得了進(jìn)一步的性能提升。
SAIN 模型細(xì)節(jié)
圖 2 SAIN 模型圖。
SAIN 模型的架構(gòu)如上圖所示,主要進(jìn)行了以下四個(gè)方面的改進(jìn):
1. 整體框架方面?;谥虚g層特征的相似性,將一組可逆網(wǎng)絡(luò)模塊解耦成重采樣和壓縮模擬兩個(gè)部分,形成自非對(duì)稱 (self-asymmetric) 架構(gòu),避免使用兩組完整的可逆網(wǎng)絡(luò)。在測(cè)試階段,使用正變換
取得高質(zhì)量的 LR 圖像,先使用逆變換
進(jìn)行壓縮恢復(fù),再使用逆變換
進(jìn)行上采樣。
2. 網(wǎng)絡(luò)結(jié)構(gòu)方面?;趬嚎s損失能夠借助高頻信息恢復(fù)的假設(shè)提出 E-InvBlock,在模塊中增加一個(gè)加性變換,使得共享大量運(yùn)算的前提下可以高效地建模壓縮前 / 壓縮后的兩組 LR 圖像。
3. 信息損失建模方面?;陔[變量的真實(shí)分布,提出使用可學(xué)習(xí)的混合高斯分布建模下采樣和有損壓縮帶來(lái)的聯(lián)合信息損失,通過(guò)重參數(shù)化技巧端到端優(yōu)化分布參數(shù)。
4. 目標(biāo)函數(shù)方面。設(shè)計(jì)了多個(gè)損失函數(shù),用于約束網(wǎng)絡(luò)的可逆性、提高重建精度,同時(shí)在損失函數(shù)中引入真實(shí)壓縮操作以增強(qiáng)對(duì)真實(shí)壓縮方案的魯棒性。
實(shí)驗(yàn)與效果評(píng)估
評(píng)估數(shù)據(jù)集為 DIV2K 驗(yàn)證集和 Set5、Set14、BSD100 和 Urban100 四個(gè)標(biāo)準(zhǔn)測(cè)試集。
量化評(píng)估指標(biāo)為:
- PSNR:Peak Signal-to-Noise Ratio,峰值信噪比,反映重建圖像與原始圖像的均方誤差,越高越好;
- SSIM:Structural Similarity Image Measurement,衡量重建圖像與原始圖像的結(jié)構(gòu)相似度,越高越好。
在表 1 和圖 3 的對(duì)比實(shí)驗(yàn)中,SAIN 在所有數(shù)據(jù)集上的 PSNR 和 SSIM 分?jǐn)?shù)都大幅領(lǐng)先 SOTA 的圖像重采樣模型。在比較低的 QF 下,現(xiàn)有方法普遍出現(xiàn)嚴(yán)重的性能下降,而 SAIN 模型仍然保持最優(yōu)性能。
表 1 對(duì)比實(shí)驗(yàn),在 DIV2K 數(shù)據(jù)集上比較不同 JPEG 壓縮質(zhì)量 (QF) 下的重建質(zhì)量 (PSNR / SSIM)。
圖 3 對(duì)比實(shí)驗(yàn),在四個(gè)標(biāo)準(zhǔn)測(cè)試集上比較不同 JPEG QF 下的重建質(zhì)量 (PSNR)。
在圖 4 的可視化結(jié)果中,可以明顯看出 SAIN 還原的 HR 圖像更加清晰、準(zhǔn)確。
圖 4 不同方法在 JPEG 壓縮下的可視化結(jié)果對(duì)比 (×4 倍率)。
在表 2 的消融實(shí)驗(yàn)中,研究人員還比較了其他幾種結(jié)合真實(shí)壓縮進(jìn)行訓(xùn)練的候選方案。這些候選方案相較于完全對(duì)稱的現(xiàn)有模型 (IRN) 具有更高的抗壓縮性能,但在參數(shù)量和準(zhǔn)確率上仍不如 SAIN 模型。
表 2 針對(duì)整體框架和訓(xùn)練策略的消融實(shí)驗(yàn)。
在圖 5 的可視化結(jié)果中,研究人員對(duì)比了 WebP 壓縮失真下不同圖像重采樣模型的重建結(jié)果??梢园l(fā)現(xiàn),SAIN 模型在 WebP 壓縮方案下同樣表現(xiàn)出了最高的重建分?jǐn)?shù),能夠清晰準(zhǔn)確地還原圖像細(xì)節(jié),證明了 SAIN 對(duì)于不同壓縮方案的兼容性。
圖 5 不同方法在 WebP 壓縮下的定性和定量對(duì)比 (×2 倍率)。
除此之外,該研究還針對(duì)混合高斯分布、E-InvBlock 和損失函數(shù)等部分進(jìn)行了消融實(shí)驗(yàn),證明了這些改進(jìn)對(duì)于結(jié)果的正向貢獻(xiàn)。
總結(jié)與展望
火山引擎多媒體實(shí)驗(yàn)室針對(duì)抗壓縮圖像重采樣提出了一個(gè)基于非對(duì)稱可逆框架的模型:SAIN。該模型包含重采樣和壓縮模擬兩個(gè)部分,使用混合高斯分布建模分辨率下降和壓縮失真造成的聯(lián)合信息損失,結(jié)合可微的 JPEG 算子進(jìn)行端到端訓(xùn)練,提出 E-InvBlock 來(lái)增強(qiáng)模型擬合能力,大幅提升了對(duì)于常見(jiàn)壓縮算法的魯棒性。
火山引擎多媒體實(shí)驗(yàn)室是字節(jié)跳動(dòng)旗下的研究團(tuán)隊(duì),致力于探索多媒體領(lǐng)域的前沿技術(shù),參與國(guó)際標(biāo)準(zhǔn)化工作,其眾多創(chuàng)新算法及軟硬件解決方案已經(jīng)廣泛應(yīng)用在抖音、西瓜視頻等產(chǎn)品的多媒體業(yè)務(wù),并向火山引擎的企業(yè)級(jí)客戶提供技術(shù)服務(wù)。實(shí)驗(yàn)室成立以來(lái),多篇論文入選國(guó)際頂會(huì)和旗艦期刊,并獲得數(shù)項(xiàng)國(guó)際級(jí)技術(shù)賽事冠軍、行業(yè)創(chuàng)新獎(jiǎng)及最佳論文獎(jiǎng)。
未來(lái),研究團(tuán)隊(duì)會(huì)繼續(xù)優(yōu)化圖像重采樣模型在有損壓縮下的性能,并且進(jìn)一步探索抗壓縮視頻重采樣、任意倍率重采樣等更加復(fù)雜的應(yīng)用情景。