Models where Stan outperforms Nutpie/Walnuts

@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?