NAX: Neural Architecture & Crossbar Explorer

• NAX is a co-design framework that jointly optimizes neural architectures and their mapping onto memristive crossbars to achieve Pareto-optimal trade-offs in accuracy, energy, area, latency, and robustness.
• It employs ILP-based techniques for spiking neural networks and differentiable NAS for deep neural networks, leveraging heterogeneous crossbar geometries and peripheral co-search for enhanced efficiency.
• Experimental benchmarks demonstrate significant reductions in area and energy-delay-area product, along with improved adversarial robustness across diverse neural network and hardware scenarios.

The Neural Architecture and crossbar eXplorer (NAX) is a class of algorithm-hardware co-design frameworks that systematically search or optimize both neural network topologies and the mapping of those networks onto physical memristive crossbar arrays, with the primary objective of achieving Pareto-optimal trade-offs among accuracy, energy, area, latency, and robustness under non-ideal hardware constraints. NAX methodologies encompass Integer Linear Programming (ILP)-based compilers for spiking neural networks (SNNs) as well as differentiable Neural Architecture Search (NAS) engines that simultaneously explore deep neural network (DNN) architectures and hardware design spaces, supporting heterogeneity across both network layers and hardware modules. The evolution of NAX approaches has been driven by the emergence of novel neuromorphic hardware and increasing recognition of the need to match neural topology, crossbar geometry, and peripheral circuit design for maximal efficiency, particularly as both DNNs and SNNs become sparser and more structurally diverse.

1. System-Level Problem Definition and Motivation

NAX frameworks address the coupled optimization of neural architecture and crossbar-based in-memory computing hardware. This co-design is motivated by the interplay between the topological properties of neural networks—such as channel size, kernel shape, and sparse connectivity—and the rigid, resource-constrained dimensions of memristor crossbars. For SNNs, this involves partitioning highly sparse, directed graphs onto crossbars with strict row and column capacities; for DNNs, it encompasses mapping convolutions and activations onto crossbars whose non-idealities (e.g., device mismatch, line resistance) can degrade inference accuracy when neglected.

The central insight underlying NAX approaches is that neural network sparsity and variable layer parameters interact significantly with hardware resource utilization: smaller or heterogeneously dimensioned crossbars can be better exploited if the structural statistics of the input network are known and utilized. As a result, NAX frameworks attempt to map or search for neural architectures in a manner aware of both computational topology and hardware constraints, seeking to minimize composite metrics such as energy-delay-area product (EDAP), reduce communication overhead, and increase adversarial and hardware robustness (Pohl et al., 3 Mar 2025, Bhattacharjee et al., 2023, Negi et al., 2021, Moitra et al., 2023).

2. Mathematical Foundations and Optimization Formulations

The NAX approach spans both strictly mathematical ILP mapping for SNNs and differentiable NAS for DNNs.

For SNN mapping (Pohl et al.), the formalism is an ILP with binary decision variables:

• $x_{i,j}$: neuron placement, $y_j$: crossbar activation, $S_{k,j}$: input axon activation, $b_{k,j}$: local route indicator, $s_{i,j}$: inter-crossbar route indicator.
• Constraints enforce unique neuron placement (∑$x_{i,j}=1$), crossbar output/input capacity (∑$x_{i,j}\le y_j N_j$, ∑$S_{k,j}\le y_j A_j$), and inter-/intra-crossbar routing definitions.
• Objectives are staged:
  1. Minimize area (total memristor count): $\min_{z} \sum_j C_j y_j$.
  2. Minimize inter-crossbar routing: $\min_{z} \sum_{i,j} (s_{i,j} - b_{i,j})$ at fixed area.
  3. Minimize spike traffic using profile-guided weighting: $y_j$0 (Pohl et al., 3 Mar 2025).

For DNN/DNN-on-crossbar co-search (NAX, XploreNAS), the approach is based on differentiable NAS with architecture logits $y_j$1 indexing candidate kernel/crossbar options per layer, and hardware cost regularization folded into the NAS objective. The training alternates between updating network weights and architecture parameters, using surrogate models (e.g., GENIEx) to account for hardware non-idealities (Negi et al., 2021). For adversarially robust design, XploreNAS incorporates explicit crossbar-noise regularization into both supernet training and fine-tuning (Bhattacharjee et al., 2023).

3. Integration of Heterogeneity in Crossbar and Architecture Design

A defining feature of advanced NAX formulations is the encoding and exploitation of heterogeneity—both in crossbar geometry across the chip and in neural architecture across network layers. For SNNs, crossbar $y_j$2 can have arbitrary output and input capacities $y_j$3 (e.g., $y_j$4, $y_j$5, $y_j$6, etc.), with the ILP mapping methodically packing clusters to optimally utilize this heterogeneity. This approach leverages “axon sharing,” so multiple neurons in a cluster that share the same pre-synaptic source consume only one column input, maximizing spatial efficiency (Pohl et al., 3 Mar 2025).

For DNNs, NAX and XploreNAS search spaces include operation tuples of $y_j$7, where $y_j$8 and $y_j$9 are both layer-variable. Differentiable or path-binarized selection ensures each layer independently selects its optimal (kernel, crossbar) configuration, resulting in final architectures with layerwise heterogeneous mappings (Negi et al., 2021, Bhattacharjee et al., 2023).

This generalization is often extended to peripheral circuit co-design: e.g., XPert co-searches ADC type, precision, column sharing, and input bitwidth per layer, achieving larger hardware efficiency gains than crossbar size adaptation alone (Moitra et al., 2023).

NAX Variant	Crossbar Heterogeneity	Architecture Heterogeneity	Peripheral Circuit Co-Search
ILP SNN Mapper (Pohl et al., 3 Mar 2025)	Yes	Topology-driven	No
NAS DNN Co-Design (Negi et al., 2021)	Yes	Layerwise	No
XploreNAS (Bhattacharjee et al., 2023)	Yes	Layerwise + adversarially robust	No
XPert (Moitra et al., 2023)	Yes	Layerwise	Yes (ADC, precision, etc.)

4. Experimental Results and Quantitative Benchmarks

NAX approaches consistently demonstrate statistically significant improvements in area, energy, delay, and adversarial/hardware robustness relative to prior homogeneous or uncoupled baselines.

SNN ILP Mapping (Pohl et al., 3 Mar 2025)

• Area reduction: Modeling axon sharing within homogeneous $S_{k,j}$0 crossbar architectures achieves 16.7–27.6% reduction over previous ILPs.

• Heterogeneous architectures (mix from $S_{k,j}$1): 66.9–72.7% additional area decrease.
• Routing reduction: Static Network Utilization (SNU) minimization yields 11.9–26.4% reduction in inter-crossbar routes at fixed area.
• Traffic reduction: Profile-based optimization produces inter-crossbar spike reductions of 0.5–14.8% with 1–6 orders of magnitude lower solve time.
• All experiments use Google OR-Tools CP-SAT. 800-neuron networks converge in minutes for area; profile-guided stages run in seconds.

DNN Co-Design Benchmarks

• On CIFAR-10, layerwise heterogeneous NAX (NAX-ni2) achieves 0.8% higher non-ideal accuracy and 17% lower EDAP compared to homogeneous ResNet-20@64×64 baseline (Negi et al., 2021).
• On Tiny ImageNet, NAX-ni2 gives 0.2% higher accuracy and ~4% lower EDAP relative to the best homogeneous baseline.
• XploreNAS models: Subnets show 8–16% higher adversarial robustness and 1.5–1.6× lower EDAP than ResNet-18 on CIFAR-10/CIFAR-100/SVHN under crossbar noise ($S_{k,j}$2); MultiXbar models retain 48–55% adversarial accuracy across $S_{k,j}$3 from 10–50% without retraining (Bhattacharjee et al., 2023).
• XPert: Full five-dimensional co-search achieves up to $S_{k,j}$4 EDAP reduction on CIFAR-10 at matched accuracy (e.g., 2.90 mJ·ms·mm$S_{k,j}$5 for XPertNet$S_{k,j}$6 vs. 29.7 mJ·ms·mm$S_{k,j}$7 for VGG16 baseline) (Moitra et al., 2023).

5. Algorithmic and Architectural Insights

Empirical and ablation analyses across the NAX literature reveal several design patterns:

• SNNs benefitting from heterogeneous crossbars utilize clusters matched to their structural fan-in/out, preventing needless area overhead inherent in square or uniform crossbar arrangements (Pohl et al., 3 Mar 2025).
• For DNNs, early layers with modest channel depth favor smaller or medium-sized crossbars to avoid peripheral under-utilization, while deep layers with high channel count leverage larger crossbars (Negi et al., 2021).
• Joint energy+latency regularization is required; focusing on a single hardware metric degrades others substantially (Negi et al., 2021).
• Co-searching peripheral and architectural parameters (XPert) is necessary for attaining Pareto-optimal hardware efficiency; e.g., tuning ADC precision higher in shallow layers, then scaling it down in deeper layers proportionally to accuracy loss (Moitra et al., 2023).

6. Extensions: Robustness and Multi-Objective Co-Optimization

A salient trend within NAX frameworks is explicit robustness to both hardware non-idealities and adversarial attacks.

• XploreNAS introduces a two-phase training regime enforcing crossbar-aware loss regularization alongside adversarial training, sampling Subnets by thresholding architecture logits, followed by fine-tuning under Projected Gradient Descent (PGD) with crossbar noise (Bhattacharjee et al., 2023).
• Profile-guided ILP for SNNs reduces inter-crossbar activity on the expected, not worst-case, spike traces, rendering solver times orders of magnitude lower at negligible accuracy cost (Pohl et al., 3 Mar 2025).
• For robust NAS, incorporating crossbar simulation (e.g., via GENIEx) in the forward path during architecture search directly penalizes architectures that are sensitive to device mismatch or variation (Negi et al., 2021).

Multi-objective formulation is central: area, delay, energy, EDAP, and robustness targets can be encoded as (weighted) loss terms or as sequential hard constraints in the optimization procedure. This design permits hardware-aware neural compilers/users to select trade-offs appropriate for target deployment scenarios.

7. Significance, Applicability, and Future Directions

NAX constitutes a generalizable methodology for co-designing and mapping neural networks onto heterogeneously structured, non-ideal crossbar-based hardware, supporting both SNN and DNN paradigms. Its contributions span staged ILP mapping, differentiable NAS co-optimization, and, more recently, full-stack circuit/peripheral co-search. NAX closes the loop between neural network training, structured sparsity exploitation, hardware mapping, and peripheral adaptation.

A plausible implication is the increasing relevance of such co-search frameworks as device-level and circuit non-idealities become more prominent at advanced technology nodes, and as neural architectures continue to diversify in scale and sparsity. Lessons from XPert further suggest that beyond crossbar size, peripheral circuit parameters should be included in the design search for systematic, scalable, and robust in-memory computing systems (Moitra et al., 2023).

Key results demonstrate that by integrating all relevant architectural and hardware parameters in the search space, NAX methodologies can yield up to an order-of-magnitude improvement in hardware efficiency (EDAP) at matched accuracy or robustness, substantially outperforming prior art limited to parameter-homogeneous designs or decoupled compiler flows. This positions NAX as a cornerstone strategy for the design and optimization of next-generation neuromorphic hardware and efficient AI accelerators (Pohl et al., 3 Mar 2025, Negi et al., 2021, Bhattacharjee et al., 2023, Moitra et al., 2023).