AutoFormer——首个Vision Transformer的One-shot NAS

Paper地址：https://arxiv.org/abs/2107.00651

GitHub地址：GitHub - microsoft/Cream: This is a collection of our NAS and Vision Transformer work.

前言

Bert/Transformer结构，已成为自然语言处理、语音识别与多模态匹配等模型的主流结构，视觉模型也逐渐往Transformer结构收敛。本文针对视觉Transformer提出了AutoFormer方法，是一种基于weight entanglement的One-shot NAS，充分预训练的超网络、能满足多种资源约束的部署需求。

如下图所示，传统的基于Block choice并排训练的可微分形式（左图）、或SPOS形式（右图）的One-shot NAS，通过权重共享方式、完成超网络预训练，并进一步通过子网络采样/搜索、完成模型结构优化。但是多个Block的并排训练，一方面限制了搜索空间的扩大，另一方面不利于超网络的充分预训练。并且，将传统One-shot NAS直接应用于视觉Transformer，训练收敛速度与精度效果都不理想（见论文分析）。

本文提出的Weight entanglement，在不额外增加Block choice的前提下，通过通道宽度、网络深度、Attention head数目等多个维度的调整，实现了视觉Transformer模型的动态可伸缩预训练与结构搜索。Weight entanglement的做法，类似于BigNAS、FBNet-v2的通道搜索，都不会额外增加通道维度的权重参数量。

相比于手工设计的CNN模型（ResNet、ResNext、DenseNet）与视觉Transformer模型（ViT、Deit），AutoFormer模型在相同资源开销条件下、能够获得最好的识别精度：

AutoFormer原理 

AutoFormer总体思路如上图所示，是一种基于Weight entanglement的动态可伸缩搜索方法，其搜索维度包括Attention heads、通道宽度与网络深度。下面从搜索空间、搜索策略与搜索效率这些维度加以分析：

• 搜索空间：Embedding dimension、Q-K-V dimension、Attention heads、MLP ratio与Network depth；根据不同的资源约束需求，分别设置Supernet-tiny、Supernet-small与Supernet-base三个基本的Template model；

  

• 搜索策略：基于Weight entanglement的One-shot NAS，完成超网络预训练之后，通过进化算法执行子网络搜索。包含l个网络层的子网络结构可按下式表示，表示第i层的Block结构，表示相应的权重参数：

  

   而每个Block结构，都是按照Weight entanglement原则从超网络采样获得，下式表示n个动态选择范围：

  

  

  

•  搜索效率：按照权重共享形式、完成超网络预训练，训练效率高、且收敛速度快；另外，由于没有引入额外的Block choice，因此训练时的Memory cost较低；

实验结果