Pyramid attention networks for image restoration

paper:https://arxiv.org/abs/2004.13824

code: https://github.com/SHI-Labs/Pyramid-Attention-Networks

1. 基本思想

作者指出，當前基於深度學習的方法只是在單個尺度上利用了self-similarity（do not take full advantage of self-similarities by relying on self-attention neural modules that only process information at the same scale）。在該論文里，作者提出了 pyramid attention module 用於圖像修復，可以利用多尺度特征金字塔捕獲 long-range 的特征關聯。

在傳統圖像去噪方法中，NLmean、BM3D都是采用自相似性進行圖像去噪的佼佼者，圖像的自相似性非常重要。同時，作者指出，基於深度學習的方法大多使用non-local方法來獲取 long-range 的信息，這樣的模塊早期被用於識別和分類任務，但是在 low-level 的視覺任務中也被證明非常有效。

但是，直接將 self-attention 應用於圖像修復有一定的局限性：

• multi-scale non-local prior has not been explorered.
• pixel-wise matching used in the self-attention module is usually noisy for image restoration.

基於上述分析，作者提出了non-local pyramid attention 模塊來捕獲long-range信息。

如圖所示，該方法可以挖掘 multi-scale self-examplars 來提高重建性能，主要有三個優點：

• 將現有non-local方法的擴展到多尺度空間（現有non-local方法的搜索空間僅限於金字塔的最底端）
• 不同大小特征之間的 long-range 依賴被較好的建模
• 可以降低特征中的噪聲，進而提升性能。

2. 方法

如圖2a所示，non-local attention 在全圖進行響應特征平均，可以這樣描述：

\[y^i=\frac{1}{\sigma(x)}\sum\phi(x^i, x^j)\theta(x^j) \]

其中\(\phi\) 計算相似性，\(\theta\)表示特征變換。上述公式僅在統一尺度空間進行處理，難以利用跨尺度的特征。

圖2b為 scale agnostic attention，是 2a 的一種擴展。給定一個尺度\(s\)，如何評價\(x^j\)和\(x^j_{\delta(s)}\)之間的相似性，以及如何把\(x^j_{\delta(s)}\)聚合到\(y^i\)是兩個關鍵問題。為了解決這兩個問題，作者將\(x^j_{\delta(s)}\)下采樣為像素特征\(z^j\)，由於需要在整個特征圖上搜索，作者對原始輸入\(x\)進行下采樣得到特征圖\(z\)，然后相似性可以這樣計算：

\[y^i=\frac{1}{\sigma(x,z)}\sum_j\phi(x^i,z^i)\theta(z^i) \]

同時，作者將該計算方式擴展到金字塔形式，金字塔注意力具有跨尺度特征自相似建模能力。

在具體實現時，相似性函數\(\phi\)使用了embeded gaussian。在特征變換函數\(\theta\)方面，使用簡單的線性變換。同時，為了改進匹配的魯棒性，添加了額外的近鄰相似約束，類似於經典的非局部均值濾波，添加了一個強約束信息：近鄰且相似。這樣有助於學習相關特征，同時抑制不相關特征。

圖3 為 Pyramid attention 模塊的一個示意圖，在損失函數方面，作者選擇了L1損失。

3. 實驗與結論

作者使用了一個5層的金字塔，網絡包含80個殘差塊，提出的PA模塊插入在第40個殘差塊之后，第41個殘差塊之前。該方法在圖像降噪、去馬賽克、去壓縮偽影、超分辨率方面都進行了實驗，效果都比較好。具體細節和ablation study 可以參考論文，不再細說。

該論文使用了一個 simple backbone，然后只加了一個 pyramid attention 模塊，就取得了非常好的效果。作者認為，該模塊可以在以后的圖像修復網絡中，作為基礎模塊來使用。