[论文解读] 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 高效计算系统发育似然的梯度，实现了基于梯度的可扩展系统发育树推断。其性能在复杂模型下与 specialized 库如 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 的自动微分功能，精确计算对分支长度和模型参数的对数似然梯度。
• 该库采用树拓扑与分支长度的分层参数化方法，支持对连续与离散组件的联合推断。
• 它支持通过自动微分实现的基于梯度的推断方法，如哈密顿蒙特卡洛和变分推断。
• TreeFlow 的实现针对计算效率进行了优化，在 GTR/Weibull 模型下性能与 specialized 库如 bito/BEAGLE 相当。
• 该框架通过基于 Python 的概率编程支持可扩展的模型定义，实现新演化模型的快速原型开发。

实验结果

研究问题

• RQ1 自动微分能否有效应用于涉及递归和控制流密集操作的系统发育似然计算？
• RQ2 通用可微分系统发育库的性能与 specialized 高性能库（如 bito/BEAGLE）相比如何？
• RQ3 TreeFlow 能否高效扩展到大规模数据集？其经验计算复杂度如何？
• RQ4 TreeFlow 在支持涉及演化速率和时变种群动态的复杂模型方面的能力如何？
• RQ5 平均场变分推断在捕捉系统发育参数后验相关性方面存在哪些局限性？结构化近似方法能在多大程度上改进这一问题？

主要发现

• 在 GTR/Weibull 模型下，TreeFlow 的计算性能仅比 bito/BEAGLE 低 2–3 倍，表明其与专用库具有很强的竞争力。
• TreeFlow 的计算扩展性接近线性，表现为对数-对数运行时间与数据大小拟合的斜率低于 2，表明其随序列数量增长的计算开销接近线性。
• 该库利用 TensorFlow 的控制流与执行模型，高效计算递归似然计算的梯度，在此场景下优于 JAX 的急切执行模式。
• TreeFlow 中的平均场变分推断无法捕捉分歧时间的后验相关性，尤其在如流感等复杂数据集中表现明显，凸显了结构化近似方法的必要性。
• 通过其灵活的概率编程接口，TreeFlow 支持新模型的快速原型开发，包括具有时变演化速率和种群动态的模型。
• 该框架具有可扩展性，非常适合未来集成通过可微分树搜索实现的拓扑推断，尽管这仍是重大的计算挑战。