Question on simplex bijector implementation

[Red-Portal]

Hi,

It appears that torch.probability simply uses softmax for the simplex bijector.
Is there a reason our simplex transform is much more complicated?
I was also thinking about a GPU-friendly implementation, which the current implementation appears hard do.

[torfjelde]

Softmax isn't bijective. The one we have now is (maps from d to d-1 dimensional)

[devmotion]

See also #51.

[Red-Portal]

Betanalpha does discuss a bijective softmax by arbitrarily setting the endpoint logits. Any experience with this?

[devmotion]

That's supported e.g. in GPLikelihoods (see maybe also the discussion in JuliaGaussianProcesses/GPLikelihoods.jl#55).

[Red-Portal]

Good to know thanks. Though, back to my original intention, I really wish that our simplex bijector could play nicely with GPUs out of the box. Among non-NF bijectors, it seems the simplex bijector is really going to be the big challenge going in that direction. Do we have any plans on how to pursue this? It does seem to me that the softmax approach would be much easier to get this done.

[Red-Portal]

Actually, nevermind. I just wrote a stick-breaking bijector using array operations based on the implementations of numpyro and tensorflow. If this were to be added to Bijectors.jl we'll probably have to add a CUDA array specialization. Let me know how to proceed on this.

[devmotion]

On Julia >= 1.9, a CUDA specialization could be put in an extension (possibly could even just be an extension with GPUArrays).

[Red-Portal]

I do have the feeling that this will have to wait until the batch operation interface is finalized. @torfjelde Do we have an expectation on when that would be?

[sethaxen]

There are three main ways to use softmax for simplex transforms. One uses parameter expansion to retain bijectivity: f(y) = [softmax(y); logsumexp(y)]. The other two come from compositional data analysis literature are called additive log-ratio f(y) = softmax(vcat(y, 0)) and isometric log-ratio f(y) = softmaxx(V * y) for a particular choice of semi-orthogonal matrix V. I'm currently testing performance of each of these versus stick-breaking.

[penelopeysm]

I tried implementing isometric log-ratio + softmax (because Stan uses that now — I directly copied the Stan implementation line for line), and IIRC it used less memory but was a tiny bit slower than stick-breaking.

Happy to reopen if somebody would like to investigate further

[sethaxen]

I'm currently testing performance of each of these versus stick-breaking.

Coming back to this, the experiment results are at https://sethaxen.github.io/simplex_transforms_results/, and I believe these results were the motivation for Stan to switch to ILR. It takes some deselection of bad transforms and browsing of different metrics to see, but in most cases the top 3 for most targets are StickbreakingLogistic (Stan's old stick-breaking transform; what we do here), ILR, and ExpandedSoftmax. There's no clear winner, but it generally seems that ILR is a better trade-off by most metrics.

In particular, min_rel_ess_per_chain uses an RMSE-based ESS estimate that computes the minimum ESS per gradient across all coordinates in the model (RMSE-based estimate, so will be low if the chains converge but to the wrong target). By that metric, ILR typically performs better than stick-breaking. Note that stick-breaking is expected to perform best for Dirichlet, since the stick-breaking transform decorrelates Dirichlet. Dirichlet is unique with all kinds of symmetries, and real posteriors are unlikely to look like Dirichlet, so it's more interesting to note that ILR matches stick-breaking's performance for all but the most extreme Dirichlet targets while it sometimes significantly outperforms it on extreme multi-logit-normal targets (these consist of multivariate-normal passed through an ALR transform, which is quite different from ILR, so by construction we wouldn't expect ILR to have an unfair advantage here).

When looking at the n_divergent, even on Dirichlet targets, stick-breaking tends to diverge for harder targets whereas ILR produced no divergences even for the hardest targets.

For me the take-home is that stick-breaking is fine, but ILR is a little better; if I was implementing a new library I would probably implement ILR first and default to it.

I tried implementing isometric log-ratio + softmax (because Stan uses that now — I directly copied the Stan implementation line for line), and IIRC it used less memory but was a tiny bit slower than stick-breaking.

Is that raw computation time for log-density or gradient, or is that ESS/wallclock time? I didn't measure wallclock in any of the experiments, though I would expect wallclock time for gradient evaluation to be quite similar for both of these transforms.

[penelopeysm]

Is that raw computation time for log-density or gradient, or is that ESS/wallclock time?

It was just the time for the transform + logjac itself. I vaguely remember doing some sampling tests but it wasn't scientific enough for me to report back. These are very useful plots! I am guessing that the implementation is in Stan? I am definitely happy to open a new issue again if you're interested in translating to Julia. It probably wouldn't be my main focus, but happy to provide support with any benchmarks, etc.