摘要
利用注意力機制的視覺表示模型在追求大感受野時面臨著巨大的計算開銷。在本研究中,我們通過引入基于物理熱傳導原理的熱傳導算子(Heat Conduction Operator, HCO)來緩解這一挑戰這么高級咩(⊙o⊙)!
。HCO將圖像塊視為熱源,并通過自適應的熱能擴散來模擬它們之間的相關性,從而實現魯棒的視覺表示。HCO的計算復雜度為O(N^1.5),因為它可以通過離散余弦變換(DCT)操作實現。HCO是一種即插即用的模塊,與深度學習主干網絡結合后可以產生具有全局感受野的視覺表示模型(稱為vHeat)。在各種視覺任務的實驗中表明,除了性能更強之外,與Swin-Transformer相比,vHeat實現了高達3倍的吞吐量、減少80%的GPU內存分配以及減少35%的計算浮點運算量(FLOPs)。
代碼可在以下鏈接獲取:https://github.com/MzeroMiko/vHeat 和 https://openi.pcl.ac.cn/georgew/vHeat。
1、引言
卷積神經網絡(CNNs)自深度學習興起以來一直是視覺表示的基石,在各種視覺任務中展現出卓越的性能。然而,依賴于局部感受野和固定的卷積操作限制了其捕捉圖像中長距離和復雜依賴關系的能力。這些局限性激發了開發替代視覺表示模型的濃厚興趣,包括基于視覺變換器(ViTs)和狀態空間模型的架構。盡管這些模型表現出色,但它們仍面臨挑戰,包括相對較高的計算復雜度和缺乏可解釋性。
-----------------------------------------------------------------------------CNN和VIT架構解析---------------------------------------------------------------------------------
- CNN
- 卷積操作。基于卷積核在輸入圖像上滑動,提取局部特征,這種局部感受野特性使得CNN能夠很好地捕捉圖像中的局部紋理、邊緣等信息。CNN的特征提取是先提取局部信息,在逐步抽象到整體特征。
- CNN結構。CNN結構通常包含卷積層、池化層和全連接層。池化層用于降低特征圖的空間維度,減少計算量,同時保留重要特征。全連接層則用于將提取到的特征進行組合,輸出最終的分類結果或回歸值。
- 歸納偏置。歸納偏置其實就是一種先驗知識,一種提前做好的假設。在CNN中的歸納偏置一般包括兩類: locality(局部性)和translation equivariance(平移不變性):
- locality:假設相鄰的區域會有相鄰的特征,靠得越近的東西相關性能也就越強。局部性可以控制模型的復雜度。
- translation equivariance:由于卷積核是一樣的所以不管圖片中的物體移動到哪里,只要是同樣的輸入進來遇到同樣的卷積核,那么輸出就是一樣的。利用平移等變形可以很好的提高模型的泛化能力。
- ViT
- 對比。先前將Transformer應用于CV領域的工作,都和CNN相關。要么時將CNN卷積后的特征圖應用自注意力機制,要么是將注意力機制當成卷積核,應用于局部圖像。這樣做的原因是,假設一張圖像的大小是224*224,那么對每個像素點進行自注意力操作后,計算量會變得很大。因此Transformer的最大改動是將圖像分塊,這樣大大減少了計算量。
- 結構。ViT通過自注意力機制同時考慮圖像塊之間的全局關系。在處理圖像時,每個圖像塊的特征更新都會受到其他所有圖像塊的影響。例如,在識別一張風景圖片時,ViT能夠同時考慮天空、山脈、河流等不同區域之間的關系,從而更好地理解整個圖像的語義信息。不具有CNN那樣的明確的層次化特征提取過程。雖然Vi也可以T通過堆疊多層Transformer模塊來提取更復雜的特征,但它的每一層都是在全局范圍內進行特征更新,而不是像CNN那樣逐層從局部到整體提取特征。
- 在大規模數據集上表現出色。當有足夠大的數據集進行訓練時,ViT能夠學習到更豐富的特征表示。例如在ImageNet - 21k(一個包含21000個類別的大規模圖像數據集)上預訓練的ViT模型,在遷移到其他任務時往往能夠取得更好的性能。
- 總結
使用基于CNN的方法存在感受野有限的問題,不能很好的建模長遠的依賴關系(全局信息),而基于transformer的方法可以很好的建模全局信息但是transformer反而缺乏類似于CNN的歸納偏置,這些先驗信息必須通過大量的數據來進行學習,所以小的數據在CNN上取得的效果一般優于基于transformer的方法。訓練基于CNN的方法通常只需要一個較小的數據集,而訓練基于transformer的方法一般需要再大的數據集上進行預訓練。
------------------------------------------------------------------------------結束---------------------------------------------------------------------------------------------------
在解決這些限制時,我們從熱傳導領域汲取靈感,其中空間局部性對于熱能的傳遞至關重要,因為相鄰粒子的碰撞導致了熱能的傳遞。值得注意的是,熱傳導原理與視覺語義在空間域內的傳播之間存在類比關系,因為特定尺度下相鄰的圖像區域往往包含相關的信息或具有相似的特征。基于這些聯系,我們引入了vHeat,這是一種受物理啟發的視覺表示模型,它將圖像塊視為熱源,并將它們之間相關性的計算建模為熱能的擴散。
你小子嘰里咕嚕說什么呢,一點都聽不懂,努力理解一下:
- 熱傳導原理:熱傳導是熱量從一個地方傳到另一個地方的過程,熱量傳導過程其實是鐵棒里的粒子相互碰撞,熱能就從熱的地方傳導冷的地方。
- 圖像中的熱傳導類比:在圖像里,相鄰的區域往往也很相似。比如,一張照片里的一塊草地,草地上的每一小塊區域(比如幾根草)看起來都差不多,它們的顏色、紋理等特征很相似。這就像是熱傳導里的“相鄰粒子”,它們之間有很強的關聯。
圖2. 信息傳導機制的比較:自注意力與熱傳導。(a)自注意力算子將信息從一個像素“傳導”到所有其他像素,導致復雜度為O(N2)。(b)熱傳導算子(HCO)將中心像素視為熱源,并通過離散余弦變換(DCT,記為F)和逆離散余弦變換(IDCT,記為F?1)來傳導信息傳播,該方法具有可解釋性、全局感受野以及O(N1?)的復雜度。
為了將熱傳導原理整合到深度網絡中,我們首先推導出二維空間中熱傳導的一般解,并將其擴展到多個維度,對應于特征通道。基于這一一般解,我們設計了熱傳導算子(HCO),它模擬了視覺語義在圖像塊之間的多維傳播。值得注意的是,我們證明了HCO可以通過二維離散余弦變換(DCT)和逆離散余弦變換(IDCT)近似實現,有效將計算復雜度降低到O(N^1.5)。這一改進由于DCT和IDCT操作的高度并行性,顯著提高了訓練和測試的效率。此外,由于DCT獲得的頻域中的每個元素都包含了圖像空間中所有塊的信息,因此vHeat能夠建立長距離的特征依賴關系,并實現全局感受野。為了增強vHeat的表示適應性,我們提出了可學習的頻率值嵌入(FVEs),以表征頻率信息并預測視覺熱傳導的熱擴散率。
我們開發了一系列vHeat模型(即vHeat-Tiny/Small/Base),并通過廣泛的實驗驗證了它們在多樣化視覺任務中的有效性。與各種架構的基準視覺主干網絡(例如ConvNeXt、Swin和Vim)相比,vHeat在圖像分類、目標檢測和語義分割任務中始終表現出更優的性能。具體而言,vHeat-Base在ImageNet-1K上實現了84.0%的top-1準確率,比Swin高出0.5%,并且其吞吐量比Swin高出超過40%(661對比456)。為了探索vHeat的泛化能力,我們還在魯棒性評估基準和低級視覺任務中驗證了其優越性。此外,由于HCO的O(N^1.5)復雜度,vHeat相比基于ViT的模型具有顯著更低的計算成本,表現出更少的浮點運算量(FLOPs)和GPU內存需求,以及隨著圖像分辨率增加而更高的吞吐量。特別是當輸入圖像分辨率增加到768×768時,vHeat-Base的吞吐量比Swin高出3倍,GPU內存分配減少80%,計算FLOPs減少35%。
本研究的貢獻總結如下:
- 我們提出了vHeat,這是一種受熱傳導物理原理啟發的視覺主干模型,同時實現了全局感受野、低計算復雜度和高可解釋性。
- 我們設計了熱傳導算子(HCO),這是一個物理上合理的模塊,將圖像塊視為熱源,通過頻率值嵌入(FVEs)預測自適應的熱擴散率,并按照熱傳導原理傳遞信息。
- vHeat在視覺任務(包括圖像分類、目標檢測和語義分割)中取得了令人鼓舞的性能,同時在高分辨率圖像上實現了更高的推理速度、更少的FLOPs和更低的GPU內存使用。
2、相關工作
2.1 卷積神經網絡(CNNs)
CNNs在視覺感知的歷史中一直是里程碑式的模型[30, 31]。CNNs的獨特特性體現在卷積核上,這些卷積核在特定設計的GPU上具有很高的計算效率。借助強大的GPU和大規模數據集[14],人們提出了越來越深[24, 29, 52, 57]、更高效的模型[27, 46, 58, 73],以在各種視覺任務中實現更高的性能。為了提升卷積操作的能力[10]、效率[28, 74]和適應性[11, 69],人們對其進行了大量改進。然而,CNNs的局部感受野這一固有限制依然存在。最近開發的大型卷積核[16]在擴大感受野方面邁出了一步,但在處理高分辨率圖像時遇到了困難。
2.2 視覺變換器(Vision Transformers)
基于自注意力算子[63]構建的ViTs在構建全局特征依賴方面具有天然優勢。憑借自注意力在所有圖像塊之間的學習能力,ViTs在有大規模數據集用于預訓練的情況下,已成為迄今為止最強大的視覺模型[18, 45, 61]。引入分層架構[12, 15, 17, 37, 41, 60, 67, 79, 80]進一步提升了ViTs的性能。然而,ViTs是O(N2)的計算復雜度,這意味著在處理高分辨率圖像時會帶來巨大的計算開銷。人們通過引入窗口注意力、線性注意力和交叉協方差注意力算子[1, 6, 37, 66]來提高模型效率,但代價是縮小了感受野或降低了非線性能力。還有研究通過在ViTs中引入卷積操作[12, 64, 68]或設計混合架構將CNN與ViT模塊結合起來[12, 41, 54],提出了混合網絡。
2.3 狀態空間模型和循環神經網絡(RNNs)
狀態空間模型(SSMs)[22, 43, 65]具有線性復雜度的長序列建模能力,也被從自然語言領域遷移過來(例如Mamba[21])。通過將選擇性掃描機制適應于二維圖像,設計了視覺SSMs[36, 82]。然而,基于選擇性掃描機制的SSMs由于并行性有限,限制了其整體潛力。最近的受控加權鍵值(RWKV)和RetNet模型[44, 56]在保持線性復雜度的同時提高了并行性。它們結合了Transformer的高效并行化訓練和RNN的高效推理,利用線性注意力機制,使模型既可以表示為Transformer,也可以表示為RNN,從而在訓練期間并行化計算,并在推理期間保持恒定的計算和內存復雜度。盡管具有這些優勢,但將二維圖像建模為序列會降低可解釋性。
2.4 生物和物理啟發模型
生物和物理原理一直是創造視覺模型的源泉。擴散模型[26, 49, 53]受到非平衡熱力學[13]的啟發,能夠通過為擴散步驟定義馬爾可夫鏈來生成圖像。QB-Heat[8]利用物理熱傳導方程作為掩碼圖像建模任務的監督信號。脈沖神經網絡(SNNs)[20, 32, 59]聲稱能更好地模擬生物神經元的信息傳遞,為簡單的視覺任務構建模型[4]。這些模型的成功鼓勵我們探索物理熱傳導原理,以開發視覺表示模型。(○′・д・)ノ
3、 Methodology
4、實驗
1. 圖像分類
圖像分類的結果總結在表1中。在相似的計算量(FLOPs)下,vHeat實現了82.2%的top-1準確率,分別超過了Swin-T和Vim-S 0.9%和0.8%。值得注意的是,vHeat在Small和Base尺度上也表現出優越性。具體來說,vHeat-B在僅11.2 GFLOPs和68M模型參數的情況下,實現了84.0%的top-1準確率,分別超過了Swin-B和Vim-B 0.5%和0.8%。
在計算效率方面,vHeat在Tiny/Small/Base模型尺度上相比基準模型享有顯著更高的推理速度。例如,vHeat-T實現了每秒1514張圖像的吞吐量,比Vim-S高出87%,比ConvNeXt-T高出26%,比Swin-T高出22%,同時保持了性能優勢。
2. 目標檢測和實例分割
作為骨干網絡,vHeat在MSCOCO2017數據集上進行了目標檢測和實例分割的測試。我們加載了分類預訓練的vHeat權重進行下游評估。由于輸入圖像大小與分類任務不同,需要將FVEs或k的形狀對齊到下游任務的目標圖像大小。更多細節請參考補充材料中的D.1節。
目標檢測的結果總結在表2中,vHeat在兩種訓練計劃(12或36個epoch)中均在box/mask平均精度(APb和APm)方面表現出優越性。例如,在12個epoch的微調計劃下,vHeat-T/S/B模型分別實現了45.1%/46.8%/47.7%的目標檢測mAP,分別超過了Swin-T/S/B 2.4%/2.0%/0.8% mAP,以及ConvNeXt-T/S/B 0.9%/1.4%/0.7% mAP。在相同的配置下,vHeat-T/S/B分別實現了41.2%/42.3%/43.0%的實例分割mAP,超過了Swin-T/S/B和ConvNeXt-T/S/B。在36個epoch(3×)的微調計劃下,vHeat的優勢依然存在,且在多尺度訓練下表現穩定。此外,vHeat相比Swin和ConvNeXt享有更高的推理速度(FPS)。例如,vHeat-B實現了每秒20.2張圖像,比Swin-B/ConvNeXt-B(13.8/14.1張圖像/秒)高出46%/43%。這些結果突出了vHeat在密集預測下游任務中提供強大性能和效率的潛力。
3. 語義分割
在ADE20K數據集上的結果總結在表3中,vHeat在Tiny/Small/Base尺度上一致地超過了其他基線模型。例如,vHeat-B分別超過了NAT-B和ViL-B 1.1%/0.8% mIoU。
4. 魯棒性評估
為了驗證vHeat的魯棒性,我們在ObjectNet和ImageNet-A等分布外分類數據集上評估了vHeat-B。我們在表4中測量了這兩個基準的Top-1準確率(%)。結果表明,vHeat在分布外數據上一致地超過了Swin和ConvNeXt(更好的結果用粗體標記)。這些實驗突出了vHeat在處理分布外數據(如旋轉對象、不同視角(ObjectNet)和自然對抗性樣本(ImageNet-A))時的魯棒性。
5. 低級視覺任務
為了進一步評估我們提出的vHeat模型的泛化能力,我們將熱傳導算子(HCO)集成到SwinIR模型中,替換了自注意力模塊,從而得到了vHeatIR架構。然后我們在幾個標準的低級視覺任務上進行了一系列實驗,以評估vHeatIR的性能。這些任務包括在Set12和McMaster數據集上的灰度和彩色圖像去噪,以及在LIVE1數據集上的JPEG壓縮偽影減少。在這些實驗中,我們使用了與SwinIR相同的設置,以確保公平比較。結果總結在表5中,表明vHeatIR一致地超過了其他基線模型。這一改進主要歸功于HCO在頻率域中高效運行的能力,這增強了模型在處理低級圖像細節方面的性能。在訓練了15,000次迭代后,我們在圖6中可視化了噪聲水平為σ=15的彩色圖像的去噪結果。如圖所示,與SwinIR相比,vHeatIR產生了明顯更清晰的圖像,表明其在恢復圖像質量方面的優越能力。這些結果不僅突出了提出的vHeat模型的有效性,還驗證了其在低級視覺任務中的強大泛化能力。