17 天前

CATE:基于Transformer的计算感知神经架构编码

Shen Yan, Kaiqiang Song, Fei Liu, Mi Zhang
CATE:基于Transformer的计算感知神经架构编码
摘要

近期研究(White 等,2020a;Yan 等,2020)揭示了神经架构搜索(Neural Architecture Search, NAS)中架构编码的重要性。这些编码用于表征神经架构的结构信息或计算信息。相较于结构感知型编码,计算感知型编码能够将具有相似性能的架构映射至相近区域,从而提升下游架构搜索的效率(Zhang 等,2019;White 等,2020a)。在本工作中,我们提出一种基于 Transformer 的计算感知编码方法——CATE(Computation-Aware Transformer-based Encoding)。与现有基于固定变换(如路径编码)的计算感知编码不同,CATE 采用成对预训练策略,利用具备交叉注意力机制的 Transformer 模型学习计算感知编码,所获得的编码能够蕴含神经架构中密集且上下文相关的计算信息。我们在小型和大型搜索空间下,针对三种主要依赖编码的 NAS 子流程,将 CATE 与十一种现有编码方法进行了对比实验。实验结果表明,CATE 显著提升了下游搜索性能,尤其在大规模搜索空间中优势更为突出。此外,跨搜索空间的实验进一步验证了 CATE 在训练以外搜索空间中仍具备优异的泛化能力。相关代码已开源,地址为:https://github.com/MSU-MLSys-Lab/CATE。