Deep Bayesian Active Learning Using In-Memory Computing Hardware

With the rapid development of artificial intelligence (AI) technologies, deep learning has made significant progress in complex tasks. However, the success of deep learning largely relies on massive amounts of labeled data, and the data labeling process is not only time-consuming and labor-intensive but also requires specialized domain knowledge, making it costly. This challenge is particularly pronounced in professional fields such as robotic skill learning, catalyst discovery, drug discovery, and protein production optimization, where obtaining labeled data is especially difficult and expensive. To address this issue, Deep Bayesian Active Learning (DBAL) has emerged. DBAL significantly improves labeling efficiency by actively selecting the most informative data for annotation, enabling high-quality learning with limited labeled data.

However, the implementation of DBAL faces a significant technical challenge: it requires handling a large number of random variables and high-bandwidth data transfer, which places high demands on traditional deterministic hardware. Traditional Complementary Metal-Oxide-Semiconductor (CMOS) hardware often results in substantial energy consumption and latency when processing these stochastic tasks. To tackle this problem, a research team proposed a Memristor-based Computation-in-Memory (CIM) framework, leveraging the intrinsic stochastic properties of memristors to enable efficient DBAL.

This research was conducted by Yudeng Lin, Bin Gao, Jianshi Tang, Qingtian Zhang, He Qian, and Huaqiang Wu from the School of Integrated Circuits and the Beijing National Research Center for Information Science and Technology at Tsinghua University. The findings were published in Nature Computational Science in January 2025.

Research Background and Problem

The success of deep learning relies heavily on large amounts of labeled data. However, in many real-world application scenarios, acquiring such data is not only costly but also time-consuming and requires substantial expertise. For example, in robotic skill learning, a robot needs to repeatedly attempt and adjust actions to learn how to perform specific tasks, and each attempt requires resetting the experimental scenario, significantly increasing time and resource overhead. DBAL addresses this by actively selecting the most informative data for annotation, significantly reducing the amount of labeled data required and thus improving learning efficiency while lowering costs.

However, the implementation of DBAL faces hardware challenges. DBAL involves a large number of random variables and high-bandwidth data transfer, and traditional CMOS hardware often results in substantial energy consumption and latency when processing these tasks. Additionally, DBAL requires the generation of a large number of Gaussian random numbers, which is a computationally intensive task, further increasing the hardware burden.

Research Methods and Innovations

To address these issues, the research team proposed a memristor-based CIM framework. Memristors are a new type of non-volatile memory device whose conductance can be modulated by external voltage and exhibit intrinsic stochastic properties. Leveraging these properties, memristors can efficiently generate random numbers and enable parallel computation, significantly reducing data transfer latency and energy consumption.

Specifically, the research team proposed a Memristor Stochastic Gradient Langevin Dynamics (MSGLD) method, utilizing the stochastic modulation properties of memristors to implement DBAL within the CIM framework. To validate the feasibility and effectiveness of this method, the team implemented DBAL on a memristor-based stochastic CIM system and successfully demonstrated a robotic skill learning task. Experimental results showed that, compared to traditional CMOS hardware, the memristor-based CIM system achieved a 44% improvement in speed and saved 153 times more energy.

Research Results and Conclusions

The team first analyzed the stochastic characteristics of memristors and found that memristors exhibit Gaussian-distributed random fluctuations during reading and modulation processes, providing a basis for efficient random number generation. Subsequently, they proposed the MSGLD method, leveraging the stochastic properties of memristors to achieve efficient updates of network weights. Through this method, the memristor Bayesian Deep Neural Network (BDNN) can efficiently learn from uncertain samples and accurately capture prediction uncertainty.

In the robotic skill learning experiment, the team used an 11×50×50×2 memristor BDNN to successfully train a robot to learn a pouring skill using the DBAL method. The results showed that, compared to passive learning methods, the active learning method significantly improved the model’s classification performance and task success rate with the same amount of labeled data. Furthermore, the memristor-based CIM system demonstrated exceptional energy efficiency and speed in handling these tasks, highlighting its potential in applications such as robotic skill learning.

Research Value and Significance

This study proposes a memristor-based CIM framework combined with the MSGLD method to achieve efficient DBAL. This innovation not only significantly reduces the cost and time of data labeling but also provides an efficient learning method for resource-constrained scenarios such as edge computing. Additionally, the intrinsic stochastic properties of memristors offer a new hardware implementation pathway for probabilistic computations in Bayesian methods, with broad application prospects.

Future research could further explore the potential of memristor technology in broader application scenarios, such as drug discovery, catalyst design, and protein production optimization. Moreover, optimizing the manufacturing and operating conditions of memristors to reduce device-to-device variations is an important direction for future research.

This study opens new avenues for the efficient implementation of deep learning, showcasing the vast potential of memristor technology in the field of artificial intelligence.