NAS新视角:图神经网络驱动的通用架构空间,混合卷积与Transformer性能跃升!

机器学习算法人工智能与算法

点击下方卡片,关注「AI视界引擎」公众号

( 添加时备注:方向+学校/公司+昵称/姓名 )

picture.image

picture.image

作者提出通用神经架构空间,一个用于神经架构搜索(NAS)的通用搜索空间,它在一个单一、灵活的框架下统一了卷积网络、transformers及其混合架构。yinqin-Transformers能够在通用框架下发现新颖的架构以及分析现有架构。

作者还提出了一种新的搜索算法来遍历所提出的搜索空间,并证明了该空间包含有趣的架构,在采用相同的训练设置时,这些架构的性能优于最先进的手工设计架构。

最后,作者引入了一个统一的工具包,其中包含标准化的训练和评估协议,以促进可复现性并能在NAS研究中实现公平比较。总而言之,这项工作为通过一个统一的、基于图的NAS视角来系统地探索神经架构的全谱系开辟了一条道路。

unsetunset1. 引言unsetunset

尽管神经架构搜索(NAS)在为预定义架构[37, 51]识别最优超参数配置或在改善边缘设备[15, 36]上的推理延迟方面取得了不可否认的成功,但据作者所知,它至今未能产生一种能够显著优于基于ResNet[13, 34, 35]或Vision Transformer[7, 10, 39]的流行但手工设计的网络架构的新型网络架构。

与计算机视觉研究的许多领域一样,NAS的进展主要由可用的基准测试驱动,因此近年来,大多数NAS方法[11, 17, 31, 42]都集中于NAS-Bench系列[9, 46]的表化NAS基准测试。表化基准测试是一个固定的网络架构数据集,其中每个网络的精度和其他参数都已被预先计算,因此NAS算法无需训练网络即可获得其精度。尽管表化数据集通过降低所需的计算成本来促进NAS研究,但根据定义,它们也阻止了NAS方法发现新的、更好的架构——搜索空间中的最佳架构已是已知的,并且只需查阅预先计算的精度表就能找到它。此外,这些表化数据集的搜索空间通常相对较小,并且局限于多次重复相同的构建单元,这是因为在创建基准测试期间,需要对每个网络至少训练一次。

得益于免训练NAS方法的出现[5, 19, 22, 23, 40, 48],现在可以在不训练网络架构的情况下估计其准确性,从而消除了NAS的计算约束,并能够探索更大的搜索空间和新型架构。为此目的引入的搜索空间——即Zen-NAS [23]和AutoFormer [3]——是通过改变所选手工设计架构块的超参数(如深度、宽度或扩展比率)创建的,这些架构块分别是MobileNetV2 [33]和Vision Transformer [10]。由于每个搜索空间都是通过改变给定架构块的超参数创建的,因此使用这些搜索空间仍然无法找到具有不同拓扑结构或构建块的网络,随后,手工设计的架构(如CoAtNet [7])仍然优于AutoFormer搜索空间中已知的最佳架构[2]。

在本文中,作者旨在通过引入一种名为通用神经架构空间(UniNAS)的新颖搜索空间来解决上述局限性。该搜索空间以通用方式设计,不偏向于任何特定的现有架构,同时确保所有最先进但手工设计的架构都包含在该搜索空间内。因此,该搜索空间不仅允许探索现有架构的各种组合和全新的架构,还允许在一个统一的框架内对最先进的网络、其拓扑结构和设计选择进行系统性分析。

此外,作者提出了一种新的架构搜索算法来遍历所提出的UniNAS搜索空间。当与最先进的免训练NAS Agent [40]结合时,在遍历搜索空间的少量步骤后,作者发现了一种新颖的网络架构——UniNAS-A,该架构在使用相同的训练协议和大小约束下,优于当前最先进的架构,这表明该搜索空间包含有趣的架构,值得进一步探索。

最后但同样重要的是,作者提供了一个工具包(作为一个易于安装的Python包),该工具包允许遍历所提出的搜索空间,为搜索空间中的任何点创建PyTorch网络模块,最重要的是提供了一致且明确的训练协议和训练代码来创建和评估最终模型。这对于可复现性和NAS方法的公平比较极其重要,因为尽管作者通常会报告在标准数据集(如ImageNet)上的最终准确率,但他们使用不同的训练计划、不同的超参数设置和数据增强技术,或者他们甚至会使用一个在显著更多数据上预训练的现有更大模型作为教师网络——这使得所提出架构之间的公平比较几乎不可能。

综上所述,作者对社区做出了如下贡献:

  1. 作者提出了一种名为UniNAS的新型通用NAS搜索空间。该搜索空间以通用方式设计,因此包含了许多新颖的网络拓扑结构,与此同时,它也包含了已知的人工设计架构,从而允许在统一框架内对网络拓扑结构进行系统分析。
  2. 作者引入了一种新的搜索算法来遍历所提出的设计空间,并表明该空间包含有趣的新颖架构——例如UniNAS-A——它在多种任务(分类、检测和分割)上优于手工设计的最先进架构。
  3. 作者提供了一个统一的工具包,包括标准化的训练和评估协议,以帮助采用所提出的通用搜索空间,增强NAS方法的可复现性,并促进未来的NAS研究。

unsetunset2. UNINASunsetunset

本节介绍了UniNAS搜索空间,首先介绍作为其核心组件的UniNAS块,然后是整体网络架构,最后通过与现有搜索空间的比较,突出关键差异和改进。

A. UniNAS块

a) 基本操作:作者的目标是创建一个跨越多种不同架构的搜索空间,因此作者在UniNAS块的定义中采用通用设置,在架构搜索过程中实现最大灵活性,同时保留与现代分层网络兼容的实际约束。具体而言,UniNAS块被定义为具有一个输入节点和一个输出节点的任意有向无环图(DAG),其中每个中间节点代表1个基本操作。基本操作包括卷积层(深度可分离、逐点、标准等)、池化、 Mask 、正则化元素、各种形式的矩阵乘法和点积、具有多输入或多输出边的操作以及非线性函数(见表1)。这使得该块能够表示多样化的局部计算模式,同时保持一致的接口属性。

虽然表1中的大多数基本操作确实很基础,因此无需进一步解释,但作者将具体介绍其中两个较为复杂的操作:对于张量

,作者定义

对于

以及对于

,

对于

。Matmul1和Matmul2表示一种乘法关系,但与逐元素乘法不同,它结合了所有通道信息,并与Softmax结合,这些可以组合产生注意力机制。然而,作者再次注意到,可变性远不止于此,因为任何图都是可能的,只有维度才是重要的(!)。

表1: UniNAS 块基本操作 作者报告了输入大小为

的单次前向传播的参数数量 (Params) 和浮点运算次数 (FLOPs)。某些节点会改变形状和/或在多输入或多输出模式下运行,

表示

个相同形状的张量。

b) 块计算图:作者对这个通用公式施加的唯一约束是:1) 块的输入和输出维度保持相同,以及 2) DAG 内相邻节点之间的维度匹配,以确保张量在图中正确传播。这一约束简化了在更大架构中的集成,其中跨块的一致特征图形状显著减少了前向传播过程中动态形状处理的复杂性,并能够在各种搜索到的块上实现稳定训练。通过强制执行一致的维度,作者避免了自动插入额外投影的需要,否则这些投影可能会干扰对搜索到的拓扑结构的分析。

作者强调,现代深度学习架构中的许多常用模块,例如ResNet [13]及其变体的残差块、基于transformer的模型[7, 10]中带或不带相对位置偏置的自注意力层、用于自适应通道重新校准的squeeze-and-excitation模块[16],以及EfficientNets [34, 35]中使用的inverted mobile bottleneck结构,都可以表示为作者的UniNAS块的实例。关于如何将这些模块制作成UniNAS块格式的示例,包括精确的图结构和节点操作,可以在图3中看到。

picture.image

然而,作者鼓励读者不要局限于这些经典例子;人们当然可以想象出更多样化的网络。搜索空间不仅包括简单地在链中堆叠卷积和非线性,还涵盖了树状结构、具有选择性注意力合并的并行路径,以及卷积和自注意力层的混合组合,所有这些都封装为一个UniNAS块。这展示了块设计的表达能力,同时在整个搜索空间中保持统一和连贯的表示。

B. UniNAS网络

最终的网络结构将不同的UniNAS块顺序堆叠成单链(见图2)。这遵循了当前最先进网络的一般设计,其分层 Backbone 网络类似于[7, 13, 35, 39]。茎阶段(S0)使用卷积层对输入进行下采样,其后是多个阶段,每个阶段包含多个不同的UniNAS块,最后是一个由全局平均池化和全连接层组成的分类头。由于UniNAS块保持维度不变,在每个阶段的开始,空间维度通过步长为2的标准最大池化减小,而通道数通过通道投影1增加。如前所述,UniNAS块的结构(即图表示)在整个网络中是变化的,这增强了网络的拓扑可变性。阶段数、块数以及空间和通道维度作为缩放超参数,允许探索不同的网络规模模式。

picture.image

表2:NAS搜索空间比较

UniNAS涵盖了最多种类的网络,它不限制块拓扑结构,并允许探索新颖的架构。作者还报告了每个空间中已知最佳网络在ImageNet-1k上的分类准确率——在未找到可比训练设置下的报告结果时,作者使用与表3相同的训练协议训练给定搜索空间的已知最佳架构(用†表示)。

picture.image

C. 与现有搜索空间的比较

现有的搜索空间存在两个主要限制:受限的拓扑可变性和较差的可扩展性。因此,尽管该领域已有广泛的前期研究,但在这些空间中发现的性能最佳的模型,与使用 UniNAS 获得的架构相比,其性能仍然不及(见表2)。

诸如[26, 49]中的DARTS方法依赖于权重共享超网络,这些网络计算成本高昂且产生有偏的梯度估计,导致不可靠的架构排名。当扩展到transformers [4]时,自注意力层的成本迫使这些搜索空间受到严重限制——通常将搜索简化为"使用注意力或不使用注意力"的二元选择。即使有这些限制,基于DARTS的网络仍然局限于小规模设置,ImageNet-1k准确率低于

(表2)。

诸如NAS-Bench [9, 46]等基准测试面临一个根本的可扩展性问题:可能网络的数量随操作数量呈指数级增长,使得穷举训练仅在玩具空间上可行。因此,这些基准测试仅限于在CIFAR [18]或ImageNet16-120 [6]等小型数据集上评估的简化卷积网络。在ImageNet-1k上,要么不存在结果,要么报告的性能远低于最先进水平,由于架构空间是有限的且已被完全探索,因此没有进一步改进的空间。

Zen-NAS [23] 和 AutoFormer (V1/V2,也称为 S3) [3, 4] 将搜索空间限制在具有重复 MobileNetV2 块 [33] 或 Vision Transformer 块 [10, 27] 的网络上。生成的架构仅在超参数上有所不同,如扩展比例、通道数量或头数。然而,目前无法公平比较这些空间中报告的顶级模型,因为它们是通过从一个更大的教师模型进行知识蒸馏来训练的,并且使用了 substantially 更多的数据 [19],因此尚不清楚报告的性能提升是由于搜索本身还是蒸馏过程。事实上,当应用蒸馏时,[12] 使用标准的 EfficientNet-B2 [34] 在相同的 FLOPs 和训练预算下实现了更高的准确率。类似地,当在相同的参数预算下训练时,[43] 中的手工设计的 transformer 在 ImageNet-1k 上的性能超过了 AutoFormer 的最佳 transformer [2, 3] 超过

总结来说,有两个主要的改进领域:网络多样性和公平可复现的比较。这两个问题都由UniNAS解决。1) UniNAS的灵活性允许进行真正拓扑感知的架构搜索,其中注意力机制可以插入到选定位置,与卷积结合,或在某些分支中完全替换,这种灵活性在之前受限于单一操作选择的框架中是不可行的;2) yinqin-Transformers允许对网络拓扑进行公平比较,因为它涵盖了所有提到的空间以及诸如ResNet [13]、EfficientNet [34]、ViT [10]和CoAtNet [7](见图3)等拓扑结构,因此它能够在一个框架内直接比较它们的特性和准确性,以及相互之间和与新型架构的比较。

与之前的搜索空间不同,在作者的UniNAS中,在每个阶段使用不同的块是可行的(参见[28]对NAS-Bench框架的批评,该批评源于以下事实:简单的模型参数数量对于估计最终网络精度具有最高的预测强度,因此阻碍了研究行人探索新的有趣拓扑结构)。

unsetunset3. 架构搜索unsetunset

在本节中,作者提出了一种基于给定标准在UniNAS搜索空间中查找网络的算法。具体而言,作者假设设计选择如阶段数、块数和基础通道维度是已知的,并且以Params和FLOPs边界形式的约束是给定的。然后作者的目标是在上述约束条件下找到最大化特定目标(如最高准确率)的网络。

a) 搜索步骤:更正式地说,任何UniNAS网络由一系列图

标识,每个图对应一个UniNAS块,其中

是总块数。这使作者能够将UniNAS搜索表述为一种基于图的算法,涉及节点的添加和消除。为了保证搜索能够高效地遍历1)仅在网络大小约束范围内的网络,并且

  1. 可行图(就节点维度和奇偶性而言),作者采用以下方法。

作者将每个节点

(其中

)与可训练参数数量Params和近似浮点运算数量FLOPs相关联。这些值可以很容易地作为每个基本操作节点的输入张量形状的函数获得(可以在表1中找到),这使作者能够在添加/删除给定节点后高效计算整体网络成本。

作者注意到,每个基本操作的Params值计算起来非常直接。然而,估计FLOPs并非易事,作者选择将每个操作的所有乘法和加法都包含在其最终的FLOPs值中。这样做是为了给所有操作赋予非零成本,以防止出现不可控的复杂度发散。因此,这个值可能与运行时FLOPs估算器(如PyTorch profiler)返回的值不同,而且,该值还可能因底层硬件的不同而有所差异。

  1. 作者为搜索空间定义了可行的节点添加和消除操作。这些操作将确保UniNAS块在一次搜索步骤后仍然是一个可行的计算图,同时允许作者在大型UniNAS空间中自由探索。请记住,图的可行性是由表1中的输入/输出节点形状决定的。只需确保:a) 仅当输入空间维度具有整数平方根时才添加RelPosBias,b) 仅当通道维度可被2或3整除时才分别添加Chunk2、Chunk3,c) 仅当通道维度可被4整除时才添加ConvRed4,d) 改变维度和/或具有多个输出的节点与其耦合节点一起添加。当消除时,这两个节点之间的分支也会被消除。

算法1 UniNAS搜索步骤

picture.image

要求:UniNAS网络,即一系列图

,可能的节点类型列表(参见表1,包含节点成本FLOPs,Params),搜索边界

,

, Paramsmax, Params

),消除概率 peliminate

,最大尝试次数

。 1: 当

时执行

2: 随机选择一个 graph

和一个 node

3:采样

4: 如果

5: {节点

消除}

6: 确定在 v 和它的耦合节点之间的最小子图

(对于只有一个输出且不改变维度的节点则为空)

7: 计算在潜在消除

后 FLOPs 和 Params 的变化

8: 如果变化相对于

和 Params

是可信的,那么

9: 消除

10: 休息

11: end if

12: 否则

13: {在

之后的节点添加}

14: 计算在潜在添加

及其耦合节点

之后 FLOPs 和 Params 的变化(对于不改变维度的单输出节点,此项为空)

15: 如果变化相对于

和 Params

是合理的,则

17: 中断

18: end if

19: 结束如果

20:

21: end while最后,作者构建了(单个)搜索步骤。在每个步骤中,作者选择一个将被调整的UniNAS块,决定是添加还是消除一个节点。如果作者添加,则在随机选择的节点后添加,前提是这根据上述1)和2)是可行的。确定添加/消除选项的参数

通常设置为低于0.5,因为作者希望先添加然后消除节点,因为对于某些类型的节点,消除它们也意味着消除该节点与其配对对应节点之间的整个分支,这可能导致网络大小的大幅削减。在实验中,作者选择

的值,如图1所示,通过这种选择,随机游走避免了网络大小向无限制增长或缩小的退化。更正式的描述请参见算法1。

picture.image

b) 训练与评估协议:当免训练NAS算法完成时,其识别为最有希望的候选架构仍需进行训练以获得最终(真实)准确率。遗憾的是,NAS文献中一直缺乏精确的训练方案,不同作者使用不同的训练数据、epoch数量、训练批次大小等来报告结果,这使得不同NAS方法的直接比较变得不可能。在UniNAS搜索空间中,因此作者也提供了训练最终网络的详细训练协议(见表3),以便不同架构和NAS方法(包括未来的工作)可以以公平且可复现的方式进行比较。

c) UniNAS工具包:作者提供了一个可通过pip安装的软件包uninas,使研究行人能够轻松访问使用所提出的UniNAS空间所需的所有组件。具体而言,作者提供了将任何UniNAS网络生成为PyTorch模型的工具,以及构建自定义UniNAS网络的直观界面。此外,作者还包括一个用于计算图图形可视化的模块,该模块专门为复杂的网络结构设计。作者还提供了算法1中描述的搜索算法的实现,该实现考虑了计算预算约束,封装在单个函数调用中,该函数调用可以触发在UniNAS空间中的随机游走或针对可从PyTorch网络计算的任何目标的优化算法。最后,作者还包括了表3中描述的训练协议,并可选支持分布式训练,以促进任何未来UniNAS网络的可复现性和公平比较。

unsetunset4. 结果unsetunset

a) UniNAS中的随机游走:作者使用作者的搜索步骤算法1和参数con执行了随机游走

表3:UniNAS评估协议

作者通过首先在ImageNet-1k上训练图像分类,然后使用上述hyper-parameters在COCO和ADE20K数据集上进行fine-tuning,从而在UniNAS空间中对classification、detection和segmentation任务评估NAS方法。batch size的选择以适合单个A100 GPU为准,并且是按每个GPU报告的,因此当使用更多GPU时,需要相应地调整learning rate。

将参数量限制在22-28M,FLOPs限制在6-20G,以评估作者的架构搜索在巨大的UniNAS空间中的导航能力。在图1中,作者可以看到50万个采样网络的Params和FLOPs,表明作者的搜索可以轻松跨越不同的网络大小和配置,还涵盖了基于transformer和卷积的非常不同的网络架构。在图4中,作者提供了UniNAS基本操作的进一步细分,显示网络大小在特定操作之间高度可变,表明作者的探索有效地覆盖了非常多样化的配置。

picture.image

b) 架构搜索:在作者的UniNAS中搜索最佳架构时,作者使用免训练的NAS Agent VKDNW[40]来评估网络性能,而不是对其进行训练。作者定义

其中,

表示在随机输入批次下初始化时,Empirical Fisher Information Matrix (FIM) 特征值的第k个十分位数,作为FIM谱的一种表示。具体而言,作者分别为每个块

计算公式(3),然后对所有这些块的结果取平均,由此得到一个单一标量值,作者在搜索中对该值进行最大化。作者之所以选择这个免训练 Agent ,是因为它与网络规模正交(见[40]中的图3),而作者的目标是在特定的网络规模预算内搜索最优拓扑结构。

然后作者在UniNAS中迭代搜索:作者从初始网络开始,并执行算法1中的1024步,

表4:UniNAS搜索空间中的模型在ImageNet-1k上的分类准确率。所有模型具有相似数量的参数,并采用表3中的训练协议以相同方式进行训练。训练在8块A100 GPU上耗时2天。

picture.image

其中只保留前64个网络(种群大小)。在搜索过程中,网络规模被限制为27M参数和20G FLOPs,在图2中作者采用了4个阶段,初始输出大小为64,各阶段分别有2、3、5和2个UniNAS块,隐藏维度分别为96、192、384、768,与现代架构3一致。搜索在单个A100 GPU上耗时12小时,最终作者根据VKDNW分数选择了最佳网络——作者将其表示为UniNAS-A(见图5)。

picture.image

c) 图像分类:首先,作者将UniNAS-A与UniNAS搜索空间中的其他网络进行比较:EfficientNet [34]、ResNet [13]、CoAtNet和具有相对位置偏置的ViT [7]。为了在公平的设置下比较这些网络,作者将网络扩展到相同数量的阶段和块,并通过调整通道数量,作者也使网络规模保持相同。每个网络都在ImageNet1k [8]上进行训练,使用与表3相同的训练协议,训练在8个A100 GPU上花费了2天时间。如表4所示,UniNAS-A以显著优势优于标准手工设计的网络。

d) 下游任务

接下来,作者对上一步的每个网络在两个下游任务上进行了微调——在 MS-COCO [24] 上的目标检测和在 ADE20K [50] 上的语义分割,使用与表3 相同的设置。在4个 A100 GPU 上,分割任务的训练耗时约10小时,检测任务的训练耗时21小时4。再一次,UniNAS-A 的表现显著优于现有网络(见表5)。

e) Agent 消融

最后,作者通过使用不同的 Agent 运行搜索算法1,并评估每个 Agent 所找到的最佳架构,来对选择 VKDNW [40] 作为 Agent 进行消融。在表6中,作者表明,对UniNAS空间的搜索对特定 Agent 的选择基本保持鲁棒性,但VKDNW在相同的搜索预算(12小时)下仍然能够找到最佳网络架构。

picture.image

unsetunset5. 相关工作unsetunset

由于NAS搜索空间的讨论已在第II-C节中呈现,在此作者专注于NAS搜索方法。

a) One-shot NAS:这些方法基于离散架构空间的松弛,通常使用一个超网,该超网包含了给定搜索空间中所有可能的节点和边。该方法首次在DARTS [26]中被提出,其中超网由所有可能的操作组成,每个操作都被分配一个权重,该权重在超网训练期间通过梯度下降进行调整。一旦超网训练完成,就保留权重最高的操作,从而构建出最终的网络架构。Robust-DARTS [47]通过在超网训练中引入数据增强来改善测试时泛化能力,而SGAS [21]则旨在通过在每个阶段显式地选择允许的操作来提高超网训练的稳定性。One-shot NAS的主要挑战是内存消耗,因为超网必须包含整个搜索空间中所有可能的操作;以及排名紊乱,即在超网中评估的架构的性能可能与其作为独立网络时的性能不同。

unsetunset6. 结论unsetunset

作者提出了UniNAS,一个通用的神经架构搜索空间,旨在一个统一的框架内系统地探索、分析和比较网络拓扑。与以往的工作相比,作者进一步将计算模块分解为基本操作,从而能够表达和扩展手工设计与NAS发现的架构,同时支持对网络设计中拓扑可变性的系统性研究。

作者提出了一种高效的架构搜索算法,该算法直接在UniNAS内运行,实现了对FLOPs和参数预算的精确控制,同时支持细粒度修改以遍历多样化的架构家族,

unsetunset参考unsetunset

[1]. Universal Neural Architecture Space: Covering ConvNets, Transformers and Everything in Between

点击上方卡片,关注「AI视界引擎」公众号

0
0
0
0
关于作者
关于作者

文章

0

获赞

0

收藏

0

相关资源
火山引擎大规模机器学习平台架构设计与应用实践
围绕数据加速、模型分布式训练框架建设、大规模异构集群调度、模型开发过程标准化等AI工程化实践,全面分享如何以开发者的极致体验为核心,进行机器学习平台的设计与实现。
相关产品
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论