[Paper Review] NAX: Co-Designing Neural Network and Hardware Architecture for Memristive Xbar based Computing Systems

Shubham Negi, Indranil Chakraborty+2 more|arXiv (Cornell University)|Jun 23, 2021

Advanced Memory and Neural Computing24 references7 citations

TL;DR

NAX is a neural architecture search framework that co-designs deep neural networks and memristive crossbar-based in-memory computing (IMC) hardware by jointly optimizing kernel sizes and crossbar dimensions to balance accuracy and hardware efficiency. It achieves 17% and 4% lower EDAP, and 0.8% and 0.2% higher accuracy on CIFAR-10 and Tiny ImageNet compared to baseline ResNet models, respectively.

ABSTRACT

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.

Motivation & Objective

• Address the co-dependence between DNN model parameters (e.g., kernel size, depth) and hardware parameters (e.g., crossbar size) in memristive crossbar-based in-memory computing (IMC) systems.
• Overcome the trade-off between hardware efficiency (energy, latency, area) and model accuracy due to device and circuit non-idealities in memristive crossbars.
• Develop a NAS framework that co-optimizes neural network architecture and IMC hardware configuration to achieve optimal performance under non-ideal conditions.
• Demonstrate that heterogeneous crossbar sizes across network layers can improve both hardware efficiency and inference accuracy in analog IMC systems.

Proposed method

• Propose a co-design search space that includes variable kernel sizes (e.g., 3x3, 5x5) mapped to different crossbar sizes (e.g., 64x64, 128x128) for each DNN layer.
• Integrate non-idealities in memristive crossbars—such as parasitic resistance, non-linearity from access transistors, and I-V non-idealities—into the accuracy estimation during NAS.
• Use a differentiable architecture search (DARTS-style) approach with a super-network that includes mixed layers with multiple kernel and crossbar size options.
• Apply a two-stage search: i-search (ignoring non-idealities) and ni-search (including non-idealities) to evaluate accuracy under realistic hardware constraints.
• Employ GENIEx inference to estimate the impact of non-idealities during architecture parameter updates, improving accuracy-aware search.
• Train the architecture parameters using Adam optimizer with a step learning rate schedule, and validate the final compact network on test sets.

Experimental results

Research questions

• RQ1 How does crossbar size affect the hardware efficiency (EDAP) and accuracy of DNNs mapped to memristive crossbar-based IMC systems?
• RQ2 What is the optimal trade-off between hardware efficiency and model accuracy when co-optimizing DNN architecture and crossbar size under non-ideal device and circuit effects?
• RQ3 Can a NAS framework that jointly explores DNN and hardware parameters outperform conventional DNNs mapped to fixed hardware configurations in terms of EDAP and accuracy?
• RQ4 How do non-idealities in larger crossbars impact inference accuracy, and can they be mitigated through intelligent crossbar size selection per layer?
• RQ5 To what extent does considering non-idealities during the search process improve the final model’s accuracy on real IMC hardware?

Key findings

• NAX achieves 17% lower EDAP and 0.8% higher accuracy on CIFAR-10 compared to the baseline ResNet-20 model when mapping to memristive crossbar hardware.
• On Tiny ImageNet, NAX achieves 4% lower EDAP and 0.2% higher accuracy than the baseline ResNet-18 model under the same hardware constraints.
• The optimal crossbar size varies per layer: larger crossbars (e.g., 128x128) are used when output channels exceed 128 to avoid underutilization, while smaller crossbars (e.g., 64x64) are preferred for smaller operations.
• The ni-search variant, which accounts for non-idealities during search, yields models with significantly better accuracy than i-search, especially on Tiny ImageNet, where the accuracy gap is 18.57% for the best-performing model.
• The search cost for ni-search is higher (350 GPU-hours on Tiny ImageNet) than i-search (5 GPU-hours), but the resulting models show superior accuracy and hardware efficiency.
• The results demonstrate that heterogeneous crossbar sizing across layers is essential for balancing hardware efficiency and accuracy, with 128x128 crossbars being optimal for high-channel operations and 64x64 for smaller ones.