[论文解读] K-shot NAS: Learnable Weight-Sharing for NAS with K-shot Supernets

Xiu Su, Shan You+5 more|arXiv (Cornell University)|Jun 11, 2021

Machine Learning and Data Classification参考文献 42被引用 6

一句话总结

本文提出 K-shot NAS，一种可学习的权重共享框架，用每种操作的 K 个超网络字典替代单个超网络权重共享，通过单纯形网络实现路径特定的权重凸组合。该方法通过更有效地逼近独立训练权重，提升了网络架构评估的准确性，在仅使用 412M FLOPs 的情况下实现了 ImageNet 上 77.9% 的 Top-1 准确率。

ABSTRACT

In one-shot weight sharing for NAS, the weights of each operation (at each layer) are supposed to be identical for all architectures (paths) in the supernet. However, this rules out the possibility of adjusting operation weights to cater for different paths, which limits the reliability of the evaluation results. In this paper, instead of counting on a single supernet, we introduce $K$-shot supernets and take their weights for each operation as a dictionary. The operation weight for each path is represented as a convex combination of items in a dictionary with a simplex code. This enables a matrix approximation of the stand-alone weight matrix with a higher rank ($K>1$). A extit{simplex-net} is introduced to produce architecture-customized code for each path. As a result, all paths can adaptively learn how to share weights in the $K$-shot supernets and acquire corresponding weights for better evaluation. $K$-shot supernets and simplex-net can be iteratively trained, and we further extend the search to the channel dimension. Extensive experiments on benchmark datasets validate that K-shot NAS significantly improves the evaluation accuracy of paths and thus brings in impressive performance improvements.

研究动机与目标

• 解决单次 NAS 的局限性，即所有路径共享权重会阻碍性能评估的准确性。
• 通过允许路径自适应学习不同的操作权重，提升基于超网络的 NAS 的可靠性。
• 通过 K 次超网络字典和路径特定的编码，实现对独立训练权重的更高秩矩阵逼近。
• 通过为通道宽度引入非参数正则化，扩展搜索范围，以提升编码生成的判别能力。
• 验证可学习、架构定制的权重共享可带来更优的最终模型性能。

提出的方法

• 引入 K-shot 超网络，其中每个操作维护一个包含 K 个权重向量的字典，而非单一共享权重。
• 使用单纯形码向量 λ，将每条路径的操作权重表示为 K 个字典项的凸组合。
• 设计一种单纯形网络，基于架构和通道宽度编码生成路径特定的码 λ。
• 将权重共享建模为更高秩矩阵逼近问题，突破单次 NAS 的秩-1 限制。
• 引入非参数正则化，以增强不同通道宽度下编码的判别能力。
• 采用迭代优化方案，联合训练 K-shot 超网络和单纯形网络。

实验结果

研究问题

• RQ1 与单次 NAS 相比，采用可学习权重共享的 K-shot 超网络是否能提升路径评估的准确性？
• RQ2 使用路径特定的单纯形码如何影响对独立训练权重的逼近效果？
• RQ3 通过专用正则化将搜索扩展至通道宽度，是否能提升最终模型性能？
• RQ4 在搜索效率与准确性之间，超网络数量（K）与训练轮数之间存在何种权衡？
• RQ5 所提方法是否能在低 FLOP 预算下实现最先进性能？

主要发现

• K-shot NAS 仅用 412M FLOPs 即在 ImageNet 上达到 77.9% 的 Top-1 准确率，显著优于基线单次 NAS 方法。
• 该方法通过实现路径特定的权重自适应，提升了评估准确性，缩小了超网络排序与 Oracle 性能之间的差距。
• 在 K-shot 设置下，使用学习的单纯形网络生成编码相比固定或随机编码分配，准确率提升 0.28%。
• 学习到的单纯形码直方图显示分布广泛且多样，表明路径特定的权重自适应效果良好。
• 权衡分析表明，K=8 在训练成本与性能之间提供了良好平衡，更大的 K 值需要更多轮次才能收敛。
• 非参数正则化提升了不同通道宽度下编码的判别能力，有助于联合搜索中的性能提升。