[论文解读] CATE: Computation-aware Neural Architecture Encoding with Transformers

Yan Shen, Kaiqiang Song+2 more|arXiv (Cornell University)|Feb 14, 2021

Neural Networks and Applications参考文献 87被引用 8

一句话总结

CATE 提出了一种基于 Transformer 的计算感知神经架构编码方法，采用成对预训练方案，以学习密集且上下文相关的神经架构表征。通过在计算相似的架构对上进行掩码语言建模，并结合因果注意力机制，CATE 在训练分布内和分布外的神经架构搜索（NAS）基准测试中均实现了更优的性能与泛化能力，显著优于邻接矩阵等固定编码基线方法。

ABSTRACT

Recent works (White et al., 2020a; Yan et al., 2020) demonstrate the importance of architecture encodings in Neural Architecture Search (NAS). These encodings encode either structure or computation information of the neural architectures. Compared to structure-aware encodings, computation-aware encodings map architectures with similar accuracies to the same region, which improves the downstream architecture search performance (Zhang et al., 2019; White et al., 2020a). In this work, we introduce a Computation-Aware Transformer-based Encoding method called CATE. Different from existing computation-aware encodings based on fixed transformation (e.g. path encoding), CATE employs a pairwise pre-training scheme to learn computation-aware encodings using Transformers with cross-attention. Such learned encodings contain dense and contextualized computation information of neural architectures. We compare CATE with eleven encodings under three major encoding-dependent NAS subroutines in both small and large search spaces. Our experiments show that CATE is beneficial to the downstream search, especially in the large search space. Moreover, the outside search space experiment demonstrates its superior generalization ability beyond the search space on which it was trained. Our code is available at: https://github.com/MSU-MLSys-Lab/CATE.

研究动机与目标

• 为解决基于固定变换的计算感知编码在神经架构搜索（NAS）中因泛化能力差和预训练效果不佳而带来的局限性。
• 通过在 Transformer 中利用自注意力与交叉注意力机制，提升计算感知编码的表征能力。
• 通过成对学习方案实现架构编码器的有效预训练，以捕捉长距离依赖关系与计算局部性。
• 证明所学习的上下文编码在分布外搜索空间中优于固定编码。
• 提供一种可扩展、通用的架构编码方法，适用于多种 NAS 子程序与搜索空间。

提出的方法

• CATE 采用成对预训练方案，对计算相似的架构对进行操作掩码，模型被训练以预测被掩码的操作符。
• 该方法使用共享权重的 Transformer 编码器分别编码每对架构，随后通过一个交叉注意力 Transformer 融合两者的联合信息。
• 通过基于邻接矩阵的因果掩码确保模型尊重神经架构的有向无环计算流，替代 BERT 中使用的标准双向注意力机制。
• 模型通过掩码语言建模（MLM）目标进行预训练，以学习深层、上下文相关的计算表征。
• 通过引入 Floyd-Warshall 算法结果增强邻接矩阵，以捕捉间接连接，从而增强长距离依赖建模。
• 最终的架构编码由交叉注意力输出生成，并用于下游 NAS 子程序，如基于预测器的搜索、强化学习与多臂赌博机搜索。

实验结果

研究问题

• RQ1 与固定变换方法相比，基于 Transformer 的成对预训练方案是否能学习到更有效的计算感知架构编码？
• RQ2 使用交叉注意力与因果掩码是否能提升对架构计算流建模与泛化能力？
• RQ3 在分布外搜索空间泛化方面，CATE 相较于邻接矩阵等固定编码方法表现如何？
• RQ4 架构对采样策略、超参数与掩码类型对下游 NAS 性能有何影响？
• RQ5 CATE 是否能在更大、未见过的搜索空间中有效泛化，同时保持高精度与高效性？

主要发现

• 在 NAS-Bench-101 上，CATE 在 150 次查询预算下，于训练分布外搜索空间评估中显著优于邻接矩阵编码，验证误差仅为 5.94%。
• 间接依赖掩码优于直接掩码，在 NAS-Bench-101 上将测试误差降低 0.09%，在 NAS-Bench-301 上降低 0.05%，表明长距离依赖建模能有效提升编码质量。
• 更多的交叉注意力块（$L_c$）与更宽的前馈网络层（$d_{ff}$）可提升 NAS 性能，其中 $L_c=24$ 与 $d_{ff}=64$ 在 NAS-Bench-101 上实现最低误差 5.95%。
• 较小的邻域大小 $K$ 与更强的计算局部性（$\delta$）带来更优性能，$\delta = 1 \times 10^6$ 与 $K=2$ 表现最佳。
• CATE 在训练分布外的搜索空间中表现出优异的泛化能力，在小型与大型 NAS 基准测试中均展现出卓越的鲁棒性与可扩展性。
• 该方法在三种依赖编码的 NAS 子程序中均达到当前最优性能：基于预测器的搜索、强化学习与多臂赌博机搜索。