來自Transformer的降維打擊！北京大學等最新發布論文，聯合提出圖像處理Transformer。通過對low-level計算機視覺任務，如降噪、超分、去雨等進行研究，提出了一種新的預訓練模型IPT，占領low-level多個任務的榜首。

本文首發自極市平臺，作者@Happy，轉載需獲授權。

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

該文是清華大學&華為諾亞等聯合提出的一種圖像處理transformer。Transformer自提出之日起即引起極大的轟動，BERT、GPT-3等模型迅速占用NLP各大榜單；后來Transformer被用于圖像分類中同樣引起了轟動；再后來，transformer在目標檢測任務中同樣引起了轟動。現在Transformer再出手，占領了low-level多個任務的榜首，甚至它在去雨任務上以1.6dB超越了已有最佳方案。

Abstract

隨機硬件水平的提升，在大數據集上預訓練的深度學習模型(比如BERT，GPT-3)表現出了優于傳統方法的有效性。transformer的巨大進展主要源自其強大的特征表達能力與各式各樣的架構。

在這篇論文中，作者對low-level計算機視覺任務（比如降噪、超分、去雨）進行了研究并提出了一種新的預訓練模型：IPT(image processing transformer)。為最大挖掘transformer的能力，作者采用知名的ImageNet制作了大量的退化圖像數據對，然后采用這些訓練數據對對所提IPT(它具有多頭、多尾以適配多種退化降質模型)模型進行訓練。此外，作者還引入了對比學習以更好的適配不同的圖像處理任務。經過微調后，預訓練模型可以有效的應用不到的任務中。僅僅需要一個預訓練模型，IPT即可在多個low-level基準上取得優于SOTA方案的性能。

上圖給出了所提方案IPT與HAN、RDN、RCDNet在超分、降噪、去雨任務上的性能對比，IPT均取得了0.4-2.0dB不等的性能提升。

Method

為更好的挖掘transformer的潛力以獲取在圖像處理任務上的更好結果，作者提出了一種ImageNet數據集上預訓練的圖像處理transformer，即IPT。

IPT architecture

上圖給出了IPT的整體架構示意圖，可以看到它包含四個成分：

• 用于從輸入退化圖像提取特征的Heads；
• encoder與decoder模塊用于重建輸入數據中的丟失信息；
• 用于輸出圖像重建的Tails。

Heads

為適配不同的圖像處理任務，作者提出采用多頭(multi-head每個頭包含三個卷積層)架構以分別處理不同的任務。假設輸入圖像表示為$x\in R^{3\times H\times W}$，每個頭可以生成特征$f_H\in R^{C\times H\times W}$(注：C默認64)。故多頭架構的計算過程可以描述為：
$$f_H=H^i(x),i=(1,?,N_t)$$
其中$N_t$表示圖像處理任務個數。

Transformer encoder

在將前述特征送入到transformer之前，需要將輸入特征拆分成塊，每個塊稱之為“word”。具體來講，輸入特征$f_H\in R^{C\times H\times W}$將被拆分為一系列塊$f_{p_i}\in R^{p^2\times C},i=\{1,?,N\}$，其中$N=\frac{HW}{P^2}$表示特征塊的數量，P表示特征塊尺寸。

為編碼每個塊的位置信息，作者還在encoder里面添加了可學習的位置編碼信息$E_{p_i}\in R^{P^2\times C}$。這里的encoder延續了原始transformer，采用了多頭自注意力模塊和前向網絡。

encoder的輸出表示為$f_{E_i}\in R^{P^2\times C}$,它與輸入塊尺寸相同，encoder的計算過程描述如下：

$$y_0=[E_{p_1}+f_{p_1},E_{p_2}+f_{p_2},?,E_{p_N}+f_{p_N}]$$

$$q_i=k_i=v_i=LN(y_{i?1})$$

$$y_i^{{}^{\prime }}=MSA(q_i,k_i,v_i)+y_{i?1}$$

$$y_i=FFN(LN(y_i^{{}^{\prime }}))+y_i^{{}^{\prime }},i=1,?,l$$

$$[f_{E_1},f_{E_2},?,f_{E_N}]=y_l,$$

其中$l$表示encoder的層數，MSA表示多頭自注意力模塊，FFN表示前饋前向網絡(它僅包含兩個全連接層)。

Transformer decoder

decoder采用了與encoder類似的架構并以encoder的輸出作為輸入，它包含兩個MSA與1個FFN。它與原始transformer的不同之處在于：采用任務相關的embedding作為額外的輸入，這些任務相關的embedding$E_t^i\in R^{P^2\times C},i=\{1,?,N_t\}$用于對不同任務進行特征編碼。decoder的計算過程描述如下：

$$z_0=[f_{E_1},f_{E_2},?,f_{E_N}]$$

$$q_i=k_i=LN(z_{i?1})+E_t,v_i=LN(z_{i?1})$$

$$z_i^{{}^{\prime }}=MSA(q_i,k_i,v_i)+z_{i?1}$$

$$q_i^{{}^{\prime }}=LN(z_i^{{}^{\prime }})+E_t,k_i^{{}^{\prime }}=v_i^{{}^{\prime }}=LN(z_0)$$

$$z_i^{{}^{\prime \prime }}=MSA(q_i^{{}^{\prime }},k_i^{{}^{\prime }},v_i^{{}^{\prime }})+z_i^{{}^{\prime }}$$

$$z_i=FFN(LNR(z_i^{{}^{\prime \prime }}))+z_i^{{}^{\prime \prime }},i=1,?,l$$

$$[f_{D_1},f_{D_2},?,f_{D_N}]=z_l$$

其中$f_{D_i}\in R^{P^2\times C}$表示decoder的輸出。decoder輸出的N個尺寸為$P^2\times C$的塊特征將組成特征$f_D\in R^{C\times H\times W}$。

Tails 這里的Tails屬性與Head相同，作者同樣采用多尾以適配不同的人物，其計算過程可以描述如下：
$$f_T=T^i(f_D),i=\{1,?,N_t\}$$
最終的輸出$f_T$即為重建圖像，其尺寸為$3\times H^{{}^{\prime }}\times W^{{}^{\prime }}$。輸出圖像的尺寸受任務決定，比如x2超分而言，$H^{{}^{\prime }}=2H,W^{{}^{\prime }}=2W$。

Pre-training on ImageNet

除了transformer的自身架構外，成功訓練一個優化transformer模型的關鍵因素為：大數據集。而圖像處理任務中常用數據集均比較小，比如圖像超分常用數據DIV2K僅僅有800張。針對該問題，作者提出對知名的ImageNet進行退化處理并用于訓練所提IPT模型。

這里的退化數據制作采用了與圖像處理任務中相同的方案，比如超分任務中的bicubic下采樣，降噪任務中的高斯噪聲。圖像的退化過程可以描述如下：

$$I_{corrupted}=f(I_{clean})$$

其中f表示退化變換函數，它與任務相關。對于超分任務而言，$f_{sr}$表示bicubic下采樣；對于降噪任務而言，$f_{noise}(I)=I+\eta$。IPT訓練過程中的監督損失采用了常規的$L_1$損失，描述如下：
$${\mathcal{L}}_{supervised}=\sum _{i=1}^{N_t}{L}_1(IPT(I_{corrupted}^i),I_{clean})$$

上式同樣表明：所提方案IPT同時對多個圖像處理任務進行訓練。也就說，對于每個batch，隨機從多個任務中選擇一個進行訓練，每個特定任務對應特定的head和tail。在完成IPT預訓練后，我們就可以將其用于特定任務的微調，此時可以移除掉任務無關的head和tail以節省計算量和參數量。

除了上述監督學習方式外，作者還引入了對比學習以學習更通用特征以使預訓練IPT可以應用到未知任務。對于給定輸入$x_j$(隨機從每個batch中挑選)，其decoder輸出塊特征描述為$f_{D_i}^j\in R^{P^2\times C,i=\{1,?,N\}}$。作者期望通過對比學習最小化同一圖像內的塊特征距離，最大化不同圖像的塊特征距離，這里采用的對比學習損失函數定義如下：

$$l(f_{D_{i_1}}^j,f_{D_{i_2}}^j)=?log\frac{exp(d(f_{D_{i_1}}^j,f_{D_{i_2}}^j))}{{\sum }_{k=1}^B{\mathbb{I}}_{k=j}exp(d(f_{D_{i_1}}^j,f_{D_{i_2}}^j))}$$

$${\mathcal{L}}_{constrastive}=\frac{1}{B{N}^2}\sum _{i_1=1}^N\sum _{i_2=1}^N\sum _{j=1}^{B}l(f_{D_{i_1}}^j,f_{D_{i_2}}^j)$$

其中$d(a,b)=\frac{a^{T}b}{\Vert a\Vert \Vert b\Vert }$表示cosine相似性。為更充分的利用監督與自監督信息，作者定義了如下整體損失：
$${\mathcal{L}}_{IPT}=\lambda ?{\mathcal{L}}_{constrastive}+{\mathcal{L}}_{supervised}$$

Experiments

Datasets

作者采用ImageNet數據制作訓練數據，輸入圖像塊大小為$48\times 48$，大約得到了10M圖像數據。采用了6中退化類型：x2、x3、x4、noise-30、noise-50以及去雨。

Training&Fine-tuning

作者采用32個NVIDIA Tesla V100顯卡進行IPT訓練，優化器為Adam，訓練了300epoch，初始學習率為$5e^{?5}$，經200epoch后衰減為$2e^{?5}$，batch=256。在完成IPT預訓練后，對特定任務上再進行30epoch微調，此時學習率為$2e^{?5}$。

Super-resolution

下表&下圖給出了超分方案在圖像超分任務上的性能與視覺效果對比。可以看到：

• IPT取得了優于其他SOTA超分方案的效果，甚至在Urban100數據集上以0.4dB優于其他超分方案；
• IPT可以更好重建圖像的紋理和結構信息，而其他方法則會導致模糊現象。

Denoising

下表&下圖給出了所提方法在降噪任務上的性能與視覺效果對比，可以看到：

• 在不同的噪聲水平下，IPT均取得了最佳的降噪指標，甚至在Urban100數據上提升高達2dB。
• IPT可以很好的重建圖像的紋理&結構信息，而其他降噪方法則難以重建細節信息。

Deraining

下表&下圖給出了所提方法在圖像去雨任務上的性能與視覺效果對比。可以看到：

• 所提方法取得了最好的指標，甚至取得了1.62dB的性能提升；
• IPT生成的圖像具有與GT最詳盡，且具有更好的視覺效果。

Generalization Ability

為說明所提方法的泛化性能，作者采用了未經訓練的噪聲水平進行驗證，結果見下表。可以看到：盡管未在該其噪聲水平數據上進行訓練，所提IPT仍取得了最佳的指標。

Ablation Study

下圖對比了IPT與EDSR在不同數量訓練集上的性能對比，可以看到：當訓練集數量較少時，EDSR具有更好的指標；而當數據集持續增大后，EDSR很快達到飽和，而IPT仍可持續提升并大幅超過了EDSR。

下表給出了對比損失對于模型性能影響性分析(x2超分任務)。當僅僅采用監督方式進行訓練時，IPT的指標為38.27；而當引入對比學習機制后，其性能可以進一步提升0.1dB。這側面印證了對比學習對于IPT預訓練的有效性。

全文到此結束，對比感興趣的同學建議查看原文。

作者介紹：Happy，一個愛“胡思亂想”的AI行者。歡迎大家聯系極市小編（微信ID:fengcall19）加入極市原創作者行列

總結

以上是生活随笔為你收集整理的Transformer再下一城！low-level多个任务榜首被占领，北大华为等联合提出预训练模型IPT的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。