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

[palday]

Some current timings:

model	optimizer	storage	model (MiB)	L (MiB)	grad ws (MiB)	peak RSS (MiB)	time (s)	feval	ms/feval	objective
insteval	LN_NEWUOA	dense	17.5	10.7	—	1178.2	0.6	43	14.0	237730.6154
insteval	LN_NEWUOA	rfp	12.7	5.8	—	1183.4	0.57	45	12.8	237730.6154
insteval	LD_LBFGS	dense	17.5	10.7	36.6	1210.5	1.45	14	103.6	237730.6154
insteval	LD_LBFGS	rfp	12.7	5.8	31.8	1190.6	1.44	14	102.5	237730.6154
ml1m	LN_NEWUOA	dense	197.5	116.5	—	1338.9	43.95	52	845.2	2.6639720116e6
ml1m	LN_NEWUOA	rfp	145.1	64.1	—	1392.2	42.64	52	819.9	2.6639720116e6
ml1m	LD_LBFGS	dense	197.5	116.5	279.6	1812.1	45.76	13	3520.1	2.6639720116e6
ml1m	LD_LBFGS	rfp	145.1	64.1	227.2	1646.1	35.81	13	2754.8	2.6639720116e6

[codecov]

Codecov Report

❌ Patch coverage is 82.14600% with 203 lines in your changes missing coverage. Please review.
✅ Project coverage is 92.82%. Comparing base (1c01802) to head (bfbbba9).

Files with missing lines	Patch %	Lines
src/gradient.jl	80.09%	201 Missing ⚠️
ext/MixedModelsForwardDiffExt.jl	97.82%	1 Missing ⚠️
src/MixedModelsNLoptExt.jl	96.87%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #903      +/-   ##
==========================================
- Coverage   94.89%   92.82%   -2.08%     
==========================================
  Files          38       39       +1     
  Lines        3958     4947     +989     
==========================================
+ Hits         3756     4592     +836     
- Misses        202      355     +153     

Flag	Coverage Δ
current	92.50% <81.80%> (-2.09%)	⬇️
minimum	92.74% <82.14%> (-2.06%)	⬇️
nightly	92.50% <81.81%> (-2.09%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dmbates]

In the comments at the beginning of src/gradient.jl it states that the weight for the logarithm of the last element on the diagonal of L is pwrss / σ² when σ is fixed. That may be a formula that produces the desired result but I don't see how it can be a weight when pwrss depends on the current value of θ. @palday Am I missing something subtle here?

I would be quite happy to have a formula that only works for ML and REML estimation.