NASPTE ASPLOS 2023

Neural Architecture Search as Program Transformation Exploration

Posted by Treaseven on December 6, 2024

Background

编译器优化专注于重组底层张量计算以利用硬件特性,但受限于必须保持计算正确性
NAS则利用神经网络的鲁棒性,通过变换网络架构(如分组卷积、瓶颈层)来优化性能

Our Approach

将NAS中的网络架构操作重新表达为程序转换,使其可以与现有编译器转换结合

Overview

Code Transformation:不影响最终计算的值,只改变内存访问模式,如interchange Model Transformation:改变了权重张量的大小和外层循环的范围,如Bottlenecking

基本原理

  • 从整体网络骨架开始,尝试设计可以插入不同位置的单元
  • 每个单元被描述为一个DAG,节点是中间特征图,边表示可能的操作
  • 目标是找到最佳的DAG结构,插入骨架并在给定数据集上训练

传统的NAS从预先定义的操作来进行组合,本文提出从程序转换来生成新的操作,以来发现预定义列表中没有的新操作类型

Unified space

Extending the Polyhedral Model

  • Bottlenecking
  • Grouping
  • Depthwise

Fisher Potential as a Legality Check

  • 传统程序转换必须保证语义等价,但神经网络允许一定的变形而不影响功能
  • Fisher Potential提供了一个无需训练就能评估转换合法性的方法

Evaluation

Thinking

(1) 本文主要针对是单个卷积层的优化,转换操作都是在层级别进行,没有考虑层与层之间的相互影响,缺乏对整个网络架构的全局视角
(2) 可能带来的问题局部最优不一定是全局最优,可能错过跨层优化机会,如skip connection、Feature Reuse
(3) 可能改进的方向构建计算图依赖关系,分析数据流模式,识别关键瓶颈;多层联合优化,考虑层间依赖,优化数据传输,平衡计算资源
(4) 本文针对cnn架构,对于transformer架构难以扩展

Reference

Neural Architecture Search as Program Transformation Exploration