An Algorithm-Hardware Co-Design for Bayesian Neural Network Utilizing SOT-MRAM’s Inherent Stochasticity

Abstract

Probabilistic machine learning plays a central role in the domains such as decision-making and autonomous control benefitting from its ability of representing and manipulating uncertainty about models and predictions. Until now, there are few hardware considerations to address the intensive computation and true random number generation for Bayesian neural network (BayesNN), whose weights are represented by probability distributions. In this article, we propose to apply the local reparameterization trick to alleviate the burden of random number generators (RNGs), which could be implemented by utilizing the inherent random noise of spin-orbit torque magnetic random access memory (SOT-MRAM). Sampling strategies are discussed to significantly reduce the number of operations and parameters of BayesNN. A device-circuit-system benchmark framework is then developed to evaluate the effects of device nonidealities such as the bias and variation of switching probability. The evaluation on the CIFAR-10 dataset suggests that BayesNN could achieve comparable accuracy as conventional deep neural network (DNN) with acceptable hardware overhead but provide much better uncertainty calibration with respect to out-of-distribution (OOD) inputs (rotated images as the example).