Transformer再下一城!low-level多个任务榜首被占领,北大华为等联合提出预训练模型IPT
來自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∈R3×H×Wx \in R^{3\times H \times W}x∈R3×H×W,每個頭可以生成特征fH∈RC×H×Wf_{H} \in R^{C \times H \times W}fH?∈RC×H×W(注:C默認64)。故多頭架構的計算過程可以描述為:
fH=Hi(x),i=(1,?,Nt)f_H = H^i(x), i=(1,\cdots,N_t) fH?=Hi(x),i=(1,?,Nt?)
其中NtN_tNt?表示圖像處理任務個數。
Transformer encoder
在將前述特征送入到transformer之前,需要將輸入特征拆分成塊,每個塊稱之為“word”。具體來講,輸入特征fH∈RC×H×Wf_{H} \in R^{C \times H \times W}fH?∈RC×H×W將被拆分為一系列塊fpi∈Rp2×C,i={1,?,N}f_{p_i} \in R^{p^2 \times C}, i=\{1, \cdots, N\}fpi??∈Rp2×C,i={1,?,N},其中N=HWP2N=\frac{HW}{P^2}N=P2HW?表示特征塊的數量,P表示特征塊尺寸。
為編碼每個塊的位置信息,作者還在encoder里面添加了可學習的位置編碼信息Epi∈RP2×CE_{p_i} \in R^{P^2 \times C}Epi??∈RP2×C。這里的encoder延續了原始transformer,采用了多頭自注意力模塊和前向網絡。
encoder的輸出表示為fEi∈RP2×Cf_{E_i} \in R^{P^2 \times C}fEi??∈RP2×C,它與輸入塊尺寸相同,encoder的計算過程描述如下:
y0=[Ep1+fp1,Ep2+fp2,?,EpN+fpN]y_0 = [E_{p_1} + f_{p_1}, E_{p_2} + f_{p_2}, \cdots, E_{p_N} + f_{p_N}] y0?=[Ep1??+fp1??,Ep2??+fp2??,?,EpN??+fpN??]
qi=ki=vi=LN(yi?1)q_i = k_i = v_i = LN(y_{i-1}) qi?=ki?=vi?=LN(yi?1?)
yi′=MSA(qi,ki,vi)+yi?1y_i^{'} = MSA(q_i, k_i, v_i) + y_{i-1} yi′?=MSA(qi?,ki?,vi?)+yi?1?
yi=FFN(LN(yi′))+yi′,i=1,?,ly_i = FFN(LN(y_i^{'})) + y_i^{'} , i = 1, \cdots, l yi?=FFN(LN(yi′?))+yi′?,i=1,?,l
[fE1,fE2,?,fEN]=yl,[f_{E_1}, f_{E_2}, \cdots, f_{E_N}] = y_l, [fE1??,fE2??,?,fEN??]=yl?,
其中lll表示encoder的層數,MSA表示多頭自注意力模塊,FFN表示前饋前向網絡(它僅包含兩個全連接層)。
Transformer decoder
decoder采用了與encoder類似的架構并以encoder的輸出作為輸入,它包含兩個MSA與1個FFN。它與原始transformer的不同之處在于:采用任務相關的embedding作為額外的輸入,這些任務相關的embeddingEti∈RP2×C,i={1,?,Nt}E_t^i \in R^{P^2 \times C},i=\{1, \cdots, N_t\}Eti?∈RP2×C,i={1,?,Nt?}用于對不同任務進行特征編碼。decoder的計算過程描述如下:
z0=[fE1,fE2,?,fEN]z_0 = [f_{E_1}, f_{E_2}, \cdots, f_{E_N}] z0?=[fE1??,fE2??,?,fEN??]
qi=ki=LN(zi?1)+Et,vi=LN(zi?1)q_i = k_i = LN(z_{i-1}) + E_t, v_i = LN(z_{i-1}) qi?=ki?=LN(zi?1?)+Et?,vi?=LN(zi?1?)
zi′=MSA(qi,ki,vi)+zi?1z_i^{'} = MSA(q_i,k_i, v_i) + z_{i-1} zi′?=MSA(qi?,ki?,vi?)+zi?1?
qi′=LN(zi′)+Et,ki′=vi′=LN(z0)q_i^{'} = LN(z_i^{'}) + E_t, k_i^{'} = v_i^{'} = LN(z_0) qi′?=LN(zi′?)+Et?,ki′?=vi′?=LN(z0?)
zi′′=MSA(qi′,ki′,vi′)+zi′z_i^{''} = MSA(q_i^{'}, k_i^{'}, v_i^{'}) + z_i^{'} zi′′?=MSA(qi′?,ki′?,vi′?)+zi′?
zi=FFN(LNR(zi′′))+zi′′,i=1,?,lz_i = FFN(LNR(z_i^{''})) + z_i^{''}, i=1,\cdots,l zi?=FFN(LNR(zi′′?))+zi′′?,i=1,?,l
[fD1,fD2,?,fDN]=zl[f_{D_1}, f_{D_2}, \cdots, f_{D_N}] = z_l [fD1??,fD2??,?,fDN??]=zl?
其中fDi∈RP2×Cf_{D_i} \in R^{P^2 \times C}fDi??∈RP2×C表示decoder的輸出。decoder輸出的N個尺寸為P2×CP^2 \times CP2×C的塊特征將組成特征fD∈RC×H×Wf_D \in R^{C\times H \times W}fD?∈RC×H×W。
Tails 這里的Tails屬性與Head相同,作者同樣采用多尾以適配不同的人物,其計算過程可以描述如下:
fT=Ti(fD),i={1,?,Nt}f_T = T^i(f_{D}) , i=\{1, \cdots, N_t\} fT?=Ti(fD?),i={1,?,Nt?}
最終的輸出fTf_TfT?即為重建圖像,其尺寸為3×H′×W′3\times H^{'} \times W^{'}3×H′×W′。輸出圖像的尺寸受任務決定,比如x2超分而言,H′=2H,W′=2WH^{'} = 2H, W^{'} = 2WH′=2H,W′=2W。
Pre-training on ImageNet
除了transformer的自身架構外,成功訓練一個優化transformer模型的關鍵因素為:大數據集。而圖像處理任務中常用數據集均比較小,比如圖像超分常用數據DIV2K僅僅有800張。針對該問題,作者提出對知名的ImageNet進行退化處理并用于訓練所提IPT模型。
這里的退化數據制作采用了與圖像處理任務中相同的方案,比如超分任務中的bicubic下采樣,降噪任務中的高斯噪聲。圖像的退化過程可以描述如下:
Icorrupted=f(Iclean)I_{corrupted} = f(I_{clean}) Icorrupted?=f(Iclean?)
其中f表示退化變換函數,它與任務相關。對于超分任務而言,fsrf_{sr}fsr?表示bicubic下采樣;對于降噪任務而言,fnoise(I)=I+ηf_{noise}(I) = I + \etafnoise?(I)=I+η。IPT訓練過程中的監督損失采用了常規的L1L_1L1?損失,描述如下:
Lsupervised=∑i=1NtL1(IPT(Icorruptedi),Iclean)\mathcal{L}_{supervised} = \sum_{i=1}^{N_t} L_1(IPT(I_{corrupted}^i), I_{clean}) Lsupervised?=i=1∑Nt??L1?(IPT(Icorruptedi?),Iclean?)
上式同樣表明:所提方案IPT同時對多個圖像處理任務進行訓練。也就說,對于每個batch,隨機從多個任務中選擇一個進行訓練,每個特定任務對應特定的head和tail。在完成IPT預訓練后,我們就可以將其用于特定任務的微調,此時可以移除掉任務無關的head和tail以節省計算量和參數量。
除了上述監督學習方式外,作者還引入了對比學習以學習更通用特征以使預訓練IPT可以應用到未知任務。對于給定輸入xjx_jxj?(隨機從每個batch中挑選),其decoder輸出塊特征描述為fDij∈RP2×C,i={1,?,N}f_{D_i}^j \in R^{P^2 \times C, i=\{1, \cdots, N\}}fDi?j?∈RP2×C,i={1,?,N}。作者期望通過對比學習最小化同一圖像內的塊特征距離,最大化不同圖像的塊特征距離,這里采用的對比學習損失函數定義如下:
l(fDi1j,fDi2j)=?logexp(d(fDi1j,fDi2j))∑k=1BIk≠jexp(d(fDi1j,fDi2j))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\ne j} exp(d(f_{D_{i_1}}^j, f_{D_{i_2}}^j))} l(fDi1??j?,fDi2??j?)=?log∑k=1B?Ik?=j?exp(d(fDi1??j?,fDi2??j?))exp(d(fDi1??j?,fDi2??j?))?
Lconstrastive=1BN2∑i1=1N∑i2=1N∑j=1Bl(fDi1j,fDi2j)\mathcal{L}_{constrastive} = \frac{1}{BN^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) Lconstrastive?=BN21?i1?=1∑N?i2?=1∑N?j=1∑B?l(fDi1??j?,fDi2??j?)
其中d(a,b)=aTb∥a∥∥b∥d(a,b) = \frac{a^Tb}{\|a\| \|b\|}d(a,b)=∥a∥∥b∥aTb?表示cosine相似性。為更充分的利用監督與自監督信息,作者定義了如下整體損失:
LIPT=λ?Lconstrastive+Lsupervised\mathcal{L}_{IPT} = \lambda \cdot \mathcal{L}_{constrastive} + \mathcal{L}_{supervised} LIPT?=λ?Lconstrastive?+Lsupervised?
Experiments
Datasets
作者采用ImageNet數據制作訓練數據,輸入圖像塊大小為48×4848 \times 4848×48,大約得到了10M圖像數據。采用了6中退化類型:x2、x3、x4、noise-30、noise-50以及去雨。
Training&Fine-tuning
作者采用32個NVIDIA Tesla V100顯卡進行IPT訓練,優化器為Adam,訓練了300epoch,初始學習率為5e?55e^{-5}5e?5,經200epoch后衰減為2e?52e^{-5}2e?5,batch=256。在完成IPT預訓練后,對特定任務上再進行30epoch微調,此時學習率為2e?52e^{-5}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的全部內容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: 何恺明团队最新力作SimSiam:消除表
- 下一篇: 首次摆脱对梯度的依赖,CMU等开源Sco