物理傳熱啟發的視覺表徵模型vHeat來了，嘗試突破注意力機制，兼具低複雜度、全域性感受野

机器之心發表於2024-06-03
視覺模型複雜度
物理傳熱啟發的視覺表徵模型vHeat來了，嘗試突破注意力機制，兼具低複雜度、全域性感受野
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文作者為VMamba的原班人馬，其中第一作者王兆植是中國科學院大學和鵬城實驗室的2022級聯合培養博士生，共同一作劉悅是中國科學院大學2021級直博生。他們的主要研究方向是視覺模型設計和自監督學習。

如何突破 Transformer 的 Attention 機制？中國科學院大學與鵬城國家實驗室提出基於熱傳導的視覺表徵模型 vHeat。將圖片特徵塊視為熱源，並透過預測熱傳導率、以物理學熱傳導原理提取影像特徵。相比於基於Attention機制的視覺模型， vHeat 同時兼顧了：計算複雜度（1.5次方）、全域性感受野、物理可解釋性。

vHeat-base 模型在高解析度影像輸入時，throughput、GPU 視訊記憶體佔用、flops 分別是 Swin-base 模型的3倍、1/4、3/4，在影像分類、目標檢測、語義/例項分割等基礎下游任務上達到了先進的效能表現。

圖片

  • 論文地址: https://arxiv.org/pdf/2405.16555

  • 程式碼地址: https://github.com/MzeroMiko/vHeat

  • 論文標題：vHeat: Building Vision Models upon Heat Conduction

Overview

CNN 和視覺 Transformer（ViT）是當前最主流的兩類基礎視覺模型。然而，CNN的效能表現受限於區域性感受野和固定的卷積核運算元。ViT 具有全域性依賴關係的表徵能力，然而代價是高昂的二次方級別計算複雜度。我們認為 CNN 和 ViT 的卷積運算元和自注意力運算元都是特徵內部的畫素傳播過程，分別是一種資訊傳遞的形式，這也讓我們聯想到了物理領域的熱傳導。於是我們根據熱傳導方程，將視覺語義的空間傳播和物理熱傳導建立聯絡，提出了一種 1.5 次方計算複雜度的視覺熱傳導運算元（Heat Conduction Operator, HCO），進而設計出了一種兼具低複雜度、全域性感受野、物理可解釋性的視覺表徵模型 vHeat。HCO 與 self-attention 的計算形式和複雜度對比如下圖所示。實驗證明了 vHeat 在各種視覺任務中表現優秀。例如 vHeat-T 在 ImageNet-1K 上達到 82.2% 的分類準確率，比 Swin-T 高 0.9%，比 Vim-S 高1.7%。效能之外，vHeat 還擁有高推理速度、低 GPU 視訊記憶體佔用和低 FLOPs 這些優點。在輸入影像解析度較高時，base 規模的 vHeat 模型相比於 Swin 達到 3 倍吞吐量、1/4 的GPU視訊記憶體佔用和 3/4 的 FLOPs。

圖片

方法介紹

用圖片表示點圖片在 t 時刻下的溫度， 物理熱傳導方程為圖片，其中 k>0，表示熱擴散率。給定 t=0 時刻下的初始條件圖片，該熱傳導方程可以採用傅立葉變換求得通解，表示如下：

圖片

其中圖片和圖片分別表示傅立葉變換和逆傅立葉變換，圖片 表示頻域空間座標。

我們利用 HCO 來實現視覺語義中的熱傳導，先將物理熱傳導方程中的圖片擴充套件為多通道特徵圖片，將圖片視為輸入，圖片視為輸出，HCO 模擬了離散化形式的熱傳導通解，如下公式所示：

圖片

其中圖片和圖片分別表示二維離散餘弦變換和逆變換，HCO 的結構如下圖 (a) 所示。

圖片

此外，我們認為不同影像內容應該對應不同的熱擴散率，考慮到圖片的輸出在頻域中，我們根據頻率值來決定熱擴散率，圖片。由於頻域中不同位置表示了不同的頻率值，我們提出了頻率值編碼（Frequency Value Embeddings, FVEs）來表示頻率值資訊，與 ViT 中的絕對位置編碼的實現和作用類似，並用 FVEs 對熱擴散率 k 進行預測，使得 HCO 可以進行非均勻、自適應的傳導，如下圖所示。

圖片

vHeat 採用多層級的結構實現，如下圖所示，整體框架與主流的視覺模型類似，其中的 HCO layer 如圖 2 (b) 所示。

圖片

實驗結果

ImageNet分類

圖片

透過對比實驗結果不難看出，在相似的引數量和 FLOPs 下:

  1. vHeat-T 取得了 82.2% 的效能，超過 DeiT-S 達 2.4%、Vim-S 達 1.7%、Swin-T 達 0.9%。
  2. vHeat-S 取得了 83.6% 的效能，超過 Swin-S 達 0.6%、ConvNeXt-S 達 0.5%。
  3. vHeat-B 取得了 83.9% 的效能，超過 DeiT-B 達 2.1%、Swin-B 達 0.4%。

同時，由於 vHeat 的 O (N^1.5) 低複雜度和可平行計算性，推理吞吐量相比於 ViTs、SSM 模型有明顯的優勢，例如 vHeat-T 的推理吞吐量為 1514 img/s，比 Swin-T 高 22%，比 Vim-S 高 87%，也比 ConvNeXt-T 高 26%，同時擁有更好的效能。

下游任務

圖片

在 COCO 資料集上， vHeat 也擁有效能優勢：在 fine-tune 12 epochs 的情況下，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。在 ADE20K 資料集上，vHeat-T/S/B 分別達到 46.9/49.0/49.6 mIoU，相比於 Swin 和 ConvNeXt 依然擁有更好的效能表現。這些結果驗證了 vHeat 在視覺下游實驗中完全 work，展示出了能平替主流基礎視覺模型的潛力。

分析實驗

有效感受野

圖片

vHeat 擁有全域性的有效感受野，視覺化對比的這些主流模型中只有 DeiT 和 HiViT 也具備這個特性。但是值得注意的是，DeiT 和 HiViT 的代價是平方級的複雜度，而 vHeat 是 1.5 次方級的複雜度。

計算代價

圖片

上圖從左到右分別為 vHeat-B 與其他 base 規模下的 ViT-based 模型的推理吞吐量 / GPU 視訊記憶體佔用 / 計算量 FLOPs 對比。可以明顯看出，由於 O (N^1.5) 的計算複雜度，vHeat 相比於對比的模型有更快的推理速度、更低的視訊記憶體佔用以及更少的 FLOPs，並且在影像解析度越大時，優勢會更為明顯。在輸入影像為 768*768 解析度時，vHeat-B 的推理吞吐量為 Swin-B 的 3 倍左右，GPU 視訊記憶體佔用比 Swin-B 低 74%，FLOPs 比 Swin-B 低 28%。vHeat 與 ViT-based 模型的計算代價對比，展示出其處理高解析度影像的優秀潛質。

相關文章