[論文レビュー] 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.

研究の動機と目的

• SBML モデルを Python/ML パイプラインでシームレスに再利用できるようにする。
• JAX を活用して ODE モデルの効率的な並列シミュレーションを実現する。
• モデルパラメータと介入の微分可能で勾配対応の最適化を提供する。
• SBMLtoODEpy に匹敵するシンプルさと拡張性を維持しつつ JAX 加速を追加する。

提案手法

• SBML を自動的に解析し、JAX で実装された Python モデルへ変換する。
• jit、vmap、grad などの JAX 変換を用いてシミュレーションを加速し、並列化を可能にする。
• 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 SBML ベースの ODE モデルを JAX エコシステム内で効率的にシミュレートし、最適化できるか。
• RQ2 JAX 変換は SBML 派生モデルのスケーラブルな並列シミュレーションを可能にするか。
• RQ3 SBML モデルのロールアウトを微分することの実用的な利点と制約は何か。
• RQ4 SBMLtoODEjax は最適化タスクのための Optax のような ML ツールにどのように統合されるか。

主な発見

• SBMLtoODEjax は jit、vmap、grad を活用して最小コードで並列化された ODE シミュレーションを可能にする。
• 生成されたモデルは JAX ツールリングと互換性があり、大規模実験のバッチ処理が可能。
• Optax との統合により、動力学パラメータと外部介入の勾配ベースの最適化をサポート。
• このアプローチは ML ワークフロー内で生物学的ネットワーク挙動の勾配ベースの発見と最適化を促進する。
• ベンチマークにより、ノートブック内の複数ロールアウトで 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