Why not gradient descent

Why not gradient descent#

Scope. Why the supported fitters are exponential-family MLE and EM, not Adam or L-BFGS on the observed-data negative log-likelihood (NLL). Why log_prob is still differentiable, and how to plug in an external optimiser if a surrounding model requires it.

Where things live. fit_mle on exponential families and BatchEMFitter for mixtures. The η→θ solve is in Solvers and Bessel Functions. EM structure is in EM Framework.

1. The supported fitters#

For a regular exponential family the MLE is the mean sufficient statistic followed by a constrained Bregman inversion:

\[ \hat\eta = n^{-1}\sum_{i=1}^n t(x_i),\qquad \hat\theta = \arg\min_\theta\bigl[\psi(\theta)-\theta\cdot\hat\eta\bigr]. \]

GIG implements that as fit_mle. The inversion is a 3-D convex problem after one \(O(n)\) reduction; it does not re-evaluate \(K_\nu\) on every observation at every step. The GIG Fisher information can have condition number \(10^{30}\) at extreme \((a,b)\); the solve η-rescales before walking in \(\theta\) (Solvers and Bessel Functions § 2).

Normal variance-mean mixtures (GeneralizedHyperbolic and special cases) are not exponential families in \(x\) alone. Fitting is EM (Dempster1977): the E-step returns \(\mathbb{E}[t(Y)\mid X]\), the M-step is the same from_expectation map. Complete-data likelihood is monotone; the normal M-step is closed form.

That is the supported path. The rest of this page is why a GPJax / FlowJAX-style loop — store unconstrained parameters, apply softplus, run Adam or L-BFGS on the marginal NLL — is not a second public fitter.

2. What was compared#

On GIG and Gamma, NLL methods store \(\phi=(p,\,\varphi_a,\,\varphi_b)\) with \(a=\mathrm{softplus}(\varphi_a)\), \(b=\mathrm{softplus}(\varphi_b)\), and minimise \(-\frac1n\sum_i\log p(x_i\mid p,a,b)\). GH uses an analogous softplus on \((a,b)\) and a Cholesky factor of \(\Sigma\).

Method

What it optimises

fit_mle / EM

\(\hat\eta\to\theta\) (GIG: η-rescaled Bregman; GH: EM)

L-BFGS + softplus

observed NLL in unconstrained \(\phi\) (JAXopt L-BFGS)

Adam + softplus

same NLL, first-order (hand-rolled; Optax is not a dependency)

L-BFGS-B box NLL

GIG only: SciPy L-BFGS-B bounds \(a,b>0\), no softplus

Reproduction: uv run python benchmarks/bench_gradient_fitting.py. CPU, JAX 0.9.1, \(n=2000\) (Gamma, GIG), \(n=600\) (GH, \(d=2\)).

3. What the comparison shows#

Gamma \((\alpha,\beta)=(2,1.5)\) — no Bessel, well-conditioned control. L-BFGS + softplus matches fit_mle in NLL and in \((\hat\alpha,\hat\beta)\). Adam reaches the same NLL in 1500 steps. The optimiser harness works; Gamma does not need it (fit_mle is a 2-D convex solve after one mean of \(t(x)=(\log x,\,x)\)).

Interior GIG \((p,a,b)=(1,2,1)\). L-BFGS + softplus reproduces fit_mle to working precision (\(\Delta\)NLL \(= -4.7\times 10^{-4}\) for both; parameters agree to printed digits). Wall-clock: fit_mle \(0.19\,\mathrm{s}\), L-BFGS \(2.57\,\mathrm{s}\) (\(\sim 14\times\)). Adam after 1500 steps is still \(0.275\) in hybrid parameter error and slightly worse NLL than the MLE.

Degenerate GIG. At large \(\sqrt{ab}\) or \(a/b\) extremes, the NLL ridge is nearly flat in \((p,a,b)\) and the MLE in classical coordinates is weakly identified. η-rescaled Bregman stays in a reasonable neighbourhood. L-BFGS + softplus matches that neighbourhood when the point is only mildly asymmetric, and drives \(b\) to underflow (\(10^{-37}\)) at \((p,a,b)=(1,10^4,10^{-3})\). Box-constrained NLL collapses \(b\) onto the bound. Adam can flip the sign of \(p\). Softplus cannot represent the Variance-Gamma boundary \(b=0\) exactly; EM’s degenerate-GIG branches can.

GH, \(d=2\). EM and L-BFGS reach the same NLL to \(\sim 10^{-4}\) (\(\Delta\)NLL \(\approx -0.013\) vs the truth on the interior draw). EM: \(0.45\)–\(1.0\,\mathrm{s}\) (25 iterations). L-BFGS: \(11\)–\(13\,\mathrm{s}\) (40 iterations), \(\sim 10\)–\(25\times\) slower. Adam at 200 steps (already \(4\)–\(8\times\) EM wall-clock) is \(0.09\)–\(0.19\) nats worse than the truth and has not left the initialisation neighbourhood. Subordinator \((p,a,b)\) differs across EM and L-BFGS at matched NLL (GH scale gauge).

The ML recipe is a valid MLE on well-conditioned interiors. It is the wrong default here: cost, first-order lag, GIG-boundary collapse, and the exact \(b=0\) boundary that VG occupies.

4. What this does not show#

A historical reason for avoiding gradient descent was \(\partial_\nu\log K_\nu\): older JAX Bessel wrappers (TensorFlow Probability’s log_besselk) returned a zero \(\nu\)-tangent, so Adam could not move \(p\). That bug is gone. log_kv is an ordinary JAX function (frozen whole-line quadrature; no custom_jvp). On the same GIG grid, jax.grad of the observed NLL with respect to \((p,a,b)\) matches a CPU finite difference of the same kernel, relative error \(\sim 10^{-9}\) (tests/test_gig_properties.py). Hessians of \(\log K_\nu\) are now \(\mathrm{Var}(u)>0\) (S10); the first-order NLL check was never the blocker.

Quasi-Newton on the GIG NLL can recover the interior MLE. We still do not lead with it, for the reasons in § 3, not because autodiff through \(K_\nu\) is broken.

5. Bring your own optimiser#

log_prob is JIT-able and differentiable. If a surrounding model must co-optimise GIG parameters with non-normix parameters, the density is the primitive — not a method="adam" fitter. Interior \((p,a,b)\) only; the caveats in § 3 still apply.

import jax
import jax.numpy as jnp
from jax.nn import softplus
from normix import GIG

def unpack(phi):
    return phi[0], softplus(phi[1]), softplus(phi[2])

def nll(phi, X):
    p, a, b = unpack(phi)
    dist = GIG(p=p, a=a, b=b)
    return -jnp.mean(jax.vmap(dist.log_prob)(X))

# jax.grad(nll)(phi, X)  — plug into JAXopt L-BFGS, Optax, or a custom loop

Adam and L-BFGS remain unsupported as public fitters: there is no grid point where they beat fit_mle / EM, and shipping them means owning \(b\to 0\) collapse, the GH gauge, and step-count / learning-rate knobs whose only documented advice would be “prefer fit_mle”.

Reopen the decision if a concrete caller needs one of:

  • streaming / minibatch where a full E-step over \(n\) is infeasible;

  • joint NLL with parameters that are not a normix exponential family.

The first abstraction in either case is a documented nll(params, X) helper, not an in-tree Adam.

6. Cross-references#