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