Redefining Parameter Estimation and Covariate Selection via Variational Autoencoders: One run is all you need
Jan Rohleff1, Dominic Bräm2, Freya Bachmann1, Dr. Uri Nahum3, Dr. Britta Steffens2, Prof. Dr. Marc Pfister2, Dr. Gilbert Koch2, Prof. Dr. Johannes Schropp1
1Department of Mathematics and Statistics, University of Konstanz, 2Pediatric Pharmacology and Pharmacometrics, University Children's Hospital Basel (UKBB), 3Institute of Biomedical Engineering and Medical Informatics, University of Applied Sciences and Arts
Objectives: Classical software packages for nonlinear mixed-effects (NLME) modeling are either based on linearization techniques (FO, FOCE, La Place) or, more stochastically, on maximizing expectation values (SAEM). Covariate selection with these methods can be laborious, as numerous runs are needed to test and estimate covariate model parameters. Our objective is to provide an alternative by using variational autoencoders (VAEs) [1], a powerful generative machine learning technique. VAEs have shown their potential across various fields. One notable example is AlphaFold [2], which won the Nobel Prize in 2024 for significantly improving our understanding of protein structures. VAEs are highly effective to model complex, uncertain data structures, making them particularly suitable for tasks such as NLME modeling. We apply VAEs to simultaneously estimate the population model parameters and perform covariate selection in one single run. Methods: We developed a VAE framework that is specifically designed to solve NLME problems. The highly flexible VAE structure allows for the additional integration of different modeling tasks, such as covariate selection. The VAE starts with a 'family of models', namely any possible combination of covariates and optimizes an information criterion (e.g., AIC, BIC, or BICc). Covariate selection then simply consists of determining the best model from the chosen model family. As a consequence, maximizing the likelihood and selection of the optimal covariate combinations can be done simultaneously. This contrasts with other methods that first optimize the likelihood for a given covariate model and afterwards compute the information criterion in an iterative process. Using mixed-integer optimization procedures, which are well adapted for such selection problems [3], our method determines the optimal covariate combination by solving a sparse regression problem and both, population parameter and optimal covariates, are determined in one single run. The VAE is implemented in Python. Results: To assess the performance of our VAE, we applied a slightly updated version of the structural model from [4] which was developed to characterize weight progression measurements over the first 7 days of life using 2425 neonates [5]. The structural model consists of 7 parameters and the clinically relevant covariates parity, sex of neonate, gestational age, mother age and delivery mode were applied. Our VAE results were compared to the results from the automated covariate selection methods SAMBA [6], COSSAC and SCM available in Monolix 2024. As previously mentioned our VAE estimates all model parameters and performs the entire covariate selection simultaneously in a single run. Number of runs for SAMBA, COSSAC and SCM were 4, 24, and > 400, respectively. We stopped the SCM at 400 runs and 2.5 days of computation, and excluded it from further comparison. The VAE selected 14 covariate effects, similarly to SAMBA with 14, while COSSAC selected 13. Among all three methods, 7 covariate effects were similar. Information criteria could not be compared among the three methods since VAEs apply linearization techniques whereas the other two methods use importance sampling. Although linearization can be applied in Monolix, covariate selection with COSSAC failed in this setting. Nevertheless, a comparison of goodness-of-fit plots and estimated model parameters suggest a strong agreement of final results. To summarize, the quality of the results achieved by COSSAC, SAMBA, and the VAE are on a comparable level. Conclusion: VAEs are a powerful machine learning tool for solving NLME models efficiently. The VAE-based approach enables the simultaneous estimation of the population model parameters and covariate selection in one single run, i.e., one run is all you need. Hence, VAEs provide a method for automating covariate selection. Furthermore, we established a VAE-NLME framework, which allows the integration of all the machine learning tools available in Python. This admits limitless integration of machine learning into NLME. Consequently, one can gain deeper insights into complex datasets and it opens new opportunities for model design.