Skip to content

A Gauss–Newton mass matrix for NUTS

gauss_newton_mass turns a fit into a mass matrix for NUTS, the No-U-Turn Sampler of Hoffman & Gelman (2014) as implemented in numpyro (Betancourt 2017 is a good introduction). Use it to sample any posterior with one peak: images, and orbits and binaries with fitted error scales and gains. Sampling the posterior shows it on an image.

The matrix

Near the maximum a posteriori, a posterior with one peak is close to a Gaussian whose precision, in the unconstrained coordinates \(z\) that numpyro samples, is \(F = J^\top J + \mathrm{diag}(P)\). \(J\) is the Jacobian, with respect to \(z\), of a vector of curvature rows: the whitened residuals of every dataset (with diagonal errors held fixed), and the residuals of likelihood terms and angle-vector priors. Since \(\delta^\top J^\top J\delta = |J\delta|^2\), \(F\) is positive semi-definite. \(P\) is each coordinate's own curvature from its elementwise prior and transform, including the Jacobian of numpyro's bijection, as NUTS's potential has it (a Normal prior on the real line gives \(1/\sigma^2\)).

The inverse mass matrix is the covariance \(F^{-1}\). It whitens that Gaussian, so directions the data fix tightly take the same step as those left to the prior; without it, the step size shrinks to suit the tightest direction. On a 62² Gaussian-field image fitted to 588 AMI observables it cut NUTS from its full tree depth (1023 leapfrog steps) per draw to 63; on an 8-parameter orbit, from 87 steps to about 6.

The matrix is a preconditioner: errors in it cost efficiency, not correctness.

Fitted error terms

With noise= the matrix covers the error terms too, under their numpyro site names ("noise.vis_scale", "noise[1].phi_scale"). The likelihood then has a second term, \(\sum_k\log\sigma_k\), that is not a sum of squares, but its expected curvature is known for Gaussian data: a scale \(s\) multiplying \(N\) errors adds \(2N\) in \(\ln s\), whatever the model, giving \(\mathrm{blockdiag}(J^\top J, 2N)\). virgil builds this, the Fisher information, with extra rows \(\sqrt2\log\sigma_k\) beside the residuals, and added errors (vis_error_rel) get their cross terms with the model parameters the same way. For phases with errors above about 0.5 rad (a von Mises rather than Gaussian likelihood) the scale's curvature is overstated by up to a few times, which costs only efficiency.

Gains

Gains and closure-phase offsets correlated across channels are integrated out analytically: their covariance whitens the residuals, and \(F\) differentiates through that whitening. That gives a good curvature for the model parameters; for fitted gain widths it is a positive semi-definite heuristic, not the Fisher information.

What is dropped

The exact Hessian also has \(\sum r\,\nabla^2 r\), each residual times the model's curvature. It is small when the residuals look like noise or the model is nearly linear over the error bars. It matters when the fit is poor (\(\chi^2/N \gg 1\)) or the model bends on the scale of the errors; \(F\) is then still a safe direction but a poorer guide to the error bars. A likelihood term's log_norm (a fitted RV jitter's, or the Jacobian of position_angle_prior) has no rows, so its curvature is left out.

When not to use it

  • Several modes (orbit aliases, the 180° flip of a symmetric binary): the matrix describes one mode, so sample each separately.
  • Nearly circular orbits (\(e \lesssim 0.05\)): \(\omega\) and the periastron time are degenerate, and NUTS reports divergences.
  • Companions too faint to detect, whose posterior fills the prior.
  • Closure phases that can pass through ±180°, where the likelihood has sharp walls.
  • \(\chi^2/N \gg 1\) (a misspecified model).

Hold field hyperparameters (\(\sigma\) and \(\ell\), chosen for example by log_evidence) fixed: the curvature depends on them, so a matrix computed at one \(\sigma\) and \(\ell\) is wrong once they move. Put every sampled parameter in priors: a tightly constrained one left out (such as an image's flux) keeps its identity mass and a tiny step size.

Size

The dense matrix has \(n^2\) entries, so above 4096 coordinates (images larger than 64²) you get a diagonal mass and a UserWarning. Its entries are the diagonal of \(F^{-1}\), computed without forming the matrix, when there are fewer rows than coordinates (an image), and \(1/\mathrm{diag}(F)\) otherwise.