本文作者為VMamba的原班人馬，其中第一作者王兆植是中國(guó)科學(xué)院大學(xué)和鵬城實(shí)驗(yàn)室的2022級(jí)聯(lián)合培養(yǎng)博士生，共同一作劉悅是中國(guó)科學(xué)院大學(xué)2021級(jí)直博生。他們的主要研究方向是視覺(jué)模型設(shè)計(jì)和自監(jiān)督學(xué)習(xí)。

如何突破 Transformer 的 Attention 機(jī)制？中國(guó)科學(xué)院大學(xué)與鵬城國(guó)家實(shí)驗(yàn)室提出基于熱傳導(dǎo)的視覺(jué)表征模型 vHeat。將圖片特征塊視為熱源，并通過(guò)預(yù)測(cè)熱傳導(dǎo)率、以物理學(xué)熱傳導(dǎo)原理提取圖像特征。相比于基于Attention機(jī)制的視覺(jué)模型， vHeat 同時(shí)兼顧了：計(jì)算復(fù)雜度（1.5次方）、全局感受野、物理可解釋性。

vHeat-base 模型在高分辨率圖像輸入時(shí)，throughput、GPU 顯存占用、flops 分別是 Swin-base 模型的3倍、1/4、3/4，在圖像分類(lèi)、目標(biāo)檢測(cè)、語(yǔ)義/實(shí)例分割等基礎(chǔ)下游任務(wù)上達(dá)到了先進(jìn)的性能表現(xiàn)。

• 論文地址: https://arxiv.org/pdf/2405.16555
• 代碼地址: https://github.com/MzeroMiko/vHeat
• 論文標(biāo)題：vHeat: Building Vision Models upon Heat Conduction

Overview

CNN 和視覺(jué) Transformer（ViT）是當(dāng)前最主流的兩類(lèi)基礎(chǔ)視覺(jué)模型。然而，CNN的性能表現(xiàn)受限于局部感受野和固定的卷積核算子。ViT 具有全局依賴關(guān)系的表征能力，然而代價(jià)是高昂的二次方級(jí)別計(jì)算復(fù)雜度。我們認(rèn)為 CNN 和 ViT 的卷積算子和自注意力算子都是特征內(nèi)部的像素傳播過(guò)程，分別是一種信息傳遞的形式，這也讓我們聯(lián)想到了物理領(lǐng)域的熱傳導(dǎo)。于是我們根據(jù)熱傳導(dǎo)方程，將視覺(jué)語(yǔ)義的空間傳播和物理熱傳導(dǎo)建立聯(lián)系，提出了一種 1.5 次方計(jì)算復(fù)雜度的視覺(jué)熱傳導(dǎo)算子（Heat Conduction Operator, HCO），進(jìn)而設(shè)計(jì)出了一種兼具低復(fù)雜度、全局感受野、物理可解釋性的視覺(jué)表征模型 vHeat。HCO 與 self-attention 的計(jì)算形式和復(fù)雜度對(duì)比如下圖所示。實(shí)驗(yàn)證明了 vHeat 在各種視覺(jué)任務(wù)中表現(xiàn)優(yōu)秀。例如 vHeat-T 在 ImageNet-1K 上達(dá)到 82.2% 的分類(lèi)準(zhǔn)確率，比 Swin-T 高 0.9%，比 Vim-S 高1.7%。性能之外，vHeat 還擁有高推理速度、低 GPU 顯存占用和低 FLOPs 這些優(yōu)點(diǎn)。在輸入圖像分辨率較高時(shí)，base 規(guī)模的 vHeat 模型相比于 Swin 達(dá)到 3 倍吞吐量、1/4 的GPU顯存占用和 3/4 的 FLOPs。

方法介紹

用

表示點(diǎn)

在 t 時(shí)刻下的溫度， 物理熱傳導(dǎo)方程為

，其中 k>0，表示熱擴(kuò)散率。給定 t=0 時(shí)刻下的初始條件

，該熱傳導(dǎo)方程可以采用傅里葉變換求得通解，表示如下：

其中

和

分別表示傅里葉變換和逆傅里葉變換，

 表示頻域空間坐標(biāo)。

我們利用 HCO 來(lái)實(shí)現(xiàn)視覺(jué)語(yǔ)義中的熱傳導(dǎo)，先將物理熱傳導(dǎo)方程中的

擴(kuò)展為多通道特征

，將

視為輸入，

視為輸出，HCO 模擬了離散化形式的熱傳導(dǎo)通解，如下公式所示：

其中

和

分別表示二維離散余弦變換和逆變換，HCO 的結(jié)構(gòu)如下圖 (a) 所示。

此外，我們認(rèn)為不同圖像內(nèi)容應(yīng)該對(duì)應(yīng)不同的熱擴(kuò)散率，考慮到

的輸出在頻域中，我們根據(jù)頻率值來(lái)決定熱擴(kuò)散率，

。由于頻域中不同位置表示了不同的頻率值，我們提出了頻率值編碼（Frequency Value Embeddings,  FVEs）來(lái)表示頻率值信息，與 ViT 中的絕對(duì)位置編碼的實(shí)現(xiàn)和作用類(lèi)似，并用 FVEs 對(duì)熱擴(kuò)散率 k 進(jìn)行預(yù)測(cè)，使得 HCO 可以進(jìn)行非均勻、自適應(yīng)的傳導(dǎo)，如下圖所示。

vHeat 采用多層級(jí)的結(jié)構(gòu)實(shí)現(xiàn)，如下圖所示，整體框架與主流的視覺(jué)模型類(lèi)似，其中的 HCO layer 如圖 2 (b) 所示。

實(shí)驗(yàn)結(jié)果

ImageNet分類(lèi)

通過(guò)對(duì)比實(shí)驗(yàn)結(jié)果不難看出，在相似的參數(shù)量和 FLOPs 下:

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

同時(shí)，由于 vHeat 的 O (N^1.5) 低復(fù)雜度和可并行計(jì)算性，推理吞吐量相比于 ViTs、SSM 模型有明顯的優(yōu)勢(shì)，例如 vHeat-T 的推理吞吐量為 1514 img/s，比 Swin-T 高 22%，比 Vim-S 高 87%，也比 ConvNeXt-T 高 26%，同時(shí)擁有更好的性能。

下游任務(wù)

在 COCO 數(shù)據(jù)集上， vHeat 也擁有性能優(yōu)勢(shì)：在 fine-tune 12 epochs 的情況下，vHeat-T/S/B 分別達(dá)到 45.1/46.8/47.7 mAP，超過(guò)了 Swin-T/S/B 達(dá) 2.4/2.0/0.8 mAP，超過(guò) ConvNeXt-T/S/B 達(dá) 0.9/1.4/0.7 mAP。在 ADE20K 數(shù)據(jù)集上，vHeat-T/S/B 分別達(dá)到 46.9/49.0/49.6 mIoU，相比于 Swin 和 ConvNeXt 依然擁有更好的性能表現(xiàn)。這些結(jié)果驗(yàn)證了 vHeat 在視覺(jué)下游實(shí)驗(yàn)中完全 work，展示出了能平替主流基礎(chǔ)視覺(jué)模型的潛力。

分析實(shí)驗(yàn)

有效感受野

vHeat 擁有全局的有效感受野，可視化對(duì)比的這些主流模型中只有 DeiT 和 HiViT 也具備這個(gè)特性。但是值得注意的是，DeiT 和 HiViT 的代價(jià)是平方級(jí)的復(fù)雜度，而 vHeat 是 1.5 次方級(jí)的復(fù)雜度。

計(jì)算代價(jià)

上圖從左到右分別為 vHeat-B 與其他 base 規(guī)模下的 ViT-based 模型的推理吞吐量 / GPU 顯存占用 / 計(jì)算量 FLOPs 對(duì)比?？梢悦黠@看出，由于 O (N^1.5) 的計(jì)算復(fù)雜度，vHeat 相比于對(duì)比的模型有更快的推理速度、更低的顯存占用以及更少的 FLOPs，并且在圖像分辨率越大時(shí)，優(yōu)勢(shì)會(huì)更為明顯。在輸入圖像為 768*768 分辨率時(shí)，vHeat-B 的推理吞吐量為 Swin-B 的 3 倍左右，GPU 顯存占用比 Swin-B 低 74%，F(xiàn)LOPs 比 Swin-B 低 28%。vHeat 與 ViT-based 模型的計(jì)算代價(jià)對(duì)比，展示出其處理高分辨率圖像的優(yōu)秀潛質(zhì)。

本文轉(zhuǎn)自 機(jī)器之心 ，作者：機(jī)器之心

原文鏈接:??https://mp.weixin.qq.com/s/SR_yMehYn9Zmh79Gm3BZFg??