CoNAX: Towards Comprehensive Co-Design Neural Architecture Search Using HW Abstractions

HW-aware neural architecture search (HW-NAS) aims to yield high-accuracy neural network (NN) architectures by automatically exploring multiple architectural parameters of potential network candidates. In most HW-NAS approaches, the HW parameter search space is limited. Hence, HW awareness is tied to only a few degrees of design freedom, leading to the following sub-optimalities - First, it restricts exploration of HW parameters, which can potentially lead to better network candidates; Second, HW-NAS is still entirely a software-centric process where HW-awareness is taken care by an HW function exposed to the NAS process and is oblivious to the actual deployment. To tackle the above challenges, this paper proposes a Co-Design Neural Architecture Search (Co-NAS) approach that simultaneously explores hardware and neural architecture variations, thus allowing for full system optimization. By connecting the mutual impact of variable neural networks and HW parameters on the network's prediction accuracy and on-device efficiency in a shared optimization loop, Co-NAS finds designs of optimum performance and enables HW/SW Co-Design. This work aims to enable more diverse HW search spaces (higher degrees of design freedom) for ML accelerators (such as using a virtual prototype) and efficient exploration by integrating abstract ML accelerator and NN architecture modeling into a comprehensive Co-NAS environment. In our experiments, we explore hardware variations of a baseline accelerator architecture to demonstrate how our work can help find designs with better hardware latency and comparable network accuracy. Designs yielded by our framework provide a speed-up of 1.4x compared to the baseline on a restricted SW search space at the same HW resources.