self-contained-math: true
anchor-sections: true
theme: default
toc: true
toc-depth: 4
toc-expand: 2
toc-location: left
toc-title: Contents
number-sections: true
code-summary: Show/Hide Code
fig-format: svg
fig-width: 8
fig-height: 6
license: CC BY-SA 4.0
---
:::{.callout-warning}
## Experimental Feature
VEM is an **experimental** estimation method introduced in Pumas 2.8.
Its API may change in future releases without deprecation.
:::
In this tutorial,
we introduce **Variational Expectation Maximization (VEM)**,
a new estimation method available in `Pumas.Experimental`.
We will cover the conceptual background,
fit a warfarin PK model with VEM,
compare the results to FOCE,
and provide a comprehensive reference of VEM options.
```{julia}
using Pumas
using Pumas.Experimental: VEM
using PumasUtilities
using PharmaDatasets
using DataFramesMeta
using CairoMakie
using AlgebraOfGraphics
using SummaryTables
using Random: default_rng, seed!
seed!(123)
```
## Conceptual Background
### The Marginal Likelihood Challenge
In nonlinear mixed-effects (NLME) modeling,
we want to estimate population parameters while accounting for between-subject variability through random effects.
The marginal likelihood requires integrating over all possible random effect values for every subject,
and these integrals are generally intractable.
Different estimation methods handle this challenge in different ways:
- **FOCE/LaplaceI** approximate the integral using a second-order Taylor expansion around the mode of the random effects.
- **SAEM** uses Markov Chain Monte Carlo (MCMC) sampling within an Expectation-Maximization framework.
- **VEM** maximizes a lower bound on the log-likelihood using variational inference.
### The Evidence Lower Bound (ELBO)
VEM works by maximizing the **Evidence Lower Bound (ELBO)**,
which is a lower bound on the log marginal likelihood.
The ELBO equals the log-likelihood minus a KL divergence term that measures how far the approximate posterior is from the true posterior.
A higher ELBO typically means a better fit.
Intuitively,
VEM simultaneously finds:
1. **Population parameters** (fixed effects, variance components, residual error), and
2. **An approximation to each subject's random effect posterior distribution**.
This joint optimization is performed in a single level,
unlike FOCE/LaplaceI which alternates between inner (random effects) and outer (population parameters) optimization loops.
One key advantage of VEM is that it can handle any response data distribution, including discrete, censored,
truncated, time-to-event and even custom distributions. FOCE by comparison is limited to a restricted set of distributions, and while LaplaceI can handle any response distribution, it can be very slow for complex models.
### Presampling vs. Resampling
VEM has two optimization variants controlled by the `presample` keyword:
| | Presampling (`presample=true`, default) | Resampling (`presample=false`) |
|:------------- |:----------------------------------------------------- |:------------------------------ |
| **Optimizer** | L-BFGS (deterministic) | Adam (stochastic) |
| **Phases** | 1 (default) | 10 (default) |
| **Samples** | Pre-drawn at start of phase, reused across iterations | Re-drawn at every iteration |
In the resampling variant, the number of samples increases (by default, unless overridden) and the learning rate of Adam decreases from one phase to the next.
### Diagonal vs. Full Posterior Covariance
VEM approximates each subject's random effect posterior with a multivariate normal distribution.
The `diagonal` keyword controls the structure of this approximation:
- **`diagonal=true`** (default): assumes independent random effects in the posterior. Faster, fewer parameters to optimize.
- **`diagonal=false`**: captures correlations between random effects in the posterior. Slower but may be more accurate when random effects are correlated.
## Warfarin PKPD Model with VEM
We will fit an oral PKPD model to the warfarin dataset, first using VEM and then comparing with FOCE.
### Model Definition
```{julia}
warfarin_model = @model begin
@param begin
# PK
pop_CL ∈ RealDomain(lower = 0.0, init = 0.134)
pop_V ∈ RealDomain(lower = 0.0, init = 8.11)
pop_tabs ∈ RealDomain(lower = 0.0, init = 0.523)
pop_lag ∈ RealDomain(lower = 0.0, init = 0.1)
pk_Ω ∈ PSDDomain([0.01, 0.01, 0.01])
σ_prop ∈ RealDomain(lower = 0.0, init = 0.00752)
σ_add ∈ RealDomain(lower = 0.0, init = 0.0661)
# PD
pop_e0 ∈ RealDomain(lower = 0.0, init = 100.0)
pop_emax ∈ RealDomain(init = -1.0)
pop_c50 ∈ RealDomain(lower = 0.0, init = 1.0)
pop_tover ∈ RealDomain(lower = 0.0, init = 14.0)
pd_Ω ∈ PSDDomain([0.01, 0.01, 0.01])
σ_fx ∈ RealDomain(lower = 0.0, init = 0.01)
end
@random begin
pk_η ~ MvNormal(pk_Ω)
pd_η ~ MvNormal(pd_Ω)
end
@covariates begin
FSZV
FSZCL
end
@pre begin
# PK
CL = FSZCL * pop_CL * exp(pk_η[1])
Vc = FSZV * pop_V * exp(pk_η[2])
tabs = pop_tabs * exp(pk_η[3])
Ka = log(2) / tabs
# PD
e0 = pop_e0 * exp(pd_η[1])
emax = pop_emax
c50 = pop_c50 * exp(pd_η[2])
tover = pop_tover * exp(pd_η[3])
kout = log(2) / tover
rin = e0 * kout
end
@dosecontrol begin
lags = (Depot = pop_lag,)
end
@init begin
Turnover = e0
end
@vars begin
cp := Central / Vc
ratein := Ka * Depot
pd := 1 + emax * cp / (c50 + cp)
end
@dynamics begin
Depot' = -ratein
Central' = ratein - CL * cp
Turnover' = rin * pd - kout * Turnover
end
@derived begin
conc ~ @. CombinedNormal(cp, σ_add, σ_prop)
pca ~ @. Normal(Turnover, σ_fx)
end
end
```
:::{.callout-tip}
## Same Model for VEM and FOCE
The `@model` definition is identical regardless of whether you fit with VEM or FOCE/LaplaceI.
VEM works with standard `@model` definitions -- no special model macro is required.
:::
:::{.callout-note}
## Random effect on `emax`
Note that no random effect is used on `emax`. We tried adding a random effect on `emax` and estimating the model parameters. However, that resulted in the row and column of `pd_Ω` corresponding to that random effect to be all close to 0 which caused numerical issues for the asymptotic esimtator of the confidence interval. So no random effect is used on `emax` here for simplicity.
:::
### Data Loading
```{julia}
warfarin_data = @chain dataset("pumas/warfarin_pumas") begin
@rtransform begin
:FSZV = :wtbl / 70
:FSZCL = (:wtbl / 70)^0.75
end
end
first(warfarin_data, 5)
```
```{julia}
pop = read_pumas(
warfarin_data;
id = :id,
time = :time,
amt = :amt,
cmt = :cmt,
evid = :evid,
covariates = [:FSZV, :FSZCL],
observations = [:conc, :pca],
)
```
### Fitting with Default VEM
Fitting with VEM follows the same `fit` interface as all other estimation methods.
Simply pass `VEM()` as the estimation algorithm:
```{julia}
#| output: false
#| echo: false
# warm-up
@time fpm_vem = fit(
warfarin_model,
pop,
init_params(warfarin_model),
VEM(; optim_options = (; iterations = 40, show_every = 20)),
)
```
```{julia}
@time fpm_vem = fit(
warfarin_model,
pop,
init_params(warfarin_model),
VEM(; optim_options = (; show_every = 20)),
)
```
During fitting,
VEM displays a convergence trace showing the negative ELBO (the objective being minimized).
The objective value is the negative of the lower bound on the log marginal likelihood,
so lower values indicate a better fit.
### Extracting Results
Population parameter estimates can be extracted with `coef`:
```{julia}
simple_table(coeftable(fpm_vem))
```
The posterior approximation modes of the random effects for each subject are available via `empirical_bayes`:
```{julia}
ebes = empirical_bayes(fpm_vem)
first(ebes, 3)
```
Note that these generally may not correspond to modes of the true posterior.
### Log-Likelihood Computation
VEM does not compute the log-likelihood during fitting.
To evaluate the log-likelihood after fitting,
you must specify an approximation method:
```{julia}
ll_foce = loglikelihood(fpm_vem, FOCE())
```
:::{.callout-warning}
## Log-Likelihood Approximation Starting Points
When using `FOCE()` (or `LaplaceI()`) for log-likelihood evaluation,
the empirical Bayes estimates (EBEs) are computed for each subject
starting from the mode of the prior
(i.e., zero random effects in the Gaussian case)
rather than the VEM posterior approximation's mode.
For well-behaved models,
this is unlikely to cause a significant difference in the log-likelihood approximation.
However, this behavior will likely change in a future release
to use the VEM posterior approximation's mode as the starting point
for the Laplace approximation.
:::
### Diagnostics
Model diagnostics are available through `inspect`.
For models with Gaussian error models (like our proportional normal error),
you must explicitly specify how to compute weighted residuals using `wres_approx`:
```{julia}
res = inspect(fpm_vem; wres_approx = FOCE())
```
:::{.callout-note}
## Weighted Residuals Linearization Point
When computing the weights of the weighted residuals, the linearization of the mean and variance are performed at the mode of the approximate posterior for each subject, learned by VEM, rather than the mode of the true posterior.
The predictions used for the residuals are also computed using the mode of the VEM posterior approximation, which may differ from the EBEs or the modes of the true posterior.
:::
Goodness-of-fit plots work as other estimation methods in Pumas:
```{julia}
goodness_of_fit(res, observations = [:conc])
```
```{julia}
goodness_of_fit(res, observations = [:pca])
```
:::{.callout-warning}
## `wres_approx` Required for Gaussian Errors
When using VEM with Gaussian-like error models (`Normal`, `ProportionalNormal`, etc.),
you must pass `wres_approx = FOCE()` or `wres_approx = LaplaceI()` to `inspect`.
Omitting it will raise an error.
For other error models (Beta, Gamma, Poisson, Bernoulli, TimeToEvent), this is not needed.
:::
### Predictions and VPCs
Predictions and visual predictive checks work as other estimation methods in Pumas:
```{julia}
preds = predict(fpm_vem)
first(preds, 3)
```
:::{.callout-note}
## VEM Posterior Approximation's Mode in Predictions
The individual predictions are computed using the mode of the VEM posterior approximation for each subject, which may differ from the EBEs or the modes of the true posterior.
:::
```{julia}
#| output: false
conc_vpc_res = vpc(fpm_vem, observations = [:conc]);
```
```{julia}
vpc_plot(conc_vpc_res)
```
```{julia}
#| output: false
pca_vpc_res = vpc(fpm_vem, observations = [:pca]);
```
```{julia}
vpc_plot(pca_vpc_res)
```
### Inference Workaround
VEM does not currently support `infer` for computing standard errors and confidence intervals:
```{julia}
#| error: true
infer(fpm_vem)
```
The recommended workaround is to use the VEM estimates as a starting point for a FOCE/LaplaceI fit,
then run inference on the FOCE/LaplaceI result.
This typically converges very quickly since VEM provides excellent initial values:
```{julia}
#| output: false
#| echo: false
# warm-up
@time fpm_foce_from_vem = fit(
warfarin_model,
pop,
coef(fpm_vem),
FOCE();
init_randeffs = empirical_bayes(fpm_vem),
optim_options = (; iterations = 40, show_every = 20),
)
```
```{julia}
#| warning: false
@time fpm_foce_from_vem = fit(
warfarin_model,
pop,
coef(fpm_vem),
FOCE();
init_randeffs = empirical_bayes(fpm_vem),
optim_options = (; show_every = 20),
)
```
```{julia}
loglikelihood(fpm_foce_from_vem)
```
```{julia}
infer(fpm_foce_from_vem)
```
:::{.callout-note}
## True Posterior Mode vs Posterior Approximation Mode
Initializing a FOCE fit with VEM estimates and running for 0 or more iterations is also a way to get a fitted Pumas model with standard EBEs (true posterior mode, instead of the VEM posterior approximation's mode) that can be used for downstream functions, such as `predict`, `inspect` and `loglikelihood`.
:::
:::{.callout-tip}
## VEM as an Initializer
Using VEM to initialize FOCE is a powerful workflow.
VEM provides good population parameter estimates _and_ per-subject random effect posteriors,
which together allow FOCE to converge quickly and reliably.
:::
## Comparing VEM and FOCE
### Direct FOCE Fit
Let us also fit the model directly with FOCE from default initial values for comparison:
```{julia}
#| warning: false
@time fpm_foce_direct = fit(
warfarin_model,
pop,
init_params(warfarin_model),
FOCE();
optim_options = (; show_every = 20),
)
```
```{julia}
loglikelihood(fpm_foce_direct)
```
### Compare Estimates
```{julia}
compare_estimates(; VEM = fpm_vem, VEM_FOCE = fpm_foce_from_vem, FOCE = fpm_foce_direct)
```
### Compare Log-Likelihoods
```{julia}
DataFrame(
Method = [
"Deterministic VEM",
"FOCE initialized from deterministic VEM",
"FOCE direct",
],
LogLikelihood = [
loglikelihood(fpm_vem, FOCE()),
loglikelihood(fpm_foce_from_vem),
loglikelihood(fpm_foce_direct),
],
) |> simple_table
```
## VEM Options Reference
This section provides a comprehensive reference for all VEM keyword arguments.
All options below are keyword arguments to `VEM()` unless otherwise noted.
### Presampling vs. Resampling (`presample`)
The `presample` keyword controls which VEM variant is used.
Many other defaults change depending on this choice:
| | Option | Presampling (`presample=true`, default) | Resampling (`presample=false`) |
|:-------------------- |:-------------------- |:--------------------------------------- |:---------------------------------- |
| **Optimizer** | `optim_alg` | L-BFGS (deterministic) | Adam (stochastic) |
| **Phases** | `nphases` | 1 | 10 |
| **Initial samples** | `init_nsamples` | 15 | 3 |
| **Sample increment** | `nsamples_increment` | half the initial samples (7 by default) | the initial samples (3 by default) |
| **Max samples** | `max_nsamples` | Unlimited | 25 |
Presampling draws samples once per phase and reuses them across iterations,
which allows a deterministic optimizer (L-BFGS) to be used.
Resampling draws fresh samples at every iteration,
which requires a stochastic optimizer (Adam)
but can be more appropriate for DeepPumas models with neural networks.
However, generally the stochastic optimizer needs many more iterations to converge than the deterministic one, so presampling is a good default for most models.
Note that all of the above options can be overridden regardless of the `presample` setting.
```{julia}
# Resampling VEM
@time fpm_vem_resample = fit(
warfarin_model,
pop,
init_params(warfarin_model),
VEM(;
presample = false,
nphases = 5,
optim_options = (; iterations = 1000, show_every = 20),
),
)
```
```{julia}
loglikelihood(fpm_vem_resample, FOCE())
```
The number of iterations and phases have been restricted in the above example for demonstration purposes.
:::{.callout-note}
## Reading the VEM Convergence Trace
When presampling is enabled, the trace is deterministic and should show steady convergence. The trace printed during fitting has three columns:
| Column | Meaning |
|:------------------ |:--------------------------------------------------------------------------------------------------------------------- |
| **Iter** | Current iteration within the phase. |
| **Function value** | The **negative ELBO** — the objective being minimized. Lower values indicate a better fit. |
| **Gradient norm** | The **gradient norm** — how steep the objective landscape is. Decreasing values indicate the optimizer is converging. |
When resampling is enabled, the trace is stochastic and may show more fluctuation. In the stochastic case, the trace printed during fitting has five columns:
| Column | Meaning |
|:--------- |:----------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| **phase** | Current phase / total phases. Each phase increases the sample count. |
| **iter** | Current iteration / max iterations within this phase. |
| **obj** | The **negative ELBO** — the objective being minimized. Lower values indicate a better fit. |
| **∇obj** | The **gradient norm** — how steep the objective landscape is. Decreasing values indicate the optimizer is converging. |
| **s∇obj** | The **scaled gradient norm** — gradient norm relative to the first iteration of the current phase (normalized to 1.0). Shows relative convergence progress within each phase. |
A phase may terminate early with "Terminating phase due to sgtol"
when the scaled gradient norm drops below the stopping tolerance (0.05), triggering the next phase, typically with an increase in the sample size and a decrease in the learning rate.
For stochastic optimizers like Adam (resampling VEM),
some fluctuation in `∇obj` is normal
because the gradients are estimated from random samples.
:::
```{julia}
compare_estimates(;
VEM_presample = fpm_vem,
VEM_resample = fpm_vem_resample,
VEM_FOCE = fpm_foce_from_vem,
FOCE = fpm_foce_direct,
)
```
```{julia}
DataFrame(
Method = [
"Deterministic VEM",
"Stochastic VEM",
"FOCE initialized from deterministic VEM",
"FOCE direct",
],
LogLikelihood = [
loglikelihood(fpm_vem, FOCE()),
loglikelihood(fpm_vem_resample, FOCE()),
loglikelihood(fpm_foce_from_vem),
loglikelihood(fpm_foce_direct),
],
) |> simple_table
```
Note that stochastic VEM took much longer to converge and it only got close to the best parameters. Deterministic VEM will generally converge faster. Stochastic VEM is mostly useful for models with neural networks and large populations where `subject_ratio < 1` (to be demonstrated below) can be used to get quick estimates by running the algorithm for a few iterations only.
### Sample Size Control
VEM draws samples from a base distribution to estimate the ELBO gradient.
More samples improve accuracy but increase computation time.
When `nsamples` is set to a nonzero value,
it overrides `init_nsamples` and sets `nsamples_increment` to 0,
so the sample count stays fixed across all phases:
```{julia}
# Fixed 30 samples per subject across all phases
@time fpm_vem_nsamples = fit(
warfarin_model,
pop,
init_params(warfarin_model),
VEM(; nsamples = 30, optim_options = (; show_every = 20)),
)
```
```{julia}
loglikelihood(fpm_vem_nsamples, FOCE())
```
Notice that increasing the number of samples improved the quality of the ELBO however that does not guarantee an improvement in the FOCE-approximated log likelihood of the solution found. In this case, it improved it. Heuristically, setting `nsamples` to 15-30 seems to be a good starting point for most models, but the best value can vary based on the number of random effects in the model.
For multi-phase runs, the number of samples in phase `p` is:
`min(init_nsamples + (p - 1) * nsamples_increment, max_nsamples)`.
More samples in later phases refine the estimate as the parameters approach their final values.
```julia
# Resampling with more phases and custom sample progression
fit(
model,
pop,
params,
VEM(;
presample = false,
nphases = 15,
init_nsamples = 5,
nsamples_increment = 5,
max_nsamples = 50,
),
)
```
### Optimizer Options (`optim_options`)
Fine-grained control over the optimizer is passed through the `optim_options` named tuple.
The available options depend on which variant is used.
For **presampling** (L-BFGS):
| Option | Default | Description |
|:------------ |:------- |:----------------------------------- |
| `iterations` | 3000 | Maximum iterations per phase |
| `show_trace` | `true` | Print convergence information |
| `show_every` | 5 | Print every N iterations |
| `f_reltol` | 1e-5 | Relative tolerance on the objective |
| `g_tol` | 1e-3 | Gradient norm tolerance |
```julia
fit(model, pop, params, VEM(; optim_options = (; iterations = 5000, show_every = 10)))
```
For **resampling** (Adam):
| Option | Default | Description |
|:--------------- |:------- |:-------------------------------------------------- |
| `iterations` | 20000 | Maximum iterations per phase |
| `show_trace` | `true` | Print convergence information |
| `show_every` | 10 | Print every N iterations |
| `subject_ratio` | 1.0 | Fraction of subjects per iteration (mini-batching) |
Setting `subject_ratio < 1.0` enables **mini-batching**,
where only a random subset of subjects is used per iteration.
This can speed up each iteration for large datasets at the cost of noisier gradients:
```julia
fit(
model,
pop,
params,
VEM(; presample = false, optim_options = (; iterations = 30000, subject_ratio = 0.5)),
)
```
### Prior on Population Parameters (`prior`)
By default (`prior=true`), VEM adds the log prior on the population parameters,
if defined in the `@param` block, to the ELBO.
This allows for maximum a-posteriori (MAP) estimation of the population parameters,
after marginalizing the random effects.
Setting `prior=false` ignores the prior contribution,
which means the ELBO depends only on the likelihood:
```julia
fit(model, pop, params, VEM(; prior = false))
```
If the `@param` block has no prior distributions and only includes domains,
this option is not used.
### Base Distribution Standard Deviation (`base_dist_std`)
VEM's variational family uses a reparameterization trick:
samples are drawn from a base distribution
and transformed to approximate the posterior.
The `base_dist_std` parameter controls the standard deviation
of this base distribution (default: 0.1):
```julia
fit(model, pop, params, VEM(; base_dist_std = 0.05))
```
Smaller values start the optimization
with a tighter initial approximation to the posterior,
which can help stability for sensitive models.
Larger values allow broader initial exploration.
The default of 0.1 works well for most models.
### Posterior Covariance Structure (`diagonal`)
VEM approximates each subject's random effect posterior
with a multivariate normal distribution.
The `diagonal` keyword controls the covariance structure:
- **`diagonal=true`** (default):
assumes independent random effects in the posterior.
Faster and requires fewer variational parameters.
- **`diagonal=false`**:
captures correlations between random effects in the posterior.
Slower but may be more accurate when random effects are correlated.
```{julia}
fpm_full_cov = fit(
warfarin_model,
pop,
init_params(warfarin_model),
VEM(; diagonal = false, optim_options = (; show_every = 20)),
)
```
```{julia}
loglikelihood(fpm_full_cov, FOCE())
```
Note that setting `diagonal = false` does not always lead to a better fit. It is generally recommended to start with `diagonal = true` and only switch to `diagonal = false` if there is reason to believe that the random effects are highly correlated.
### Holding Parameters Constant (`constantcoef`)
You can hold specific parameters constant during estimation
using the `constantcoef` keyword.
This is a keyword to `fit`, **not** to `VEM()`:
```julia
fit(model, pop, params, VEM(); constantcoef = (:σ_prop,))
```
:::{.callout-note}
`constantcoef` is a keyword argument to `fit`, not to `VEM()`.
This is consistent with how it works for FOCE and other algorithms.
:::
### Reproducibility (`rng`)
VEM uses random sampling internally,
so results vary across runs by default.
Set a specific random number generator for reproducible results:
```julia
rng = default_rng()
seed!(rng, 123)
fit(model, pop, params, VEM(; rng))
```
:::{.callout-note}
Even with presampling,
different seeds produce different results
because the pre-drawn samples differ.
For exact reproducibility across runs,
always set the `rng`.
:::
### Threading (`ensemblealg`)
Control parallelism across subjects:
- **`EnsembleThreads()`** (default): uses multiple threads for gradient computation.
- **`EnsembleSerial()`**: single-threaded, useful for debugging.
```julia
fit(model, pop, params, VEM(; ensemblealg = EnsembleSerial()))
```
### ODE Solver Options (`diffeq_options`)
Pass differential equation solver options through `diffeq_options`:
```julia
fit(model, pop, params, VEM(; diffeq_options = (; abstol = 1e-8, reltol = 1e-8)))
```
This is useful when the default ODE tolerances
are insufficient for a stiff or sensitive model.
## When to Use VEM vs FOCE/LaplaceI
**Choose VEM when:**
- You have a response model that FOCE does not support and LaplaceI is too slow, e.g. truncated, censored or time-to-event data.
- The model has many random effects and FOCE is too slow or struggles to converge, and you want a faster alternative, e.g. to initialize FOCE or as a standalone method.
- You have a very large dataset and want a fast initial fit to use as a starting point for FOCE/LaplaceI, then stochastic VEM with mini-batching can help.
**Choose FOCE when:**
- You have a model without too many random effects and the response model is supported by FOCE.
- You need a well-established (non-experimental) fitting method, diagnostics and inference tools.
**Choose LaplaceI when:**
- Your response model is not supported by FOCE.
- You need a well-established (non-experimental) fitting method, diagnostics and inference tools.
## Conclusion
In this tutorial,
we introduced VEM as a flexible estimation method in Pumas.
We covered its conceptual foundations (ELBO maximization, presampling vs. resampling, and posterior covariance structure),
demonstrated its use with a warfarin PK model,
showed how to compare results with FOCE,
and provided a reference for the main VEM options.
