PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural ODE enhanced linear mixed effect models for estimating complex association patterns of time-varying covariates with the marker trajectory

Zhe Aurore Li, Quentin Clairon, Cécilia Samieri, Rodolphe Thiébaut, Mélanie Prague, Cécile Proust-Lima

arXiv:2608.29714v1stat.MLcs.LG

TL;DR

Longitudinal LMMs support inference for irregular, partially observed repeated data but require pre-specified covariate–outcome forms. The Neural ODE-LMM learns continuous-time trajectory effects within the LMM framework and uses counterfactual contrasts for interpretation; simulations recover instantaneous and cumulative effects, while the 3C analysis finds trajectory-dependent associations with cognitive decline.

  • Problem

    LMMs require analysts to pre-specify how exposure histories relate to outcomes, limiting flexible modelling of complex and potentially cumulative time-varying associations.

  • Method

    The Neural ODE-LMM encodes covariate trajectories into a continuous-time latent state that drives fixed- and random-effect designs, with grouped regularisation encouraging cumulative representations.

  • Results

    The model recovered instantaneous and cumulative effects in simulations and detected trajectory-dependent associations between BMI trajectory shape and cognitive decline in the 3C cohort.

  • Takeaways & Limitations

    Counterfactual prediction contrasts provide interpretable association measures for complex covariate trajectories while retaining the LMM observation framework.

  • Takeaways & Limitations

    The method treats covariates as measured without error and uses within-subject linear interpolation, while delta-method variance estimation was demonstrated for architectures with fewer than 2,000 parameters.

Abstract

from arXiv · show

Longitudinal cohort studies produce repeated data that enable the assessment of time-varying association patterns between exposures and health outcomes. Classical linear mixed-effects models (LMMs) can accommodate a large variety of association patterns while accounting for the irregularly spaced, partially observed measurement. But they require the analyst to pre-specify the functional form linking the exposure history to the outcome. We propose the Neural ODE-LMM, which embeds a Neural Ordinary Differential Equation (Neural ODE) within the linear mixed-effects framework: a learned vector field encodes covariate trajectories into a continuous-time latent state that drives both the fixed- and random-effect design, while preserving the standard LMM observation model. This retains classical likelihood-based inference while learning complex, potentially cumulative, covariate effects flexibly. All parameters are estimated by maximising a penalised marginal likelihood. To quantify covariate effects, we introduce contrasts of counterfactual predictions that compare the expected outcome under alternative covariate trajectories with variance estimated via the delta method. In simulations, the model recovers both instantaneous and cumulative-burden effects without prior specification of the functional form. Applied to the Trois-Cités (3C) cohort, a population-based study of 7{,}324 participants, the method reveals trajectory-dependent associations of BMI and fasting glucose with cognitive decline.

Fundings

The work was supported by Inserm’s Exposome Booster Program, project ANR EDyLES, and the University of Bordeaux’s France 2030 program/RRI PHDS. It addresses flexible and interpretable modelling of time-varying covariate associations in longitudinal epidemiology.

  • Support came from Inserm through the Exposome Booster Program, project ANR EDyLES, and the University of Bordeaux’s France 2030 program/RRI PHDS.
  • Longitudinal epidemiology requires modelling repeated, irregularly spaced, and partially missing measurements of time-varying exposures and health outcomes.
  • Classical LMMs provide interpretable population-level associations but require analysts to specify covariate–outcome functional forms in advance.
  • The Neural ODE-LMM encodes covariate trajectories into a continuous-time latent state driving both fixed- and random-effect designs within an LMM observation model.
  • A grouped skip-connection penalty encourages covariate effects to be represented through the ODE pathway rather than only through instantaneous inputs.

3 Interpretability and inference

The paper quantifies covariate associations by comparing counterfactual predictions under alternative covariate trajectories. These contrasts can represent time-varying effects and are accompanied by delta-method uncertainty estimates.

  • Counterfactual contrasts compare model predictions under alternative covariate trajectories while holding other covariates at their observed values.
  • Population-level counterfactual means marginalise predictions over the joint distribution of remaining covariates, following the partial-dependence principle.
  • Contrasts between intervention paths can use constant or time-varying trajectories, enabling detection of trajectory-dependent effects.
  • The contrast is evaluated over a time grid spanning the follow-up window.
  • The delta method estimates contrast variance from the gradient and regularised empirical Fisher information, supporting pointwise 95% confidence intervals.

4 Simulation study

The simulation evaluates whether Neural ODE-LMM recovers both instantaneous and cumulative BMI effects while retaining valid uncertainty quantification. It uses realistic 3C-based data-generating settings and counterfactual trajectory contrasts.

  • Study design: The study compares recovery of instantaneous BMI–age associations with cumulative BMI-burden effects across two simulation regimes.The cumulative regime is unavailable to a standard LMM unless the correct accumulation functional is specified in advance.
  • Study design: 100 replications simulate BMI and cognitive outcomes using parameters and follow-up structures calibrated from the 3C cohort.The design retains observed covariates and irregular follow-up windows to produce realistic within-subject variability.
  • Scenarios: S1 uses an instantaneous age-modulated BMI association with βBMI = −0.30 and βint = −0.05, providing a calibration benchmark for inferential validity.The correctly specified LMM serves as the benchmark when the true association is linear.
  • Scenarios: S2 specifies a cumulative BMI effect with no instantaneous BMI effect or age interaction, which standard LMMs cannot recover without an explicit cumulative definition.The Neural ODE-LMM routes BMI through the ODE pathway and uses group-lasso regularization to encourage cumulative representation learning.
  • Estimand: The primary estimand compares population-averaged predictions under BMI values at the 25th and 75th observed percentiles across follow-up time.Variance is reported alongside the contrast on a time grid.
  • Results: In S1, relative bias is approximately 3–6% with correct 95% confidence-interval coverage at every time point.Bias is larger than with the oracle LMM, as expected, while remaining bounded when the true linear structure is representable.
  • Results: In S2, estimates track the true contrast from −0.50 at t = 2 to −2.50 at t = 10 with approximately 3–5% relative bias and 0.92–0.98 coverage.The delta-method variance closely matches empirical variance at later times, and late decline produces worse outcomes than late spike despite equal terminal BMI.

5 Application to cardiometabolic health and cognitive decline

The 3C application examines cardiometabolic trajectories and cognitive decline using Neural ODE-LMM alongside classical LMM benchmarks. The Neural ODE-LMM fits the observed cognitive trajectory closely and identifies BMI associations that depend on exposure history.

  • Data: The 3C cohort enrolled 9,294 older adults, and the analytical sample contained 7,324 participants with 2–7 visits.The study includes BMI, blood pressure, fasting glucose, and HDL cholesterol as time-varying cardiometabolic covariates.
  • Models: Classical LMM benchmarks use natural cubic splines for time and progressively add time-varying covariates and their interactions.All benchmark models share a random intercept and random slopes on a two-degree-of-freedom time spline.
  • Models: The Neural ODE vector field receives five time-varying covariates and their observation masks, while static covariates enter both encoder and skip pathways.Group lasso regulates the dynamic-covariate skip pathway and encourages ODE routing when cumulative representations improve fit.
  • Goodness of fit: The Neural ODE-LMM achieves slightly higher marginal log-likelihood on training and test data with comparable subject-level MSE.Its population-averaged predictions track the observed mean closely, whereas the classical LMM underestimates early time points.
  • Counterfactual analysis: Counterfactual analyses compare six BMI, glucose, and HDL trajectory profiles defined by the 25th and 75th covariates’ percentiles.Profiles include stable, late-spike, late-decline, gradual-rise, and gradual-decline patterns.
  • BMI results: Gradually rising BMI trajectories have the highest predicted cognitive scores, while declining trajectories show the steepest cognitive decline.Late decline is worse than late spike at later visits despite both profiles crossing the same BMI value at midpoint; the classical LMM instead responds to current BMI and reverses their ordering.

6 Discussion

The discussion presents Neural ODE-LMM as a continuous-time extension of the LMM that learns trajectory associations while retaining likelihood-based structure. Simulations and the 3C application support recovery of cumulative and trajectory-dependent effects, but uncertainty quantification remains constrained by measurement and model complexity.

  • Contribution: Neural ODE-LMM encodes covariate trajectories into a continuous-time latent state that drives fixed- and random-effect designs within the LMM framework.The approach preserves the LMM observation model and its statistical structure.
  • Inference: Counterfactual prediction contrasts provide interpretable association measures, with delta-method variance based on a shrinkage-stabilized inverse empirical Fisher information.The contrast compares population-averaged predictions under alternative covariate paths.
  • Findings: Simulations recover both instantaneous and cumulative associations with negligible bias and well-calibrated confidence intervals.The 3C analysis additionally identifies trajectory-dependent BMI associations and distinct temporal patterns for fasting glucose and HDL cholesterol.
  • Findings: The method recapitulates cumulative BMI association patterns without pre-specifying a parametric trajectory effect and supplies delta-method confidence intervals.The interpretation is linked to prior evidence that BMI decline in older populations can mark preclinical neurodegeneration.
  • Limitations: Measurement-error uncertainty is not modeled: interpolated covariates are treated as observed, and errors may accumulate through long ODE integration paths.The delta-method variance was demonstrated for architectures with fewer than 2,000 parameters, while larger-network uncertainty quantification remains open.

7 Data and code availability

The paper provides access conditions for anonymized 3C data and makes scripts for reproducing the application and simulations available online.

  • Data availability: Anonymized 3C data may be shared upon a reasonably justified request to the 3C scientific committee.The request contact is e3C.CoordinatingCenter@u-bordeaux.fr.
  • Code availability: Scripts reproducing the application and simulations, with documentation, are available at the NodeLmm GitHub repository.The models were trained with PyTorch 2.6.0 and Python 3.10+ without GPU acceleration.

Supplementary Materials

The simulations were calibrated to the 3C cohort and evaluated instantaneous versus cumulative BMI effects using repeated longitudinal data. Neural ODE-LMM contrasts were assessed for bias, variance estimation, and confidence-interval coverage across replicated datasets.

  • Simulation design: The reference outcome model used IST scores with spline time effects, baseline covariates, time-varying BMI, random intercepts and time slopes, and residual error.Parameters and spline knots were calibrated from a preliminary LMM fit to the 3C data.
  • Scenario S1: Scenario S1 represented an instantaneous linear BMI effect through BMI and BMI × centred-age fixed-effect terms, with no history function.The oracle specification was compared with a spline × age interaction specification to assess robustness to mild misspecification.
  • Scenario S2: Scenario S2 represented cumulative BMI burden through the time-integrated exposure history, which cannot be recovered from contemporaneous BMI alone.Under constant BMI, the cumulative component becomes αvt, producing a contrast whose magnitude increases with follow-up time.
  • Simulation design: 100 independent datasets were generated for each simulation scenario using 3C-like marginal distributions, visit schedules, and variance components.Covariate trajectories were drawn from auxiliary mixed-effects models, and simulated visits followed subject-specific follow-up windows.

S2 Trois-Cit´es cohort: data description

The 3C cohort analysis used an IST-based longitudinal LMM framework with spline time effects, baseline covariates, time-varying cardiometabolic factors, and subject-specific random effects. Five nested classical LMM specifications differed in their fixed-effect means and were compared using BIC.

  • Classical LMM specifications: The five classical LMMs differed only in the fixed-effect mean, progressively adding time-varying covariates and interactions with time, age, or BMI.Time-varying terms included BMI, systolic and diastolic blood pressure, fasting glucose, and HDL cholesterol.
  • Model selection: Model TV3 was selected by BIC among the five classical LMM specifications.The comparison used the 5,859-participant analytic sample.
  • Neural ODE-LMM implementation: The Neural ODE-LMM architecture and training hyperparameters for the cohort analysis were documented separately after a two-tier selection procedure.The supplementary architecture table reports the selected configuration details.

S5 Model selection results

The cohort architecture was selected with subject-level held-out validation and a two-tier search over regularisation, skip pathways, and decoder capacity. Group lasso was preferred because it supported data-driven covariate routing, while the selected second-tier model used a larger latent dimension.

  • Selection protocol: The held-out cohort selection used 4,687 training, 1,172 validation, and 1,465 test subjects, with early stopping based on validation marginal NLL.The test set was used once for unbiased performance reporting, and standardisation statistics came from training data only.
  • Search grid: Tier 1 varied regularisation, penalty weight, and dynamic skip-path inclusion before Tier 2 varied latent and fixed-effect output dimensions.Tier 2 fixed the winning Tier-1 regularisation and skip choices while sweeping decoder capacity.
  • Tier 1 selection: Config 5 with group lasso λGL = 0.1 achieved NLL = 15,841 versus NLL = 15,843 for the comparable unregularised Config 2 and was preferred for retaining the skip pathway.Near-zero skip-weight norms provided empirical evidence that covariate effects were mediated through the ODE latent state.
  • Tier 2 selection: Config 7 with latent dimension d = 16 and fixed-effect output dimension p = 2 was selected in Tier 2 because it had the lowest NLL.Tier 2 validation NLLs differed by at most 8 units across 1,172 validation subjects.
  • Simulation selection: In simulations, skip-only routing won 9/10 replicates for instantaneous effects, while group lasso won 8/10 for cumulative effects; unregularised models were never selected.These results matched the intended distinction between instantaneous and cumulative covariate pathways.

S7 Stationarity and Fisher diagnostics

Stationarity and Fisher diagnostics assessed whether the empirical Fisher supported delta-method variance estimation for counterfactual contrasts. The reported diagnostics indicated small mean scores, no detected null-space contribution, and minimal shrinkage adjustment.

  • 3C cohort diagnostics: For the selected 3C Neural ODE-LMM, the mean penalised score norm was 1.45, the typical individual score norm was 41.3, and their ratio was 0.035.The ratio was used to assess the first-order stationarity condition required by the empirical Fisher variance estimator.
  • Stationarity diagnostics: Stationarity ratios below 0.05 indicated that the mean penalised score was small relative to individual score magnitudes.This supported the first-order condition underlying the empirical Fisher variance estimator.
  • Fisher diagnostics: The null-space fraction of the counterfactual contrast gradient was exactly zero at every evaluation time.The diagnostic projected the gradient onto Fisher eigenvectors below the relative threshold λk < 10^-10λmax.
  • Fisher regularisation: Ledoit–Wolf shrinkage intensities were 0.0009 ± 0.0001 for S2 and 0.0030 ± 0.0005 for S1.The shrinkage target contributed less than 0.3% to the regularised estimator.
  • Subject-level predictions: BLUP-based subject-level trajectory predictions were displayed for 25 randomly selected participants from both training and test sets.These figures provided subject-level prediction views complementary to the aggregate diagnostics.

S9 Glucose counterfactual predictions

The glucose counterfactual analysis shows a cumulative association with cognition: trajectory history matters beyond terminal glucose values, and higher glucose profiles are associated with worse predicted cognition.

  • Trajectory-profile counterfactual predictions: Sustained low-glucose profiles yield the highest predicted cognitive scores, whereas sustained high-glucose profiles yield the lowest.The late-spike profile remains above the stable-high profile late in follow-up despite both reaching the same terminal glucose value.
  • Trajectory-profile counterfactual predictions: Approximately 5 IST points separate the most favorable and unfavorable glucose trajectory profiles.
  • Pairwise contrasts: Under the classical LMM, contrasts depend only on current glucose, becoming zero when profiles converge; Neural ODE-LMM contrasts persist after convergence.
  • Pairwise contrasts: All four Neural ODE-LMM glucose contrasts are significantly negative throughout follow-up, with confidence intervals excluding zero.
  • Cumulative versus baseline-and-current diagnostic: Temporary glucose perturbations produce non-zero endpoint differences, confirming cumulative exposure encoding, although effects are smaller than for BMI.The spike and dip comparisons yield Δt=12 values of −0.47 and −0.28, respectively.
  • Skip-connection diagnostic: All five time-varying covariate skip pathways have norms below 10^-3, so trajectory-profile effects are routed almost entirely through the ODE pathway.The latent-state decoder columns have Frobenius norm 6.27 under group-lasso penalty λGL = 0.1.

S14 Latent state visualisation

The latent-state visualisation evaluates how counterfactual covariate trajectories are represented over time. Divergence analyses show that the ODE state retains trajectory history, with responsive latent dimensions differing across BMI and glucose.

  • Counterfactual latent trajectories: Counterfactual profiles span Q25–Q75 and include stable, spiking, declining, and gradual trajectory shapes.Each target covariate is replaced across subjects, the ODE is re-integrated, and population-averaged latent trajectories are computed.
  • Latent state divergence: Profiles with identical endpoints but different paths produce non-zero latent-state divergence, indicating sensitivity to trajectory shape.
  • Latent state divergence: Profiles sharing the same current value after crossover show the largest divergence, consistent with accumulated historical information in the ODE state.
  • Dimension-level latent responses: BMI-responsive latent dimensions include Λ9 and Λ3, while glucose-responsive dimensions include Λ14 and Λ9.For BMI, Λ14 and Λ7 show more moderate differentiation; the glucose results are reported for the four most responsive dimensions.
Loading 2608.29714v1…