Variational autoencoder for inference of nonlinear mixed effect models based on ordinary differential equations

We propose a variational autoencoder (VAE) approach for parameter estimation in nonlinear mixed-effects models based on ordinary differential equations (NLME-ODEs) using longitudinal data from multiple subjects. In moderate dimensions, likelihood-based inference via the stochastic approximation EM algorithm (SAEM) is widely used, but it relies on Markov Chain Monte-Carlo (MCMC) to approximate subject-specific posteriors. As model complexity increases or observations per subject are sparse and irregular, performance often deteriorates due to a complex, multimodal likelihood surface which may lead to MCMC convergence difficulties. We instead estimate parameters by maximizing the evidence lower bound (ELBO), a regularized surrogate for the marginal likelihood. A VAE with a shared encoder amortizes inference of subject-specific random effects by avoiding per-subject optimization and the use of MCMC. Beyond pointwise estimation, we quantify parameter uncertainty using observed-information-based variance estimator and verify that practical identifiability of the model parameters is not compromised by nuisance parameters introduced in the encoder. We evaluate the method in three simulation case studies (pharmacokinetics, humoral response to vaccination, and TGF-$β$ activation dynamics in asthmatic airways) and on a real-world antibody kinetics dataset, comparing against SAEM baselines.

💡 Research Summary

This paper introduces a variational auto‑encoder (VAE) framework for estimating parameters in nonlinear mixed‑effects (NLME) models that embed ordinary differential equations (ODEs). Traditional inference for NLME‑ODEs relies on the stochastic approximation EM (SAEM) algorithm, which uses Markov chain Monte Carlo (MCMC) to sample subject‑specific random effects. While SAEM works well in moderate dimensions, it becomes computationally burdensome and unstable when data are sparse, irregularly sampled, or when the dimension of the random effects grows, because the MCMC step may mix poorly and get trapped in local modes of a highly non‑convex likelihood surface.

The authors replace the MCMC‑based E‑step with amortized variational inference. They define a joint evidence lower bound (ELBO) on the marginal likelihood of the population parameters ϕ = (θ̄, LΩ, LΣ). For each subject i, a shared neural encoder parameterized by ψ maps the observed longitudinal measurements Y_i to the mean μ_i and Cholesky factor L_i of a Gaussian variational posterior qψ(b_i|Y_i) = N(μ_i, L_iL_iᵀ) over the latent random effects b_i. The decoder takes a sampled b_i, plugs it into the subject‑specific ODE vector field f_{θ_i} (θ_i = g(θ̄, b_i)), solves the ODE with an adaptive solver (Diffrax), and reconstructs the measurements via the observation function h_{θ_i}. The reconstruction term together with the KL divergence between qψ and the prior p(b_i) (which itself has a learnable covariance Ω) yields the ELBO, which is maximized jointly over ϕ and ψ using Adam.

Key methodological contributions include:

1. Amortized inference: A single encoder learns a smooth mapping from sparse, possibly irregular time series to latent random effects, eliminating per‑subject optimization and dramatically reducing computational cost.
2. Learnable prior covariance: By estimating Ω jointly, the model avoids posterior collapse—a common failure mode where the variational posterior collapses to the prior—because the prior adapts to the data.
3. ODE‑conditioned decoder: The decoder is a mechanistic ODE solver, ensuring that latent variables retain a clear biological interpretation and that the learned representations are not arbitrary.
4. Uncertainty quantification: After convergence, the observed Fisher information matrix for ϕ is computed via automatic differentiation; its inverse provides a covariance estimator for population parameters, offering frequentist‑style confidence intervals without resorting to costly MCMC.
5. Identifiability analysis: The authors examine practical identifiability by inspecting the condition number of the Fisher matrix and performing sensitivity analyses, demonstrating that the additional encoder parameters do not compromise the ability to uniquely recover the underlying ODE parameters.

The framework is evaluated on three simulated case studies—pharmacokinetic one‑compartment model, humoral response to vaccination, and TGF‑β activation in asthmatic airways—and on a real dataset of antibody kinetics following COVID‑19 vaccination. Across all scenarios, the VAE‑NLME method achieves estimation accuracy comparable to or better than SAEM, especially when each subject contributes only 3–5 observations. Computationally, the VAE approach on a GPU is 5–10× faster than SAEM run on CPU, and it provides near‑instantaneous posterior estimates for new subjects after training.

Limitations are acknowledged: the current implementation uses a simple Gaussian variational family and restricts the random‑effects dimension to ≤5. Extending to richer families (e.g., normalizing flows or mixture of Gaussians) and scaling to higher‑dimensional random effects are identified as future work. Moreover, encoder architecture choices may need to be tailored to specific domains, and systematic hyper‑parameter tuning remains an open challenge.

In summary, the paper presents a principled, scalable alternative to MCMC‑based SAEM for NLME‑ODE inference. By maximizing the ELBO with an amortized encoder‑decoder architecture, it delivers fast, accurate parameter estimates, robust uncertainty quantification, and maintains practical identifiability, making it a valuable tool for pharmacometrics, immunology, and other fields that rely on mechanistic longitudinal modeling.