【ICML2021】 One-shot 权重共享神经网络结构搜索算法

为了设计出高性能低计算量的网络结构，研究人员往往需要采用神经网络结构搜索（NAS）算法来获得突破手工设计上限的结构。为了提高搜索效率，目前方法通常采用一个 One-shot 的超网络，使得在搜索空间中的结构都能从超网络中获得相互共享的网络权重（weights）。在此类方法中，权重共享的假设十分重要。特别地，对于超网络中的任意 operation，所有的结构具有同样的网络权重。尽管这种假设能有效提高网络的训练效率，然而这也限制了不同网络结构在 Supernet 中选择不同 weights 的自由度。事实上，在单个的网络结构训练中，不同结构对应的同一 operation 的 weights (stand-alone weights) 应该有一定的不同。Supernet 作为一个网络评估器，应尽可能满足其中包含任意子结构 weights 与 stand-alone 训练的 weights 的一致性。显然 One-shot 框架下的 Supernet 构建方式很难满足这种需求并实现准确地评估不同结构的性能。

在本篇论文中，我们提出了一种 K-shot NAS 的 Supernet 构建框架以释放子网络结构在 Supernet 中的潜力。不同于 One-shot 框架中每个 operation 仅保留一套独立的网络权重的做法，在 K-shot NAS 方法中，我们对于每一个 operation 构建了一个权重的字典，对应于 K 个Supernets 的 weights。所有包含特定 operation 的子网络结构都将共享这 K 个 Supernets 的全部 weights， 因此我们致力于探索每个子网络结构对于 K 个 Supernets 的 weights 的定制（组合）方式即可。具体而言，对于任意子结构，每个 operation 的 weight 都将通过一个 simplex 上字典中的凸组合进行表示。一个 simplex-net 用来生成不同子结构对应的 operation 的 K 个 Supernets 凸组合的权重。基于每个结构对应的凸组合，其独有的 weight sharing 方式将被构造和使用。简单而言，One-shot 框架利用 rank-1 的矩阵来近似所有子结构对应的单独训练的权重，而我们的 K-shot NAS 方法则鼓励相应更高阶的近似关系。在此基础上，我们进一步探索了联合搜索 operations 和channels 的情况，channel 分支可以直接编码进入 simplex-net 来探索不同层之间的上下文信息或差异。我们进一步提出了一个非参数化的正则化方法来提高不同 path 之间 weight sharing 编码的区分性。本文算法在 ImageNet 和 NAS-Bench-201 上都取得了 SOTA (state-of-the-art) 的结果。