Couple of questions regarding IVON

[dimarkov]

Thank you for excellent work on IVON, really enjoyed the related papers.

I am interested in using it for my own research on Bayesian sparstfication and pruning of DNNs,
and I have an early implementation in JAX/Optax, you can find the repo (here)[https://github.com/dimarkov/blrax].

I would like to know if you have some general recommendations on tuning ivon's parameters. I did some testing
on deep MLPs, like MLP mixers, with CIFAR10 and CIFAR100. It is difficult to find a set of parameter values that work well
on both datasets. Also, to achieve comparable performance to adamw and lion, I need to use Cosine annealing for the
learning rate. Does this matches with your experience, or is this potentially an indication of a mistake in the implementation?

Also, I am curies if you have done some comparison with algorithms based on Lie-Groups BLR. Any idea on how do they compare in terms of results quality and convergence speed. I expect lie-group based ones
to be slower per iteration.

[moellenh]

Hi! Thanks for the interest in IVON. We have an (experimental) optax version that you may want to check out: https://github.com/ysngshn/ivon-optax

We haven't tried IVON on MLP mixers or large MLPs, but the main principles are:

• Setting learning rate as large as possible until it gets unstable.
• Setting hess_init (h0) and ess (lambda in the paper) such that not too much Gaussian noise is added to the weights during training. If too much noise is added to the weights, the training diverges. In the paper, we used a larger h0 with beta2 close to one and a smaller ess (equal to number of data points) so that the posterior is initialized close to a point mass and slowly changed during training. This strategy seems to give the right amount of noise and learning dynamics to get good results.
• An alternative strategy is to try h0 close to zero, use very large ess (e.g., data size times 100 or 1000). This one also works with smaller beta2 such as 0.9999 and the Hessian is adapting faster, but we found the results to be slightly worse.
• In our experience, we found cosine annealing to be necessary to match the performance of SGD/Adam.
• For problems where Adam(W) works well, typically IVON can benefit from gradient clipping.

The Lie-Group BLR algorithms require a more memory (one extra momentum variable) and have roughly the same speed as IVON. The quality should be a bit higher, since there is more flexibility in the choice of posterior. For example, they can use non-Gaussian posteriors such as Laplace. A high quality implementation is however lacking so far, and we hope to look into this in the future. There are many small details which may matter in practice, such as scale of the variables, initializations, etc.

[dimarkov]

Thank you! This all matches well with what I have been able to figure out. I will check out the optax repo.