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 by profiled restricted maximum likelihood (REML). 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. Check that each proposed update does not materially worsen its draw’s profiled REML objective; refit rejected draws with exact lme4.
5. Solve the fixed effects exactly for each accepted draw conditional on the updated covariance parameters.

If the shared anchor optimisation fails or its Hessian is unusable, ALDEx3 falls back to exact lme4 for the entire feature. If only particular draw-specific updates fail, ALDEx3 refits those draws exactly and retains the accepted BLMM draws. In either case, it emits one consolidated warning.

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 generalized least-squares (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 \in \mathbb{R}^{N \times p}\) is the fixed-effects design matrix, \(\beta_s \in \mathbb{R}^{p}\) is the draw-specific fixed-effects vector, \(Z \in \mathbb{R}^{N \times q}\) is the random-effects design matrix, and \(u_s \in \mathbb{R}^{q}\) is the draw-specific random-effects vector. Thus \(p\) is the number of fixed-effect coefficients and \(q\) is the number of random-effect coefficients. The residual standard deviation is \(\sigma_s>0\), \(I_N\) denotes the \(N \times N\) identity matrix, and the superscript \(T\) denotes matrix transpose.

The parameter vector \(\theta_s\) determines the random-effects covariance relative to the residual variance. As in lme4, BLMM represents this relative covariance by a lower-triangular factor \(\Lambda_{\theta_s} \in \mathbb{R}^{q \times q}\): \[ \Sigma_{\theta_s} = \Lambda_{\theta_s}\Lambda_{\theta_s}^T, \qquad \operatorname{Var}(u_s) = \sigma_s^2\Lambda_{\theta_s}\Lambda_{\theta_s}^T. \] For a random-intercept term with \(m\) grouping levels, \(\Lambda_{\theta_s}=\theta_s I_m\), where \(I_m\) is the \(m \times m\) identity matrix. In that special case, \(\theta_s\) is the ratio of the random-intercept standard deviation to the residual standard deviation.

BLMM optimises an unconstrained parameter vector \(\phi_s\) rather than \(\theta_s\) directly. For each covariance-factor component constrained to be nonnegative, BLMM uses \[ \phi_{sj}=\log(\theta_{sj}), \qquad \theta_{sj}=\exp(\phi_{sj}). \] Components without a lower bound, such as off-diagonal covariance-factor entries, are not log-transformed. Accordingly, the scalar relation \(\theta_s=\exp(\phi_s)\) applies directly to a random-intercept-only model but not necessarily to every component of a correlated random-effects model.

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 \(I_q\) is the \(q \times q\) identity matrix and \(L_{\theta_s}\) is the lower Cholesky factor defined by \[ A_{\theta_s}=L_{\theta_s}L_{\theta_s}^T. \] Thus \(\Lambda_{\theta_s}\) describes the random-effects covariance, whereas \(L_{\theta_s}\) is a computational factor derived from that covariance and the random-effects design matrix.

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), \] where \(L_{X,\theta_s}\) is the lower Cholesky factor defined by \(M_{\theta_s}=L_{X,\theta_s}L_{X,\theta_s}^T\). Here \(\|\cdot\|\) denotes the Euclidean norm, and \(|\cdot|\) below denotes a matrix determinant.

Let \(\ell_s\) denote the profiled REML objective that BLMM minimises. Up to an additive constant that does not depend on the covariance parameters, it is \[ \ell_s(\theta_s) = \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\}. \] Here expressions such as \(A_\phi\) and \(M_\phi\) mean that \(\phi\) is first mapped to \(\theta\), after which the corresponding covariance factor \(\Lambda_\theta\) and profiled matrices are constructed. The anchor is the minimiser \[ \bar{\phi}=\operatorname*{arg\,min}_{\phi}\bar{\ell}(\phi). \]

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=\nabla^2\bar{\ell}(\bar{\phi}) \] denote the observed Hessian of the averaged objective at that anchor. For draw \(s\), BLMM computes the score correction \[ g_s = \nabla\bar{\ell}(\bar{\phi})-\nabla\ell_s(\bar{\phi}), \] the difference between the average and draw-specific profiled gradients. It then applies the local Newton update \[ \phi_s \approx \bar{\phi} + H^{-1} g_s. \]

This yields one proposed covariance parameter vector per draw without rerunning a full REML optimisation from scratch. Because the shared Hessian may be a poor local approximation for a particular draw, BLMM accepts the proposal only when \[ \ell_s(\phi_s) \leq \ell_s(\bar{\phi}) + \tau_s, \] where \(\tau_s\) is a scale-aware numerical tolerance. A non-finite proposal or an objective increase larger than \(\tau_s\) triggers exact lme4 fallback for that draw. This check rejects demonstrably harmful local updates without requiring a damped Newton or trust-region algorithm.

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 the marginal covariance implied by the updated parameters is \[ V_s = \sigma_s^2 \left( Z\Sigma_{\theta_s}Z^T+I_N \right) = \sigma_s^2 \left( Z\Lambda_{\theta_s}\Lambda_{\theta_s}^T Z^T+I_N \right). \]

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 when shared-anchor optimisation or Hessian validation fails, and for rejected draw-specific updates, instead of silently returning an unreliable approximation.

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
)
#> Warning: blmm: 31 draw(s) across 1 feature(s) fell back to exact lme4 because
#> their local covariance updates failed validation. Examples: feature 7: draw 10:
#> proposed covariance update worsened the profiled REML objective by 0.00210645
#> (tolerance 1.37528e-06); draw 21: proposed covariance update worsened the
#> profiled REML objective by 4.7827 (tolerance 1.44495e-06); draw 24: proposed
#> covariance update worsened the profiled REML objective by 0.788883 (tolerance
#> 1.25464e-06) ...

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

In this example, clr.sm centers the scale model on the scales implied by the centered log-ratio (CLR) transformation, and gamma = 0 adds no scale uncertainty around that center. These choices reproduce a fixed CLR scale assumption; they are not required by the BLMM engine.

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,
    " timing run(s). 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 maximum absolute posterior-mean coefficient difference was 0.024, and the median absolute difference was 0. The maximum absolute posterior standard-error difference was 0.029, and the median standard-error difference was 0.003. Across 1 timing replicate, blmm was approximately 31.51 times as fast as lme4 on this subset. These comparisons illustrate the practical target for BLMM: on identical posterior draws, it should stay close to exact lme4 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.