Infrastructure level code to support various experiments.

[edmundlth]

We need quick and easy way to run various experiments and ways to get measurements from them.

Desiderata:

• Need to be able to run models of different sizes d and with different data set size n.
• Need efficient method to do posterior sampling from neural network models.
• Measure free energy, Bayesian generalisation error, all as a function of n.
• Include SGD method of training in addition to Bayesian predictive. Should also include "ensembling" in the sense of doing multiple SGD training with different initialisation.

Tools:

• Pytorch
• Numpyro for MCMC and other Bayesian computation support.

[edmundlth]

First step:

• [ ] Need program that takes in a model and data and generate samples from its posterior distribution using MCMC
  • model can mean: torch.module if that is convenient. One possible alternative is to just think of "model" as being completely specified paramterised regression function (e.g. neural net) so that $y = f(x, w) + $ gaussian noise.
  • data might just be torch.Dataset
  • the MCMC code should parameterise on different MCMC algorithm (HMC, NUTS, etc...), number of samples, MCMC burn in length etc.
• [ ] Once MCMC is implemented, we should try to compute estimation for Free energy $F_n$, Bayesian generalisation error $B_g$ and training error $B_t$. Once we are happy this works reasonably well (can be challenging for large model size, $d$), we can try to plot the learning curve for different model of interest.

Thoughts:

• Alternative include using variational inference, but that open up a whole other can of worm (RLCT for variational posterior is not the same + realisability / renormalisability issues).
• If we struggle with "global" MCMC (due to high dimension and complicated local minimums etc), we might be able to do something with "local" MCMC where we SGD to a good parameter and sample near it. We might want to be able to do this anyway since it relates to the $P(w^) = p(w^ | D_n) Z(w^*)$ decomposition mentioned in Section 7.6 of the Grey book.