SteinVI: Recompute Score Function* For Each Particle Interaction in the Attractive Force

[OlaRonning]

Previously

The score function(1) in the attractive force was computed once for each particle and reused. While this is cheap, it also biases the Stein force.

Changes

• The score function is recomputed with fresh samples for every particle interaction.
• Added comments with equation references.
• Factored out the loss function in preparation for making it a class parameter that allows SteinLoss and ELBOs.
• ruff format v0.7.1 -> v0.9

The behavior-changing code is in einstein/steinvi.py and example/stein_bnn.py. The rest is from using the latest ruff.

Effect

• More stable optimization (Stein force fluctuates less at minimum)
• BNN example no longer needs tight initialization.
• Increased computation time, but outweighed by improvement in stability

(1) It is not actually a score function (grad(log p(x|D))(x)) as we are taking the gradient of an ELBO, not the log posterior log p(x|D)).

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB