MobileNetV4-移动设备生态系统的通用模型学习笔记

代码地址：models/official/vision/modeling/backbones/mobilenet.py at master · tensorflow/models · GitHub

论文地址：https://arxiv.org/pdf/2404.10518.pdf

这篇论文介绍了最新一代MobileNet卷积神经网络，被称为MobileNetV4(MNv4)，旨在为移动设备提供通用高效的架构设计。主要贡献包括:

引入了Universal Inverted Bottleneck(UIB)搜索块，这是一种统一灵活的结构，合并了Inverted Bottleneck(IB)、ConvNext、前馈网络(FFN)和一种新的额外深度卷积(Extra Depthwise)变体。

提出了针对移动加速器优化的移动版多头注意力(Mobile MQA)，相比传统多头自注意力(MHSA)提供39%的推理加速。介绍了一种优化的神经架构搜索(NAS)方法，提高了MNv4搜索的有效性。

结合UIB、Mobile MQA和改进的NAS，设计出一系列在移动CPU、DSP、GPU及专用加速器(如苹果Neural Engine和谷歌Pixel EdgeTPU)上表现极优的MNv4模型，在多种硬件平台上达到mostly Pareto optimal。

提出了一种新的知识蒸馏技术，通过数据集混合和增加平衡类内数据来增强泛化能力，实现了87%的ImageNet-1K分类精度。

---

Introduction

强调在移动设备上部署高效神经网络的重要性，既能提供实时交互体验，又能避免将私人数据传输到公共互联网。提出了两个创新的构建模块-Universal Inverted Bottleneck(UIB)和Mobile MQA，并通过改进的NAS方法将它们集成,创建了一系列在多种移动硬件平台上表现优异的模型。

简要介绍了Universal Inverted Bottleneck(UIB)模块的功能，它通过可选的深度卷积改进了Inverted Bottleneck，统一了多种微架构如Inverted Bottleneck、ConvNext和FFN。提及Mobile MQA是一种针对加速器优化的注意力块，相比多头自注意力(MHSA)提供39%的推理加速。指出该工作提出了一种优化的神经架构搜索(NAS)方案，以提高MNv4的搜索有效性。

提及通过一种新的蒸馏技术，进一步提高了MNv4模型的精度。

作者的模型系列涵盖了一系列计算范围，从具有3.8M参数和0.2G MAC的极其紧凑的MNv4-Conv-S设计，在Pixel 6 CPU上在2.4毫秒内实现73.8%的顶级ImageNet-1K精度，到为移动模型精度建立新参考的MNv4-Hybrid-L高端变体，在Pixel8 EdgeTPU上以3.8毫秒的速度运行。

---

Hardware-Independent Pareto Efficiency

介绍了Roofline Model，它通过考虑工作负载的运算强度(operational intensity)来估计给定工作负载在特定硬件上的性能，并预测它是受内存带宽还是计算能力的瓶颈所限制。

提出了Ridge Point的概念，它是硬件的峰值运算能力与峰值内存带宽的比率，代表了达到最大性能所需的最小运算强度。

通过分析不同Ridge Point下的延迟与精度权衡，发现MobileNetV4在Ridge Point从0到500时大多数情况下都是Pareto Optimal的。
解释了在低Ridge Point硬件(如CPU)上，减少MAC数量有利于降低延迟；而在高Ridge Point硬件(如加速器)上，增加MAC数量可提高模型容量而不显著影响延迟。

分析了MobileNetV4模型在不同阶段如何平衡MAC和内存带宽，使得昂贵的层不会同时导致计算和内存两方面的瓶颈，从而达到Mostly Pareto Optimal的效果。
通过理论分析和实际测量结果说明,MobileNetV4在多种移动硬件(CPU、GPU、DSP和加速器)上都能取得Mostly Pareto Optimal的性能表现,展现了其通用性。

---

Universal Inverted Bottlenecks

作者介绍了一种新的构建模块-Universal Inverted Bottleneck，UIB)。

UIB是在传统的Inverted Bottleneck块的基础上，通过引入两个可选的深度卷积层进行扩展和改进。虽然改动简单，但UIB成功地统一了几种流行的架构,包括Inverted Bottleneck、ConvNext和ViT中的前馈网络(FFN)。此外，UIB还引入了一种新的变体-Extra Depthwise Inverted Bottleneck, ExtraDW。
UIB具有四种可能的实例化形式，提供了不同的空间和通道混合权衡，扩大感受野的选择，以及最大化计算利用率的能力。四种实例包括:标准Inverted Bottleneck(IB)、ConvNext、ExtraDW和FFN。
作者解释了每种实例的工作原理和特点。例如ConvNext允许在扩张之前进行更便宜的空间混合，ExtraDW结合了ConvNext和IB的优点等。
在每个网络阶段，UIB块都可以灵活地在空间/通道混合、感受野调整和计算利用率之间进行权衡。

---

Mobile MQA

介绍了一种针对移动加速器优化的新型注意力块设计，相比传统多头自注意力(MHSA)提供39%的推理加速。

指出在移动加速器上,内存访问往往比算力更成为性能瓶颈,因此简单减少MAC操作数无法获得更好的性能，需要考虑操作强度(运算与内存访问的比率)。提出在混合模型中使用Multi-Query Attention (MQA)而不是标准的Multi-Head Self-Attention (MHSA)。

MQA通过在所有head共享keys和values的方式大幅减少内存访问开销。实验表明,相比MHSA，MQA在EdgeTPU和三星S23 GPU上实现了超过39%的推理加速，同时参数量和MACs也减少了25%以上，而精度只下降0.03%。

在MQA的基础上，作者进一步融入了Spatial Reduction Attention (SRA)的思想，对keys和values进行非对称的空间下采样，从而进一步提升效率。具体做法是在MQA中用3x3深度可分离卷积代替平均池化，对keys和values的空间分辨率进行2倍下采样。

实验表明，加入空间下采样后，相比不使用下采样，Samsung S23 GPU上的推理时间减少了25%，CPU时间减少了23%，而精度仅下降0.06%。Mobile MQA注意力块的数学表达式如下：

---

Design of MNv4 Models

作者阐述了设计MobileNetV4系列模型时采用的理念和方法-"Simplicity Meets Efficiency"。

通过广泛的相关性分析，作者总结出一些确保高相关性和接近Pareto前沿性能的关键组件和参数。总结了一些重要发现:

1、多路径设计如组卷积可能因内存访问复杂性而降低效率。

2、硬件支持很关键，一些高级模块如Squeeze and Excite（SE）、GELU等在DSP上表现不佳。

3、传统组件如深度可分离卷积、ReLU、BatchNorm和简单注意力具有良好的效率和硬件兼容性。

基于上述发现，提出了一系列设计原则:

1、优先使用标准广泛支持的组件，确保无缝部署和硬件高效。

2、采用灵活的UIB构建块，允许在空间/通道混合、感受野调整和计算利用率之间权衡。

3、使用直接简单的Mobile MQA注意力机制，追求最佳性能。

介绍了用于UIB模型搜索的改进NAS配方:采用两阶段搜索策略，分别专注于确定最佳滤波器尺寸和UIB块配置。

提出了一种增强的两阶段搜索策略，旨在缓解TuNAS中由于参数共享而导致的对小滤波器和扩张系数的偏向。

粗粒度搜索阶段: 专注于确定最优滤波器大小，同时保持固定参数(默认扩张系数4，3x3深度卷积核)。

细粒度搜索阶段: 在粗粒度结果的基础上，搜索UIB块中两个深度卷积层的配置(3x3或5x5卷积核)，保持扩张系数为4。

与传统的一次性联合搜索空间相比，两阶段搜索策略显著提高了搜索效率和模型质量。为了提高TuNAS对架构质量的评估准确性(对于奖励计算和策略学习至关重要)，作者引入了一个离线蒸馏数据集，消除了对数据增强、正则化和优化设置的敏感性。选择JFT蒸馏数据集作为TuNAS的训练集，相比在ImageNet-1K上训练，可显著提高效率(参数减少20.9%，MACs减少13.9%)。将TuNAS的训练时长延长至750个epoch，以产生更深层次、更高质量的模型,因为深层模型在延长训练时间后往往能超越宽度扩展模型。

---

Results

ImageNet classification

 COCO Object Detection