[논문 리뷰] SBMLtoODEjax: Efficient Simulation and Optimization of Biological Network Models in JAX

Mayalen Etcheverry, Michael Levin외 2명|arXiv (Cornell University)|2023. 07. 17.

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 모델의 매끄러운 재사용을 가능하게 한다.
• ODE 모델의 효율적 병렬 시뮬레이션을 위해 JAX를 활용한다.
• 모델 매개변수와 개입에 대한 미분 가능하고 기울기 기반의 최적화를 제공한다.
• JAX 가속을 추가하되 SBMLtoODEpy에 비해 단순성과 확장성을 유지한다.

제안 방법

• SBML을 자동으로 구문 분석하고 JAX로 구현된 Python 모델로 변환한다.
• 시뮬레이션을 가속하고 병렬화를 가능하게 하기 위해 JAX 변환(jit, vmap, grad)을 사용한다.
• SBMLtoODEpy 규칙을 확장하여 JAX 친화적인 함수형 프로그래밍 스타일과 Equinox PyTree 통합에 맞춘다.
• 매개변수와 개입의 기울기 기반 최적화를 위해 Optax와의 통합을 허용한다.
• ML 파이프라인 내에서 병렬 롤아웃 벡터화와 기울기 기반 최적화를 시연한다.

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

실험 결과

연구 질문

• RQ1 Can SBML-based ODE models be efficiently simulated and optimized inside the JAX ecosystem?
• RQ2 Do JAX transformations enable scalable parallel simulations of SBML-derived models?
• RQ3 What are the practical benefits and limitations of differentiating through SBML model rollouts?
• RQ4 How does SBMLtoODEjax integrate with ML tooling like Optax for optimization tasks?

주요 결과

• SBMLtoODEjax enables parallelized ODE simulations with minimal code, leveraging jit, vmap, and grad.
• Generated models are compatible with JAX tooling and can be batched for large-scale experiments.
• Integration with Optax supports gradient-based optimization of kinetic parameters and external interventions.
• The approach facilitates gradient-based discovery and optimization of biological network behaviors within ML workflows.
• Benchmarking shows speed advantages over SBMLtoODEpy for multiple rollouts in their notebooks.

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