pyro-ppl / numpyro / 39422435324

93%

push

github

fix(contrib.hsgp): numerically stable spectral densities with finite gradients for large m (#2365)

* fix(contrib.hsgp): compute SE and Matern spectral densities in log space

In float32, exp(-w^2 length^2 / 2) underflows to exactly 0 for most of the
first m eigenvalues once m is a few dozen (92/100 zeros at length=1, ell=1).
The Matern density also overflows through gamma(nu + dim/2) and (2 nu)^nu
for nu >~ 30. Evaluate both densities in log space and expose
log_diag_spectral_density_{squared_exponential,matern}; the existing
functions become exp() of the log-space result.

* fix(contrib.hsgp): stable log-space rational quadratic spectral density

log(exp(-z) * kve(nu, z)) underflows to log(0) for large z, which gives
NaN gradients once sqrt() is applied; use log_bessel_kve(nu, z) - z
instead. The regular branch of the jnp.where at w=0 is now evaluated at a
safe argument so its log(0) can no longer leak NaN into the gradient of
the S(0) branch. Expose log_diag_spectral_density_rational_quadratic and
drop the unused modified_bessel_second_kind helper.

* fix(contrib.hsgp): periodic coefficients via exponentially scaled Bessel

exp(a) * ive(j, a) / exp(a) with a = 1 / length^2 overflows to inf / inf
= nan in float32 for length <~ 0.1. Use ive(j, a) = I_j(a) exp(-a)
directly and drop the unused modified_bessel_first_kind helper.

* fix(contrib.hsgp): finite gradients for HSGP approximations with large m

sqrt(S) has an infinite derivative at S = 0, so once the spectral density
underflows (most basis functions for m ~ 100 in float32) the gradient of
the approximation w.r.t. the kernel hyperparameters is inf * 0 = nan and
initialize_model cannot find valid initial parameters, in both the
centered and the non-centered parameterization (the clip in the centered
path zeroes the cotangent but 0 * inf is still nan). Build the square
root of the density as exp(0.5 * log S) instead, whose gradient at an
underflowed entry is exactly 0. Document the remaining limi... (continued)

53 of 54 new or added lines in 1 file covered. (98.15%)

15770 of 16984 relevant lines covered (92.85%)

1.7 hits per line