What？UFO！ | UFO-ViT用X-Norm讓你的Transformer模型回歸線性復(fù)雜度
在本文中提出了一種簡單的方法，使自注意力具有線性復(fù)雜度而不損失性能。通過替換softmax函數(shù)，利用矩陣乘法的關(guān)聯(lián)律消除二次運算。在早期的研究中，這種類型的因子分解通常會導(dǎo)致性能下降。在圖像分類和目標(biāo)檢測方面，UFO-ViT模型優(yōu)于大多數(shù)現(xiàn)有的基于Transformer和CNN模型的結(jié)果。
1簡介
Vision transformers已成為計算機視覺任務(wù)的重要模型之一。雖然它們優(yōu)于早期的卷積網(wǎng)絡(luò)，但使用傳統(tǒng)的自注意力算法時，其復(fù)雜度是。
在這里，本文提出了UFO-ViT(Unit Force Operated Vision Trnasformer)，通過修改少數(shù)行自注意力消除一些非線性來減少自注意力的計算復(fù)雜度，UFO-ViT實現(xiàn)線性復(fù)雜性而不降低性能。該模型在圖像分類和密集預(yù)測任務(wù)上優(yōu)于大多數(shù)基于Transformer的模型。
本文方法的貢獻:
提出了新的約束方案XNorm，它生成單元超球來提取關(guān)系特征，防止自注意力過度依賴于初始化方式。此外，它具有O(N)復(fù)雜度，能有效地處理高分辨率輸入。 證明了UFO-ViT可以用于一般用途。模型在圖像分類和密集預(yù)測任務(wù)上進行了實驗。在較低的容量和FLOPS，UFO-ViT使用線性自注意力方案依然優(yōu)于大多數(shù)基于Transformer模型。特別是，模型在輕量狀態(tài)下表現(xiàn)良好。 提出了一種動態(tài)處理輸入圖像的新模塊。這意味著權(quán)重的大小與輸入分辨率無關(guān)。它適用于密集的預(yù)測任務(wù)，如目標(biāo)檢測或語義分割，這通常需要比訓(xùn)練前更高的分辨率。對于基于MLP的結(jié)構(gòu)，如果模型在不同分辨率上預(yù)訓(xùn)練，則需要進行后處理。
2本文方法
模型的結(jié)構(gòu)如圖2所示。它是一個混合的卷積層，UFO模塊和一個簡單的前饋MLP層。接下來將指出所提出的方法是如何取代softmax函數(shù)并確保線性復(fù)雜度的。
2.1 理論解釋
對于輸入，傳統(tǒng)的自注意力機制表示為:
A表示注意力操作符。
由于softmax的非線性導(dǎo)致不能分解為O(N×h+h×N)。本文方法是先用結(jié)合律消去softmax來計算。由于應(yīng)用identity反而會導(dǎo)致退化，所以作者建議使用簡單的約束來防止它退化。
本文提出方法叫XNorm，定義如下：
式中，γ為可學(xué)習(xí)參數(shù)，h為嵌入維數(shù)。它是簡單的l2范數(shù)，但它可以適用于的空間維數(shù)和q的通道維數(shù)，這就是為什么它被稱為交叉歸一化。
使用關(guān)聯(lián)法則，key和value先相乘，query再相乘。如圖3所示。這兩個乘法的復(fù)雜度都是O(hNd)，所以這個過程與n是線性的。
2.2 XNorm
將softmax替換為XNorm
在過程中，key和value直接相乘。與線性核方法相同生成h個聚類。
XNorm應(yīng)用于query和輸出。
其中x為輸入。最后，投影權(quán)重標(biāo)度，加權(quán)求和。
在這個公式中，關(guān)系特征是由patches和聚類之間的余弦相似度定義的。XNorm將query和clusters限制為單位向量。它通過規(guī)范化它們以縮小范圍，防止它們的value抑制關(guān)系屬性。如果它們有任意值，則注意力區(qū)域太依賴于初始化。
然而，這種解釋不足以解釋原因XNorm必須是l2標(biāo)準(zhǔn)化形式。在這里介紹另一種理論觀點，考慮簡單的物理模擬。
考慮殘差連接，任意模塊的輸出公式為:
其中n和x表示當(dāng)前深度和輸入圖像。假設(shè)x為某物體的位移，n為時間，則上式重新定義為:
大部分神經(jīng)網(wǎng)絡(luò)是離散的，所以?t是常數(shù)。(簡便起見，設(shè)?t=1)殘差項表示速度，與力相同，?t=1。
在物理學(xué)中，胡克定律被定義為彈性向量k與位移向量x的點積。彈性力產(chǎn)生調(diào)和勢U是的函數(shù)。從物理上講，勢能干擾了粒子通過它的運動路徑。(為了便于理解，想象一個球沿著拋物線形狀的軌道移動)
以上公式一般用來近似分子在x≈0處的勢能。對于多分子，利用彈性的線性，
除了上面的公式?jīng)]有歸一化。假設(shè)k的個數(shù)很小，那么k就太大了。如果某個粒子繞著它們運動，它就不容易逃脫。在這種情況下，只存在兩種情況。1. 崩潰了。粒子沿最大k方向排列。關(guān)系特征被忽略了。一些x大的粒子在坍縮中幸存下來。
Collapsed。粒子沿最大k的方向排列。 Relational features are neglected。一些x大的粒子在坍縮中幸存下來。
這與前一部分提到的情況完全相同。XNorm強制所有向量為單位向量，以防止這種情況發(fā)生。換句話說，XNorm不是標(biāo)準(zhǔn)化，而是約束。
在注意力模塊中，F(xiàn)FN是不可忽視的。然而，這里有一個簡單的解釋。FFN是靜態(tài)的，不依賴于空間維度。在物理上，這類函數(shù)表示諧振子方程的驅(qū)動力。
這就起到了放大器或減速器的作用，增強或消除與空間關(guān)系無關(guān)的特征。
3UFO-ViT
構(gòu)建UFO-ViT模型采用了早期視覺轉(zhuǎn)換器模型中的一些架構(gòu)策略。在本節(jié)將介紹幾種架構(gòu)優(yōu)化技術(shù)。總體結(jié)構(gòu)如圖2所示。
1、Patch embedding with convolutions
最近的一些研究表明Patch embedding with convolutions可以較好地訓(xùn)練ViT，而不是線性投影。
2、Positional encoding
使用位置編碼作為可學(xué)習(xí)參數(shù)。
3、Multi-headed attention
本文所設(shè)計的模塊是多方向的，以便更好的正則化。將公式3中的γ參數(shù)應(yīng)用于所有Head以衡量每個Head的重要性。
4、Local patch interaction
為自注意力模塊設(shè)計額外的模塊來提取局部特征，目前已不是什么特別的想法。作者選擇其中最簡單的方法，3×3深度可分離卷積。采用早期在xCiT中提出的LPI，使用堆疊的2個conv-bn-GELU層。
5、Feed-forward network
與傳統(tǒng)基于Transformer的模型一樣，本文的模型使用4維隱藏MLP的point-wise feed-forward。
6、Class attention
在ImageNet1k實驗中使用了CaiT中提供的類注意力層。它幫助Class Token收集空間信息。為了減少計算，只在Class Token上計算Class attention。這與原始文件是一樣的。注意，Class Attention由UFO模塊組成，而CaiT使用正常的自注意力模塊來進行Class attention。
4實驗
4.1 Image Classification
消融研究重點是XNorm的重要性，以及在3.3節(jié)中解釋過的架構(gòu)優(yōu)化。作者試驗了各種歸一化方法。大多數(shù)其他歸一化方法都不能減少損失。它可以是證明本文的理論解釋是合理的隱性證據(jù)之一。有趣的是，應(yīng)用single L2Norm的性能也很差。
用DeiT對三種具有相同結(jié)構(gòu)設(shè)計方案的模型進行了實驗。如圖1所示，所有的模型都顯示出比大多數(shù)基于Transformer的并發(fā)模型更高的性能和參數(shù)效率。此外，本文提出的模型在復(fù)雜性和數(shù)據(jù)效率方面有優(yōu)勢，而原始的ViT需要更大的額外數(shù)據(jù)集，如JFT-300M或ImageNet21k。
4.2 Object Detection with Mask R-CNN
比較CNN和基于Transformer的目標(biāo)檢測和實例分割任務(wù)。為了公平比較，所有結(jié)果的實驗環(huán)境都是相同的。所有模型都是在ImageNet1k數(shù)據(jù)集上預(yù)訓(xùn)練的。
本文模型明顯優(yōu)于基于CNN的模型。而且，與容量較低的最先進的視覺Transformer相比，它們?nèi)〉昧烁呋蚋懈偁幜Φ慕Y(jié)果。但是使用相同的設(shè)計空間，XCiT顯示的結(jié)果稍微好一些?？赡苁荴CiT模型有更大的嵌入空間d。UFO-ViT-B在bbox檢測任務(wù)上的表現(xiàn)略低于UFO-ViT-M，但在較小的bbox和整體實例分割得分上更好。
5參考
[1].UFO-ViT: High Performance Linear Vision Transformer without Softmax
6推薦閱讀
Transformer-Unet | 如何用Transformer一步一步改進 Unet？
高效Transformer | 85FPS！CNN + Transformer語義分割的又一境界,真的很快！
【書童的學(xué)習(xí)筆記】集智小書童建議你這么學(xué)習(xí)Transformer，全干貨?。?！
長按掃描下方二維碼添加小助手。
可以一起討論遇到的問題