SpinBayes: Algorithm-Hardware Co-Design for Uncertainty Estimation Using Bayesian In-Memory Approximation on Spintronic-Based Architectures

Recent development in neural networks (NNs) has led to their widespread use in critical and automated decision-making systems, where uncertainty estimation is essential for trustworthiness. Although conventional NNs can solve many problems accurately, they do not capture the uncertainty of the data or the model during optimization. In contrast, Bayesian neural networks (BNNs), which learn probabilistic distributions for their parameters, offer a sound theoretical framework for estimating uncertainty. However, traditional hardware implementations of BNNs are expensive in terms of computational and memory resources, as they (i) are realized with inefficient von Neumann architectures, (ii) use a significantly large number of random number generators (RNGs) to implement the distributions of BNNs, and (iii) have a substantially greater number of parameters than conventional NNs. Computing-in-memory (CiM) architectures with emerging resistive non-volatile memories (NVMs) are promising candidates for accelerating classical NNs. In particular, spintronic technology, which is distinguished by its low latency and high endurance, aligns very well with these requirements. In the specific context of Bayesian neural networks (BNNs), spintronics technologies are very valuable, thanks to their inherent potential to act as stochastic or as deterministic devices. Consequently, BNNs mapped on spintronic-based CiM architectures could be a highly efficient implementation strategy. However, the direct implementation on CiM hardware of the learned probabilistic distributions of BNN may not be feasible and can incur high overhead. In this work, we propose a new Bayesian neural network topology, named SpinBayes, that is able to perform efficient sampling during the Bayesian inference process. Moreover, a Bayesian approximation method, called in-memory approximation, is proposed that approximates the original probabilistic distributions of BNN with a distribution that can be efficiently mapped to spintronic-based CiM architectures. Compared to state-of-the-art methods, the memory overhead is reduced by 8× and the energy consumption by 80×. Our method has been evaluated on several classification and semantic segmentation tasks and can detect up to 100% of various types of out-of-distribution data, highlighting the robustness of our approach, without any performance sacrifice.