BrainUnit: Integrating Physical Units into High-Performance AI-Driven Scientific Computing

Chaoming Wang,Sichao He,Shouwei Luo,Yuxiang Huan,Si Wu
DOI: https://doi.org/10.1101/2024.09.20.614111
2024-09-25
Abstract:Artificial intelligence (AI) is revolutionizing scientific research across various disciplines. The foundation of scientific research lies in rigorous scientific computing based on standardized physical units. However, current mainstream high-performance numerical computing libraries for AI generally lack native support for physical units, significantly impeding the integration of AI methodologies into scientific research. To fill this gap, we introduce BrainUnit, a unit system designed to seamlessly integrate physical units into AI libraries, with a focus on compatibility with JAX. BrainUnit offers a comprehensive library of over 2000 physical units and more than 300 unit-aware mathematical functions. It is fully compatible with JAX transformations, allowing for automatic differentiation, just-in-time compilation, vectorization, and parallelization while maintaining unit consistency. We demonstrate BrainUnit's efficacy through several use cases in brain dynamics modeling, including detailed biophysical neuron simulations, multiscale brain network modeling, neuronal activity fitting, and cognitive task training. Our results show that BrainUnit enhances the accuracy, reliability, and interpretability of scientific computations across scales, from ion channels to whole-brain networks, without significantly impacting performance. By bridging the gap between abstract computational frameworks and physical units, BrainUnit represents a crucial step towards more robust and physically grounded AI-driven scientific computing.
Neuroscience
What problem does this paper attempt to address?
The paper attempts to address the issue that current mainstream high-performance numerical computing libraries (such as PyTorch, TensorFlow, and JAX) lack native support for handling physical units, which severely hinders the application of AI methods in scientific research. Specifically, the paper points out: 1. **Importance of Physical Units**: The foundation of scientific research lies in rigorous calculations based on standardized physical units. The precise application of physical units not only signifies scientific rigor but is also crucial for preventing dimensional errors and avoiding catastrophic failures. 2. **Limitations of Existing Computing Libraries**: Although these computing libraries perform excellently in machine learning and deep learning, they generally do not support the handling of physical units. This leads researchers to manually perform unit conversions during complex calculations, increasing operational complexity and the risk of errors, while also consuming a significant amount of time and effort. 3. **Barriers to Interdisciplinary Collaboration**: Differences in unit conventions across various research fields increase communication costs, hinder the sharing of data and models, and consequently affect the speed of innovation and discovery. To fill this gap, the paper introduces **BrainUnit**, a unit system designed to seamlessly integrate physical units into high-performance computing libraries, with a particular emphasis on compatibility with JAX. BrainUnit provides over 2000 physical units and more than 300 mathematical functions that support units, and it is fully compatible with JAX's automatic differentiation, just-in-time compilation, vectorization, and parallelization features, thereby enhancing computational performance while maintaining unit consistency. Through several use cases in brain dynamics modeling, the paper demonstrates the effectiveness of BrainUnit in improving computational accuracy, reliability, and interpretability.