Uncertainty quantification via a memristor Bayesian deep neural network for risk-sensitive reinforcement learning

Abstract

Many advanced artificial intelligence tasks, such as policy optimization, decision making and autonomous navigation, demand high-bandwidth data transfer and probabilistic computing, posing great challenges for conventional computing hardware. Since digital computers based on the von Neumann architecture are good at precise and deterministic computing, their computing efficiency is limited by the high cost of both data transfer between memory and computing units and massive random number generation. Here we develop a stochastic computation-in-memory computing system that can efficiently perform both in situ random number generation and computation based on the nanoscale physical behaviour of memristors. This system is constructed based on a hardware-implemented multiple-memristor-array system. To demonstrate its functionality and efficiency, we implement a typical risk-sensitive reinforcement learning task, namely the storm coast task, with a four-layer Bayesian deep neural network. The computing system efficiently decomposes aleatoric and epistemic uncertainties by exploiting the inherent stochasticity of memristor. Compared with the conventional digital computer, our memristor-based system achieves a 10 times higher speed and 150 times higher energy efficiency in uncertainty decomposition. This stochastic computation-in-memory computing system paves the way for high-speed and energy-efficient implementation of various probabilistic artificial intelligence algorithms.