Efficient analytic gradient of the LMM objective + gradient-based optimizers

[palday]

Summary

This PR adds an efficient analytic gradient of the profiled objective (deviance / REML criterion) for LinearMixedModel, exposes it as the exported function objective_gradient!, and wires it into the NLopt backend so that the gradient-based optimizers :LD_LBFGS, :LD_MMA, and :LD_SLSQP can be selected via fit(MixedModel, form, data; optimizer=:LD_LBFGS). The default optimizer remains the derivative-free :LN_NEWUOA.

It also reworks the ForwardDiff extension to reuse the core linear-algebra pipeline instead of maintaining parallel fd_* implementations, and adds a gradient keyword to fit/fit! (:analytic default, :forwarddiff opt-in) so that ForwardDiff can serve as an alternative gradient source for the LD optimizers — see the sections below the BLAS-3 one.

Approach

The objective is an affine function of the logarithms of the diagonal elements of the blocked Cholesky factor, obj = 2 Σⱼ wⱼ log Lⱼⱼ + const, with weights 1 on the random-effects rows, 1 (REML) or 0 (ML) on the fixed-effects rows, and n, n − p, or pwrss/σ² on the ℓ_yy element for ML, REML, and fixed-σ objectives respectively. Differentiating Ω = LLᵀ (Murray 2016, arXiv:1602.07527) gives

∂obj/∂θₚ = tr(W L⁻¹ Ω̇ₚ L⁻ᵀ) = ⟨S, Ω̇ₚ⟩, S = L⁻ᵀWL⁻¹,

where S is shared by all parameters. A single blocked computation of the lower blocks of L⁻¹ (stored in a reusable GradientWorkspace), followed by weighted Gram products contracted against the sparse structure of the A blocks, therefore yields the entire gradient: Ω̇ₚ is never formed and no solve is repeated per parameter. Blocks between scalar terms with sparse A, and Diagonal/UniformBlockDiagonal diagonal blocks, use selected-entry kernels that evaluate only the entries of S that the contraction touches.

This replaces the earlier per-parameter draft in src/gradient.jl, which performed two full blocked triangular solves per parameter and disagreed with ForwardDiff for vector-valued terms (Λ where Λᵀ was needed in Omega_dot_diag_block!, and a missing Λᵣᵀ premultiplication of the off-diagonal blocks).

API

• objective_gradient!(g, m) — gradient at the current parameter values; returns the objective value.
• objective_gradient!(g, m, θ) — installs θ via setθ!/updateL! first, mirroring objective!.
• MixedModels.GradientWorkspace(m) (internal) — preallocated storage for allocation-free repeated evaluation inside the optimizer loop.
• Profiling a model fitted with an LD optimizer falls back to a derivative-free optimizer, since the profiling objectives do not evaluate gradients.
• The LD optimizers internally work on a per-observation scaling of the objective and gradient: on the deviance scale the gradient norm grows with n, making the identity a poor initial inverse-Hessian estimate, and the first LBFGS line-search step overshoots wildly for large data sets. fitlog, the progress display, and the reported fmin stay on the deviance scale; ftol_rel is invariant under the scaling and ftol_abs acts as a per-observation tolerance.

Performance

Gradient evaluation vs the ForwardDiff extension: ~60× faster on the maximal kb07 model, ~9× on insteval, with far less allocation.

Fits, gradient-based vs derivative-free:

dataset	nθ	LN_NEWUOA evals / time / fmin	LD_LBFGS evals / time / fmin	Δfmin (LN − LD)
sleepstudy	3	72 / 0.00 s / 1751.9393444641	17 / 0.00 s / 1751.9393444632	+8e-10
penicillin	2	37 / 0.09 s / 332.1883486689	25 / 0.25 s / 332.1883486684	+5e-10
mrk17_exp1 (maximal)	36	2537 / 16.7 s / 7147.5506	91 / 3.5 s / 7147.5201	+0.031
ml1m	2	52 / 43.8 s / 2663972.0116	13 / 48.1 s / 2663972.0116	−2e-9
kb07 (maximal)	20	845 / 0.45 s / 28637.12273	49 / 0.06 s / 28637.12269	+4e-5

Gradient-based optimization wins decisively on many-parameter models: mrk17 fits ~5× faster and reaches an objective 0.031 lower; kb07 fits ~7× faster. On ml1m (n = 10⁶, 2 parameters) LD_LBFGS needs only 13 evaluations but each gradient costs ~4× an objective evaluation, so wall time is on par with the derivative-free default; LD_MMA and LD_SLSQP also converge cleanly there (19 and 25 evaluations). Without the per-observation scaling, LBFGS failed outright on ml1m — see the second commit.

Dense-crossed cross term (BLAS-3)

For two scalar random-effects terms whose Cholesky factor has dense fill-in (the classic crossed-effects case, e.g. ml1m), the dominant cost of the gradient is the cross term S[r,b] on the sparsity pattern of A[r,b], previously one BLAS-1 column dot product per nonzero — memory-bandwidth bound (~2 Gflop/s measured). When the crossing is dense enough, this is now evaluated with a BLAS-3 matrix product X[r,r]ᵀX[r,b] a column-panel at a time (transient product bounded to q_r × 128, never fully materialized), running at ~80 Gflop/s. The path is gated on crossing density (nnz > 0.03·q_r·q_b) and dense fill, so sparse crossings (e.g. insteval) keep the BLAS-1 path; results are identical up to floating-point reassociation. On ml1m the gradient is ~1.24× faster.

Note on the original selected-inverse follow-up: a Takahashi-style selected inversion does not help these models. With dense Cholesky fill the factor pattern is dense, so forming X[r,b] directly is already cheaper than the selected-inverse alternative (which needs the full dense Z[r,r]), and the dominant selected-entry extraction cost is unchanged. The measurement that establishes this is in the commit message; the achievable win was the BLAS-3 restructuring above.

ForwardDiff extension rework

The extension previously duplicated the core objective pipeline in out-of-place fd_* variants (fd_setθ!, fd_updateL!, fd_pwrss, fd_logdet, fd_cholUnblocked!, fd_rankUpdate!). Every divergence from core was one of: array arguments instead of model fields, a BLAS/LAPACK call requiring BlasFloat, or a ::Float64/::T assertion blocking dual numbers. All six are now deleted in favor of:

• generic (non-BLAS) fallback methods for rankUpdate! (dense syrk! and UBD + BlockedSparse syr! paths; the BLAS methods are now T<:BlasFloat specializations and the dense-C signatures are narrowed to HermOrSym{T,Matrix{T}} to avoid ambiguity with the Diagonal/UniformBlockDiagonal-storage methods) and for cholUnblocked! (generic cholesky!(Hermitian(A, :L)); the Diagonal bound is relaxed from AbstractFloat to Real);
• array-based methods of setθ!, updateL!, pwrss, and a _logdet helper, with the model-based methods as thin wrappers.

Float64 dispatch still reaches the BLAS/LAPACK methods (verified with which); updateL! time and allocations on kb07 are identical to main, and no new method ambiguities are introduced.

Semantics change: fd_deviance now profiles σ (or holds it at optsum.sigma when fixed) and includes the constant weights term, exactly matching objective. ForwardDiff.gradient/hessian therefore refer to the profiled objective and agree with objective_gradient! at any θ (≤1e-12 on sleepstudy) without syncing the model state first — previously agreement relied on the envelope theorem at synced state. The Hessian reference values in test/forwarddiff.jl changed accordingly.

Gradient source selection

fit(MixedModel, form, data; optimizer=:LD_LBFGS, gradient=:forwarddiff) selects forward-mode AD as the gradient source; the default gradient=:analytic uses objective_gradient!. The option is stored in OptSummary (serialization-aware) and validated in fit!; requesting :forwarddiff without ForwardDiff.jl loaded throws an informative error. The ForwardDiff path caches the dual-promoted copies of A/L/reterms together with a GradientConfig and a DiffResults buffer in a reusable workspace (fd_gradient_workspace), so repeated gradient evaluations do not re-promote; value and gradient come from a single dual sweep. Both sources use the same per-observation scaling and take identical optimization paths (same evaluation counts, same optima).

Analytic vs ForwardDiff gradient source

From gradients/fd_vs_analytic.jl (one subprocess per configuration so peak RSS is per-configuration; measured on a refit! after a warm-up fit, so time/bytes/allocs exclude compilation):

model	optimizer	gradient	nθ	feval	objective	time (s)	alloc (MiB)	# allocs	peak RSS (MiB)
kb07	LN_NEWUOA	—	20	970	28637.1228	0.422	1.8	57331	767
kb07	LD_LBFGS	analytic	20	53	28637.1234	0.074	2.7	25555	776
kb07	LD_LBFGS	forwarddiff	20	53	28637.1234	4.036	11.7	20501	794
ml1m	LN_NEWUOA	—	2	52	2663972.0116	46.9	0.9	18563	1075
ml1m	LD_LBFGS	analytic	2	13	2663972.0116	45.9	280.5	18346	1253
ml1m	LD_LBFGS	forwarddiff	2	13	2663972.0116	279.2	522.8	17985	1567

The analytic gradient dominates on both axes: 55× faster than ForwardDiff on kb07 and 6× on ml1m, and — contrary to the intuition that the one-shot GradientWorkspace is the memory-heavy option — ForwardDiff also allocates more (1.9× the bytes on ml1m) and has a ~310 MiB higher peak RSS there, because the cached dual arrays carry nθ partials per entry. :forwarddiff is thus primarily a correctness oracle and fallback, not a memory-saving alternative.

Possible follow-ups

• Stream the dense fill block itself out of _invL! to bound peak memory (currently O(q_r·q_b) for the fill block); invasive to the hot path, memory-only benefit at scales beyond ml1m.
• Heuristics for when a gradient-based optimizer should be preferred (many covariance parameters, moderate n).
• Update profiling code to accept any optimizer

🤖 Generated with Claude Code

[dmbates]

When running the tests I keep getting a test failure of the form

Test Summary:           | Pass  Total   Time
gradient vs ForwardDiff |   87     87  11.9s
┌ Warning: NLopt optimization failure: FAILURE
└ @ MixedModels.MixedModelsNLoptExt ~/.julia/dev/MixedModels/src/MixedModelsNLoptExt.jl:182
LD_LBFGS penicillin: Test Failed at /Users/dmbates/.julia/dev/MixedModels/test/grad.jl:147
  Expression: m.optsum.returnvalue in (:SUCCESS, :FTOL_REACHED, :XTOL_REACHED)
   Evaluated: FAILURE in (:SUCCESS, :FTOL_REACHED, :XTOL_REACHED)

but if I rerun that example by hand I get a returnvalue of :FTOL_REACHED, not FAILURE.

Is anyone else encountering this? For reference this is fitting the model

julia> fit(MixedModel, @formula(diameter ~ 1 + (1|plate) + (1|sample)), dataset(:penicillin); optimizer=:LD_LBFGS)
Linear mixed model fit by maximum likelihood
 diameter ~ 1 + (1 | plate) + (1 | sample)
   logLik   -2 logLik     AIC       AICc        BIC    
  -166.0942   332.1883   340.1883   340.4761   352.0676

Variance components:
            Column   Variance Std.Dev. 
plate    (Intercept)  0.714993 0.845572
sample   (Intercept)  3.135188 1.770646
Residual              0.302425 0.549932
 Number of obs: 144; levels of grouping factors: 24, 6

  Fixed-effects parameters:
─────────────────────────────────────────────────
               Coef.  Std. Error      z  Pr(>|z|)
─────────────────────────────────────────────────
(Intercept)  22.9722    0.744596  30.85    <1e-99
─────────────────────────────────────────────────

julia> m.optsum.returnvalue
:FTOL_REACHED