论文题目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 等扩展实验均以该版本为准。文章目录[toc]五分钟回顾模型实验结果消融实验全局与多尺度缺一不可语义分割图像超分辨率EfficientViT-SAMImageNet 分类本文方法一、Softmax Attention 的问题二、ReLU 线性注意力为什么还不够三、Multi-Scale Linear Attention四、EfficientViT Module五、整体架构与特征融合总结与思考近期高分辨率视觉模型的发展越来越快但真正把模型放到手机 CPU、边缘 GPU 或云端 GPU 上运行时我们经常会遇到一个熟悉的问题论文中的 FLOPs 很低实际推理却不一定快。对于密集预测任务输入分辨率高、token 数量多。标准 Softmax Attention 需要处理 token 两两之间的关系计算和内存开销会随分辨率快速增长大核卷积虽然可以扩大感受野其真实速度又依赖硬件和算子支持。本篇内容本文介绍 EfficientViT一种面向高分辨率密集预测的高效视觉模型。它的核心是Multi-Scale Linear Attention多尺度线性注意力使用 ReLU Linear Attention 建立全局感受野再利用轻量小卷积聚合邻域信息弥补线性注意力局部建模与多尺度学习能力不足的问题。五分钟回顾模型EfficientViT 想解决的并不是“如何再造一个更小的分类网络”而是一个更具体的问题如何在高分辨率密集预测中同时获得全局感受野、多尺度特征和真实硬件上的高效率论文观察到高分辨率密集预测通常需要两种能力全局感受野理解远距离像素之间的语义关系。多尺度学习兼顾目标整体、局部边界和小目标细节。此前的方法通常沿着两条路线展开Transformer 路线借助 Softmax Attention 建立全局关系。标准 Softmax Attention 关于 token 数量具有二次复杂度SegFormer 虽使用 Spatial-Reduction Attention 缓解开销高分辨率下仍存在较大的计算和访存压力。卷积路线利用多分支和大卷积核扩大感受野但大核卷积是否高效依赖硬件与底层算子的支持。EfficientViT 的思路可以概括为一句话用 ReLU 线性注意力负责全局关系用轻量卷积补充局部信息和多尺度特征。从宏观结构看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 和输出层得到最终的密集预测结果。EfficientViT Module 主要包含两个部分Multi-Scale Linear Attention提取多尺度的全局上下文信息。FFN DWConv利用前馈网络中的深度卷积增强局部信息提取能力。这里有一个很重要的设计取舍Backbone 已经具有较强的上下文建模能力因此 Head 没有继续堆叠复杂模块而是选择了简单的加法融合与 MBConv。换句话说EfficientViT 的效率不是来自某一个孤立算子而是来自注意力、局部卷积、特征融合和网络宏观结构的共同约束。实验结果论文在语义分割、图像超分辨率、Segment Anything 和 ImageNet 分类上进行了验证。速度数据覆盖移动 CPU、边缘 GPU 与云端 GPU。由于不同表格的输入尺寸、数值精度和推理框架不同下面只在相同实验条件内比较。消融实验全局与多尺度缺一不可论文首先在 Cityscapes 上验证 Multi-Scale Linear Attention 的两个核心组成部分。输入分辨率为1024×2048各模型通过调整宽度被控制在相同的4.4G MACs并且均从随机初始化开始训练。在参数量和 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/sSegNeXt-T79.84.3M51G93.2 ms95 image/sEfficientViT-B180.54.8M25G24.3 ms175 image/sSegNeXt-B82.628M276G228 ms41 image/sEfficientViT-B383.040M179G81.8 ms70 image/sSegNeXt-L83.249M578G374 ms26 image/sEfficientViT-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/sSegNeXt-T41.14.3M6.6G12.4 ms842 image/sEfficientViT-B142.84.8M3.1G4.0 ms1142 image/sSegNeXt-S44.314M16G17.2 ms592 image/sEfficientViT-B245.915M9.1G7.3 ms846 image/sSegFormer-B450.364M96G44.9 ms212 image/sEfficientViT-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 msSwinIR43.49 dB61.2 ms32.31 dB9.7 msEfficientViT w0.7543.54 dB14.3 ms32.31 dB2.8 msEfficientViT43.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 mAPSAM-ViT-H641M2973G11 image/s46.544.2EfficientViT-SAM-L261M69G538 image/s46.642.7EfficientViT-SAM-XL0117M185G278 image/s47.543.9EfficientViT-SAM-XL1203M322G182 image/s47.844.4EfficientViT-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 倍加速。本文方法EfficientViT 的创新并不是简单删除 Softmax而是围绕 ReLU Linear Attention 的优势和缺点组合出一套适合高分辨率密集预测的模块。一、Softmax Attention 的问题高分辨率输入会产生大量 token。标准 Softmax Attention 需要显式计算 token 两两之间的关系token 数量为 N 时注意力矩阵的尺寸为 N×N。当分辨率提高时计算量和内存占用会快速增长。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 的分布通常更平滑对局部细节的捕获能力较弱。在语义分割、超分辨率等密集预测任务中物体边界、小目标和纹理信息都依赖局部建模因此这一缺点尤其值得重视。为此作者采用两项补偿设计在 FFN 中插入 DWConv增强局部信息提取能力。聚合 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。论文 v6 的默认实现采用两分支设计一条分支保留原始 Q/K/V另一条分支聚合5×5邻域 token在多尺度能力和实际效率之间做平衡。这里的5×5是论文实验配置而不是该模块只能使用的固定定义。多尺度聚合采用轻量的小核卷积而不是直接堆叠大卷积核。实现层面作者将原本针对各 Q/K/V 和各注意力头的独立聚合融合为一次 DWConv 与一次分组1×1卷积从而减少零散算子调用。timm源码中通过ConvNormAct使用1x1卷积输出 3 倍维度刚好对应 Q/K/V的生成。self.qkvConvNormAct(in_channels,3*total_dim,1,biasuse_bias[0],norm_layernorm_layer[0],act_layeract_layer[0],**dd,)多尺度聚合的实现对qkv做多尺度空间加权融合。用深度卷积 分组 1×1 卷积实现多尺度上下文聚合替代普通注意力单尺度融合。第一个卷积kernelscale, groups3*total_dim。逐通道深度卷积用于空间多尺度上下文提取groups输入通道数 →深度卷积 DepthwiseConv等价于对 Q 整组、K 整组、V 整组分别做多尺度空间邻域信息聚合。cale 越大 → 抓取更远距离空间关系实现多尺度上下文第二个卷积1×1, groups3*heads。多头内通道交互打散分组、融合头内特征self.aggregnn.ModuleList([nn.Sequential(nn.Conv2d(3*total_dim,3*total_dim,scale,paddingget_same_padding(scale),groups3*total_dim,biasuse_bias[0],**dd,),nn.Conv2d(3*total_dim,3*total_dim,1,groups3*heads,biasuse_bias[0],**dd),)forscaleinscales])最后线性注意力层舍弃传统 Softmax 缩放点积用矩阵聚合 归一化实现全局注意力无指数运算、速度极快也是你 TRT FP16 出 0 输出的高危算子FP16 下矩阵相乘极易炸 NaN先升 FP32 计算最后切回原精度。def_attn(self,q,k,v):dtypev.dtype q,k,vq.float(),k.float(),v.float()kvk.transpose(-1,-2) v outq kv outout[...,:-1]/(out[...,-1:]self.eps)returnout.to(dtype)四、EfficientViT ModuleMulti-Scale Linear Attention 与FFN DWConv共同构成 EfficientViT ModuleInput │ Multi-Scale Linear Attention ← 多尺度全局上下文 │ FFN DWConv ← 通道变换与局部信息 │ Output两部分的分工非常清楚:组成部分主要作用解决的问题Multi-Scale Linear Attention建立全局关系并融合多个尺度高分辨率下全局建模开销大FFN DWConv变换通道并提取局部特征ReLU 注意力对局部细节不够敏感五、整体架构与特征融合以论文图示中的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 的核心贡献可以归纳为三点ReLU Linear Attention利用可分解的相似度形式和矩阵乘法结合律使注意力关于 token 数量由二次增长降为线性增长。Multi-Scale Linear Attention使用轻量卷积生成多尺度 Q/K/V将全局感受野与多尺度学习结合起来。卷积与注意力协同设计利用 FFN 中的 DWConv、后置注意力 stage 和轻量 Head兼顾局部细节与硬件效率。从实验看EfficientViT 的价值主要体现在性能—效率折中它不只追求更低的 MACs也不只追求最高精度而是努力让计算量下降转化为论文所测试硬件和软件栈上的真实延迟与吞吐量优势。需要注意MACs 描述乘加操作数量不能在没有统一计数约定时直接等同于 FLOPs不同设备上的延迟也会受到算子实现、数值精度和推理框架影响。论文中的速度结论只适用于相应测试条件不应无条件外推到所有部署环境。