# NAX: Co-Designing Neural Network and Hardware Architecture for Memristive Xbar based Computing Systems

arXiv: https://arxiv.org/abs/2106.12125 · HTML: https://mlsanity.com/paper/2106.12125/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 23 references that are in the map.
- First posted on arXiv 23 Jun 2021.

## Builds on (most cited first, 13 of 13)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [MnasNet: Platform-Aware Neural Architecture Search for Mobile](https://mlsanity.com/paper/1807.11626/) (2018)
- [ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware](https://mlsanity.com/paper/1812.00332/) (2018)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [PUMA: A Programmable Ultra-efficient Memristor-based Accelerator for Machine Learning Inference](https://mlsanity.com/paper/1901.10351/) (2019)
- [Co-Exploration of Neural Architectures and Heterogeneous ASIC Accelerator Designs Targeting Multiple Tasks](https://mlsanity.com/paper/2002.04116/) (2020)
- [Progressive DARTS: Bridging the Optimization Gap for NAS in the Wild](https://mlsanity.com/paper/1912.10952/) (2019)
- [Device-Circuit-Architecture Co-Exploration for Computing-in-Memory Neural Accelerators](https://mlsanity.com/paper/1911.00139/) (2019)
- [Best of Both Worlds: AutoML Codesign of a CNN and its Hardware Accelerator](https://mlsanity.com/paper/2002.05022/) (2020)
- [DANCE: Differentiable Accelerator/Network Co-Exploration](https://mlsanity.com/paper/2009.06237/) (2020)
- [Rethinking Co-design of Neural Architectures and Hardware Accelerators](https://mlsanity.com/paper/2102.08619/) (2021)
- [PABO: Pseudo Agent-Based Multi-Objective Bayesian Hyperparameter Optimization for Efficient Neural Accelerator Design](https://mlsanity.com/paper/1906.08167/) (2019)

## Built on by (most cited first, 3 of 3)

- [Examining the Role and Limits of Batchnorm Optimization to Mitigate Diverse Hardware-noise in In-memory Computing](https://mlsanity.com/paper/2305.18416/) (2023)
- [On the Noise Stability and Robustness of Adversarially Trained Networks on NVM Crossbars](https://mlsanity.com/paper/2109.09060/) (2021)
- [HERCULES: Hardware-Efficient, Robust, Continual Learning Neural Architecture Search](https://mlsanity.com/paper/2605.04103/) (2026)

## Abstract (arXiv)

In-Memory Computing (IMC) hardware using Memristive Crossbar Arrays (MCAs) are gaining popularity to accelerate Deep Neural Networks (DNNs) since it alleviates the "memory wall" problem associated with von-Neumann architecture. The hardware efficiency (energy, latency and area) as well as application accuracy (considering device and circuit non-idealities) of DNNs mapped to such hardware are co-dependent on network parameters, such as kernel size, depth etc. and hardware architecture parameters such as crossbar size. However, co-optimization of both network and hardware parameters presents a challenging search space comprising of different kernel sizes mapped to varying crossbar sizes. To that effect, we propose NAX -- an efficient neural architecture search engine that co-designs neural network and IMC based hardware architecture. NAX explores the aforementioned search space to determine kernel and corresponding crossbar sizes for each DNN layer to achieve optimal tradeoffs between hardware efficiency and application accuracy. Our results from NAX show that the networks have heterogeneous crossbar sizes across different network layers, and achieves optimal hardware efficiency and accuracy considering the non-idealities in crossbars. On CIFAR-10 and Tiny ImageNet, our models achieve 0.8%, 0.2% higher accuracy, and 17%, 4% lower EDAP (energy-delay-area product) compared to a baseline ResNet-20 and ResNet-18 models, respectively.