[论文解读] FNAS: Uncertainty-Aware Fast Neural Architecture Search

Jihao Liu, Ming Zhang+5 more|arXiv (Cornell University)|May 25, 2021

Advanced Neural Network Applications参考文献 42被引用 5

一句话总结

FNAS 通过引入三个模块——不确定性感知评论家（UAC）、架构知识库（AKP）和架构经验回放缓冲区（AEB）——实现了基于强化学习的神经架构搜索（NAS）速度提升10倍。这些模块实现了跨任务和搜索过程的知识迁移，将搜索时间从20,000 GPU小时减少至2,000 GPU小时，同时在ImageNet、人脸识别和COCO目标检测任务上达到最先进准确率。

ABSTRACT

Reinforcement learning (RL)-based neural architecture search (NAS) generally guarantees better convergence yet suffers from the requirement of huge computational resources compared with gradient-based approaches, due to the rollout bottleneck -- exhaustive training for each sampled generation on proxy tasks. In this paper, we propose a general pipeline to accelerate the convergence of the rollout process as well as the RL process in NAS. It is motivated by the interesting observation that both the architecture and the parameter knowledge can be transferred between different experiments and even different tasks. We first introduce an uncertainty-aware critic (value function) in Proximal Policy Optimization (PPO) to utilize the architecture knowledge in previous experiments, which stabilizes the training process and reduces the searching time by 4 times. Further, an architecture knowledge pool together with a block similarity function is proposed to utilize parameter knowledge and reduces the searching time by 2 times. It is the first to introduce block-level weight sharing in RLbased NAS. The block similarity function guarantees a 100% hitting ratio with strict fairness. Besides, we show that a simply designed off-policy correction factor used in "replay buffer" in RL optimization can further reduce half of the searching time. Experiments on the Mobile Neural Architecture Search (MNAS) search space show the proposed Fast Neural Architecture Search (FNAS) accelerates standard RL-based NAS process by ~10x (e.g. ~256 2x2 TPUv2 x days / 20,000 GPU x hour -> 2,000 GPU x hour for MNAS), and guarantees better performance on various vision tasks.

研究动机与目标

• 为解决基于强化学习的神经架构搜索（NAS）所面临的高计算成本问题，该方法通常需要20,000 GPU小时的训练时间，以充分训练所采样的架构。
• 通过利用先前搜索过程中的可迁移知识，提升基于强化学习的NAS在样本效率和训练速度方面的表现。
• 通过在不同任务和搜索设置之间重用架构与参数知识，实现更快收敛和更优性能。
• 证明在NAS中实现架构与参数层面的知识迁移，可显著减少搜索时间而不损失准确率。

提出的方法

• 在近端策略优化（PPO）中引入不确定性感知评论家（UAC），以联合建模架构-性能分布并估计不确定性，提升训练稳定性，将所需样本数减少50%。
• 提出架构知识库（AKP），用于存储先前训练检查点中的模块级权重，并利用模块相似性函数初始化新架构，实现约2倍的收敛加速。
• 采用结合了离策略校正因子的架构经验回放缓冲区（AEB），用于存储并重放先前训练好的模型，通过高效复用经验，将搜索时间减少一半。
• 使用多目标奖励函数，直接在ImageNet、COCO和人脸识别数据集上进行搜索，实现端到端性能评估，无需剪枝。
• 通过迁移学习，利用一个任务（如人脸识别）的知识初始化UAC和AKP，以加速另一任务（如ImageNet）的搜索，验证了跨任务知识迁移的有效性。
• 将UAC、AKP和AEB三个模块整合，实现NAS整体流程10倍的速度提升，同时保持或提升最终模型的准确率。

实验结果

研究问题

• RQ1 不确定性感知的价值函数能否提升基于强化学习的神经架构搜索中的样本效率？
• RQ2 能否重用先前搜索过程中获得的模块级参数知识，以加速新架构的收敛？
• RQ3 结合校正因子的离策略经验回放能否减少NAS中的训练周期数？
• RQ4 从一个搜索过程（如人脸识别）中获得的知识，能在多大程度上用于加速另一任务（如ImageNet）的搜索？
• RQ5 整合多种知识迁移机制，能否在不造成性能下降的前提下，实现基于强化学习的NAS 10倍速度提升？

主要发现

• FNAS将搜索时间从基线MNAS的20,000 GPU小时减少至2,000 GPU小时，在Mobile NAS搜索空间中实现10倍加速。
• 该方法在ImageNet上实现了66.25%的Top-1准确率，计算量为85M MACs，优于MBv3，并在相同计算约束下超越EfficientNet-B0。
• 在细粒度人脸识别任务中，FNAS在计算量减少24%的前提下，将验证准确率提升4个百分点，优于MBv2。
• 使用在人脸识别数据上预训练的UAC，FNAS仅用25%的样本即在ImageNet上达到73.91%的Top-1准确率，较基线提升0.67%。
• 使用从人脸识别搜索中初始化的AKP，FNAS在仅使用一半样本的情况下，ImageNet上的Top-1准确率达到75.22%，较基线提升0.6%。
• 消融实验证实，结合UAC、AKP和AEB可实现最高10倍加速和最佳性能，且各模块独立贡献于效率与准确率的提升。