stefanradev93 / BayesFlow

A Python library for amortized Bayesian workflows using generative neural networks.
https://bayesflow.org/
MIT License
297 stars 45 forks source link

Coin Sampling Optimizer #122

Open marvinschmitt opened 7 months ago

marvinschmitt commented 7 months ago

Coin sampling (Sharrock & Nemeth, 2023, [1]) is a learning-rate-free optimization scheme that performs on par with perfectly LR-tuned SVGD up to O(log N).

According to @LouisSharrock, coin sampling is currently implemented in JAX. If there was a TensorFlow implementation, we could directly use it within BayesFlow if it meets the tf.keras.optimizer.Optimizer API.

Within the BayesFlow pipeline, the optimizer is passed when we initiate the training loop. Here's a pointer for online training via train_online: https://github.com/stefanradev93/BayesFlow/blob/f5f7a6a2c7a8a5ef69f46fc3a3030be1166d7f63/bayesflow/trainers.py#L371

[1] https://proceedings.mlr.press/v202/sharrock23a

louissharrock commented 7 months ago

The COntinuous COin Betting (COCOB) algorithm does actually have a tensorflow implementation here: https://www.tensorflow.org/addons/api_docs/python/tfa/optimizers/COCOB

Unfortunately it is implemented in the soon-to-be deprecated TensorFlow Addons, but this might provide a starting point for an updated implementation.

Also, a minor correction to your attribution. Coin betting as an optimisation method was introduced by Orabona and Pal in [1], and extended by Orabona and Tommasi for training deep neural nets in [2]. Our innovation in [3] was to adapt this approach for sampling problems.

[1] https://proceedings.neurips.cc//paper/by-source-2016-320 [2] https://proceedings.neurips.cc/paper_files/paper/2017/hash/7c82fab8c8f89124e2ce92984e04fb40-Abstract.html [3] https://proceedings.mlr.press/v202/sharrock23a