物理傳熱啟發(fā)的視覺表征模型vHeat來了，嘗試突破注意力機制，兼具低復雜度、全局感受野
本文作者為VMamba的原班人馬，其中第一作者王兆植是中國科學院大學和鵬城實驗室的2022級聯(lián)合培養(yǎng)博士生，共同一作劉悅是中國科學院大學2021級直博生。他們的主要研究方向是視覺模型設計和自監(jiān)督學習。
如何突破 Transformer 的 Attention 機制？中國科學院大學與鵬城國家實驗室提出基于熱傳導的視覺表征模型 vHeat。將圖片特征塊視為熱源，并通過預測熱傳導率、以物理學熱傳導原理提取圖像特征。相比于基于Attention機制的視覺模型， vHeat 同時兼顧了：計算復雜度（1.5次方）、全局感受野、物理可解釋性。
vHeat-base 模型在高分辨率圖像輸入時，throughput、GPU 顯存占用、flops 分別是 Swin-base 模型的3倍、1/4、3/4，在圖像分類、目標檢測、語義/實例分割等基礎下游任務上達到了先進的性能表現(xiàn)。
- 論文地址: https://arxiv.org/pdf/2405.16555
- 代碼地址: https://github.com/MzeroMiko/vHeat
- 論文標題：vHeat: Building Vision Models upon Heat Conduction
Overview
CNN 和視覺 Transformer（ViT）是當前最主流的兩類基礎視覺模型。然而，CNN的性能表現(xiàn)受限于局部感受野和固定的卷積核算子。ViT 具有全局依賴關(guān)系的表征能力，然而代價是高昂的二次方級別計算復雜度。我們認為 CNN 和 ViT 的卷積算子和自注意力算子都是特征內(nèi)部的像素傳播過程，分別是一種信息傳遞的形式，這也讓我們聯(lián)想到了物理領(lǐng)域的熱傳導。于是我們根據(jù)熱傳導方程，將視覺語義的空間傳播和物理熱傳導建立聯(lián)系，提出了一種 1.5 次方計算復雜度的視覺熱傳導算子（Heat Conduction Operator, HCO），進而設計出了一種兼具低復雜度、全局感受野、物理可解釋性的視覺表征模型 vHeat。HCO 與 self-attention 的計算形式和復雜度對比如下圖所示。實驗證明了 vHeat 在各種視覺任務中表現(xiàn)優(yōu)秀。例如 vHeat-T 在 ImageNet-1K 上達到 82.2% 的分類準確率，比 Swin-T 高 0.9%，比 Vim-S 高1.7%。性能之外，vHeat 還擁有高推理速度、低 GPU 顯存占用和低 FLOPs 這些優(yōu)點。在輸入圖像分辨率較高時，base 規(guī)模的 vHeat 模型相比于 Swin 達到 3 倍吞吐量、1/4 的GPU顯存占用和 3/4 的 FLOPs。
方法介紹
用
表示點
在 t 時刻下的溫度， 物理熱傳導方程為
，其中 k>0，表示熱擴散率。給定 t=0 時刻下的初始條件
，該熱傳導方程可以采用傅里葉變換求得通解，表示如下：
其中
和
分別表示傅里葉變換和逆傅里葉變換，
表示頻域空間坐標。
我們利用 HCO 來實現(xiàn)視覺語義中的熱傳導，先將物理熱傳導方程中的
擴展為多通道特征
，將
視為輸入，
視為輸出，HCO 模擬了離散化形式的熱傳導通解，如下公式所示：
其中
和
分別表示二維離散余弦變換和逆變換，HCO 的結(jié)構(gòu)如下圖 (a) 所示。
此外，我們認為不同圖像內(nèi)容應該對應不同的熱擴散率，考慮到
的輸出在頻域中，我們根據(jù)頻率值來決定熱擴散率，
。由于頻域中不同位置表示了不同的頻率值，我們提出了頻率值編碼（Frequency Value Embeddings, FVEs）來表示頻率值信息，與 ViT 中的絕對位置編碼的實現(xiàn)和作用類似，并用 FVEs 對熱擴散率 k 進行預測，使得 HCO 可以進行非均勻、自適應的傳導，如下圖所示。
vHeat 采用多層級的結(jié)構(gòu)實現(xiàn)，如下圖所示，整體框架與主流的視覺模型類似，其中的 HCO layer 如圖 2 (b) 所示。
實驗結(jié)果
ImageNet分類
通過對比實驗結(jié)果不難看出，在相似的參數(shù)量和 FLOPs 下:
- vHeat-T 取得了 82.2%的性能，超過 DeiT-S 達 2.4%、Vim-S 達 1.7%、Swin-T 達 0.9%。
- vHeat-S 取得了 83.6%的性能，超過 Swin-S 達 0.6%、ConvNeXt-S 達 0.5%。
- vHeat-B 取得了 83.9%的性能，超過 DeiT-B 達 2.1%、Swin-B 達 0.4%。
同時，由于 vHeat 的 O (N^1.5) 低復雜度和可并行計算性，推理吞吐量相比于 ViTs、SSM 模型有明顯的優(yōu)勢，例如 vHeat-T 的推理吞吐量為 1514 img/s，比 Swin-T 高 22%，比 Vim-S 高 87%，也比 ConvNeXt-T 高 26%，同時擁有更好的性能。
下游任務
在 COCO 數(shù)據(jù)集上， vHeat 也擁有性能優(yōu)勢：在 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 數(shù)據(jù)集上，vHeat-T/S/B 分別達到 46.9/49.0/49.6 mIoU，相比于 Swin 和 ConvNeXt 依然擁有更好的性能表現(xiàn)。這些結(jié)果驗證了 vHeat 在視覺下游實驗中完全 work，展示出了能平替主流基礎視覺模型的潛力。
分析實驗
有效感受野
vHeat 擁有全局的有效感受野，可視化對比的這些主流模型中只有 DeiT 和 HiViT 也具備這個特性。但是值得注意的是，DeiT 和 HiViT 的代價是平方級的復雜度，而 vHeat 是 1.5 次方級的復雜度。
計算代價
上圖從左到右分別為 vHeat-B 與其他 base 規(guī)模下的 ViT-based 模型的推理吞吐量 / GPU 顯存占用 / 計算量 FLOPs 對比?？梢悦黠@看出，由于 O (N^1.5) 的計算復雜度，vHeat 相比于對比的模型有更快的推理速度、更低的顯存占用以及更少的 FLOPs，并且在圖像分辨率越大時，優(yōu)勢會更為明顯。在輸入圖像為 768*768 分辨率時，vHeat-B 的推理吞吐量為 Swin-B 的 3 倍左右，GPU 顯存占用比 Swin-B 低 74%，F(xiàn)LOPs 比 Swin-B 低 28%。vHeat 與 ViT-based 模型的計算代價對比，展示出其處理高分辨率圖像的優(yōu)秀潛質(zhì)。
本文轉(zhuǎn)自 機器之心 ，作者：機器之心
原文鏈接:??https://mp.weixin.qq.com/s/SR_yMehYn9Zmh79Gm3BZFg??
