---
title: "ALDEx3 Mixed-Effects Engines"
output: rmarkdown::html_vignette
vignette: >
%\VignetteIndexEntry{ALDEx3 Mixed-Effects Engines}
%\VignetteEngine{knitr::rmarkdown}
%\VignetteEncoding{UTF-8}
---
```{r, include = FALSE}
knitr::opts_chunk$set(
collapse = TRUE,
comment = "#>"
)
```
## Motivation
ALDEx3 fits a regression model to every Monte Carlo draw of the latent
log-abundance for every feature. For mixed-effects models, repeating a full
exact variance-component optimisation for each draw can dominate runtime.
The exact ALDEx3 mixed-effects engines remain available through
`method = "lme4"` and `method = "nlme"`. The `method = "blmm"` engine is an
ALDEx3-specific approximation designed to reduce repeated optimisation cost
while preserving the same fixed-effects model. When `n.cores > 1`, BLMM and
the exact mixed-effects engines parallelize across features, not across the
Monte Carlo draws within a feature.
## What BLMM Approximates
For one feature and one Monte Carlo draw, the target remains the usual Gaussian
linear mixed model fitted in profiled REML form. BLMM does **not** approximate
the fixed-effects solve once covariance parameters are chosen.
The approximation is only this:
1. Fit one feature-specific batched anchor covariance model across all draws.
2. Build draw-specific local score updates around that anchor.
3. Update the covariance parameters per draw with a local Newton step.
4. Solve the fixed effects exactly for each draw conditional on the updated
covariance parameters.
If the approximate path cannot be evaluated cleanly for a feature, ALDEx3 warns
and falls back to the exact `lme4` engine for that feature.
## Conceptual Formulation
Conceptually, each draw has its own REML objective for the same design matrices
and grouping structure but a different response vector. BLMM introduces:
1. One anchor fit per feature.
2. One shared local curvature matrix at the anchor.
3. Draw-specific local score corrections around the anchor.
4. Exact conditional GLS fixed-effect solves for each draw.
This is the cleanest way to describe the method statistically.
## Implemented Profiled Formulation
To keep the notation manageable, fix one feature \(d\). For Monte Carlo draw
\(s\), let \(y_{ds} \in \mathbb{R}^{N}\) denote the vector of latent
log-abundances across the \(N\) samples. Since the feature is fixed throughout
this section, I write \(y_s\) instead of \(y_{ds}\).
For that fixed feature and draw, BLMM targets the same Gaussian linear mixed
model as `lme4`:
\[
y_s = X\beta_s + Z u_s + \varepsilon_s,
\qquad
u_s \sim N(0, \sigma_s^2 \Sigma_{\theta_s}),
\qquad
\varepsilon_s \sim N(0, \sigma_s^2 I_N).
\]
Here \(X\) is the fixed-effects design matrix, \(Z\) is the random-effects
design matrix, \(\beta_s\) is the draw-specific fixed-effects vector, and
\(\theta_s\) denotes the top-level covariance parameters for draw \(s\). As in
`lme4`, BLMM works with a relative covariance parameterisation in which
\(\Sigma_{\theta_s}\) is represented through a lower-triangular factor
\(\Lambda_{\theta_s}\).
Write \(p = \mathrm{ncol}(X)\) for the number of fixed-effect coefficients and
\(q = \mathrm{ncol}(Z)\) for the number of random-effect coefficients.
BLMM is implemented in the profiled `lme4` formulation rather than through a
dense \(N \times N\) covariance matrix. This does not change the statistical
target. It is simply the algebraic form that makes repeated mixed-model fitting
much more efficient.
After profiling out the fixed effects and residual scale, the criterion depends
only on the covariance parameters. In this formulation, the key matrices are
\[
A_{\theta_s} = I_q + \Lambda_{\theta_s} Z^T Z \Lambda_{\theta_s}^T,
\]
\[
R_{ZX,\theta_s} = L_{\theta_s}^{-1}\Lambda_{\theta_s} Z^T X,
\qquad
M_{\theta_s} = X^T X - R_{ZX,\theta_s}^T R_{ZX,\theta_s},
\]
where \(L_{\theta_s}\) is the Cholesky factor of \(A_{\theta_s}\).
The response-dependent part of the draw-specific criterion is the profiled
weighted residual sum of squares
\[
\mathrm{PWRSS}_s(\theta_s)
=
\|y_s\|^2
- \|R_{ZY,\theta_s}\|^2
- \|C_{\beta,\theta_s}\|^2,
\]
with
\[
R_{ZY,\theta_s} = L_{\theta_s}^{-1}\Lambda_{\theta_s} Z^T y_s,
\qquad
C_{\beta,\theta_s} =
L_{X,\theta_s}^{-1}
\left(
X^T y_s - R_{ZX,\theta_s}^T R_{ZY,\theta_s}
\right),
\]
and \(L_{X,\theta_s}\) the Cholesky factor of \(M_{\theta_s}\).
Up to constants, the profiled REML criterion for one draw is
\[
\ell_s(\theta_s)
\propto
\frac{1}{2}
\left\{
\log |A_{\theta_s}|
+
\log |M_{\theta_s}|
+
(N-p)\log\left(
\frac{\mathrm{PWRSS}_s(\theta_s)}{N-p}
\right)
\right\}.
\]
Now collect the \(S\) Monte Carlo draws for feature \(d\) into the response
block
\[
Y_d = [y_{d1}, \dots, y_{dS}] \in \mathbb{R}^{N \times S}.
\]
The computational advantage of BLMM is that one anchor evaluation reuses the
same structural calculations across all \(S\) columns of \(Y_d\). In each
evaluation it performs
1. one factorisation of the profiled random-effects system,
2. one log-determinant calculation,
3. one solve against the fixed-effects design matrix,
4. one multi-right-hand-side solve against \(Y_d\),
5. one average over the draw-specific profiled residual log terms.
BLMM therefore does **not** run \(S\) separate REML optimisations inside the
anchor optimiser.
More precisely, the anchor fit targets the average of the per-draw profiled
REML criteria:
\[
\bar{\ell}(\phi)
=
\frac{1}{S}\sum_{s=1}^S \ell_s(\phi),
\]
which is, up to constants,
\[
\bar{\ell}(\phi)
\propto
\frac{1}{2}
\left\{
\log |A_{\phi}|
+
\log |M_{\phi}|
+
\frac{N-p}{S}\sum_{s=1}^S
\log\left(
\frac{\mathrm{PWRSS}_s(\phi)}{N-p}
\right)
\right\}.
\]
This point is important: the batched anchor uses the **average of the per-draw
profiled REML criteria**, not a surrogate such as
\(\log(\mathrm{mean}\,\mathrm{PWRSS}_s)\). The anchor, the draw-specific score
correction, and the curvature matrix are all defined with respect to this same
averaged profiled objective.
BLMM then takes a local second-order approximation around the anchor
\(\bar{\phi}\). Let \(H\) denote the observed Hessian of \(\bar{\ell}\) at that
anchor. For draw \(s\), BLMM computes a draw-specific score correction
\(g_s\), defined as the difference between the average profiled score and the
draw-specific profiled score at \(\bar{\phi}\), and applies the local Newton
update
\[
\phi_s \approx \bar{\phi} + H^{-1} g_s.
\]
This yields one updated covariance parameter vector per draw without rerunning a
full REML optimisation from scratch.
Conditional on \(\phi_s\), BLMM then solves the fixed effects exactly by
generalized least squares:
\[
\hat{\beta}_s
=
\left(X^T V_s^{-1}X\right)^{-1}X^T V_s^{-1} y_s,
\]
where \(V_s = Z\Sigma_{\theta_s}Z^T + \sigma_s^2 I_N\) is the covariance model
implied by the updated parameters.
The approximation is therefore confined to the repeated covariance-optimisation
step. The per-draw fixed-effects solve remains exact conditional on the updated
covariance model.
## Assumptions and Scope
BLMM targets the same Gaussian linear mixed model as the exact `lme4` engine for
each feature and Monte Carlo draw. It reuses `lme4::lFormula()` for formula
parsing and random-effects design construction, so the same model syntax can be
used for random intercepts, random slopes, and correlated random effects.
The method is designed specifically for the ALDEx3 setting, where the same
mixed-effects design is fit repeatedly across many Monte Carlo draws for each
feature. Its computational advantage comes from exploiting the fact that the
fixed-effects design, random-effects design, grouping structure, and covariance
parameterisation are shared across those draws, while only the response vector
changes.
BLMM is an approximate engine. The approximation is confined to the repeated
covariance-optimisation step: it replaces a full exact variance-component fit
for every draw with one feature-specific anchor fit and draw-specific local
updates around that anchor. Conditional on the updated covariance parameters,
the per-draw fixed-effects solve remains exact.
Accordingly, BLMM should be viewed as a fast ALDEx3-specific alternative to
repeated exact `lme4` optimisation, not as a replacement for the exact engine in
all settings. The exact `lme4` engine remains the reference implementation, and
BLMM should be validated against exact `lme4` fits on the target problem class
before large-scale use.
## Validation Strategy
The BLMM test suite focuses on algorithmic correctness first.
1. Single-draw consistency checks compare `blmm` to exact `lme4`.
2. Identical-draw checks verify that identical Monte Carlo draws produce
identical BLMM outputs.
3. Simulated-data comparisons check posterior means and posterior uncertainty
against the exact engine.
4. Random-slope and correlated-random-effect tests verify that covariance terms
are retained in the returned random-effects output.
5. Failure-mode tests verify that BLMM warns and falls back to exact `lme4`
instead of silently changing the model.
The practical target is that approximation error in posterior means should stay
materially smaller than posterior uncertainty used for inference, and that
posterior uncertainty itself should remain close enough to avoid changing
routine interpretation.
## Worked Example
The `oral_mouthwash_data` dataset provides a repeated-measures design with one
random intercept per participant. The exact and approximate engines use the
same mixed-effects formula:
```{r eval = TRUE}
set.seed(42)
library(ALDEx3)
data("oral_mouthwash_data", package = "ALDEx3")
Y <- oral_mouthwash_data$counts
keep_names <- row.names(Y[((rowSums(Y == 0)) / ncol(Y)) <= 0.75, ])
other <- colSums(Y[((rowSums(Y == 0)) / ncol(Y)) > 0.75, ])
Y <- rbind(Y[keep_names, ], other = other)
meta <- oral_mouthwash_data$metadata
fit_blmm <- aldex(
Y,
~ treat * timec + (1 | participant_id),
data = meta,
method = "blmm",
n.cores = 1,
nsample = 250,
scale = clr.sm,
gamma = 0
)
head(summary(fit_blmm))
```
Switching between exact and approximate fitting only changes the engine:
replace `method = "blmm"` with `method = "lme4"` for exact variance-component
optimisation.
## Example Comparison: `lme4` Versus `blmm`
On a small subset of the `oral_mouthwash_data` features, `blmm` tracks the
exact `lme4` results closely for most coefficients while running the
covariance-optimisation step approximately. The code below fits the same model
with both engines and compares the posterior mean coefficients and their
posterior standard errors.
```{r eval = TRUE}
set.seed(42)
data("oral_mouthwash_data", package = "ALDEx3")
Y0 <- oral_mouthwash_data$counts
keep_names <- row.names(Y0[((rowSums(Y0 == 0)) / ncol(Y0)) <= 0.75, ])
other <- colSums(Y0[((rowSums(Y0 == 0)) / ncol(Y0)) > 0.75, ])
Y <- rbind(Y0[keep_names, ], other = other)
# Keep a small subset to make the comparison easy to inspect in the vignette.
Y <- Y[c(seq_len(min(8, nrow(Y) - 1)), nrow(Y)), ]
meta <- oral_mouthwash_data$metadata
K <- 1
fit_args <- list(
data = meta,
n.cores = 1,
nsample = 16,
scale = clr.sm,
gamma = 0
)
fit_model <- function(method) {
do.call(
aldex,
c(
list(Y, ~ treat * timec + (1 | participant_id)),
fit_args,
list(method = method)
)
)
}
set.seed(42)
fit_lme4 <- fit_model("lme4")
set.seed(42)
fit_blmm <- fit_model("blmm")
time_lme4 <- replicate(K, {
set.seed(42)
system.time({
invisible(fit_model("lme4"))
})[["elapsed"]]
})
time_blmm <- replicate(K, {
set.seed(42)
system.time({
invisible(fit_model("blmm"))
})[["elapsed"]]
})
s_lme4 <- summary(fit_lme4)
s_blmm <- summary(fit_blmm)
cmp <- merge(
s_lme4,
s_blmm,
by = c("parameter", "entity"),
suffixes = c(".lme4", ".blmm")
)
cmp$abs_est_diff <- abs(cmp$estimate.lme4 - cmp$estimate.blmm)
cmp$abs_se_diff <- abs(cmp$std.error.lme4 - cmp$std.error.blmm)
cmp_tbl <- cmp[order(cmp$abs_est_diff, decreasing = TRUE), ][1:8, c(
"parameter", "entity",
"estimate.lme4", "estimate.blmm", "abs_est_diff",
"std.error.lme4", "std.error.blmm", "abs_se_diff"
)]
cat('\n')
knitr::kable(
cmp_tbl,
format = "html",
digits = 4,
caption = "Largest absolute differences between exact lme4 and blmm on the shared-draw subset."
)
cat('\n
\n')
runtime_tbl <- data.frame(
engine = c("lme4", "blmm"),
mean_elapsed_sec = c(mean(time_lme4), mean(time_blmm)),
median_elapsed_sec = c(median(time_lme4), median(time_blmm)),
speedup_vs_lme4 = c(1, mean(time_lme4) / mean(time_blmm))
)
cat('\n\n')
cat('\n')
knitr::kable(
runtime_tbl,
format = "html",
digits = 3,
caption = paste0(
"Elapsed runtime on the shared-draw vignette subset averaged over K = ",
K,
" repeated runs. Times will vary by hardware."
)
)
cat('\n
\n')
```
For an engine comparison to be meaningful, both fits must use the same Monte
Carlo draws. The duplicated `set.seed(42)` above enforces that. Without that
reset, the comparison would conflate engine differences with differences in the
posterior draw block itself.
With identical draws on this vignette subset, the comparison gave a maximum
absolute posterior-mean coefficient difference of about `0.023`, a median
absolute difference effectively `0`, a maximum absolute posterior
standard-error difference of about `0.025`, and a median standard-error
difference of about `0.003`. Averaging the runtime over `K = 10` repeated runs
gave a `blmm` speedup of about `11.95x` relative to `lme4` on this subset.
That is the practical target for BLMM on this example: when compared on the
same posterior draws, `blmm` should stay very close to the exact `lme4`
engine while avoiding repeated exact variance-component optimisation.
## Runtime Scaling
Runtime is driven by three quantities:
- `N`, the number of samples,
- `D`, the number of features, and
- `S`, the number of Monte Carlo draws.
Both engines pay the shared cost of generating and transforming the Monte Carlo
draws, which scales roughly with `D * N * S`. The difference is in the mixed
model solve.
The `n.cores` argument distributes features across workers, so it changes the
feature loop but not the per-feature draw loop. That is why the runtime gains
from parallelism and the runtime gains from BLMM are complementary.
`lme4` performs a full variance-component optimisation for every feature and
every draw. That means its expensive nonlinear step scales roughly like
`D * S * cost(full_REML_opt(N, random_effect_structure))`. As `N` grows or the
random-effects structure gets more complex, each exact optimisation becomes
more expensive.
`blmm` replaces those repeated exact optimisations with one batched anchor fit
per feature plus per-draw local updates and exact conditional fixed-effect
solves. Its costly nonlinear work is therefore closer to
`D * cost(anchor_REML_opt(N, random_effect_structure)) + D * S * cost(local_update(N))`.
The savings are largest when the repeated exact optimisation is the dominant
cost.
In practice, BLMM is most valuable when:
1. `S` is moderate to large, so repeated exact REML fits become expensive.
2. `D` is large, so the same model structure is reused across many features.
3. The random-effects structure is nontrivial, so each exact `lme4` fit is
expensive.
4. The approximation can be anchored once and reused cleanly across draws.
BLMM is less compelling when `S = 1`, when the random-effects structure is very
simple, or when the dataset is so small that exact `lme4` optimisation is already
cheap. In those settings, the extra approximation machinery may not buy much.
On the vignette subset above, the elapsed times in the table give a concrete
example of the expected pattern: `blmm` should be faster because it avoids
repeating the full covariance optimisation for every draw, while still solving
the fixed effects exactly conditional on the updated covariance parameters.