當Mamba遇上視覺壓縮：首個具有選擇性狀態空間的視覺壓縮網路

©PaperWeekly 原創 · 作者 | 秦詩雨

單位 | 清華大學碩士生

研究方向 | 影像壓縮

摘要

在多媒體領域，學習型視覺壓縮是一項重要的任務。現有的方法已經探索了各種基於 CNN 和 Transformer 的網路設計，以建模內容分佈並消除冗餘，但在平衡效率（即率失真權衡）和效能方面仍然是一個挑戰。

最近，狀態空間模型（SSM）因其長距離建模能力和效率而表現出潛力。受此啟發，我們首次探索了 SSM 在視覺壓縮中的應用。我們設計了 MambaVC，一個基於 SSM 的簡單、強大且高效的壓縮網路。MambaVC 開發了一個視覺狀態空間（VSS）模組，並在每次下采樣後使用 2D 選擇掃描（2DSS）模組作為非線性啟用函式，這有助於捕捉有資訊的全域性上下文並增強壓縮效果程式碼已開源。

論文標題：

MambaVC: Learned Visual Compression with Selective State Spaces

論文地址：

https://arxiv.org/pdf/2405.15413

程式碼地址：

https://github.com/QinSY123/2024-MambaVC

Introduction

學習型視覺壓縮的發展主要分為兩個方向，一個方向是構建更有效的熵模型用於率估計；另一個方向是設計更好的神經網路以消除冗餘資訊並準確捕捉內容分佈。

對於後者，我們可以根據網路結構的不同將現有模型分為以為 CNNs 基礎的方法和以 Transformers 基礎的方法。CNN 方法計算量更小但是感受野受限，無法捕捉更多的冗餘；Transformer 的方法透過注意機制在全域性感知方面表現出色，有效減少了冗餘，然而計算效率十分低下。像 TCM 這樣的混合方法試圖融合二者平衡壓縮效果和效率，但並不是一個可持續的長期解決方案。

最近，狀態空間模型（SSMs）得到了廣泛關注。它採用資料依賴的選擇機制來增強關鍵資訊提取，同時消除不相關的噪聲。這種有效收集全域性上下文的能力表明，基於 Mamba 的模型在壓縮方面可能具有優勢。Mamba 還結合了結構化重新引數化技術和硬體高效的並行掃描演算法，確保了在 GPU 上的快速訓練和推理。

因此，我們提出了 MambaVC，一種具有選擇性狀態空間的視覺壓縮網路。我們將視覺狀態空間模組作為神經壓縮網路中每個下采樣後的非線性啟用函式。VSS 模組集成了專門的二維選擇性掃描（2DSS）機制用於空間建模。該機制沿四條預定義的遍歷路徑並行執行選擇性掃描，捕捉全面的全域性上下文，促進有效和高效的壓縮。

總而來說，我們的貢獻如下：

1. 我們提出了 MambaVC，這是首個具有選擇性狀態空間的視覺壓縮網路。其中二維選擇性掃描機制（2DSS）改進了全域性上下文建模，有助於實現有效和高效的壓縮。

2. 在基準資料集上進行的廣泛實驗表明，MambaVC 在影像和影片壓縮方面具有卓越的效能和計算效率。這些強勁的結果突顯了在 CNN 和 Transformers 之外壓縮網路設計的新方向。

3. 我們展示了 MambaVC 在高解析度壓縮中的特殊有效性和可擴充套件性，表明其在許多重要但具有挑戰性的應用中具有潛力。

4. 我們對不同網路設計進行了全面比較和分析，展示了 MambaVC 相比於 CNNs 和 Transfromers 的方法在各方面的優勢，以驗證和理解其有效性。

方法

▲ 圖1 網路結構

模型的編碼器和解碼器主要包括兩部分，上/下采樣和視覺狀態空間模組（VSS Block）。對於下采樣和上取樣我們分別使用了卷積和反捲積操作處理。如圖 1（b）所示，每個視覺狀態空間模組由多個視覺狀態空間層組成，在 Mamba 的基礎上，VSS 層在層歸一化（LN）後採用了一個帶有兩個分支的門控結構。給定輸入特徵圖，主分支做如下處理：

類似地，門控分支計算權重向量 w：

最後，這兩個分支結合起來生成輸出特徵圖：

基礎的 Mamba 只能處理一維序列，無法直接應用於二維影像資料。為了有效地建模空間上下文，我們擴充套件了四種展開方式用於選擇性掃描。具體來說，對於特徵圖，表示特徵圖中第行和第列的標記，展開定義如下：

其中。是展開後的標記序列。我們應用 S6 操作進行選擇性掃描生成上下文標記序列。然後，我們透過以下摺疊模式對上下文標記序列應用反向操作：

最後，我們合併變換後的特徵圖以獲得輸出特徵圖：

實驗

影像壓縮

▲ 圖 2 Kodak資料集上的率失真效能

與之前最好的方法 MLIC+ 相比，我們的方法在平均 PSNR 上提高了 0.1 dB，而只需一半的計算複雜度和 60% 的記憶體開銷。

在不同解析度的四個資料集上 MambaVC 的效能表現均比 CNN 變體和 Transformer 變體表現要更好。

影片壓縮

▲ 圖3 影片壓縮效能

我們還將 MambaVC 擴充套件到影片壓縮以探索其潛力。我們選擇了學習型影片壓縮模型 scale-space flow（SSF）作為擴充套件的基礎框架，將其中的非線性變換替換為視覺狀態空間模組。同時我們也比較了 Transformer 變體，可以看出 MambaVC-SSF 的效能超過了 CNN 變體和 Transformer 變體。

高解析度場景壓縮

▲ 圖4 MambaVC 在 UHD 資料集上相對不同變體的 BD-rate

我們探究了 Mamba 在視覺壓縮任務上的特有優勢。我們將 UHD 中的高解析度影像按不同的比例下采樣，以建立具有相同分佈但不同大小的多組影像。如圖 4 所示，隨著解析度的增加，MambaVC 相比其他變體節省了更多的位元。

為了減輕特定資料集分佈的影響，我們在四個不同解析度的資料集上進行了測試。MambaVC 在高解析度 UHD 上的效能優勢明顯大於在低解析度 Kodak 上的效能優勢。

計算儲存開銷

儘管 MLIC+ 由於採用了更先進的熵模型而帶來了更高的計算成本，但它並未實現更優的效能。同時相比 CNN 和 Transformer 變體，計算量和儲存量都達到了最優。

增益來源

為了更加清楚瞭解 Mamba 在視覺壓縮領域相較於 CNN 和 Transformer 的優勢，我們從潛在分佈和相關性、有效感受野以及量化偏差方面進行了探究。

潛在分佈和相關性

▲ 圖5 潛在相關性

我們可視化了中每個空間畫素與其周圍位置之間的相關性，MambaVC 在所有距離上的相關性均低於 SwinVC 和 ConvVC。

▲ 圖6 潛在分佈

理論上，解相關的潛變數應遵循標準正態分佈（SND），我們擬合了不同方法的分佈曲線並計算了與 SND 的 KL 散度，MambaVC 的曲線明顯更接近 SND，KL 散度更小，這表明基於 Mamba 的超先驗網路可以更準確地學習 (µ,σ)。

有效感受野

▲ 圖7 有效感受野

更大的感受野使得網路能夠從更廣的區域捕捉相關資訊。這一特性與視覺壓縮中的非線性編碼器完美契合，因為它透過特徵提取和降維來減少影像中的冗餘。MambaVC 是唯一具有全域性 ERF 的模型，而 ConvVC 的感受野最小。

量化偏差

在有失真壓縮中，量化是資訊損失的主要來源。我們透過對比隱特徵和量化解碼後的之間偏差來評估這種損失。圖 8 展示了縮放後的偏差圖和具體數值。與 MambaVC 相比，SwinVC 和 ConvVC 的資訊損失平均分別增加了 3.3% 和 17%。視覺化結果還表明，在大多數位置上，MambaVC 的資訊損失更小（更深的藍色和更淺的紅色）。

總結

在本文中，我們介紹了 MambaVC，這是第一個基於狀態空間模型的視覺壓縮網路。MambaVC 透過構建具有二維選擇性掃描（2DSS）機制的視覺狀態空間（VSS）塊來改進全域性上下文建模和內容壓縮。

實驗結果表明，與 CNN 和 Transformer 變體相比，MambaVC 在率失真效能上表現出色，同時保持了計算和記憶體效率。這些優勢在高解析度影像上更為顯著，突顯了 MambaVC 在實際應用中的潛力和可擴充套件性。與其他設計相比，MambaVC 展示了更強的冗餘消除能力、更大的感受野和更低的量化損失，體現了其在壓縮方面的綜合優勢。我們希望 MambaVC 能為在壓縮中探索狀態空間模型（SSMs）提供基礎，並激發未來的研究工作。

更多閱讀