Nonparametric Modeling of Continuous-Time Markov Chains

Technische Samenvatting: Nonparametrische Modellering van Continue-Tijd Markovketens met Schaalbare Exacte en Benaderde Gradiënten

Probleemstelling
Het afleiden van de instantane snelheden van Continue-Tijd Markovketens (CTMC's) is een fundamentele uitdaging in velden zoals fylogenetica, epidemiologie en overlevingsanalyse. De primaire obstakels zijn de kwadratische groei van de parametersnelheden naarmate de toestandsruimte ($K$) groter wordt, sterke afhankelijkheden tussen deze snelheden, en incomplete informatie over transities in gedeeltelijk geobserveerde data. Traditionele benaderingen vertrouwen vaak op directe parametrisatie met structurele restricties of log-lineaire modellen die de snelheden koppelen aan covariaten. Echter, directe parametrisatie wordt combinatorisch complex naarmate de toestandsruimten groeien, en log-lineaire modellen falen om niet-lineaire, niet-monotone of drempelgebaseerde relaties tussen covariaten en transitiesnelheden te vangen. Bovendien vereist Bayesiaanse inferentie met behulp van Hamiltonian Monte Carlo (HMC) voor deze modellen herhaalde berekening van likelihood-gradiënten. Bestaande methoden voor het differentiëren van matrix-exponenten schalen als $O(NK^3)$ (waarbij $N$ het aantal observaties is), wat een ernstige computationele bottleneck creëert die de toepassing van flexibele, hoog-dimensionale modellen belemmert.

Methodologie
De auteurs stellen een Bayesiaans kader voor dat zowel de modelleringsflexibiliteit als de computationele schaalbaarheid van CTMC-inferentie aanpakt.

1. Nonparametrische Modellering via Gaussische Processen (GP's):
   In plaats van een restrictieve log-lineaire relatie tussen covariaten en getransformeerde snelheden aan te nemen, modelleren de auteurs de getransformeerde snelheden $\theta_{ij}$ als een willekeurige functie van covariaten met behulp van een Gaussisch Proces (GP) prior. Dit maakt het mogelijk om complexe niet-lineaire afhankelijkheden te vangen, zoals verzadigingseffecten of niet-monotone relaties, terwijl de interpreteerbaarheid behouden blijft. De instantane snelheidsmatrix $Q$ wordt geconstrueerd via een differentieerbare transformatie $g(\theta)$ (bijv. $g(\theta) = e^\theta$) om te garanderen dat aan de non-negativiteitsrestricties wordt voldaan.

2. Schaalbare Gradiëntberekening:
   Om efficiënte HMC-sampling mogelijk te maken, introduceren de auteurs twee complementaire strategieën voor het berekenen van gradiënten van de gedeeltelijk geobserveerde CTMC-likelihood met betrekking tot de parametersnelheden. Beide strategieën maken gebruik van een blokdiagonaal eigendecompositie van de snelheidsmatrix $\Lambda = VBV^{-1}$.

   • Likelihood Evaluatie: De likelihood op een Directed Acyclic Graph (DAG) wordt geëvalueerd met een message-passing algoritme (post-order en pre-order traversals) dat matrix-vector producten berekent in $O(K^2)$ per edge na een initiële $O(K^3)$ eigendecompositie.
   • Benaderde Gradiënt (Forward-Mode): Deze methode maakt gebruik van een eerste-orde benadering van de Fréchet-afgeleide van de matrix-exponent (waarbij alleen de eerste term van de reeksontwikkeling wordt behouden). Deze is ingebed in een HMC-schema met surrogaat-trajecten, waarbij een Metropolis–Hastings correctie ervoor zorgt dat de exacte posterior de stationaire verdeling blijft. Deze aanpak reduceert de gradiëntkosten tot ongeveer twee likelihood-evaluaties.
   • Exacte Gradiënt (Reverse-Mode): Deze methode leidt een analytische gesloten vorm af voor de reverse-mode gradiënt met behulp van de adjoint Fréchet-afgeleide van de matrix-exponent. Het omvat het oplossen van blokintegralen die volledig paralleliseerbaar zijn. Net als de benaderde methode bereikt dit een totale complexiteit van $O(K^3 + NK^2)$.

Belangrijkste Bijdragen

• Flexibele Covariaat Modellering: De introductie van GP-priors op getransformeerde snelheden versoepelt de log-lineaire assumptie, wat het mogelijk maakt om complexe, niet-lineaire drijvers van CTMC-dynamiek te ontdekken zonder de interpreteerbaarheid op te offeren.
• Computationele Schaalbaarheid: De ontwikkeling van exacte en benaderde gradiëntstrategieën die de computationele complexiteit van gradiëntevaluatie verminderen van $O(NK^3)$ naar $O(K^3 + NK^2)$. Dit wordt bereikt door de kubische kosten (eigendecompositie) te ontkoppelen van de lineaire afhankelijkheid van het aantal observaties.
• Implementatie en Benchmarking: De methoden zijn geïmplementeerd in de breed gebruikte Bayesiaanse fylogenetische software Beast X. De auteurs leveren rigoureuze benchmarks die aantonen dat hun benaderingen orden van grootte sneller zijn dan state-of-the-art autodifferentiatie-frameworks (bijv. torchtree, treeflow) en gespecialiseerde analytische reverse-mode implementaties, met name voor niet-reversibele modellen en grote toestandsruimten.

Resultaten

• Computationele Prestaties: Empirische benchmarks op synthetische data met tot 10.000 boomtips en 256 CTMC-toestanden laten zien dat de voorgestelde methoden een empirische kwadratische schaling in $K$ vertonen, zelfs wanneer $N < K$. In tegenstelling hiertoe benaderen concurrerende autodifferentiatie-methoden een kubische schaling en lijden ze onder grote constante factoren, waardoor ze honderden tot duizenden keren langzamer zijn.
• Statistische Recoverie: Simulatiestudies met gegevens over de transmissie van het rabiës-virus tonen aan dat het GP-model niet-lineaire log-snelheidsfuncties (kwadratische relaties) nauwkeurig herstelt die log-lineaire modellen niet kunnen vangen. Het GP-model behaalde een mediane Root Mean Square Error (RMSE) van 0,170 vergeleken met 0,487 voor het log-lineaire model, met een superieure dekking van de ware functie.
• Real-World Applicaties:
  • Vleermuis-Rabiës: Analyse van 372 rabiës-virus monsters van 17 vleermuissoorten bevestigde dat genetische afstand de transmissiesnelheden beïnvloedt. Het GP-model ondersteunde een ongeveer log-lineaire relatie in deze specifieke dataset, wat de capaciteit van de methode valideert om eenvoudiger structuren te herstellen wanneer deze aanwezig zijn.
  • Globale Influenza (H3N2): Analyse van de wereldwijde influenza-verspreiding toonde aan dat de log-lineaire assumptie te restrictief was. Het GP-model onthulde een niet-lineaire relatie met de bevolkingsdichtheid van het land van herkomst, waarbij log-lineaire modellen de effecten van lage dichtheden overschatten en de effecten van hoge dichtheden onderschatten.

Betekenis en Claims
Het artikel claimt een robuuste en schaalbare oplossing te bieden voor Bayesiaanse inferentie van CTMC's in hoog-dimensionale toestandsruimten. Door flexibele nonparametrische modellering te combineren met computationeel efficiënte gradiënt-evaluaties, maakt het kader dieper inzicht in stochastische processen die worden gedreven door externe covariaten mogelijk. De auteurs benadrukken dat hun werk de mogelijkheden van huidige modellen uitbreidt, waardoor de exploratie van complexe covariaat-effecten mogelijk wordt die voorheen computationeel onhandelbaar waren of statistisch vertekend waren door restrictieve parametrische assumpties. De implementatie in Beast X is bedoeld om bredere adoptie binnen de fylogenetische en filogeografische gemeenschappen te faciliteren.

Beperkingen en Toekomstige Richtingen
De auteurs merken op dat hoewel de likelihood en gradiëntkosten gunstig schalen, de GP-hyperparameter sampling een computationele kostenpost introduceert die schaalt als $O(K^6)$ door de inversie van de GP-covariantie matrix over de $K^2$ snelheden. Zij suggereren dat voor zeer grote toestandsruimten, sparse of gestructureerde GP-benaderingen, of andere nonparametrische alternatieven, noodzakelijk kunnen zijn. Daarnaast behandelt het huidige framework covariaten als vaststaand en bekend, wat ruimte laat voor toekomstig werk om de onzekerheid van covariaten te incorporeren.