Neural Architecture Search as Program Transformation Exploration

TL;DR

This paper presents a unified framework that models neural architecture search operations as program transformations, enabling the discovery of optimized DNNs with significantly reduced inference time and faster NAS search.

Contribution

It introduces a novel approach to unify NAS operations with program transformations, allowing exploration of new tensor convolutions and improving optimization efficiency.

Findings

Achieved over 3x inference time reduction in most cases.

Significantly reduced NAS search time.

Enabled exploration of new tensor convolution operations.

Abstract

Improving the performance of deep neural networks (DNNs) is important to both the compiler and neural architecture search (NAS) communities. Compilers apply program transformations in order to exploit hardware parallelism and memory hierarchy. However, legality concerns mean they fail to exploit the natural robustness of neural networks. In contrast, NAS techniques mutate networks by operations such as the grouping or bottlenecking of convolutions, exploiting the resilience of DNNs. In this work, we express such neural architecture operations as program transformations whose legality depends on a notion of representational capacity. This allows them to be combined with existing transformations into a unified optimization framework. This unification allows us to express existing NAS operations as combinations of simpler transformations. Crucially, it allows us to generate and explore new tensor convolutions. We prototyped the combined framework in TVM and were able to find optimizations across different DNNs, that significantly reduce inference time - over 3$\times$ in the majority of cases. Furthermore, our scheme dramatically reduces NAS search time. Code is available at~\href{https://github.com/jack-willturner/nas-as-program-transformation-exploration}{this https url}.