ALDEx3 Mixed-Effects Engines

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:

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))
#>                                    parameter         entity   estimate
#> 1                  treatantiseptic-mouthwash g__Haemophilus  2.3758996
#> 2                 treatalcoholfree-mouthwash g__Haemophilus  3.2202728
#> 3                                  treatsoda g__Haemophilus  1.6586761
#> 4                           timec15min-later g__Haemophilus  0.8264883
#> 5                            timec2hrs-later g__Haemophilus  2.3811563
#> 6 treatantiseptic-mouthwash:timec15min-later g__Haemophilus -2.3012944
#>   std.error  p.val.adj
#> 1 0.8196615 0.32251795
#> 2 0.8196615 0.01418721
#> 3 0.8196615 0.91176518
#> 4 0.6338824 0.99960265
#> 5 0.6338824 0.02803160
#> 6 0.8964451 0.36972657

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.

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('<div style="overflow-x:auto; width:100%;">\n')
#> <div style="overflow-x:auto; width:100%;">
knitr::kable(
  cmp_tbl,
  format = "html",
  digits = 4,
  caption = "Largest absolute differences between exact lme4 and blmm on the shared-draw subset."
)

cat('\n</div>\n')
#> 
#> </div>

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('<div style="overflow-x:auto; width:100%;">\n')
#> <div style="overflow-x:auto; width:100%;">
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</div>\n')
#> 
#> </div>

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.