[論文レビュー] TreeFlow: probabilistic programming and automatic differentiation for phylogenetics

Christiaan Swanepoel, Mathieu Fourment+5 more|arXiv (Cornell University)|Nov 9, 2022

Genomics and Phylogenetic Studies被引用数 4

ひとこと要約

TreeFlow は、TensorFlow を用いて系統樹の尤度の勾配を効率的に計算することで、系統樹上の勾配ベース推論をスケーラブルに可能にする確率的プログラミングおよび自動微分ライブラリである。複雑なモデルにおいては、bito/BEAGLE のような専用ライブラリと比較して 2–3 倍の性能を達成しており、線形の計算スケーリングを示し、新しい進化モデルの柔軟で迅速な実装を可能にしている。

ABSTRACT

Probabilistic programming frameworks are powerful tools for statistical modelling and inference. They are not immediately generalisable to phylogenetic problems due to the particular computational properties of the phylogenetic tree object. TreeFlow is a software library for probabilistic programming and automatic differentiation with phylogenetic trees. It implements inference algorithms for phylogenetic tree times and model parameters given a tree topology. We demonstrate how TreeFlow can be used to quickly implement and assess new models. We also show that it provides reasonable performance for gradient-based inference algorithms compared to specialized computational libraries for phylogenetics.

研究の動機と目的

• 一般用途の確率的プログラミングと系統樹推論の間のギャップを埋め、系統樹上の勾配ベース推論を可能にすること。
• 離散的な木構造と連続的な枝長を組み合わせた確率的モデルにおける計算課題に対処すること。
• 自動微分を用いた、新しい系統樹モデルの実装とベンチマークのための柔軟で効率的なフレームワークを提供すること。
• 進化率、集団動態、時間変動パラメータを含む複雑なモデルに対するスケーラブルな推論を可能にすること。
• 微分可能な確率的プログラミングフレームワーク内での木構造推論の今後の支援の基盤を築くこと。

提案手法

• TreeFlow は、計算グラフにおける再帰を可能にする TensorArray を用いた制御フローを備えた動的計画法を用いて、微分可能な系統樹尤度を実装している。
• TensorFlow の自動微分を活用し、枝長およびモデルパラメータに関して尤度の対数の正確な勾配を計算している。
• 木構造と枝長の階層的パrameterization を用いることで、連続的および離散的成分を同時に推論可能としている。
• 自動微分を介してハミルトニアンモンテカルロや変分ベイズ推論といった勾配ベース推論手法をサポートしている。
• TreeFlow の実装は計算効率に最適化されており、GTR/ワイブルモデルにおいて bito/BEAGLE のような専用ライブラリと同等の性能を示している。
• Python を用いた確率的プログラミングによるモデル仕様の拡張性を備えており、新しい進化モデルの迅速なプロトタイピングを可能にしている。

実験結果

リサーチクエスチョン

• RQ1 再帰的かつ制御フローが多く含まれる系統樹尤度計算に、自動微分が効果的に適用可能かどうか。
• RQ2 一般用途の微分可能な系統樹ライブラリの性能は、bito/BEAGLE のような専用高性能ライブラリと比べてどの程度か。
• RQ3 TreeFlow は大規模データセットに対しても効率的にスケーリング可能か。その実効計算複雑度は何か。
• RQ4 進化率や時間変動集団動態を含む複雑なモデルに対して、TreeFlow がどの程度サポートできるか。
• RQ5 平均場変分ベイズ推論が系統樹パラメータの事後分布相関をどの程度正しく捉えられておらず、構造的近似がどのように改善できるか。

主な発見

• GTR/ワイブルモデルにおいて、TreeFlow は bito/BEAGLE と比較して 2–3 倍の性能を達成しており、専用ライブラリと競合する強力な性能を示している。
• TreeFlow の計算スケーリングは、対数スケールの実行時間とデータサイズのフィットにおいて、2 未満の勾配推定値を示しており、配列数に対してほぼ線形であることが判明した。
• TensorFlow の制御フローと実行モデルの使用により、再帰的尤度計算の勾配計算が効率的に行えるようになり、この文脈では JAX のエージャー実行を上回った。
• TreeFlow における平均場変分ベイズ推論は、インフルエンザのような複雑なデータセットにおいて、分岐時刻の事後分布相関を捉えられていない。これは、構造的近似の必要性を示唆している。
• TreeFlow は、時間変動する進化率や集団動態を含む新しいモデルの迅速なプロトタイピングを、柔軟な確率的プログラミングインタフェースにより可能としている。
• フレームワークは拡張可能であり、今後の微分可能な木探索による構造推論の統合に適しているが、依然として大きな計算課題である。