XBTorch: A Unified Framework for Modeling and Co-Design of Crossbar-Based Deep Learning Accelerators

TL;DR

XBTorch addresses the need for a unified framework to study memristive neural networks across device models, hardware-aware training, and inference-time fault tolerance. It integrates with PyTorch, supports analytical and tabular device models (e.g., FeFET, ReRAM), and provides gradient decomposition, loss landscape analysis, and LLM evaluation with analog noise. The results demonstrate interplay between device non-idealities and training dynamics, showing how hardware-aware strategies can improve robustness and enable realistic benchmarking. This framework enables systematic algorithm–hardware co-design for crossbar-based accelerators and could accelerate progress in neuromorphic computing.

Abstract

Emerging memory technologies have gained significant attention as a promising pathway to overcome the limitations of conventional computing architectures in deep learning applications. By enabling computation directly within memory, these technologies - built on nanoscale devices with tunable and nonvolatile conductance - offer the potential to drastically reduce energy consumption and latency compared to traditional von Neumann systems. This paper introduces XBTorch (short for CrossBarTorch), a novel simulation framework that integrates seamlessly with PyTorch and provides specialized tools for accurately and efficiently modeling crossbar-based systems based on emerging memory technologies. Through detailed comparisons and case studies involving hardware-aware training and inference, we demonstrate how XBTorch offers a unified interface for key research areas such as device-level modeling, cross-layer co-design, and inference-time fault tolerance. While exemplar studies utilize ferroelectric field-effect transistor (FeFET) models, the framework remains technology-agnostic - supporting other emerging memories such as resistive RAM (ReRAM), as well as enabling user-defined custom device models. The code is publicly available at: https://github.com/ADAM-Lab-GW/xbtorch

Figures (6)

• Figure 1: XBTorch features and an example script comparison. (a) XBTorch features (in bold), corresponding library modules (in monospace), and lists of important key capabilities. (b) A code diff for converting a PyTorch script for training a simple neural network to XBTorch. XBTorch requires minimal changes for initialization and model patching.
• Figure 2: Comparison of a multi-layer perceptron network trained on MNIST using a low-precision hardware-aware training in XBTorch against a full precision software baseline trained using vanilla PyTorch. (a) Network test accuracies averaged over 10 independent runs. Loss landscapes for the (b) software baseline network and the (c) hardware-aware network. Loss landscape axes correspond to principal components over the parameter evolution matrix $\boldsymbol{M}$, with explained variances presented in parentheses. The optimization trajectories are included as red lines concluding at the circle marker. The hardware-aware network is trained using a 2-8-8-8 WAGE configuration which produces ternary weight networks, and a tabular FeFET device model with $V_{gs} = 0.9$$V$ and $1 \%$ variability (see Fig. \ref{['fig:tabular_models']}).
• Figure 3: Demonstration of gradient decomposition methods using XBTorch. (a) Representative decomposition of a hypothetical gradient matrix into a low-rank $k$ representation according to Eq. \ref{['eq:gradient_decomp']}. The decomposed form has less memory overhead compared to the full rank gradient matrix if $k$ is small. Network test accuracies averaged over 10 independent runs at different ranks for (b) SBPCA and (c) NMF algorithms.
• Figure 4: Hardware-aware inference using XBTorch. (a) Simulated $2500 \times 2500$ crossbar with realistic noise shows conductance maps after writing a trained solution (from Fig. \ref{['fig:hwa_training']}). Each weight matrix $\theta$ is differentially encoded into $G_{pos}$ and $G_{neg}$, mapped to random, non-overlapping locations. Red/orange boxes mark the first/second network layers; dark/light shades represent $G_{pos}$/$G_{neg}$. Insets show a $25 \times 25$ zoomed region. Unused devices are disabled during readout. (b) Actual conductance distributions for $G_{pos}$ from the first layer. (c) Inference-time test accuracy as a function of ADC/DAC bit precision. The dotted line corresponds to the ideal baseline with no noise sources. Bar heights are averages across 10 independent iterations. Error bars correspond to 3 standard deviations.
• Figure 5: Demonstration of layer ensemble averaging in XBTorch. (a) The overall flow of mapping an individual network layer to an emulated crossbar using XBTorch. System parameters, such as total crossbar size and a percentage of stuck devices within that crossbar can be specified to gauge network performance under realistic noise sources. (b) Network inference accuracy as a function of mapping redundancy and percentage of stuck devices. Under layer ensemble averaging, outputs are averaged layer-wise from the redundant mappings during inference. Bar heights correspond to averages across 10 independent iterations, and error bars correspond to a single standard deviation. (c) Conductance map after encoding and mapping the network solution for the case where the redundancy level is $6$ and $20 \%$ devices are stuck.