👋 Blackjax Meeting - Jan 2023

[github-actions]

Hello! This is an issue to track the next Blackjax meeting. Here's some relevant information:

• Date: Paris time at 5:00PM on 13/01/2022
• Google Calendar Link.
• Video Link
• Meeting Notes

These meetings are public and can be attended by anyone. If you'd like to discuss something at the meeting, we kindly ask you to comment on this issue beforehand.

[junpenglao]

Top of mind in terms of agenda:

• Remaining task for v1
• Refactor documentation (move to /~https://github.com/blackjax-devs/sampling-book)
• Developer doc

[rlouf]

In addition to what Junpeng said:

Blackjax

• The VI API. Rewrite everything in terms of an update function that generates optimisation paths (Add Stein Variational Gradient Descent #385, Add Mean Field Variational Inference implementation  #433, Pathfinder #194);
• Continuous benchmarking in Blackjax to prevent performance regression;
• Better testing;
• General organisation of the repo.
  • Move types and aliases to subfolders since nothing generalizes beyond the algorithm classes.
  • Move adaptation to mcmc
  • Ideally get rid of the optimiser folder
  • Add a meta folder. SMC goes there.
• Add Matrix chat room to give an easy way for users to reach out to the team

Sampling Book

• Caching executed notebooks to prevent the CI from exploding;
• Evaluation metrics
• General book structure

[AdrienCorenflos]

Distinction between samplers and meta-algorithms. Structure, API contract, etc.

[GStechschulte]

The Blackjax twitter tweeted a call for "which example models would you like to see in the documentation". Furthermore, the examples should preferably "illustrate common difficulties that are encountered when sampling, and how they can be overcome with the right choice of sampler.".

There were several responses. Summarized below:

• Classical time series
• Hierarchical GLMMs (using Bambi as front-end)
• Survival regression
• ConvNets
• Vision transformers

Do the maintainers have a pipeline of examples models they would like to see added? Furthermore, is there a preference on which PPL these models could be implemented in?

[rlouf]

Meeting notes

Rollcall

Please write your name and GitHub handle below

• Rémi Louf (@rlouf)
• Junpeng Lao (@junpenglao)
• Antoine Carnec (@antotocar34)
• Carlos Iguaran (@ciguaran)
• Adrien Corenflos (@AdrienCorenflos)
• Alberto Cabezas (@albcab)
• Rob Zinkov (@zaxtax)
• Ludger Paehler (@ludgerpaehler)
• Gabriel Stechtschulte (@GStechschulte)

Notes

Blackjax

What are the remaining tasks for V1? (@junpenglao)

v1 issue list

• SMC refactor
• Use logdensity in HMC family instead of potential
• Refactor the RMH step. The rationale is to be able to replace the traditional RMH step with Radford Neal's scheme.
• Make sure we can distribution gradient computation. To be able to follow the approach taken e.g. by this paper. Might be useful for VI later.
• Documentation (moving examples/ to sampling-book, see below)
• Organisation of the library
• Developer documentation
  • Structure of Blackjax
  • Expectations

@AdrienCorenflow: Regarding making sure that PyTrees are used everywhere, mgrad issue with the covariance matrix when wanting to use PyTree related discussion. High-level API work with PyTrees.
@ludgerpaehler: I will take a look at the gradient distribution.

Developer documentation (@junpenglao)

Carlos Iguaran: can we somewhere define what do we mean by meta-algorithm vs adaptation vs kernel?

Organisation of the library (@AdrienCorenflos,@rlouf)

• Distinction samplers/meta-algorithms
• Move types and aliases to subfolders
• /adaptation -> /mcmc/adaptation
• Get rid of /optimizers
• Add /meta folder, SMC goes there
• Upstream dual averaging (or in step-size adaptation)
• Move LBFGS logic to Pathfinder
• Make sure that

@junpenglao: MAP estimation needs non-standard optimizer. Maybe we can leave the optimizers in /optimizer;
@AdrienCorenflos: We need to make sure users cannot import them easily;
@junpenglao: In this case if we only use them in one place then we can leave them in this place

Rewrite the VI API (@rlouf)

We should adopt the step function design adopted everywhere else in the library and let users do the looping. Gives more flexibility to e.g. deep learning users.

Antoine Carnec: Choice wrt the optimizer library for VI.
@AdrienCorenflos: JaxOpt has support for convex optimizer;
@rlouf: Let's go with Optax for optimizers in the VI module, it seems like a safe choice.

Continuous benchmarking (@rlouf)

Look at what has been done in Aesara. We need to version the results and have it publicly available somewhere.

Warnings in CI for big discrepancies (x2), but graph every change somewhere.

We need better testing (@rlouf)

I generally dislike the "test by example" approach taken for many algorithms. We need to work with theoretical results on convergence.

@junpenglao: TFP have some really great convergence test: /~https://github.com/tensorflow/probability/blob/main/tensorflow_probability/python/distributions/internal/statistical_testing.py - but it needs a lot of samples for some of the convergence test

@AdrienCorenflos: we can use invariance + distribution testing.

Test lower level API (one step of each kernel and check invariance etc)

Reflect the directory structure of the repo in the tests suite

@AdrienCorenflos: I'll open an issue specifically for unittesting of SMC, these should be efficient enough. Top of my mind: (i) check that adaptive ESS gets the right tempering delta by using the Chi-square distance, (ii) check that we can move a Gaussian with tempered likelihood to another close one i.e. distribution testing, (iii) verify that target ESS is reached.

Community (@rlouf)

@junpenglao: We keep working with GitHub discussions for now, we can add a Matrix chat room later on.

Sampling book

@rlouf: Open individuals issues for the examples we want to see. From the Twitter post:

• Classical time series
• Hierarchical GLMMs (using Bambi as front-end)
• Survival regression
• ConvNets
• Vision transformers

@rlouf: Moving the examples to the Sampling Book will make our transition to ReadTheDocs (and get versioning) possible. We are currently hitting RTD's timeout when the Blackjax docs are built.