【第9篇】 EfficientViT(ICCV 2023):基于多尺度线性注意力的高效高分辨率密集预测网络

论文题目:EfficientViT: Multi-Scale Linear Attention for High-Resolution Dense Prediction
文献地址:https://arxiv.org/abs/2205.14756
源码地址:https://github.com/mit-han-lab/efficientvit

本文内容依据用户提供的 arXiv v6 版本(2024 年 2 月 6 日)整理。ImageNet 与 EfficientViT-SAM 等扩展实验均以该版本为准。
image.png

image.png
近期高分辨率视觉模型的发展越来越快,但真正把模型放到手机 CPU、边缘 GPU 或云端 GPU 上运行时,我们经常会遇到一个熟悉的问题:论文中的 FLOPs 很低,实际推理却不一定快
对于密集预测任务,输入分辨率高、token 数量多。标准 Softmax Attention 需要处理 token 两两之间的关系,计算和内存开销会随分辨率快速增长;大核卷积虽然可以扩大感受野,其真实速度又依赖硬件和算子支持。

🌠🌠🌠 本篇内容:本文介绍 EfficientViT,一种面向高分辨率密集预测的高效视觉模型。它的核心是 Multi-Scale Linear Attention(多尺度线性注意力):使用 ReLU Linear Attention 建立全局感受野,再利用轻量小卷积聚合邻域信息,弥补线性注意力局部建模与多尺度学习能力不足的问题。

五分钟回顾模型

EfficientViT 想解决的并不是“如何再造一个更小的分类网络”,而是一个更具体的问题:

如何在高分辨率密集预测中,同时获得全局感受野、多尺度特征和真实硬件上的高效率?
image.png
论文观察到,高分辨率密集预测通常需要两种能力:

  • 全局感受野:理解远距离像素之间的语义关系。
  • 多尺度学习:兼顾目标整体、局部边界和小目标细节。
    此前的方法通常沿着两条路线展开:
  • Transformer 路线借助 Softmax Attention 建立全局关系。标准 Softmax Attention 关于 token 数量具有二次复杂度;SegFormer 虽使用 Spatial-Reduction Attention 缓解开销,高分辨率下仍存在较大的计算和访存压力。
  • 卷积路线利用多分支和大卷积核扩大感受野,但大核卷积是否高效依赖硬件与底层算子的支持。
    EfficientViT各阶段结构
    EfficientViT 的思路可以概括为一句话:

用 ReLU 线性注意力负责全局关系,用轻量卷积补充局部信息和多尺度特征。

EfficientViT核心模块示意
从宏观结构看,EfficientViT 采用常见的 Backbone-Head,也可以理解为 Encoder-Decoder 结构:

  • Input Stem 对高分辨率输入进行初始特征提取和下采样。
  • Backbone 包含四个 stage,空间尺寸逐级减小、通道数逐级增加。
  • EfficientViT Module 被放在 Stage 3 和 Stage 4。
  • Stage 2、3、4 的输出分别记作 P2、P3、P4,构成多尺度特征金字塔。
  • 面向不同密集预测任务,Head 使用 1×1 卷积和标准上采样操作对齐特征,再通过逐元素相加完成融合;具体上采样方式随任务配置选择,例如双线性或双三次插值。
  • 融合后的特征经过若干 MBConv 和输出层,得到最终的密集预测结果。
    image.png
    EfficientViT Module 主要包含两个部分:
  1. Multi-Scale Linear Attention:提取多尺度的全局上下文信息。
  2. FFN + DWConv:利用前馈网络中的深度卷积增强局部信息提取能力。
    EfficientViT整体网络结构
    这里有一个很重要的设计取舍:Backbone 已经具有较强的上下文建模能力,因此 Head 没有继续堆叠复杂模块,而是选择了简单的加法融合与 MBConv。换句话说,EfficientViT 的效率不是来自某一个孤立算子,而是来自注意力、局部卷积、特征融合和网络宏观结构的共同约束
    image.png

实验结果

论文在语义分割、图像超分辨率、Segment Anything 和 ImageNet 分类上进行了验证。速度数据覆盖移动 CPU、边缘 GPU 与云端 GPU。由于不同表格的输入尺寸、数值精度和推理框架不同,下面只在相同实验条件内比较。

消融实验:全局与多尺度缺一不可

论文首先在 Cityscapes 上验证 Multi-Scale Linear Attention 的两个核心组成部分。输入分辨率为 1024×2048,各模型通过调整宽度被控制在相同的 4.4G MACs,并且均从随机初始化开始训练。
image.png
在参数量和 MACs 基本不变的条件下,单独引入多尺度学习或全局注意力,都能带来约 4 个百分点的 mIoU 提升;二者结合后达到 74.5 mIoU。这一结果支撑了论文的核心判断:高分辨率密集预测不仅需要看得“远”,也需要看得“细”。

语义分割

Cityscapes 实验统一使用 1024×2048 输入。GPU 数据通过 TensorRT、FP16 获得,bs1 表示 batch size 为 1。

模型mIoU参数量MACsJetson AGX Orin 延迟A100 吞吐量
SegFormer-B178.514M244G146 ms49 image/s
SegNeXt-T79.84.3M51G93.2 ms95 image/s
EfficientViT-B180.54.8M25G24.3 ms175 image/s
SegNeXt-B82.628M276G228 ms41 image/s
EfficientViT-B383.040M179G81.8 ms70 image/s
SegNeXt-L83.249M578G374 ms26 image/s
EfficientViT-L283.253M396G60.0 ms102 image/s
以最后一组为例,EfficientViT-L2 和 SegNeXt-L 都取得 83.2 mIoU,但前者在 Jetson AGX Orin 上的延迟为 60.0 ms,后者为 374 ms。在论文给定的测试环境中,EfficientViT-L2 的延迟约为 SegNeXt-L 的六分之一。

ADE20K 实验按照常见设置将图像短边缩放至 512。代表性结果如下:

模型mIoU参数量MACsJetson AGX Orin 延迟A100 吞吐量
SegFormer-B142.214M16G12.3 ms542 image/s
SegNeXt-T41.14.3M6.6G12.4 ms842 image/s
EfficientViT-B142.84.8M3.1G4.0 ms1142 image/s
SegNeXt-S44.314M16G17.2 ms592 image/s
EfficientViT-B245.915M9.1G7.3 ms846 image/s
SegFormer-B450.364M96G44.9 ms212 image/s
EfficientViT-L250.751M45G9.0 ms758 image/s
实验表明,EfficientViT 追求的不是单独降低理论计算量,而是让 MACs 的下降尽可能转化为论文所测试设备上的延迟下降和吞吐量提升。

图像超分辨率

论文包含轻量超分辨率和高分辨率超分辨率两种设置:轻量设置在 DIV2K 上训练并在 BSD100 上测试;高分辨率设置使用 FFHQ training split 的前 3000 张图像训练,并在 FFHQ validation split 的前 500 张图像上测试。表中速度在 A100 GPU、batch size 1 条件下测得。

模型FFHQ PSNRFFHQ 延迟BSD100 PSNRBSD100 延迟
Restormer43.43 dB92.0 ms32.31 dB15.1 ms
SwinIR43.49 dB61.2 ms32.31 dB9.7 ms
EfficientViT w0.7543.54 dB14.3 ms32.31 dB2.8 ms
EfficientViT43.58 dB17.8 ms32.33 dB3.2 ms
在 FFHQ 的 512×512 → 1024×1024 设置中,EfficientViT w0.75 相比 Restormer 将延迟从 92.0 ms 降至 14.3 ms,即论文所述的 6.4 倍加速,同时 PSNR 提高 0.11 dB。

EfficientViT-SAM

作者用 EfficientViT 替换 SAM 的图像编码器,同时保留 Prompt Encoder 和 Mask Decoder,构建 EfficientViT-SAM。训练分为两个阶段:先以 SAM 图像编码器为教师训练 EfficientViT 图像编码器,再使用 SA-1B 数据集进行端到端训练。
在 ViTDet 预测框作为提示的零样本实例分割实验中,A100 吞吐量采用 TensorRT 和 FP16 测量,并包含图像编码器与 SAM Head。

模型参数量MACsA100 吞吐量COCO mAPLVIS mAP
SAM-ViT-H641M2973G11 image/s46.544.2
EfficientViT-SAM-L261M69G538 image/s46.642.7
EfficientViT-SAM-XL0117M185G278 image/s47.543.9
EfficientViT-SAM-XL1203M322G182 image/s47.844.4
EfficientViT-SAM-XL1 在 COCO 和 LVIS 上均略高于 SAM-ViT-H,并将 A100 吞吐量从 11 image/s 提升至 182 image/s,对应 16.5 倍吞吐量提升。论文还给出了最高 48.9 倍吞吐量提升的更轻量变体,但不同变体对应不同精度,不能把“最高速度”和“最高精度”理解成同一个模型的结果。

ImageNet 分类

EfficientViT 虽然面向高分辨率密集预测设计,其 Backbone 在 ImageNet 上同样具有竞争力。EfficientViT-L2 在 384×384 输入下取得 86.0% Top-1 Accuracy,参数量为 64M、MACs 为 20G;论文指出,它相对 EfficientNetV2-L 提高 0.3 个百分点,并在 A100 上获得 2.6 倍加速。
image.png

本文方法

EfficientViT 的创新并不是简单删除 Softmax,而是围绕 ReLU Linear Attention 的优势和缺点,组合出一套适合高分辨率密集预测的模块。

一、Softmax Attention 的问题

高分辨率输入会产生大量 token。标准 Softmax Attention 需要显式计算 token 两两之间的关系,token 数量为 N 时,注意力矩阵的尺寸为 N×N。当分辨率提高时,计算量和内存占用会快速增长。
Softmax注意力在高分辨率下的开销
EfficientViT 使用 ReLU Linear Attention 替代标准 Softmax Attention。其相似度可以简写为:
Sim ⁡ ( Q , K ) = ReLU ⁡ ( Q ) ReLU ⁡ ( K ) T , \operatorname{Sim}(Q,K)=\operatorname{ReLU}(Q)\operatorname{ReLU}(K)^T, Sim(Q,K)=ReLU(Q)ReLU(K)T,
其中,Q、K 分别表示 Query 和 Key。利用矩阵乘法结合律,模型可以先聚合 \operatorname{ReLU}(K)^TV 以及用于归一化的统计量,再由 \operatorname{ReLU}(Q) 查询聚合结果,从而避免显式构造 N×N 注意力矩阵。在每个注意力头的维度固定时,计算与内存关于 token 数量的增长关系由二次降为线性。
这里不必纠结完整推导,只需抓住核心变化:

  • 标准注意力:计算 QKᵀ,经 Softmax 归一化后再与 V 聚合
  • 线性注意力:先聚合 ReLU(K)ᵀV 与归一化统计量, 再由 ReLU(Q) 查询聚合结果。

二、ReLU 线性注意力为什么还不够

速度更快并不意味着表达能力完全不受影响。Softmax 中的指数函数会放大相似度差异,因此注意力更容易集中到少数关键位置;ReLU Linear Attention 的分布通常更平滑,对局部细节的捕获能力较弱。
Softmax与ReLU线性注意力对比在语义分割、超分辨率等密集预测任务中,物体边界、小目标和纹理信息都依赖局部建模,因此这一缺点尤其值得重视。为此,作者采用两项补偿设计:

  1. 在 FFN 中插入 DWConv,增强局部信息提取能力。
  2. 聚合 Q、K、V 邻近位置的 token,生成多尺度 token。
    于是,ReLU Linear Attention 负责建立全局关系,卷积负责注入局部归纳偏置,两者形成互补。

三、Multi-Scale Linear Attention

输入首先通过 1×1 投影生成 Q、K、V。除原始 Q/K/V 外,作者还使用小核深度可分离卷积聚合邻域信息,生成额外尺度的 Q/K/V,再在不同尺度上执行 ReLU Linear Attention。
Multi-Scale Linear Attention结构
论文 v6 的默认实现采用两分支设计:一条分支保留原始 Q/K/V,另一条分支聚合 5×5 邻域 token,在多尺度能力和实际效率之间做平衡。这里的 5×5 是论文实验配置,而不是该模块只能使用的固定定义。
多尺度聚合采用轻量的小核卷积,而不是直接堆叠大卷积核。实现层面,作者将原本针对各 Q/K/V 和各注意力头的独立聚合,融合为一次 DWConv 与一次分组 1×1 卷积,从而减少零散算子调用。
timm源码中通过ConvNormAct使用1x1卷积,输出 3 倍维度,刚好对应 Q/K/V的生成。

self.qkv = ConvNormAct(
            in_channels,
            3 * total_dim,
            1,
            bias=use_bias[0],
            norm_layer=norm_layer[0],
            act_layer=act_layer[0],
            **dd,
        )

多尺度聚合的实现对qkv做多尺度空间加权融合。用深度卷积 + 分组 1×1 卷积实现多尺度上下文聚合,替代普通注意力单尺度融合。

  • 第一个卷积:kernel=scale, groups=3*total_dim。逐通道深度卷积用于空间多尺度上下文提取,groups=输入通道数 → 深度卷积 DepthwiseConv,等价于对 Q 整组、K 整组、V 整组分别做多尺度空间邻域信息聚合。cale 越大 → 抓取更远距离空间关系,实现多尺度上下文
  • 第二个卷积:1×1, groups=3*heads。多头内通道交互,打散分组、融合头内特征
self.aggreg = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(
                    3 * total_dim,
                    3 * total_dim,
                    scale,
                    padding=get_same_padding(scale),
                    groups=3 * total_dim,
                    bias=use_bias[0],
                    **dd,
                ),
                nn.Conv2d(3 * total_dim, 3 * total_dim, 1, groups=3 * heads, bias=use_bias[0], **dd),
            )
            for scale in scales
        ])

最后线性注意力层:舍弃传统 Softmax + 缩放点积,用矩阵聚合 + 归一化实现全局注意力,无指数运算、速度极快,也是你 TRT FP16 出 0 输出的高危算子
image.png
FP16 下矩阵相乘极易炸 NaN,先升 FP32 计算,最后切回原精度。

def _attn(self, q, k, v):
	dtype = v.dtype
	
	q, k, v = q.float(), k.float(), v.float()
	kv = k.transpose(-1, -2) @ v
	out = q @ kv
	out = out[..., :-1] / (out[..., -1:] + self.eps)
	return out.to(dtype)

四、EfficientViT Module

Multi-Scale Linear Attention 与 FFN + DWConv 共同构成 EfficientViT Module:

Input
  │
Multi-Scale Linear Attention  ← 多尺度全局上下文
  │
FFN + DWConv                  ← 通道变换与局部信息
  │
Output

两部分的分工非常清楚:

组成部分主要作用解决的问题
Multi-Scale Linear Attention建立全局关系并融合多个尺度高分辨率下全局建模开销大
FFN + DWConv变换通道并提取局部特征ReLU 注意力对局部细节不够敏感

EfficientViT Module组成

五、整体架构与特征融合

image.png
以论文图示中的 3×1024×2048 输入为例,Input Stem 首先完成初始下采样,随后四个 stage 逐级缩小空间尺寸、增加通道数。Stage 之间使用 stride 为 2 的 MBConv 下采样,EfficientViT Module 只放在 Stage 3 和 Stage 4。
这样安排的原因很直观:前面的特征图分辨率最高,卷积处理更经济;经过下采样后再进行全局建模,可以以更低成本获得较大的有效感受野。
Head 使用 P2、P3、P4 三个尺度的特征。P4 上采样 4 倍、P3 上采样 2 倍,使空间尺寸与 P2 一致;三者再分别通过 1×1 卷积对齐通道数,随后一次相加完成融合。由于 Backbone 已经具备较强的上下文建模能力,Head 不需要继续堆叠复杂模块。
因此,EfficientViT 的高效并不只来自线性注意力,而是由以下设计共同实现:

  • 在线性注意力中避免显式构造大规模注意力矩阵。
  • 用小核卷积弥补局部与多尺度能力。
  • 把 EfficientViT Module 放在较低分辨率的后两个 stage。
  • 使用 MBConv 完成高效下采样和局部特征提取。
  • 使用简单的多尺度加法融合,避免 Head 成为新的计算瓶颈。

总结与思考

EfficientViT 的核心贡献可以归纳为三点:

  1. ReLU Linear Attention:利用可分解的相似度形式和矩阵乘法结合律,使注意力关于 token 数量由二次增长降为线性增长。
  2. Multi-Scale Linear Attention:使用轻量卷积生成多尺度 Q/K/V,将全局感受野与多尺度学习结合起来。
  3. 卷积与注意力协同设计:利用 FFN 中的 DWConv、后置注意力 stage 和轻量 Head,兼顾局部细节与硬件效率。
    从实验看,EfficientViT 的价值主要体现在性能—效率折中:它不只追求更低的 MACs,也不只追求最高精度,而是努力让计算量下降转化为论文所测试硬件和软件栈上的真实延迟与吞吐量优势。
    需要注意,MACs 描述乘加操作数量,不能在没有统一计数约定时直接等同于 FLOPs;不同设备上的延迟也会受到算子实现、数值精度和推理框架影响。论文中的速度结论只适用于相应测试条件,不应无条件外推到所有部署环境。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

木棉知行者

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值