Thanks for jumping in @monnahc. I should have noticed this was a good example for SNUTS! Here’s a link to the paper if people want to read more:
Leveraging Sparsity to Improve No-U-Turn Sampling Efficiency for Hierarchical Bayesian Models. 2026. Cole C. Monnahan, Kasper Kristensen, James T. Thorson, Bob Carpenter.
This can also be used to initialize another mass matrix estimator, but for these kinds of models, the basic SNUTS mass matrix estimator should be really good.
Thanks @Bob_Carpenter I am tagged on this b/c it was split off from my SNUTS post so I was following it. I wasn’t trying to advertise SNUTS here, but just point out that nothing unusual popped out when integrating with the Laplace approximation that would indicate why some chains might adapt poorly. It seems like a straightforward geometry at least…
Yes, I deliberately picked this to address / eliminate points raised and try to narrow down what we’re talking about.
My naive theory is that the denominator is significantly unstable in early sampling, which affects both Nutpie and Walnuts as both use it, but Walnuts hurts more because it uses a raw denominator whereas Nutpie does try to condition it a bit. And of course Stan doesn’t have it so it neither benefits nor helps. Stan also is able to start over with better information with its windows.
I tried a few other things, like using a sort of student-t type conditioning weight on the denominator in one approach and the metric itself in another, and tried down weighting the denominator as warmup moves forward, None of those experiments were a success for me.
This got so long I had to insert subtitles.
I wish we could pin this comment somewhere. It was my refrain through early Stan development when people told me their samplers were faster in terms of iteration/second, but I tried to tell them that didn’t matter because (a) they got the wrong answer when run a long time and/or (b) they had lower ESS/second than HMC.
Me, too! That’s probably down to initialization. Because if they’re strong modes, neither Nuts nor Nutpie adaptation is likely to escape.
As @aseyboldt is pointing out, a serious problem we have in practice is that ESS and R-hat can be fooled by samplers that mix well through a subregion of the posterior. This was particularly common back in the Gibbs/Metropolis era when strong correlation structure would frustrate the samplers (there’s a great example of a highly correlated multivariate normal in the original Nuts paper). The obvious case is multimodal posteriors where you get stuck in a mode, but as @andrewgelman keeps reminding us, the funnel is just a continuous mixture. So you can get stuck in the neck and mix there with a too-low step size, or you can get stuck in the mouth with a too high step size. A good diagnostic is when the step size is high enough that it makes the occasional excursion into high curvature regions, then can’t get out because of detailed balance requirements (those pesky Metropolis rejections).
Trying to answer it with square error (a proper scoring metric) introduces a chicken-and-egg problem of not being able to get the true answer without effective sampling. The best thing I know to answer this question independently is something like kernel Stein divergence. Unlike ESS/R-hat estimates, which are necessary but not sufficient conditions for getting the right answer, kernel Stein discrepancy matching is both necessary and sufficient. I have a post recently on @andrewgelman’s blog, “Stein’s method, learning and inference -or- how to really monitor convergence and thin chains.”
This is the key insight to what’s going on. Check out this plot I (with Claude’s help, of course) made for the funnel, the density of which is
The plots are histograms of the values indicated on the column headers for the variables indicated on the row headers. The top row is for the log variance variable (\theta_1) and the bottom row for one of the coefficients (\theta_2). The columns are histograms of draws for (1) theta, (2) score(theta), (3) 1/theta, and (4) 1/score(theta).
Note the log-absolute scale x-axis. The red dashed vertical lines are at the minimum and maximum value out of 2M draws (you can exactly simulate from the funnel, but it has very long tails due to the \exp(\theta_1/2) term for the scale of the coefficients).
To get the inverse mass matrix (working on this is disorienting with all the inversion) we need for the leapfrog algorithm, Nuts and Nutpie do the following:
thetatheta with the inverse of the variance of score(theta).We also need the Cholesky factor of the mass matrix to initialize draws, which is the just the elementwise square root of the inverse of the diagonal variance matrix.
This isn’t naive—I’m pretty sure this is exactly what’s going on. Unprompted, it was also Claude’s guess as to what’s happening. Estimating the variance of theta is very stable because values almost all fall in the (0, 10) range for log variance and the (0, 100) range for the coefficients. Contrast this with estimating the variance of the score and inverting it. The plot includes 1 / score(theta), for which values can get as high as 1 million for the log scale parameter and 10 million for the coefficient. Now things aren’t quite as bad as this makes out unless the sampler gets stuck in one region, because we’re actually estimating the variance of the scores, then inverting, not trying to estimate the variance of the inverse scores.
So I need to go in and instrument exactly where our estimates are going off the rails and see if there’s some kind of simple regularization that will fix this problem. I’m curious if anyone like @aseyboldt or @ssp3nc3r know what Nutpie’s doing to regularize—I don’t recall regularization there from the paper, but I’m about to go check.
If we shift from using a mass matrix to preconditioning, then we only need the Cholesky factor of the inverse mass matrix (e.g.., the Cholesky factor of the posterior variance in Stan) to do the preconditioning and the final transformation of draws back to the original scale.
What happens in practice is that if we run 16 chains of Walnuts, a few of them get stuck and can’t move. In NUTS, getting stuck usually happens due to walking into a region of high curvature with too high of a step size and not being able to get out. Walnuts shouldn’t have the problem of not being able to get out of areas of high curvature, but it will be defeated when the mass matrix is off by several orders of magnitude, as it seems to be in these examples.
I’m wondering now if it might be as easy as replacing the geometric mean with an arithmetic mean. When we take geom_mean(a, b), if they are off by orders of magnitude, you get an order halfway in between, roughly sqrt(max(a, b)). When we take mean(a, b) and they are off by orders of magnitude the result is roughly max(a, b)/2. If the max is 10000, then sqrt(max) is 100, whereas max/2 is 5000.
The good news for Walnuts is that even if we have to back off using Nutpie adaptation, we’ll be able to do Stan-style adaptation just as robustly as Stan, but much faster because of the online adaptation and the improvement from using Adam rather than dual averaging for step size.
I’m probably not the person to comment directly, but here’s what I saw digging around in the code itself, Walnuts uses
inv_mass = sqrt(draw_var / score_var)
raw. The paper had this in it:
Minimizing KL divergence for mass matrix estimation using either scores or draws alone is suboptimal even for Gaussian targets, because it does not penalize both large and small eigenvalues of the covariance matrix equally.
and which I think in Nutpie’s code seems to compute the same ratio but clamps the result and non-finite/zero:
let val = (draw_var / grad_var).sqrt(); // same as Walnuts
if (!val.is_finite()) | (val == 0f64) {
// fall back to `fill_invalid`
} else {
let val = val.clamp(clamp.0, clamp.1); // clamp to a fixed range
}
So the conditioning seems to be a per-parameter clamp(val, clamp.0, clamp.1) plus a non-finite/zero fallback — where Walnuts feeds the raw ratio straight through. The draw-only and grad-only variants, array_update_var_inv_std_draw / _grad, apply the same clamp to their single variance.
I didn’t like the idea of a hard clamp and thought that could be contributing to the attenuated (compared with the raw form) issues I still saw in Nutpie, so I tried a Student-t discount instead of the hard clamp: down-weight each incoming draw/gradient by w = (ν+1)/(ν+z²) — z standardized against that parameter’s own running moments, with a consistency correction so the scale doesn’t self-shrink — so implausible updates are softly discounted rather than hard-bounded.
Assuming I even got both the idea and the code right (big if), it didn’t seem to help.
Thanks, @ssp3nc3r. I don’t recall the clamping from our paper. It sounds like you understand what’s going on as well as the rest of us, if not better!
I thought I’d investigate what happens with Neal’s funnel for the variance of the draws, inverse variance of the scores, and the geometric mean of the two.
For reference, the 10D standard normal is perfect. The unit mass matrix estimated by the geometric mean is perfect.
d0 d1 d2 d3 d4 d5 d6 d7 d8 d9
-----------------------------------------------------------------------------------------------------------------------------------------------------
run 1
(1) var(draws) 8.809e-01 9.700e-01 1.012e+00 1.087e+00 9.067e-01 7.358e-01 8.007e-01 8.020e-01 1.185e+00 1.017e+00
(2) 1/var(scores) 1.135e+00 1.031e+00 9.880e-01 9.198e-01 1.103e+00 1.359e+00 1.249e+00 1.247e+00 8.439e-01 9.837e-01
(3) geo-mean 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00
run 2
(1) var(draws) 1.143e+00 1.072e+00 9.109e-01 8.142e-01 1.231e+00 1.072e+00 1.166e+00 9.737e-01 1.039e+00 9.505e-01
(2) 1/var(scores) 8.750e-01 9.329e-01 1.098e+00 1.228e+00 8.122e-01 9.327e-01 8.579e-01 1.027e+00 9.629e-01 1.052e+00
(3) geo-mean 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00 1.000e+00
And here’s what it looks like for Neal’s funnel. The first argument is the log variance variable, which we know should be normal(0, 3) distributed. The rest have much wider tails. In this case, the variance of draws and inverse variance of scores are nowhere near each other. I’m not sure what the best static mass matrix is here because the eigenstructure of the target density changes radically between the mouth and neck of the funnel. The log variance should move more quickly in the neck and the coefficients should move more quickly in the mouth. So probably something close to a unit would be best.
v x0 x1 x2 x3 x4 x5 x6 x7 x8
-----------------------------------------------------------------------------------------------------------------------------------------------------
run 1
(1) var(draws) 8.3313 18.2773 35.6691 19.2342 25.6999 29.5199 15.3369 36.1177 23.9525 14.5868
(2) 1/var(scores) 0.2420 0.0652 0.0368 0.2226 0.0131 0.0555 0.0362 0.0781 0.0293 0.0836
(3) geo-mean 1.4201 1.0918 1.1462 2.0692 0.5794 1.2804 0.7450 1.6800 0.8383 1.1040
run 2
(1) var(draws) 10.3134 56.7193 165.0118 420.4187 82.3591 139.1720 142.7161 65.6530 30.1875 192.1642
(2) 1/var(scores) 0.2554 0.0363 0.0064 0.0084 0.0151 0.0128 0.0097 0.0401 0.0101 0.0114
(3) geo-mean 1.6230 1.4358 1.0251 1.8804 1.1134 1.3354 1.1781 1.6223 0.5514 1.4774
...
run 18
(1) var(draws) 11.2267 2420.5693 1188.5712 420.1101 1312.5732 1164.9731 1247.5511 852.3186 1485.4239 2333.0946
(2) 1/var(scores) 0.3063 0.0365 0.0523 0.0154 0.0289 0.0446 0.0251 0.0462 0.0306 0.0142
(3) geo-mean 1.8544 9.3968 7.8864 2.5472 6.1627 7.2058 5.5914 6.2738 6.7422 5.7520
run 19
(1) var(draws) 10.2156 32.1898 21.7779 14.2635 19.1638 29.2772 23.2764 15.6435 25.5616 20.9609
(2) 1/var(scores) 0.2149 0.0219 0.0056 0.0151 0.0073 0.0270 0.0321 0.0179 0.0302 0.0084
(3) geo-mean 1.4816 0.8400 0.3478 0.4636 0.3741 0.8896 0.8648 0.5290 0.8791 0.4205
run 20
(1) var(draws) 9.1248 63.5939 97.9800 36.0413 10.9327 20.1982 110.2366 73.7308 13.3208 20.0976
(2) 1/var(scores) 0.2523 0.0179 0.0166 0.0131 0.0531 0.0134 0.0226 0.0408 0.0197 0.0273
(3) geo-mean 1.5173 1.0661 1.2749 0.6865 0.7622 0.5195 1.5801 1.7334 0.5117 0.7408
I looked up the values of clamp.0, clamp.1 in nuts-rs/src/transform/adapt /diagonal.rs, and these aren’t going to help us. The last argument of update_diag_draw_grad is clamp, so clamp.0 = LOWER_LIMIT and clamp.1=UPPER_LIMIT.
const LOWER_LIMIT: f64 = 1e-20f64;
const UPPER_LIMIT: f64 = 1e20f64;
...
mass_matrix.update_diag_draw_grad(
math,
draw_mean,
grad_mean,
draw_var,
grad_var,
None,
(LOWER_LIMIT, UPPER_LIMIT),
);
So I don’t think this is going to save Nutpie’s version.
Here’s what the Rosenbrock distribution looks like in both mild and standard (extreme) form.
=== mild banana (a=1.0, b=1.0) ===
x1 x2
-------------------------------
run 1
(1) var(draws) 0.926 2.223
(2) 1/var(scores) 0.273 1.095
(3) geo-mean 0.503 1.560
run 2
(1) var(draws) 1.244 4.928
(2) 1/var(scores) 0.246 1.172
(3) geo-mean 0.553 2.404
and
=== sharp banana (classic Rosenbrock) (a=1.0, b=100.0) ===
x1 x2
-------------------------------
run 1
(1) var(draws) 0.962 1.876
(2) 1/var(scores) 0.003 0.011
(3) geo-mean 0.051 0.143
run 2
(1) var(draws) 1.090 1.557
(2) 1/var(scores) 0.003 0.010
(3) geo-mean 0.061 0.127
I’m only showing a couple runs here because they all looked very similar. This is another varying curvature case where it’s not clear we can do much better than a unit mass matrix.
Edit. Just adding more here rather than continuing to spew posts.
I took Rosenbrock and looked at two alternatives. One applied unit norm to variance vectors before geometric average and one after. Here’s what that looks like:
=== mild banana (a=1.0, b=1.0) ===
x1 x2
---------------------------------------------
run 1
(1) var(draws) 0.926 2.223
(2) 1/var(scores) 0.273 1.095
(3) geo-mean(norm) 0.307 0.952
(4) norm(geo-mean) 0.305 0.946
run 2
(1) var(draws) 1.244 4.928
(2) 1/var(scores) 0.246 1.172
(3) geo-mean(norm) 0.224 0.975
(4) norm(geo-mean) 0.224 0.974
=== sharp banana (classic Rosenbrock) (a=1.0, b=100.0) ===
x1 x2
---------------------------------------------
run 1
(1) var(draws) 0.962 1.876
(2) 1/var(scores) 0.003 0.011
(3) geo-mean(norm) 0.334 0.943
(4) norm(geo-mean) 0.329 0.930
run 2
(1) var(draws) 1.090 1.557
(2) 1/var(scores) 0.003 0.010
(3) geo-mean(norm) 0.433 0.901
(4) norm(geo-mean) 0.424 0.882
Question: Does anybody understand why the scales are so different between the variance of the draws and inverse variance of the scores? In a standard normal, they’re identical.
I reproduced the collapse using the model (y ~ normal(0, exp(lsig)),
2-level non-centered, G=15-25/J=10-15). With WALNUTS’ upstream default
(max_step_halvings=5), 8/8 chains on my reduced repro froze completely — draw sd ~1e-15
(machine epsilon), inv_mass collapsed to 1e-22, split-R̂ up to 5e15. Exact-replaying both the
geomean and draw-cov estimators on a frozen chain’s warmup trajectory shows the chain is frozen
from the very first transition — |Δθ|=0 at every timestep, because max_step_halvings=5
can’t find a valid first leapfrog at the stiff init. Both estimators collapse together in that
state (draw-cov drops three orders of magnitude too), so at least this specific freeze isn’t
distinguishing score-vs-draw-cov — it’s whether the chain survives its first step at all.
Raising max_step_halvings to ≥20 rescues this failure mode but is necessary, not sufficient:
on the faithful model it only recovers ~50% of chains (4 healthy / 1 collapsed / 3
thrash-to-timeout of 8 seeds), vs 16/16 on a milder mean-scale (non-log) version of the same
hierarchy. So there’s a real warmup-machinery gap here, independent of the metric question
(Stan’s find_reasonable_epsilon + windowed step/metric adaptation with step re-init) — worth
fixing regardless of how the metric-scale question below resolves.
Take Neal’s funnel: v ~ N(0, σ), x_i | v ~ N(0, e^{v/2}), so Var(x_i | v) = e^v. Two
population quantities:
Var(draws)_i — the marginal variance of x_i. Since E[x_i | v] = 0,Var(draws)_i = E[e^v].Var(score)_i — the variance of score_i = ∂/∂x_i log p = -x_i · e^{-v}. Conditional onv, Var(score_i | v) = e^{-2v} · Var(x_i|v) = e^{-2v} · e^v = e^{-v}, so marginallyVar(score)_i = E[e^{-v}].Now the punchline: v is symmetric about 0 (v =d= -v, since v ~ N(0,σ)), so
E[e^v] and E[e^{-v}] are literally the same integral under the substitution v → -v.
For any σ:
Var(draws)_i = E[e^v] = e^{σ²/2} = E[e^{-v}] = Var(score)_i
Var(draws)_i and Var(score)_i are not just similar order — they are exactly equal, for
every σ, by this symmetry alone. (3-line check in any CAS: Expectation[Exp[v], NormalDistribution[0,s]] and Expectation[Exp[-v], NormalDistribution[0,s]] both return
E^(s^2/2), identically, for symbolic s.)
So why do var(draws) and 1/var(scores) look so different in #26? Because the metric
formula sqrt(Var(draws)/Var(score)) needs 1/Var(score), i.e. the reciprocal of the shared
value — and e^{σ²/2} and its reciprocal e^{-σ²/2} are far apart in absolute terms whenever
e^{σ²/2} itself is far from 1 (at σ=3 the population values are e^{4.5}≈90 and ≈0.011, a
population gap of e^{σ²}=e^9≈8100×). But sqrt(Var(draws) · 1/Var(draws)) = 1 for any
number, trivially — that’s why the population geomean is exactly 1 (and why the 10-D
standard-normal check in #26 is perfect: that’s the funnel with σ=0, so both quantities are
already 1).
Two things to be careful about here. First, those are population values; your #26 tables
are finite-sample estimates and they scatter enormously around them — for the coefficient
coordinates, var(draws) ranges from ~15 (run 1) to ~2400 (run 18) against a population value
of 90, and the per-coordinate var(draws):(1/var(score)) ratio from ~90 to ~160,000 against a
population e^9≈8100. That scatter is not noise around the answer — it is the phenomenon:
two estimators of one and the same population number, disagreeing by 1–2 orders of magnitude in
finite samples. Second, that means the geomean metric is exactly right in expectation, and
the practical failure is purely a finite-sample estimation problem, not a formula problem. The
two sample estimators of the same population quantity (e^{σ²/2}) are dominated by opposite
tails of the underlying lognormal: Var(draws)'s sample estimate is pulled up by
large-positive-v excursions (the funnel’s mouth), while Var(score)'s sample estimate — via
E[e^{-v}] — is pulled up by large-negative-v excursions (the neck). A chain that hasn’t
symmetrically explored both regions will produce two very different-looking finite-sample
numbers for what is, in the population, one number. That’s a concrete, testable hypothesis: it predicts a
robust (trimmed / median-of-means) variance estimator in each channel should shrink the
estimation gap without changing the target, since the population identity doesn’t depend on
outlier draws — worth trying directly in nuts-rs/WALNUTS’s own variance-tracking code. (I’d
guess this is a better lever than clamping the final ratio — post #27 already shows nutpie’s
existing clamp, 1e-20/1e20, is far too loose to be doing this work.)
10-line standalone check (no dependencies beyond numpy; anyone can paste this in):
import numpy as np
rng = np.random.default_rng(0)
sigma, n = 3.0, 200_000
v = rng.normal(0, sigma, n)
x = rng.normal(0, np.exp(v / 2)) # x | v
score = -x * np.exp(-v) # d/dx log p
var_draws, var_score = x.var(), score.var()
print(var_draws, var_score, 1 / var_score, np.sqrt(var_draws / var_score))
# var_draws ~ var_score (both estimate e^{sigma^2/2} ~ 90.0); geomean ~ 1
warmup-robustness gap alone is enough to cause
total freezes regardless of the metric question. On the metric side the fix direction is robustifying the variance estimators feeding the geomean, not clamping the output ratio.
Just back from vacation, and there is a very long thread about a clear failure case for nutpie… :-)
@ssp3nc3r Thanks for the reproducer, I see the same thing locally now.
From what I’ve seen so far, I’m pretty sure this isn’t a fundamental issue with the mass matrix target, but something is clearly going wrong with the gigantic gradient values early on in some of the chains. nutpie tries to deal with those by clipping extreme mass matrix values and completely forgetting early gradient values quickly, but in this case, this clearly doesn’t do the trick as it should. I’ll have to investigate a bit more to see what to do about this.
I was also curios to see what happens if we switch the link function from exp to something that doesn’t blow up as horribly. I personally like to use `x + sqrt(1 + x ** 2) = exp(arcsinh(x)) for things like this in my models, it behaves linearly like softplus for large numbers, but also has a non-exponential left tail. (symmetric on log-scale).
With this, everything seems to be working fine. I get about 1000 grad evals per ess in stan and 500 grad evals per ess in nutpie, so about a 2x improvement. For the low rank version I had to increase the gamma regularization parameter, but with that I got 130 grad evals per ess.
Here are the resulting eigenvalue spectrums:
The diagonal fisher one doesn’t actually look that much better compared to the stan/kl one, but not any worse either. Maybe it just gets a bit lucky with the trajectory sizes?
About the arithmetic means, I don’t think that would do anything helpful. Extreme values have an even bigger impact, and also the hamiltonian matching from appendix B, which I think really motivates the fisher divergence goes completely out of the window.
Thanks for all the other replies, everyone. This is really really helpful.
I also realized that all the norming of the vectors doesn’t really do anything because it’s just constant multipliers and the ratio is what matters.
Thanks—the AI is right—I’m going to add find_reasonable_epsilon to Walnuts!
I should’ve realized that. Maybe I would have had I done it analytically rather than with Stan’s autodiff. As I noted above, I realize the norming doesn’t do anything for conditioning because it’s just a constant multiplier.
This is also a good point—it’s what the Fisher divergence is designed to do, but it depends on some level of mixing.
I’ll have to figure out how to do this, but this is what I meant by needing something that would condition, like a square root (i.e., geometric meaning with the identity).
The clipping range is very broad, (10^{-20), 10^{20}) in the Rust code. Maybe what we need instead of clipping is some kind of conditioning that will bring all the values closer to the center. That’s what geometric meaning with the identity does with the result being a sqrt().
This is an interesting statistical point we should investigate further in every context. A huge part of the problem with the funnel is the non-linearity of exp() in the tails. It’s not clear that’s desirable statistically, and we know it’s bad computationally. The folk theorem, as @andrewgelman would say. With
we have
With Adrian’s inverse link function,
we have
That gradient should be less prone to underflow/rounding than inverse logit, which underflows to 0 quickly and rounds to 1 even more quickly.
We could actually start with how Stan represents positive-constrained parameters as we currently do that with \exp(), too.
I guess this is a bit of a detour from the main topic here, but let’s look at that function a bit more. I’m bad at naming things, but I’ll call it bisoftplus(x) = x + sqrt(1 + x ** 2) = exp(arcsinh(x)) for now. (because it behaves a bit like softplus for the inverse as well, so “bi”-directionally. Other ideas are welcome. Claude tells me someone proposed it as an activation function under the name of squareplus some time back). The inverse is (x^2 - 1) / 2x
We use maps from \mathbb{R} to \mathbb{R}^+ in two prominent places: As inverse link functions in regressions, and in the inverse unconstraining map. Those might look similar, but I think they often require different properties.
The bisoftplus would be a pretty bad constraining map. If we wanted to sample a simple half-normal, the positive limit of the density as x goes to 0 turns into a cauchy-like left tail after an inverse bisoftplus transform. You can see how that happens if you think of the inverse of bisoftplus as sinh(log(x)). First, we take the log of the half-normal as we currently do, and then we exponentially blow up both tails with sinh. Not what we want.
softplus might have a case against exp in many applications. The left tail works like exp, but doesn’t squish the right tail of distributions. But the right tail of positive distributions is often perfectly fine as it is, there is no reason to squish it in exponentially.
But in constrast to softplus, the exp function does have some very friendly properties if we think of all of those functions as families of functions. Before we can log transform anything, we always have to remove the units, so there is a natural scaling factor we need to take into account. We also have the choice of which exponential base we want to use. That really gives us log(cx)/log(b) with the free hyperparameters c and b for the inverse-exp family. But those hyperparameters turn into a constant factor and a constant offset after the log transformation. HMC is invariant to an offset, and the diagonal mass matrix estimation will choose a scale for us anyway. So we can just use any values for c and b and completely ignore that those hyperparameters even exist. But this doesn’t work for softplus, where the scaling parameter really matters. We need to choose it somehow if we want the transformation to work properly. (We could actually learn it by minimizing the fisher divergence just as we learn the mass matrix/affine transform in nutpie, just food for thought…)
(cc @paul.buerkner from a previous discussion on this…)
So let’s say we want to run a regression, where we also predict variances based on some predictors. I’m guessing that this is roughly what’s happening in the model of @ssp3nc3r. The predictors live in \mathbb{R}, so we need an inverse link function to map that to the positive standard deviation.
I’ve seen exp and softplus in the wild for this, but they both have a lot of problems, and I think bisoftplus does something quite sensible here.
exp feels like a very natural choice, but it blows up predicted standard deviations. If our predictors have a normal distribution, then the standard deviation (and also the variance) has a log normal distribution. The log-normal has very long right tails, thicker than an exponential, but lighter than any pareto distribution (and often looks like a power law for quite a long time). This can lead to computational issues, but I don’t think it is usually what we want in a model either. “Sure, all standard deviations we’ve seen so far were between 0.5 and 2, but this out-of-sample standard deviation might just be 1000, who knows…”. I’ve never worked on anything where I thought long tails like this sound realistic.
The softplus fixes the long right tails, but it introduces new problems. We lose nice symmetry properties we had before. In the exp link, we had the property that if you swap the sign of a predictor, that the predicted standard deviation is 1/old. That’s gone. And the precision of the predicted values still has the very long tail of a log-normal, and that might be just as bad as the long tails of the standard deviation with the exp link.
The bisoftplus fixes both. We get our symmetry back, zero maps to 1, and the tails of the standard deviation and its inverse are friendly. And for large or small predictors it behaves additively on the standard deviation or its inverse. And around 0 it behaves like exp to second order.
One thing that wasn’t too important with the exp link is if we want to run the regression against the standard deviation or against the variance. That’s just a constant factor of 2 different with an exp link. But with a bisoftplus link it suddenly makes a difference. Do we want additivity of the standard deviation or the variance? Given that the variance is usually the thing that is additive, maybe it’s the better choice?
But I guess if we want to continue this, we should move it into a different thread.
Thanks, Adrian. This is super useful. I think we need to rethink our underlying transform. And then once we have an alternative for (-\infty, \infty) \mapsto (0, 1), we can also use it in place of \exp() in things like softmax.
In Stan, what we’d do is include the change of variables adjustment so that the implied distribution on (0, \infty) is uniform. Then users can place whatever priors they want on top of that. So it should lead to the exact same inferences statistically, even though the computational target density will be different due to the different change of variables.
Following the discussion of link functions for positive-constrained scales (aseyboldt’s bisoftplus, x + \sqrt{1+x^2} = \exp(\text{arcsinh}\,x)), here is a one-parameter generalization and an analytical argument that \tau = \frac{1}{2} — not \tau = 1 (bisoftplus) — is the exact value at which the funnel-neck gradient ceases to grow.
Consider substituting \sigma = g(\eta) into a Gaussian scale term (the funnel neck):
Taking the derivative with respect to the unconstrained parameter \eta:
The exploding term is not g' or (\log g)' — those stay remarkably tame for both \exp and bisoftplus. The term that detonates at the neck is \frac{d}{d\eta} (1/\sigma^2).
| Link Function | 1/\sigma^2 Behavior | \frac{d}{d\eta} (1/\sigma^2) Behavior | Verdict |
|---|---|---|---|
| \exp(\eta) | e^{-2\eta} | -2e^{-2\eta} | Exponential blowup |
| Bisoftplus (\tau=1) | \sim(2abs(\eta))^2 | \sim 8abs(\eta) | Linear growth |
| g_{\frac{1}{2}} | \sim 4abs(\eta) | \to \text{constant } (\approx 2) | Bounded |
We define the generalized family as:
For every \tau > 0, this function neatly preserves three crucial geometric properties:
Asymptotic analysis of the lower tail reveals g_\tau(\eta) \sim \left(\frac{2\eta}{\tau}\right)^\tau. Hence:
\tau = \frac{1}{2} is the exact mathematical phase transition where the neck gradient stops growing and saturates to a constant. It yields the gentlest possible tail that bounds the gradient while remaining fully expressive. It is derived from the objective function, not tuned.
| Quantity | Closed-Form Expression |
|---|---|
| Link Function (g_{1/2}) | \sigma = \sqrt{ 2\eta + \sqrt{1 + 4\eta^2} } = \exp\left(\frac{1}{2} \cdot \text{arcsinh}(2\eta)\right) |
| Inverse Link (g_{1/2}^{-1}) | \eta = \frac{1}{4} \left(\sigma^2 - \sigma^{-2}\right) |
| Log-Jacobian | \log \sigma - \frac{1}{2} \log(1 + 4\eta^2) |
Gentler tails buy numerical stability at the price of tail reach:
Tail Reach: Predictor value (η) needed to achieve scale (σ)
───────────────────────────────────────────────────────────
exp(η) │ 4.6
Bisoftplus │ 50
g_½ (τ=½) │ 2500
tau = 1/2 is the “run the regression against the variance” I mentioned.
I hope you don’t take this the wrong way, but your last two posts read a bit like I’m talking to an AI, not you, and I’d much rather talk to you than the AI bot.
Just wanted to present the math, case helpful, never mind too. Did the formatting with AI, saves so much time.
I undeleted this, but I’m not sure who deleted it the first time. If you want to delete it, @andre.pfeuffer, feel free. I hope I can’t undelete something a user wants to delete.
And for the record, I actually found the math useful and am in the process of writing the step size initialization now.
@andre.pfeuffer I really didn’t want to turn you away from this thread, or mean to imply that your posts were useless, or pure AI. And I’d much prefer it if you continue to post here if you want.
In optimization people often use gradient clipping to get past issues like this. While that doesn’t have strong theoretical reasoning that I’m aware of, this seems to work quite well in practice.
So I’ve been experimenting with something similar here. Hard clipping might actually introduce new issues however, because if several gradients are clipped in a row, their variance would suddenly be zero. So I’ve tried soft clipping via f(x) = c asinh(x / c), where I set c = 1e10. This is pretty much exactly the identity (we could make it exactly the identity if we care) until x = 1e8 or so, after that it starts to fall off logarithmically, so that f(10^{10}) \approx 10^{10}, f(10^{15}) \approx 10^{11}, f(10^{60}) \approx 10^{12} etc.I think we could also implement something like that such that it is pretty fast to compute as long as we don’t have any huge gradients with a little bit of simd magic.
I’m pretty sure this doesn’t affect the validity of hmc even in cases where we do end up changing the gradient in a meaningful way, as it only changes the proposal point in a reversible way, but I’d want to think about this a bit more before actually deploying this.
With this fix, no chains get stuck anymore, and nutpie converges fine with the exp link as well.
This time stan does end up being 1.5x more efficient, which I can no more explain than the 2x advantage for nutpie in previous run with the bisoftplus link function. The posterior eigenvalue spectrum again looks quite different, but about equally as bad. I evaluated it with 40 chains, 1000 tune, and 10_000 draws for both, so I don’t think this is ess estimation noise, but it might be “where does the power of two end up for the treesize” noise again?
I think we should have a separate discussion about etiquette with AI generated content.
To preface that conversation, since a response here gives context to those that come to this thread, I deliberately chose etiquette as the discussion point as to distinguish from a rule. I doubt we could strictly enforce a hard rule and we probably wouldn’t want to because AI can be very helpful. But it can be long-winded or used wantonly, it’s so easy to generate. In an online forum such as this, it can easily overtake a conversation which obsfucates and overwhelms, which are at odds with the spirit of the forum. I have sympathy for all here as it was helpful and it was a bit jarring and rambling. I suspect an open, honest clause about pasting AI content will help, plus editing the AI content for conciseness and correctness as a best practice.