Identifiability-aware variational autoencoder for inference of NONLINEAR mixed effect models based on ordinary differential equations

Zhe Li 1,2,3, Mélanie Prague 1,2,3,4, Rodolphe Thiébaut 1,2,3,4, Quentin Clairon 1,2,3

1 Inserm Bordeaux Population Health (Bordeaux, France), 2 INRIA (Bordeaux, France), 3 University of Bordeaux (Bordeaux, France), 4 VRI (Paris, France)

Objectives: We propose NN-NLME, a variational autoencoder framework for identifiability-aware parameter estimation with uncertainty quantification in nonlinear mixed-effects models governed by ordinary differential equations (NLME-ODEs) using longitudinal data from multiple subjects. The marginal likelihood is defined by p(Yi) = ∫p(Yi|bi)p(bi)dbi, where Yi​ represents longitudinal observations per-subject and bi the individual random effects​. NN-NLME aims to bypass identifiability issues encountered by classic SAEM-based ([2]) maximization likelihood methods for complex models, due to MCMC convergence issues and inherent multi-modality of the likelihood in such settings. We estimate parameters by maximizing the evidence lower bound (ELBO), a regularized and smoother surrogate of the marginal likelihood without relying on MCMC sampling. Beyond providing pointwise estimation escaping the mentioned identifiability issues, our method also provides a variance estimator, a feature rarely addressed in deep learning. We evaluate the method on simulated and on a real-world antibody kinetics dataset, comparing against SAEM baselines. The coupling between NLME-ODEs and neural networks has already been investigated [6, 7], but most deep or variational extensions of NLME-ODEs emphasize predictive performance and scalability over statistical interpretability and identifiability. Recent ELBO-based inference methods for NLME-ODEs close to ours have been proposed by [5]. However, they do not address practical identifiability and do not construct a variance estimator. Moreover, they do not provide a rigorous comparison with SAEM-based methods on these points.

Methods: To avoid MCMC step, we approximates the intractable posterior p(bi​∣Yi​) with a Gaussian variational distribution q(bi∣Yi), whose mean and variance are output by a shared encoder network as in [3]. The encoder is designed to be lightweight to limit overfitting and identifiability issues: it uses convolution components for regularly sampled series and a recurrent encoder for irregular sampling data, incorporating time information and padding/masking when sequences have variable length. The decoder corresponds to the mechanistic ODE solution, computed with adaptive solvers and differentiated end-to-end. Optimization is performed via gradient descent (Adam [8]) by maximizing the ELBO summed over subjects, which balances data fit (expected log-likelihood under q and a regularization term (KL divergence between q(bi∣Yi) and the Gaussian random-effects prior). Uncertainty quantification is achieved by computing the observed Fisher Information matrix. This involves approximating the marginal likelihood via Monte Carlo integration using independent samples drawn from the random effects prior distribution, enabling the derivation of standard errors for population parameters. The practical identifiability is assessed by fitting the model from multiple random initializations and checking whether the optimization consistently converges to the same solution. After training, individual random effects bi​ are obtained by a single forward pass of the subject’s observed data through the encoder, rather than via a separate post-hoc optimization step as in standard approaches.

Results: The approach is evaluated on simulation cases of increasing complexity regarding the ODE structure and the number of parameters to estimate as well as a real vaccine trial dataset composed of antibody concentration measurements. In a simple setting, our method reproduced the performance of SAEM based methods with accurate parameters and variance estimators, without adding identifiability issues. In more complex dynamics with irregular and sparse sampling, NN-NLME still provided practically identifiable estimators with more reliable uncertainty quantification contrast to SAEM. On real longitudinal antibody kinetics from a cohort of SARS-CoV-2 naive subjects who received three injections of BNT162b2 in [1, 4] (25 subjects, 2-17 observations per subject collected over 483 days after the first injection), NN-NLME and SAEM provided consistent population estimates with overlapping confidence intervals;in addition, NN-NLME delivered better-calibrated prediction intervals.

Conclusion: The proposed NN-NLME inference is a practical alternative to MCMC-driven SAEM for NLME-ODEs in sparse and complex settings. By optimizing a smoother ELBO objective and sharing an encoder across subjects, the approach preserves practical identifiability as model complexity increases, a crucial point to ensure the interpretability of inferred NLME-ODE model. It also provides a workable route to population-level uncertainty quantification through the observed-information matrix using Monte Carlo approximation and automatic differentiation. The code implemented in Python based on JAX is made publicly available, allowing for the extension to new models of interest. A dedicated software package in python is currently under development.

References:
[1] Planas, D. et al. (2021) Reduced sensitivity of sars-cov-2 variant delta to antibody neutralization, Nature 596(7871).
[2] Lavielle (2014), Mixed Effects Models for the Population Approach: Models, Tasks, Methods and Tools, Chapman & Hall/CRC Biostatistics Series.
[3] Kingma, D. P. & Welling, M. (2014), Auto-encoding variational bayes, ICLR.
[4] Clairon, Q. et al. (2023), Modeling the kinetics of the neutralizing antibody response against sars-cov-2 variants after several administrations of bnt162b2, PLoS Computational Biology
[5]Rohleff, J et al (2025). Redefining parameter estimation and covariate selection via variational autoencoders: One run is all you need. CPT: Pharmacometrics & Systems Pharmacology 14, 2232–2243.
[6] Roeder, G. et al (2019). Efficient amortised bayesian inference for hierarchical and nonlinear dynamical systems. ICML.
[7] Arruda, J. et al (2024). An amortized approach to non-linear mixed-effects modeling based on neural posterior estimation. In Proceedings of the 41st ICML.
[8] Kingma, D. P. and Ba, J. (2015). Adam: A method for stochastic optimization. ICLR.

Reference: PAGE 34 (2026) Abstr 12024 [www.page-meeting.org/?abstract=12024]

Poster: Oral: Methodology - New Tools