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

本文首發(fā)自極市平臺，作者@Happy，轉(zhuǎn)載需獲授權(quán)。

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

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

Abstract

隨機硬件水平的提升，在大數(shù)據(jù)集上預(yù)訓(xùn)練的深度學(xué)習(xí)模型(比如BERT，GPT-3)表現(xiàn)出了優(yōu)于傳統(tǒng)方法的有效性。transformer的巨大進展主要源自其強大的特征表達能力與各式各樣的架構(gòu)。

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

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

Method

為更好的挖掘transformer的潛力以獲取在圖像處理任務(wù)上的更好結(jié)果，作者提出了一種ImageNet數(shù)據(jù)集上預(yù)訓(xùn)練的圖像處理transformer，即IPT。

IPT architecture

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

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

Heads

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

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}$表示特征塊的數(shù)量，P表示特征塊尺寸。

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

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的層數(shù)，MSA表示多頭自注意力模塊，FFN表示前饋前向網(wǎng)絡(luò)(它僅包含兩個全連接層)。

Transformer decoder

decoder采用了與encoder類似的架構(gòu)并以encoder的輸出作為輸入，它包含兩個MSA與1個FFN。它與原始transformer的不同之處在于：采用任務(wù)相關(guān)的embedding作為額外的輸入，這些任務(wù)相關(guān)的embedding$E_t^i\in R^{P^2\times C},i=\{1,?,N_t\}$用于對不同任務(wù)進行特征編碼。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 }}$。輸出圖像的尺寸受任務(wù)決定，比如x2超分而言，$H^{{}^{\prime }}=2H,W^{{}^{\prime }}=2W$。

Pre-training on ImageNet

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

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

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

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

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

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

$$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相似性。為更充分的利用監(jiān)督與自監(jiān)督信息，作者定義了如下整體損失：
$${\mathcal{L}}_{IPT}=\lambda ?{\mathcal{L}}_{constrastive}+{\mathcal{L}}_{supervised}$$

Experiments

Datasets

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

Training&Fine-tuning

作者采用32個NVIDIA Tesla V100顯卡進行IPT訓(xùn)練，優(yōu)化器為Adam，訓(xùn)練了300epoch，初始學(xué)習(xí)率為$5e^{?5}$，經(jīng)200epoch后衰減為$2e^{?5}$，batch=256。在完成IPT預(yù)訓(xùn)練后，對特定任務(wù)上再進行30epoch微調(diào)，此時學(xué)習(xí)率為$2e^{?5}$。

Super-resolution

下表&下圖給出了超分方案在圖像超分任務(wù)上的性能與視覺效果對比?？梢钥吹?#xff1a;

• IPT取得了優(yōu)于其他SOTA超分方案的效果，甚至在Urban100數(shù)據(jù)集上以0.4dB優(yōu)于其他超分方案；
• IPT可以更好重建圖像的紋理和結(jié)構(gòu)信息，而其他方法則會導(dǎo)致模糊現(xiàn)象。

Denoising

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

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

Deraining

下表&下圖給出了所提方法在圖像去雨任務(wù)上的性能與視覺效果對比?？梢钥吹?#xff1a;

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

Generalization Ability

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

Ablation Study

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

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

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

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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。