Bringing uncertainty quantification to the extreme-edge with memristor-based Bayesian neural networks

Abstract Safety-critical sensory processing applications, like medical diagnosis, require making accurate decisions based on a small amount of noisy input data. For these applications, using Bayesian neural networks, able to quantify the uncertainty of the predictions, is a superior approach to using conventional artificial neural networks. However, because of the probabilistic nature of Bayesian neural networks, they can be computationally intensive to use for inference stage and thus not well suited for extreme-edge applications. An emerging idea to solve this problem is to use the intrinsic probabilistic nature of memristors to efficiently implement Bayesian neural network inference: the variability in the resistance of memristors would represent the probability distribution of weights in Bayesian neural networks. However, when using memristors, statistical effects follow the laws of device physics, whereas in Bayesian neural networks, those effects can take on arbitrary shapes. In this work, we overcome this difficulty by adopting a dedicated synapse architecture based on two memristors, and by training Bayesian neural networks with a dedicated variational inference technique that includes a “technological loss” to take into account specificities of memristor physics. This technique allowed us to program a two-layer Bayesian neural network on 75 physical crossbar arrays of 1,024 memristors, incorporating CMOS periphery circuitry to do in-memory computing, to classify arrhythmia in electrocardiograms. Our experimental neural network classified heartbeats with high accuracy, and estimated the certainty of its predictions. In the case of uncertain predictions, it differentiated between ambivalent heartbeats (aleatoric uncertainty), and heartbeats with never-seen patterns (epistemic uncertainty). We show that our technique can also be used with phase change memories, by employing a different “technological loss” term. The great advantage of this approach is its low energy consumption: we estimate an 800 times improvement in energy efficiency compared to a GPU performing the same task.