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,在图像分类、目标检测、语义/实例分割等基础下游任务上达到了先进的性能表现。
![](http://n.sinaimg.cn/spider20240603/730/w1080h450/20240603/3963-f5f376bca752e163bda3d9ee127d9cb3.png)
论文地址: 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。
![](http://n.sinaimg.cn/spider20240603/88/w1080h608/20240603/b2cb-54bb2f186ebcac5fbd45f129e932b2b6.png)
方法介绍
用
![](http://n.sinaimg.cn/spider20240603/136/w102h34/20240603/9e95-09951d57e96667490394836483d56e4d.png)
![](http://n.sinaimg.cn/spider20240603/266/w208h58/20240603/fe34-a36df724a4f65c607032bfeea8b35fa2.png)
![](http://n.sinaimg.cn/spider20240603/261/w225h36/20240603/ef93-c4aea3bb4afaee01e339a966308f86f8.png)
,该热传导方程可以采用傅里叶变换求得通解,表示如下:
,其中 k>0,表示热扩散率。给定 t=0 时刻下的初始条件
在 t 时刻下的温度, 物理热传导方程为
表示点
![](http://n.sinaimg.cn/spider20240603/561/w492h69/20240603/9a98-8bd2c18e898cfcc356a977cfa9f11e49.png)
其中
表示频域空间坐标。
分别表示傅里叶变换和逆傅里叶变换,
和
我们利用 HCO 来实现视觉语义中的热传导,先将物理热传导方程中的
![](http://n.sinaimg.cn/spider20240603/342/w300h42/20240603/b26b-a34b59b0d8abef9cd9c6ce3dece65c02.png)
![](http://n.sinaimg.cn/spider20240603/168/w132h36/20240603/ccdc-c4fd1cfbd6dd0ab4ee16c50bccbaeb55.png)
![](http://n.sinaimg.cn/spider20240603/166/w130h36/20240603/a5ca-1d0bd5fe8e37a3b77939d5a2dedcc8a3.png)
视为输出,HCO 模拟了离散化形式的热传导通解,如下公式所示:
视为输入,
,将
扩展为多通道特征
![](http://n.sinaimg.cn/spider20240603/511/w459h52/20240603/f8db-c2ffd4beb49c9df872cefba21c13ef30.png)
其中
![](http://n.sinaimg.cn/spider20240603/169/w126h43/20240603/4bc5-7d5eda285d4b80a21c4984eea05e9d7e.png)
![](http://n.sinaimg.cn/spider20240603/171/w135h36/20240603/81b1-fb71881ea6a968c3350823f7c991b796.png)
分别表示二维离散余弦变换和逆变换,HCO 的结构如下图 (a) 所示。
和
![](http://n.sinaimg.cn/spider20240603/705/w1080h425/20240603/c99e-f65a504ae5de137acfe3942872fc9b37.png)
此外,我们认为不同图像内容应该对应不同的热扩散率,考虑到
![](http://n.sinaimg.cn/spider20240603/150/w114h36/20240603/521d-3a583bc4226de65d18587fee091e0c2b.png)
![](http://n.sinaimg.cn/spider20240603/210/w171h39/20240603/f636-fbbb029f405777cc144199df965c6053.png)
。由于频域中不同位置表示了不同的频率值,我们提出了频率值编码(Frequency Value Embeddings, FVEs)来表示频率值信息,与 ViT 中的绝对位置编码的实现和作用类似,并用 FVEs 对热扩散率 k 进行预测,使得 HCO 可以进行非均匀、自适应的传导,如下图所示。
的输出在频域中,我们根据频率值来决定热扩散率,
![vHeat 采用多层级的结构实现,如下图所示,整体框架与主流的视觉模型类似,其中的 HCO layer 如图 2 (b) 所示。](http://n.sinaimg.cn/spider20240603/756/w1053h503/20240603/a0af-074ba9e5edaafaae2aba520daff1ef65.png)
![](http://n.sinaimg.cn/spider20240603/581/w1080h301/20240603/4748-a91977541c702303418a5c0c343131da.png)
实验结果
ImageNet分类
![](http://n.sinaimg.cn/spider20240603/236/w1015h821/20240603/1032-30f10cb36a21fc243bc0eb2bc365c4fe.png)
通过对比实验结果不难看出,在相似的参数量和 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 模型有明显的优势,例如 vHeat-T 的推理吞吐量为 1514 img/s,比 Swin-T 高 22%,比 Vim-S 高 87%,也比 ConvNeXt-T 高 26%,同时拥有更好的性能。
下游任务
![](http://n.sinaimg.cn/spider20240603/96/w1080h616/20240603/4855-af8b6535c7a5d586a51e1d34594005f9.png)
在 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,展示出了能平替主流基础视觉模型的潜力。
分析实验
有效感受野
![](http://n.sinaimg.cn/spider20240603/658/w1080h378/20240603/88c3-d1483c8e6ed56df1c9d1c2f9f1144387.png)
vHeat 拥有全局的有效感受野,可视化对比的这些主流模型中只有 DeiT 和 HiViT 也具备这个特性。但是值得注意的是,DeiT 和 HiViT 的代价是平方级的复杂度,而 vHeat 是 1.5 次方级的复杂度。
计算代价
![](http://n.sinaimg.cn/spider20240603/533/w1080h253/20240603/78a8-d7078bbae80ad2b4efa61270fef81cfa.png)
上图从左到右分别为 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 模型的计算代价对比,展示出其处理高分辨率图像的优秀潜质。
![© THE END](http://n.sinaimg.cn/spider20240603/685/w1080h1205/20240603/b056-a02509b0558b8b7e97bd764c1bbf7072.jpg)
![新浪科技公众号 新浪科技公众号](http://n.sinaimg.cn/tech/content/tech_qr2x.png)
“掌”握科技鲜闻 (微信搜索techsina或扫描左侧二维码关注)
![](http://n.sinaimg.cn/tech/content/tech_weixin2.png)