吴思团队推出面向AI驱动科学计算的物理单位系统SAIUnit

本文来源：北大脑科学

人工智能技术正在深刻变革各领域的科学研究。从预测复杂的蛋白质结构到模拟宏大的星系碰撞，再到优化核聚变反应堆设计，AI已成为科学家不可或缺的助手。然而，AI与传统科学研究方法的融合仍面临着根本性挑战。AI算法通常基于抽象的数值计算，缺乏直接的物理对应，且往往能在低精度条件下正常运作；相比之下，科学研究则严重依赖于标准单位体系，要求对物理世界进行精确测量，精度要求极其严格。

这种本质差异导致了一系列棘手问题：即便是微小的量纲错误也可能引发灾难性后果，就如同1999年美国火星气候轨道器任务失败的案例所示。研究者在将AI方法应用于科学研究时面临两难困境：一方面，不得不耗费大量时间进行繁琐的单位转换，以防止潜在的量纲错误；另一方面，由于缺乏统一标准，不同研究团队构建的模型和数据往往难以直接沟通和交流，这严重阻碍了跨学科合作和创新。

2025年4月16日，北京大学心理与认知科学学院、IDG麦戈文脑科学研究所、定量生物学中心、北大-清华生命科学联合中心的吴思研究团队，在《Nature Communications》期刊上发表一篇题为“Integrating physical units into high-performance AI-driven scientific computing”的文章。他们开发的SAIUnit（原BrainUnit系统），能够在强大的AI数值计算框架中引入了标准化物理单位，填补了现有AI框架的重要空白。它提供了超过2000多个常见物理单位和500多个单位感知的数学函数，并且兼容AI框架的自动微分、编译、并行、低精度计算等核心功能。

Nature Communications 评委对SAIUnit的贡献给予了充分肯定：

SAIUnit研究成果在Nature Communications上得以发表。

SAIUnit如何为科学计算注入新活力?

以下是对SAIUnit的简要介绍。

设计原则

SAIUnit的设计遵循七大原则：

• 标准化：采用国际单位制（SI），确保全球科学兼容性。
• 一致性：严格遵循物理定律，确保基本单位和复合单位之间的逻辑一致性。
• 通用性：支持已知的物理量，并具备扩展性以容纳新的物理量。
• 准确性：高精度的单位转换和计算，特别是在处理极端量级的数据时。
• 计算效率：优化实现，最小化单位检查的计算开销。
• 类型安全：防止不兼容单位之间的非法操作。
• AI兼容性：完全兼容AI功能，包括自动微分、即时编译、GPU/TPU加速、向量化和并行化。

核心数据结构

图1: 物理单位处理的核心数据结构

SAIUnit的核心创新在于三个相互关联的数据结构：saiunit.Dimension、saiunit.Unit和saiunit.Quantity。saiunit.Dimension基于国际单位制（SI）建立了严格的单位转换系统，将七个基本物理单位（长度、质量、时间、电流、温度、物质的量和发光强度）表示为七元整数向量。saiunit.Unit结合了度量尺度和维度实例，允许在不同量级上表示广泛的物理量。saiunit.Quantity通过将数值与单位结合，实现了单位感知的高性能数组计算。

优势1：兼容AI生态的物理单位、函数、和转换

图2：兼容JAX的物理单位和单位感知数学函数

SAIUnit构建了完整的单位感知数学体系，提供了科学、工程、技术等领域常见的2000多个物理单位（见图2a），以及500多个单位感知的数学函数（见图2b），其核心设计完美兼容JAX生态的四大核心变换（见图2c）：

• 自动微分(jax.grad): 系统仅对无量纲数值进行梯度计算，避免对单位尺度进行不正确的微分
• 即时编译(jax.jit): 单位转换检查仅在编译阶段进行，运行时性能媲美原始数值计算
• 向量化处理(jax.vmap): 保持批次数据的单位一致性
• 多设备并行(jax.pmap): 支持跨GPU/TPU集群的分布式单位计算

优势2：广泛科学计算适用性

图3：跨科学计算领域的物理单位集成

SAIUnit已成功集成到多个基于JAX的科学计算框架，实现了跨领域单位系统的统一。在数值微分方程求解（Diffrax软件）中，它可以直接定义各种时间单位的步长并自动检测方程的量纲一致性。在多尺度脑模拟（BrainPy软件）领域，它使不同尺度的参数可在统一量纲系统中共存，并支持实验数据的直接输入。在物理信息神经网络（PINNx软件）中，它自动执行单位一致性检查，归一化异构数据并过滤单位异常。

优势3：无额外性能开销和精度损失

图4：SAIUnit单位感知计算在CPU和GPU平台上的计算开销分析

SAIUnit的架构设计支持物理单位的检查、计算和运行都发生在编译时，使得运行时计算几乎不产生额外开销。实验证明，使用SAIUnit的单位感知计算与传统无单位计算性能相当。比如，在处理包含个神经元的网络时，单位感知模拟的运行时间基本与无单位仿真持平。

图5：带物理单位系统和不带物理单位系统的Lorenz系统模拟数值比较

同时，SAIUnit的架构设计支持低精度计算。在Lorenz系统模拟中，传统物理单位库（如Pint或Quantities）在32位或16位精度下会出现显著的轨迹偏差，而SAIUnit即使在16位低精度条件下仍保持零误差。

优势4：提高科研效率

图6：单位感知计算提高离子通道建模中的研究效率

SAIUnit通过自动化单位检查与转换，显著提高了科研效率。以计算神经科学中的离子通道建模为例：传统NEURON模拟器需要手动处理各种单位转换，而基于SAIUnit的工具可直接表达生物物理方程，减少代码量并消除人为错误风险。

优势5：促进跨学科合作

图7：SAIUnit标准化单位系统使用SI标准统一不同的测量单位

SAIUnit还促进了跨学科研究合作。不同学科使用的单位和测量标准差异，常使数据集成和协作变得复杂。例如，物理学家使用焦耳（J）来表示能量，而生态学家则倾向于使用卡路里（cal）；医学领域中的血压通常以毫米汞柱（mmHg）表示，而物理学则以帕斯卡（Pa）来衡量压强；化学反应的能量变化往往可以与机械功直接对应；生物体的新陈代谢率通常以千卡每小时（kcal/h）表示，而在物理学中则常常以瓦特（W）作为功率单位。这些差异使得跨学科的数据共享和协作变得复杂且低效。SAIUnit引入统一的物理单位系统，使用SI单位的一致维度表示不同领域的数据单位，使研究人员能够无缝理解和利用彼此的数据。

总结

SAIUnit实现了一个新的科学计算范式：

数值计算+物理单位=科学计算

这一公式带来四大优势：规避量纲失配风险、提升代码可读性、增强跨领域代码可移植性、减轻开发者负担。在脑动力学、数值积分、物理信息网络、系统生物建模、分子动力学模拟等多个领域的大量实验证明，SAIUnit提高了科学计算的准确性、可靠性和可解释性，并且不会影响性能。

SAIUnit已在PyPI和GitHub开源，并集成至多个科学计算框架，包括数值积分库Diffrax、分子动力学模拟库JAX-MD、生物系统建模库Catalax、神经元建模库BrainCell和物理信息神经网络库PINNx等专业工具。通过提供统一、严谨的物理单位体系，SAIUnit正为AI驱动的科学研究注入新的活力，推动科学计算向更精确、更可靠、更易理解的方向发展。

更多信息

SAIUnit框架前身为BrainUnit，初期通过PyPI发布（https://pypi.org/project/brainunit/）。

当前正式版已在PyPI（https://pypi.org/project/saiunit/）和GitHub（https://github.com/chaobrain/saiunit）开源，文档托管于：https://saiunit.readthedocs.io/。

本研究的实验复现代码已公开于GitHub仓库：https://github.com/chaobrain/saiunit-experiments。

北京大学心理与认知科学学院、IDG麦戈文脑科学研究所、定量生物学中心、北大-清华生命科学联合中心的吴思教授为该论文的通讯作者，北京大学心理与认知科学学院博雅博士后王超名博士和北大-清华生命科学联合中心博士生贺思超为该论文的第一作者，王超名博士为SAIUnit项目的发起和设计者，上海交通大学博士生罗守为、广东智能科学与技术研究院环宇翔研究员为该论文做出重要贡献。

原文链接

https://www.nature.com/articles/s41467-025-58626-4