[论文解读] SBMLtoODEjax: Efficient Simulation and Optimization of Biological Network Models in JAX

Mayalen Etcheverry, Michael Levin+2 more|arXiv (Cornell University)|Jul 17, 2023

Gene Regulatory Network Analysis被引用 8

一句话总结

SBMLtoODEjax 是一个轻量级库，将 SBML 模型转换为用 JAX 编写的端到端 Python 模型，从而在 ML 流水线中实现快速并行仿真与基于梯度的优化。

ABSTRACT

Advances in bioengineering and biomedicine demand a deep understanding of the dynamic behavior of biological systems, ranging from protein pathways to complex cellular processes. Biological networks like gene regulatory networks and protein pathways are key drivers of embryogenesis and physiological processes. Comprehending their diverse behaviors is essential for tackling diseases, including cancer, as well as for engineering novel biological constructs. Despite the availability of extensive mathematical models represented in Systems Biology Markup Language (SBML), researchers face significant challenges in exploring the full spectrum of behaviors and optimizing interventions to efficiently shape those behaviors. Existing tools designed for simulation of biological network models are not tailored to facilitate interventions on network dynamics nor to facilitate automated discovery. Leveraging recent developments in machine learning (ML), this paper introduces SBMLtoODEjax, a lightweight library designed to seamlessly integrate SBML models with ML-supported pipelines, powered by JAX. SBMLtoODEjax facilitates the reuse and customization of SBML-based models, harnessing JAX's capabilities for efficient parallel simulations and optimization, with the aim to accelerate research in biological network analysis.

研究动机与目标

• 在 Python/ML 流水线中实现对 SBML 模型的无缝重用。
• 利用 JAX 实现对常微分方程模型的高效并行仿真。
• 提供对模型参数和干预的可微分、带梯度的优化。
• 保持与 SBMLtoODEpy 相当的简洁性和可扩展性，同时加入 JAX 加速。

提出的方法

• 自动解析 SBML 并转换为用 JAX 实现的 Python 模型。
• 使用 JAX 转换（jit、vmap、grad）来加速仿真并实现并行化。
• 将 SBMLtoODEpy 的约定扩展为适合 JAX 的函数式编程风格和 Equinox PyTree 集成。
• 允许与 Optax 集成，实现对参数和干预的梯度优化。
• 在 ML 流水线中演示并行 rollout 向量化与基于梯度的优化。

Figure 1 : Example code (left) and output snapshot (right) reproducing original simulation results of Kholodenko 2000’s paper [ 34 ] hosted on BioModels.

实验结果

研究问题

• RQ1 SBML 基于的 ODE 模型是否可以在 JAX 生态系统内高效地仿真和优化？
• RQ2 JAX 转换是否能够实现对 SBML 派生模型的可扩展并行仿真？
• RQ3 对 SBML 模型 rollout 的微分有哪些实际收益和局限性？
• RQ4 SBMLtoODEjax 如何与 Optax 等 ML 工具结合用于优化任务？

主要发现

• SBMLtoODEjax 通过 jit、vmap 和 grad 实现了最小代码量的并行化 ODE 仿真。
• 生成的模型与 JAX 工具链兼容，可以对大规模实验进行批量处理。
• 与 Optax 的集成支持对动力学参数和外部干预的梯度优化。
• 该方法在 ML 工作流中促进对生物网络行为的基于梯度的发现与优化。
• 基准测试显示在笔记本中对多次 rollout 相比 SBMLtoODEpy 具备速度优势。

Figure 2 : (left) Example code to vectorize calls to model rollout (right) Results of a (rudimentary) benchmark comparing the average simulation time of models implemented with SBMLtoODEpy versus SBMLtoODEjax (for different number of rollouts i.e. batch size). Mean and standard-deviation results of