Batch-and-Match algorithm for minimizing the covariance-weighted Fisher divergence

[Red-Portal]

This PR adds the recently proposed batch-and-match (BaM) algorithm¹. BaM is a proximal-point-like optimization algorithm that minimizes the covariance-weighted fisher divergence:

$$\mathrm{F}(q, \pi) = \mathbb{E}_{z \sim q} {\left\lVert \nabla \log \frac{q}{\pi} (z) \right\rVert}_{\mathrm{Cov}(q)}^2 .$$

This algorithm is a measure space algorithm and is therefore restricted to the Gaussian variational family. According to my experience, BaM converges really quickly on certain low-dimensional problems (as shown in the paper¹) and requires pretty much no tuning. Therefore, this should be an interesting addition.

I expect this to be the last feature addition for v0.6.

Footnotes

1. Cai, D., Modi, C., Pillaud-Vivien, L., Margossian, C. C., Gower, R. M., Blei, D. M., & Saul, L. K. (2024). Batch and match: black-box variational inference with a score-based divergence. ICML'24. ↩ ↩²

[github-actions]

AdvancedVI.jl documentation for PR #218 is available at:
https://TuringLang.github.io/AdvancedVI.jl/previews/PR218/

[github-actions]

Benchmark Results

Benchmark suite	Current: 9b8ee3b	Previous: 40ad7e9	Ratio
normal/RepGradELBO + STL/meanfield/Zygote	3922494622.5 ns	3782348781.5 ns	1.04
normal/RepGradELBO + STL/meanfield/ReverseDiff	1018051158 ns	1134417597 ns	0.90
normal/RepGradELBO + STL/meanfield/Mooncake	917458100 ns	1190437984 ns	0.77
normal/RepGradELBO + STL/fullrank/Zygote	3939266435 ns	3771397659 ns	1.04
normal/RepGradELBO + STL/fullrank/ReverseDiff	1565912835 ns	1677575257 ns	0.93
normal/RepGradELBO + STL/fullrank/Mooncake	968276455.5 ns	1230722757 ns	0.79
normal/RepGradELBO/meanfield/Zygote	2493077502 ns	2714039964.5 ns	0.92
normal/RepGradELBO/meanfield/ReverseDiff	675764270.5 ns	788004640 ns	0.86
normal/RepGradELBO/meanfield/Mooncake	808425531 ns	1071454609 ns	0.75
normal/RepGradELBO/fullrank/Zygote	2534939838 ns	2742116219.5 ns	0.92
normal/RepGradELBO/fullrank/ReverseDiff	855604752 ns	985372676.5 ns	0.87
normal/RepGradELBO/fullrank/Mooncake	849225869.5 ns	1102447331 ns	0.77
normal + bijector/RepGradELBO + STL/meanfield/Zygote	5883218367 ns	5481930888 ns	1.07
normal + bijector/RepGradELBO + STL/meanfield/ReverseDiff	1993631184 ns	2420147924 ns	0.82
normal + bijector/RepGradELBO + STL/meanfield/Mooncake	4018512550.5 ns	4080623209 ns	0.98
normal + bijector/RepGradELBO + STL/fullrank/Zygote	6055822955 ns	5610718262 ns	1.08
normal + bijector/RepGradELBO + STL/fullrank/ReverseDiff	2740073379.5 ns	3086620384 ns	0.89
normal + bijector/RepGradELBO + STL/fullrank/Mooncake	4136432672.5 ns	4189021458 ns	0.99
normal + bijector/RepGradELBO/meanfield/Zygote	4324421989.5 ns	4218617739.5 ns	1.03
normal + bijector/RepGradELBO/meanfield/ReverseDiff	1603425162 ns	2038692481 ns	0.79
normal + bijector/RepGradELBO/meanfield/Mooncake	3812348756 ns	3928918644.5 ns	0.97
normal + bijector/RepGradELBO/fullrank/Zygote	4437575445 ns	4395929362.5 ns	1.01
normal + bijector/RepGradELBO/fullrank/ReverseDiff	1874124331 ns	2329503422 ns	0.80
normal + bijector/RepGradELBO/fullrank/Mooncake	3957099289.5 ns	4027695545.5 ns	0.98

This comment was automatically generated by workflow using github-action-benchmark.

[penelopeysm]

Looks great, thanks! :)

[Red-Portal]

Thanks as always for the review!