TPAMI 2024 | 动态网络新设计范式：LAUDNet实现时延可感知的动态网络|卷积|大模型|时延|神经网络|算子|算法

在深度学习领域，动态神经网络受到越来越多研究者的关注。这类网络能够在推理阶段根据需要选择性地激活网络层或通道，或者针对图像特征中信息丰富的区域动态分配计算资源，有效减少了每个输入样本所需的不必要计算量。然而，动态模型在实践中的效率往往与理论预期有所出入。这一现象主要由三个挑战引起：首先，由于研究社区的碎片化，各种动态推理范式尚缺乏一个统一的建模方式；其次，过分重视算法设计而忽略了调度策略的重要性，特别是在CUDA支持的GPU环境下，调度策略对优化计算性能和资源利用尤为关键；最后，由于大多数现有的库都是针对静态算子开发的，评估动态算子实际的时延变得格外复杂。

为了解决这些问题，作者提出了“时延可感知的统一动态网络”（LAUDNet），它将空间自适应计算、动态层跳过和动态通道跳过这三个主流的动态推理范式融为一体。进一步的，LAUDNet通过一个时延预测模型来准确评估不同类型动态算子的实际推理时延，并基于时延预测结果来优化动态网络的算法设计和调度优化。作者在LAUDNet框架内对多种视觉任务进行了试验验证，包括图像分类、目标检测和实例分割。结果表明，LAUDNet有效地缩减了理论与实际效率之间的差距。例如，LAUDNet能在V100、RTX3090及TX2等硬件平台上将静态模型ResNet-101的实际延迟降低超过50%。此外，LAUDNet在精确度与效率的平衡上也超过了其他具有竞争性的方法。

论文题目： Latency-aware Uniﬁed Dynamic Networks for Efﬁcient Image Recognition 论文链接： https://arxiv.org/pdf/2308.15949 项目链接： https://www.github.com/LeapLabTHU/LAUDNet

一、研究动机

深度神经网络在计算机视觉、自然语言处理、以及多模态理解与生成等多个领域展示了卓越的能力。尽管它们的性能非常出色，但这些深度网络高昂的计算成本经常限制了它们在资源受限的平台上的部署，这凸显了需要更高效的深度学习模型的重要性。

传统的静态网络在推理阶段采用相同的计算处理不同样本，忽略了样本之间以及样本内部不同时空位置的复杂度差异，从而带来大量冗余计算。动态神经网络[1]是一个近年来受到大量关注的方向，其以数据依赖的方式在推理阶段实现自适应的计算资源分配。这种自适应性涉及动态地跳过某些网络层或卷积通道，或执行空间自适应推理（集中计算努力于图像中最具信息量的区域）。随着这一领域的发展和各种动态模型显示出的潜力，一个重要的问题是：

在实际使用中，如何设计一个动态网络，以达到最优的精度-效率平衡？

解决这个问题具有相当大的挑战性，因为在不同的动态计算范式之间进行公平比较很困难。作者将这些挑战总结为三个方面：

缺乏一个统一的框架来涵盖不同的动态推理范式，因为这一领域的研究往往是碎片化的，研究者们通常在自己的应用场景中设计某个单一的方法；
过分关注算法设计，这经常导致实际效率与理论计算潜力之间的不匹配，因为调度策略和硬件属性对现实世界的推理时延也具有重大影响；
在不同硬件平台上评估动态模型的延迟是一项费力的任务，因为常见库（例如cuDNN）并不是为加速许多动态操作而构建的。

为解决这些挑战以回答上述问题，作者提出了时延可感知的统一动态网络（LAUDNet）。

首先，作者重新思考了跳层、通道跳过和空间动态卷积的算法设计，并通过“掩码与计算 (Mask-and-Compute)”的形式将它们整合在一个统一的建模框架中。
其次，作者深入探讨了将理论效率转化为实际加速的挑战，特别是在多核处理器如GPU上。传统方法通常采用与硬件无关的FLOPs（浮点运算）作为一个粗略的效率衡量标准，未能为算法设计提供延迟感知指导。在动态网络中，自适应计算结合次优调度策略加剧了FLOPs与延迟之间的差距。此外，大多数现有方法在最细粒度(Granularity)上执行自适应推理。例如，在空间维度动态推理中，每个特征像素的计算决策都是独立做出的。这种细粒度的灵活性导致了非连续的内存访问，需要专门的调度策略。
鉴于动态算子具有独特的内存访问模式和调度策略，为静态模型设计的库(如cuDNN)，并不能有效的优化动态模型。库的缺乏支持意味着每个动态算子都需要个性化的调度优化、编译和部署过程，使得在硬件平台上评估网络时延成为一个学界几乎不可能完成的工程问题。为解决这一问题，作者提出了一个新颖的时延预测模型，该模型通过考虑算法设计、调度策略和硬件属性有效地估算网络延迟。与与硬件无关的FLOPs相比，预测的延迟提供了更现实的动态模型效率表示。基于此时延预测模型，作者优化了动态模型的算法设计与底层调度策略，实现了实际效率的提升，

本文是作者此前于NeurIPS 2022发表的会议论文[2]的拓展。在会议论文中，方法仅针对空间自适应推理提出了时延预测模型，以优化该类型动态网络的实际推理效率。本文在以下方面进行了拓展与补充:

算法上，对不同动态推理范式的统一建模；
时延预测模型的增强，以支持更多动态推理范式；
Transformer架构的支持(包括算法与时延预测)；
支持动态算子的批量推理，提升了时延预测的准确性以及模型的推理效率；
在统一的框架下对不同动态范式的综合性比较分析。

二、方法

LAUDNet框架的核心在于三个关键动态推理范式的统一：空间自适应计算、动态跳层和动态通道跳过。这些范式通过一个“掩码与计算”的形式统一在一起。进一步的，LAUDNet通过引入时延预测模型，从而指导算法设计中的自适应计算粒度选择，以及通过支持批量推理、算子融合等方式优化底层调度，提升模型的实际推理效率。

2.1 算法设计

“掩码与计算"范式。如图所示，LAUDNet 采用“掩码和计算”的方式统一了三种不同的动态推理范式。
- 空间自适应计算：作者采用轻量化的一层卷积生成二值空间掩码，掩码上1元素的位置则代表该位置需要计算，否则跳过该空间区域的计算。该掩码作用于每个残差块最后一层输出特征。
- 动态跳层：与空间自适应计算共用同一个掩码生成模块，区别是只生成标量，决定整个残差块是否计算，如果不计算，则直接通过跳层连接进入下一个残差块。
- 动态通道跳过：采用另一个轻量模块生成维的通道掩码，作用于残差块中卷积的输入与输出，以保证其输入、输出特征不同通道是否被计算的结果一致，从而兼容于采用 group convolution 或 depthwise convolution 的架构。
粗粒度动态计算。与此前工作算法设计上最大的不同在于，LAUDNet采用了粗粒度的动态计算，通过在图像块(patch, 对应空间自适应计算)/组(group, 对应通道跳过)级别而非单个像素/通道级别做出“是否计算”的决策，以鼓励连续的内存访问，从而提高硬件上的实时加速。作者将空间和通道的自适应计算粒度分别定义为 . 这两个超参数合适的选取则根据后文的时延预测模型的指导，在算法灵活性和实际推理时延二者之间取得平衡。
复杂度分析。设残差块中三个卷积的计算量分别为 .
- 空间自适应计算：定义空间激活率，其等于空间掩码中1元素所占比例, 也表示后两个卷积层( 和卷积)中特征图需要计算的区域比例。而由于卷积需要被选中像素的相邻像素也参与计算，所以定义第一个卷积的激活率为 .可以得到，整个动态残差块的计算量与静态算子之比为
- 动态跳层：其计算量比例等于层激活率 .
- 动态通道跳过：设通道激活率为 , 则计算量比例为
基于Transformer的实现。在Transformer中实现几种动态推理范式则更为自然，因为视觉Transformer中图像块先天以token的形式存在。本文基于已有工作 AdaViT [3]实现，并支持了Transformer中动态算子的时延预测，以综合比较不同动态推理范式的性能。
模型训练。
- 损失函数。根据每个残差块的激活率，作者在每次前向传播中都计算出当前动态模型的理论计算量，并在损失函数中添加一个辅助项，让与对应的静态模型计算量只比接近某个数值 : 此外，作者还加入了之前方法[4]中对于激活率的上下界的约束项目 . 最后，作者采用了“自蒸馏"的训练损失，让动态模型的输出尽可能接近静态模型。最终损失函数定义为其中，是任务原本的损失，是超参数，是温度参数，是SoftMax函数。
- 离散的掩码生成模块均采用 Gumbel SoftMax 技术，与网络参数一起进行端到端优化。

2.2 时延预测模型

作者通过对GPU硬件进行三级建模，综合考虑硬件属性、网络参数、动态范式、粒度以及激活率，实现了对动态算子实际推理时延的显式估算：。最终一个动态算子的推理数据搬运、并行计算以及一个硬件相关的常数相加组成:

2.3 调度优化

LAUDNet还基于Nvidia Cutlass改进了调度策略，允许了批量推理、算子融合等加速推理的技术，进一步提高了GPU平台上的效率。

三、实验结果

LAUDNet在多个视觉任务上进行了广泛的评估，包括图像分类、目标检测和实例分割。以下按照逻辑顺序介绍主要结果与部分结论，请感兴趣的读者查阅原文获取更多实验结果与分析。

时延预测模型能准确预估不同动态算子的实际时延：

根据时延预测模型指导空间粒度的选择（一定范围内增大能有效降低实际时延）：

批量推理对于时延预测结果稳定性、提高推理效率的作用：

算子融合对于提高推理效率的作用：

动态计算粒度对时延-准确率曲线的影响（更大的粒度带来更低的时延，且空间粒度对时延影响大于通道粒度）

不同动态范式的综合比较（CNN上动态跳层在更多硬件上更加高效，而视觉Transformer上三种范式的结合最为有效）：

空间自适应计算的可视化（能够在只使用类别标签的情况下准确定位前景物体）：

下游任务：

四、总结

LAUDNet的主要贡献在于提出了一个统一的动态网络框架，并通过延迟预测模型指导了算法设计和调度优化。该工作不仅提高了深度卷积神经网络的效率，而且在多种视觉任务上超越了现有的方法。LAUDNet的研究工作为深度学习模型的效率优化提供了新的思路，特别是在资源受限的设备上的应用前景广阔。本文提出的动态网络设计范式，也可以启发未来动态神经网络在生成式任务或LLM中的运用。

参考文献

[1]: Han, Yizeng, et al. "Dynamic neural networks: A survey." TPAMI 2021.

[2]: Han, Yizeng, et al. "Latency-aware spatial-wise dynamic networks." NeurIPS 2022.

[3]: Meng, Lingchen, et al. "Adavit: Adaptive vision transformers for efficient image recognition." CVPR 2022.

[4]: Verelst, Thomas, and Tinne Tuytelaars. "Dynamic convolutions: Exploiting spatial sparsity for faster inference." CVPR 2020.

llustration From IconScout By Twiri

-The End-

扫码观看！

本周上新！

“AI技术流”原创投稿计划

TechBeat是由将门创投建立的AI学习社区（