Fitting Large Nonlinear Mixed Effects Models Using Variational Expectation Maximization

1. Probleemstelling

Niet-lineaire gemengde effecten (NLME)-modellen zijn de standaard voor het analyseren van hiërarchische en longitudinale data in farmacometrie, Quantitative Systems Pharmacology (QSP) en klinische trials. Deze modellen schatten populatieparameters ($\theta$) door de marginale likelihood te maximaliseren, wat vereist dat subject-specifieke willekeurige effecten ($\eta$) worden geïntegreerd.

De Kernuitdaging:

• Berekeningsonhaalbaarheid: Naarmate modellen complexer worden (bijvoorbeeld door het integreren van neurale netwerken, hoogdimensionale ODE's of duizenden parameters), wordt de integratie die nodig is om de marginale likelihood te berekenen, berekeningstechnisch onuitvoerbaar.
• Beperkingen van Bestaande Methoden: Traditionele algoritmen zoals First-Order Conditional Estimation (FOCE) en Stochastic Approximation EM (SAEM) vertrouwen op numerieke integratie of Markov Chain Monte Carlo (MCMC)-sampling. Deze methoden schalen slecht met het aantal parameters en willekeurige effecten, en falen vaak in convergentie of vergen excessive tijd voor modellen met duizenden parameters (bijvoorbeeld DeepNLME-modellen).

2. Methodologie: Variational Expectation Maximization (VEM)

De auteurs stellen het gebruik van Variational Expectation Maximization (VEM) voor als een schaalbaar alternatief. In plaats van te samplingen uit de posteriorverdeling van willekeurige effecten (zoals bij MCEM/SAEM), benadert VEM de posterior met een parametrische verdeling (variational family).

Belangrijke Technische Componenten:

• Variational Inference (VI) Kader:
  • Het doel is het maximaliseren van de Evidence Lower Bound (ELBO) in plaats van de marginale likelihood direct.
  • De ELBO wordt gedefinieerd als: $\text{ELBO} = \mathbb{E}_{q}[\log p(y, \eta|\theta)] - \text{KL}(q(\eta) || p(\eta|y, \theta))$.
  • Het maximaliseren van de ELBO is equivalent aan het minimaliseren van de Kullback-Leibler (KL)-divergentie tussen de variational verdeling $q(\eta)$ en de ware posterior.
• Reparameterization Trick:
  • Om gradient-based optimalisatie mogelijk te maken, gebruiken de auteurs de reparameterization trick. Willekeurige effecten worden gegenereerd als een deterministische transformatie van een vaste basisverdeling (bijvoorbeeld $\eta = T(\xi; \kappa)$ waarbij $\xi \sim \mathcal{N}(0, I)$).
  • Dit maakt het mogelijk dat gradients door het sampling-proces stromen, waardoor efficiënte optimalizers kunnen worden gebruikt.
• Gecombineerde Optimalisatie:
  • In tegenstelling tot traditionele EM, die strikt alterneert tussen E-stappen en M-stappen, voert de Pumas-implementatie gecombineerde optimalisatie uit van zowel de populatieparameters ($\theta$) als de variational parameters ($\kappa$) gelijktijdig met behulp van gradient-based methoden.
• Automatic Differentiation (AD):
  • Reverse-mode AD wordt gebruikt voor modellen met een groot aantal parameters ten opzichte van outputs (vaak bij DeepNLME). Dit is cruciaal voor het efficiënt berekenen van gradients van de ELBO met betrekking tot duizenden parameters.
  • De implementatie ondersteunt zowel forward-mode (voor kleinere modellen) als reverse-mode AD.
• Optimalisatiestrategieën:
  • Deterministische VEM: Gebruikt vooraf gesamplede basisvariabelen en de L-BFGS optimizer (quasi-Newton) voor stabiele, snelle convergentie.
  • Stochastische VEM: Gebruikt her-sampling per iteratie en de Adam optimizer, geschikt voor mini-batch verwerking van grote populaties.
• Omgaan met Beperkingen: Voor beperkte willekeurige effecten (bijvoorbeeld log-normaal) past de methode domeintransformaties toe (bijvoorbeeld log-transformatie) voordat wordt benaderd met een Gaussische variational familie.

3. Belangrijkste Bijdragen

1. Schaalbaar Algorithmus voor NLME: Het artikel introduceert een VEM-variant die specifiek is toegespitst op NLME-modellen, en in staat is om modellen te fitten met >15.000 populatieparameters en 16+ willekeurige effecten, een regime waar traditionele methoden falen.
2. Implementatie in Pumas: Het algoritme is geïmplementeerd in de statistische software Pumas, waarbij gebruik wordt gemaakt van de high-performance computing-mogelijkheden van Julia en automatic differentiation.
3. Theoretische Verbindingen: De auteurs verduidelijken de relatie tussen VEM, Variational Autoencoders (VAEs), Probabilistic PCA (PPCA) en de Laplace-benadering. Ze tonen aan dat de Laplace-methode in wezen een Taylor-benadering van tweede orde is van het VEM-doelwit.
4. Empirische Validatie: Het artikel biedt rigoureuze benchmarks die VEM vergelijken met FOCE op zowel kleine schaal (Warfarine) als grote schaal (DeepNLME) modellen.

4. Experimentele Resultaten

De auteurs voerden experimenten uit op twee modellen met behulp van de Pumas-software:

A. Warfarine Model (Validatie op Kleine Schaal)

• Opzet: Standaard PK-PD model met 31 subjecten en 464 observaties.
• Vergelijking: VEM versus FOCE.
• Resultaten:
  • Nauwkeurigheid: VEM produceerde parameterschattingen en log-likelihoods die bijna identiek waren aan FOCE (Log Likelihood: -994,5 versus -995,3).
  • Snelheid: VEM was aanzienlijk sneller (16,6 seconden versus 44,4 seconden).
  • Conclusie: VEM is een correct en efficiënt alternatief voor standaard farmacometrische modellen.

B. DeepNLME Friberg Model (Schaalbaarheid op Grote Schaal)

• Opzet: Een kunstmatig groot model dat neurale netwerken (MLP's) integreert met 15.410 populatieparameters en 16 willekeurige effecten.
• Vergelijking: VEM (5 iteraties) versus FOCE (poging tot 1 iteratie).
• Resultaten:
  • VEM: Voltooid 5 iteraties in 3092 seconden (618 seconden/iteratie) met gebruik van reverse-mode AD.
  • FOCE: Faalde om een enkele iteratie binnen 24 uur te voltooien.
  • Conclusie: VEM schaalt succesvol naar complexe, hoogdimensionale modellen waar FOCE berekeningstechnisch onuitvoerbaar wordt.

5. Betekenis en Implicaties

• Mogelijk maken van DeepNLME en QSP: Het vermogen om modellen met duizenden parameters te fitten opent de deur voor DeepNLME (het integreren van neurale netwerken in mechanistische modellen) en complexe Quantitative Systems Pharmacology (QSP)-modellen, die eerder te rekenintensief waren om als NLME's te fitten.
• Verschuiving in Optimalisatieparadigma: Het artikel pleit voor een verschuiving van MCMC-gebaseerde sampling (SAEM/MCEM) naar gradient-based variational inference voor problemen op grote schaal, waarbij gebruik wordt gemaakt van moderne deep learning-infrastructuur (reverse-mode AD, compatibiliteit met GPU/TPU).
• Praktische Beperkingen: De auteurs erkennen dat VEM gevoelig is voor hyperparameters (aantal samples, keuze variational familie) en aparte stappen vereist om Empirical Bayes Estimates (EBE's) en log-likelihood-benaderingen na de optimalisatie te berekenen, in tegenstelling tot FOCE, wat dit tijdens de schatting doet.
• Toekomstige Richtingen: De methode is gepositioneerd om een standaardtool te worden voor het fitten van grote gezamenlijke biomarkermodellen en complexe biologische systemen waar traditionele methoden op een "muur" van rekenkosten stuiten.

Samenvattend toont dit artikel aan dat Variational Expectation Maximization, in combinatie met reverse-mode automatic differentiation, een robuust, schaalbaar en nauwkeurig kader biedt voor het fitten van de volgende generatie complexe, hoogdimensionale niet-lineaire gemengde effectenmodellen in de farmacometrie.