[논문 리뷰] TreeFlow: probabilistic programming and automatic differentiation for phylogenetics

Christiaan Swanepoel, Mathieu Fourment외 5명|arXiv (Cornell University)|2022. 11. 09.

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의 자동 미분을 활용하여 분지 길이 및 모델 매개변수에 대한 로그 우도의 정확한 기울기를 계산합니다.
• 나무의 구조와 분지 길이에 대한 계층적 매개변수화를 사용하여 연속적 및 이산적 구성요소를 동시에 추론할 수 있도록 합니다.
• 자동 미분를 통해 하이퍼볼릭 몬테카를로 및 변분 추론과 같은 기울기 기반 추론 방법을 지원합니다.
• TreeFlow의 구현은 계산 효율성에 최적화되어 있으며, GTR/Weibull 모델에서 bito/BEAGLE와 같은 전용 라이브러리와 유사한 성능을 보입니다.
• 파이썬 기반의 확률적 프로그래밍을 통해 모델 사양을 확장 가능하게 하여 새로운 진화 모델의 신속한 프로토타이핑을 가능하게 합니다.

실험 결과

연구 질문

• RQ1 재귀적이고 제어 흐름이 많은 연산을 포함하는 계량적 우도 계산에 자동 미분가 효과적으로 적용될 수 있는가?
• RQ2 일반적인 기울기 가능 계량적 나무 라이브러리의 성능가 bito/BEAGLE와 같은 전용 고성능 라이브러리와 비교해 어떻게 되는가?
• RQ3 TreeFlow는 대규모 데이터셋에 대해 효율적으로 스케일링될 수 있으며, 그 실질적인 계산 복잡도는 어떠한가?
• RQ4 TreeFlow는 진화 속도와 시간에 따라 변하는 인구 역학을 포함하는 복잡한 모델을 어느 정도 지원할 수 있는가?
• RQ5 평균장 변분 추론이 계량적 매개변수의 사후 상관관계를 얼마나 잘 포착할 수 있으며, 구조적 근사가 이를 어떻게 향상시킬 수 있는가?

주요 결과

• TreeFlow는 GTR/Weibull 모델에서 bito/BEAGLE의 성능의 약 2~3배 이내로 성능을 달성하여 전용 라이브러리와의 경쟁력이 높음을 보여줍니다.
• TreeFlow의 계산 스케일링은 시퀀스 수에 대해 약 선형적임을 보이며, 로그-로그 런타임 대 데이터 크기 피팅에서 기울기 추정치가 2 이하임을 통해 확인되었습니다.
• TensorFlow의 제어 흐름 및 실행 모델을 활용함으로써 재귀적 우도 계산의 기울기 계산을 효율적으로 수행할 수 있었으며, 이는 이 맥락에서 JAX의 즉각 실행 모드보다 우수한 성능을 보였습니다.
• TreeFlow의 평균장 변분 추론은 특히 인플루엔자와 같은 복잡한 데이터셋에서 분기 시점의 사후 상관관계를 포착하지 못함을 보여주며, 이는 구조적 근사가 필요함을 시사합니다.
• TreeFlow는 유연한 확률적 프로그래밍 인터페이스를 통해 시간에 따라 변하는 진화 속도와 인구 역학을 포함한 새로운 모델의 빠른 프로토타이핑을 지원합니다.
• 이 프레임워크는 확장 가능하며, 향후 기울기 가능 나무 탐색을 통한 구조 추론 통합에 적합하지만, 이는 여전히 큰 계산적 과제로 남아 있습니다.