Exponential Family Core

Exponential Family Core#

Scope. Why every distribution lives behind one log-partition function, how the triad pattern keeps gradient/Hessian and JAX/CPU evaluation consistent, and how these primitives plug into the η→θ Bregman solver.

Where things live. Module hierarchy and the full triad table are in the API Reference. This file records the rationale.

1. One Function, Three Parametrizations#

Each distribution is described by a single convex log-partition \(\psi(\theta)\). Densities and the \(\eta\to\theta\) map are derived from it. On a full exponential family, \(\eta = \nabla\psi(\theta)\) and the Fisher information is \(\nabla^2\psi(\theta)\). Joint variance-gamma, normal-inverse-gamma, and normal-inverse-Gaussian are the exception in § 2.4. The three parametrizations still form a triangle:

classical (μ, σ², α, β, …)   ←→   natural θ   ←→   expectation η = ∇ψ(θ)
                  via from_classical / natural_params                ↑
                                                                from_expectation

from_natural, natural_params, and from_expectation are the only three constructors a user ever needs. from_expectation(η) runs the universal Bregman solver min_θ [ψ(θ) − θ·η] from fitting/solvers.py.

Why it matters: distributions never need to know about parameter conversions other than their own classical mapping. The conversion graph collapses into derivatives and Bregman minimisation of a single function.

2. The Log-Partition Triad#

Three functions × two backends:

                     JAX (JIT-able)                  CPU (numpy/scipy)
log-partition        _log_partition_from_theta        _log_partition_cpu
gradient             _grad_log_partition              _grad_log_partition_cpu
Hessian              _hessian_log_partition           _hessian_log_partition_cpu

Tier	Default	Override when
1: `_log_partition_from_theta`	abstract	always (it is the distribution)
2: `_grad_`, `_hessian_` (JAX)	`jax.grad` / `jax.hessian`	closed form, or one `log_kv_moments` bundle (GIG)
3: `_log_partition_cpu` etc.	numpy wrappers around the JAX version	Bessel-dependent distributions need the same kernel in NumPy

2.1 Why `@classmethod` for the triad#

Tier 2 needs to differentiate the subclass’s _log_partition_from_theta, not the parent’s. @classmethod gives cls, which routes to the right implementation. @staticmethod cannot — it would always call the parent. cls is not traced by JAX, so this is JIT-safe.

2.2 Why the CPU tier exists#

The CPU tier exists so solve_bregman(backend='cpu') never dispatches through JAX. GIG’s CPU triad is the same moment-quadrature kernel in NumPy (Solvers and Bessel Functions § 3). Distributions that don’t call log_kv (Gamma, InverseGamma, InverseGaussian) inherit the default np.asarray(jax_version(...)) wrappers — they pay nothing for the CPU tier.

2.3 Fusion vs coordinate separation#

Sharing Bessel work between \(\nabla\psi\) and \(\nabla^2\psi\) happens inside normix.utils.bessel.log_kv_moments(), below the triad. GIG’s _grad_log_partition and _hessian_log_partition are affine images of one BesselMoments bundle (\(\eta\) = scaled mean, \(H = D\,\mathrm{cov}\,D\)). The triad remains three θ-space classmethods; the solver still applies the φ↔θ chain rule (S2). A combined _grad_hess on the distribution would re-introduce a φ-space chain rule the distribution must understand, and is not used.

An earlier GIG gradient used Bessel ratios and a finite-difference \(\partial_p\) (ee8cf8b). PR #113 replaced that stencil with this bundle. jax.grad through the degenerate jnp.where still fails for this gradient.

2.4 Ambient moments on the curved joints#

Joint variance-gamma, normal-inverse-gamma, and normal-inverse-Gaussian keep the generalized-hyperbolic statistic \(t = [\log y,\, 1/y,\, y,\, x,\, x/y,\, \mathrm{vec}(xx^\top/y)]\). expectation_params and fisher_information return \(E[t]\) and \(\mathrm{Cov}[t]\) under the special-case law. Both are assembled from the subordinator’s moments of \((\log Y, 1/Y, Y)\) in that GIG order, plus the Gaussian conditional blocks. An entry that needs a moment of \(Y\) which does not exist is \(+\infty\).

Differentiating the restricted log-partition is the wrong source: that \(\psi\) ignores the frozen slot, so \(\nabla\psi_{\mathrm{ext}} = \nabla\psi_{\mathrm{GH}} - m v\), and the Hessian is indefinite for variance-gamma and normal-inverse-gamma. Lifting the law to joint generalized hyperbolic and differentiating there has the same defect at \(b = 0\) or \(a = 0\). Joint generalized hyperbolic itself stays on the triad. In the interior, \(\nabla\psi = E[t]\) and \(\nabla^2\psi = \mathrm{Cov}[t]\).

2.5 Valid input for joint `from_expectation`#

The canonical input is a NormalMixtureEta pytree. A flat array uses the restricted-\(\psi\) Bregman solver. For variance-gamma, normal-inverse-gamma, and normal-inverse-Gaussian, \(\nabla\psi\) omits the frozen slot, so it is not the ambient moment. Only the pytree path inverts ambient \(\eta\) for those three families. Interior joint GH is the exception: flat \(\eta = \nabla\psi\) is the ambient moment. That GH solve is still numerical.

The normal block needs \(\mathbb{E}[Y] > 0\), \(\mathbb{E}[1/Y] > 0\), and \(\mathbb{E}[Y]\,\mathbb{E}[1/Y] \ge 1\). The update sets \(D = 1 - \mathbb{E}[Y]\,\mathbb{E}[1/Y]\) and floors \(|D|\) at SAFE_DENOMINATOR. A violation of the product bound is not rejected. The scatter must stay positive definite, or the Cholesky factor fails.

Subordinator bounds on the pytree path:

Variance-gamma needs \(\mathbb{E}[Y] > 0\) and \(\mathbb{E}[\log Y] < \log\mathbb{E}[Y]\). Equality is the infinite-shape limit. A concrete Gamma.from_expectation call raises ValueError on that \(\eta\). A traced \(\eta\) skips the check.
Normal-inverse-gamma needs \(\mathbb{E}[1/Y] > 0\) and \(\mathbb{E}[\log Y] > -\log\mathbb{E}[1/Y]\). That is Jensen’s inequality for the Gamma law of \(1/Y\). The inverse-gamma map does not raise when the inequality fails.
Normal-inverse-Gaussian needs \(\mathbb{E}[Y] > 0\) and \(\mathbb{E}[1/Y] > 1/\mathbb{E}[Y]\), so \(\lambda > 0\). A smaller inverse moment is clamped at LOG_EPS.
Generalized hyperbolic needs finite moments with \(\mathbb{E}[Y] > 0\), \(\mathbb{E}[1/Y] > 0\), \(\mathbb{E}[Y]\,\mathbb{E}[1/Y] > 1\), and \(-\log\mathbb{E}[1/Y] < \mathbb{E}[\log Y] < \log\mathbb{E}[Y]\). Equality is a point-mass boundary. These conditions are necessary and not sufficient. A concrete GIG.from_expectation call raises ValueError before the solve. A traced invalid target returns a NaN model. A failed host solve raises SolverDidNotConverge. Under jit the natural parameters are NaN.

3. Bregman Solver Interface#

fitting/solvers.py exposes solve_bregman(f, η, θ₀, ...) for any convex \(f\) — not just log-partitions. The interface uses θ-space only for gradient/Hessian; bounds are handled by the solver via _setup_reparam (φ ↔ θ). Distributions never see the reparametrization.

solve_bregman(f, eta, theta0, *, backend, method, bounds,
              grad_fn, hess_fn, max_steps, tol, verbose)

Decision	Choice	Why
`grad_fn` + `hess_fn` (vs combined `grad_hess_fn`)	separate	Distributions don’t know about φ-space; solver applies chain rule via `jax.jacobian(to_theta)`
Newton implementation	hand-rolled `lax.while_loop`	No JAX library provides a Newton minimizer that accepts a user-supplied Hessian (Optimistix Newton is root-finding only, JAXopt has no Newton). The loop stops when the pre-step residual is inside `tol`.
Multi-start	JAX `lax.while_loop` over ordered starts; Python `for` for CPU	JAX stops at the first finite converged start
Result type	`BregmanResult` with `Any`-typed scalars	Survives `lax.scan` without `ConcretizationTypeError`
Cached JIT	module-level `_gig_jax_newton_jit` via `make_jit_newton_solver`	The GIG warm-start hot path otherwise retraces per call

backend × method matrix:

backend	method	bounds	What runs
`jax`	`newton`	reparam	hand-rolled `lax.while_loop`, autodiff or supplied `grad_fn`/`hess_fn`
`jax`	`lbfgs`	native (LBFGSB)	JAXopt `LBFGSB`
`jax`	`bfgs`	reparam	JAXopt `BFGS`
`cpu`	`newton`	reparam	SciPy `minimize(method='trust-exact')`
`cpu`	`lbfgs`	native	SciPy `minimize(method='L-BFGS-B')`
`cpu`	`bfgs`	reparam	SciPy `minimize(method='BFGS')` in \(\phi\)-space

GIG warm-start hot path: backend='cpu', method='lbfgs' (scipy’s L-BFGS-B

numpy log_kv) avoids GPU dispatch on this 3-D scalar problem.

BregmanResult.num_steps counts steps taken. A one-step quadratic Newton solve reports 1. A start already inside the tolerance reports 0. The pre-step check that stops the loop is not a step. reason is an integer code. Eager and traced calls store that code as the same type. The codes are 0 converged, 1 atol, 2 max steps, and 3 non-finite. kkt_residual is the free residual of \(\nabla f(\theta) - \eta\), with a bound multiplier removed.

CPU BFGS and trust-exact apply box bounds in \(\phi\)-space. The derivatives of that diagonal map are closed form. SciPy gtol is the \(\phi\)-gradient. Acceptance is the free \(\theta\) residual. CPU L-BFGS-B keeps the bounds natively.

4. Pre-1.0 Decisions Recorded Here#

4.1 D3 — `MultivariateNormal` as `ExponentialFamily`#

Promoted from a plain eqx.Module to a full ExponentialFamily. EF structure:

Component	Expression
\(t(x)\)	\([x,\;\operatorname{vec}(xx^\top)]\), shape \((d + d^2,)\)
\(\theta\)	\([\Sigma^{-1}\mu,\;-\tfrac12\operatorname{vec}(\Sigma^{-1})]\)
\(\log h(x)\)	\(0\)
\(\psi(\theta)\)	\(\tfrac12\theta_1^\top\Lambda^{-1}\theta_1 - \tfrac12\log\|\Lambda\| + \tfrac d2\log(2\pi)\), \(\Lambda = -2\,\mathrm{reshape}(\theta_2)\)

vec uses row-major (ravel()) throughout. All conversions are analytical (Tier 2 override) — no Bregman solver is ever invoked for MVN. _log_partition_from_theta uses Cholesky of \(\Lambda\) for numerical stability; log_prob overrides the inherited EF formula with a direct Cholesky path (more efficient).

4.2 D4 — Keep `jaxopt` for now#

JAXopt is unmaintained (last release 0.8.3) and emits a DeprecationWarning on import. We keep it: it is the only pure-JAX library with a native box-constrained quasi-Newton (LBFGSB). Migrating to optax.scale_by_lbfgs loses the convergence loop; optimistix.LBFGS lacks box constraints.

The deprecation warning is suppressed in normix/__init__.py. Migration recipe (when JAXopt breaks): wrap optax.scale_by_lbfgs in jax.lax.while_loop (~100 lines), then drop jaxopt.

4.3 Constraints handling#

jnp.maximum(x, LOG_EPS) (clamp), not paramax. Reasons:

The reparametrization we need is 8 lines, fully understandable.
EM does not need gradients through the constraints (it parameterises θ in the constrained space, not log-space).
No extra dependency.

jnp.where is preferred over lax.cond whenever possible: it is vmap-compatible without changing the trace, and the clamping prevents NaN gradients without branch divergence.

Natural-parameter input is validated, not clamped. GIG.from_natural raises (via eqx.error_if, compatible with jit, vmap, and grad) when a rate coordinate theta[1] or theta[2] is positive; zero is still allowed so the Gamma and InverseGamma limits remain reachable. This is separate from the M-step GIG_CLAMP_LO/GIG_CLAMP_HI window, which is unchanged.

Commit 9346ad6 selected one Bessel regime with lax.cond. PR #113 removed that dispatch and the rest of the four-regime kernel.

4.4 Module-level functions are forbidden#

Distribution behaviour lives on the class as @classmethod or @staticmethod. No _helper(self.alpha, ...) module-level functions that close over attributes — they leak the class API into module globals and are hard to override.

5. Cross-References#

η→θ optimization for GIG: Solvers and Bessel Functions.
Why not likelihood gradient descent: Why not gradient descent.
Theory: GIG distribution, EM algorithm.