High-Efficient Memristor-Based Bayesian Convolutional Neural Networks for Out-of-Distribution Detection by Uncertainty Estimation

By using the nonvolatile and stochastic properties of memristor, memristors crossbar arrays can efficiently accelerate Bayesian neural networks (BNNs). However, Bayesian convolutional neural networks (BCNNs), one of the representative algorithms of Bayesian deep learning, have not yet been implemented using memristors. The unique attribute of local weight connection of convolutional kernels poses a challenge in constructing Gaussian weights. In this work, a highly energy-efficient implementation method of probabilistic convolutional kernels in memristor-based BCNN is developed, which takes advantage of both the write and read variations of devices. Meanwhile, a Bayesian-based hardware-aware ex-situ training method for memristor-based BCNNs is proposed for further deployment on large-scale BCNNs. Incorporating HfOx-based memristor’s parameters, memristor BCNNs are used to demonstrate classification task and out-of-distribution detection (OOD) by estimating uncertainty for the first time. The two memristor BCNNs achieved accuracies of 98.67% on MNIST and 87.81% on CIFAR10 and area under the ROC curve (AUC) of 0.96 and 0.83, indicating comparable classification and OOD detection performance to the digital floating-point implementation. After analyzing the impact of various effects at the device and array level, including the read variation, write variation, retention, conductance weight level, and ADC bit widths, the proposed method shows high robustness. Moreover, compared with CMOS-based GPUs, our memristor-based system achieved a 12-fold speed increase and a 335-fold energy efficiency improvement during the prediction cycle.