Provably Scalable Black-Box Variational Inference with Structured Variational Families

Joohwan Ko, Kyurae Kim, Woo Chang Kim, Jacob R. Gardner
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:24896-24931, 2024.

Abstract

Variational families with full-rank covariance approximations are known not to work well in black-box variational inference (BBVI), both empirically and theoretically. In fact, recent computational complexity results for BBVI have established that full-rank variational families scale poorly with the dimensionality of the problem compared to e.g. mean-field families. This is particularly critical to hierarchical Bayesian models with local variables; their dimensionality increases with the size of the datasets. Consequently, one gets an iteration complexity with an explicit $\mathcal{O}(N^2)$ dependence on the dataset size $N$. In this paper, we explore a theoretical middle ground between mean-field variational families and full-rank families: structured variational families. We rigorously prove that certain scale matrix structures can achieve a better iteration complexity of $\mathcal{O}\left(N\right)$, implying better scaling with respect to $N$. We empirically verify our theoretical results on large-scale hierarchical models.

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-ko24d, title = {Provably Scalable Black-Box Variational Inference with Structured Variational Families}, author = {Ko, Joohwan and Kim, Kyurae and Kim, Woo Chang and Gardner, Jacob R.}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {24896--24931}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/ko24d/ko24d.pdf}, url = {https://proceedings.mlr.press/v235/ko24d.html}, abstract = {Variational families with full-rank covariance approximations are known not to work well in black-box variational inference (BBVI), both empirically and theoretically. In fact, recent computational complexity results for BBVI have established that full-rank variational families scale poorly with the dimensionality of the problem compared to e.g. mean-field families. This is particularly critical to hierarchical Bayesian models with local variables; their dimensionality increases with the size of the datasets. Consequently, one gets an iteration complexity with an explicit $\mathcal{O}(N^2)$ dependence on the dataset size $N$. In this paper, we explore a theoretical middle ground between mean-field variational families and full-rank families: structured variational families. We rigorously prove that certain scale matrix structures can achieve a better iteration complexity of $\mathcal{O}\left(N\right)$, implying better scaling with respect to $N$. We empirically verify our theoretical results on large-scale hierarchical models.} }
Endnote
%0 Conference Paper %T Provably Scalable Black-Box Variational Inference with Structured Variational Families %A Joohwan Ko %A Kyurae Kim %A Woo Chang Kim %A Jacob R. Gardner %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-ko24d %I PMLR %P 24896--24931 %U https://proceedings.mlr.press/v235/ko24d.html %V 235 %X Variational families with full-rank covariance approximations are known not to work well in black-box variational inference (BBVI), both empirically and theoretically. In fact, recent computational complexity results for BBVI have established that full-rank variational families scale poorly with the dimensionality of the problem compared to e.g. mean-field families. This is particularly critical to hierarchical Bayesian models with local variables; their dimensionality increases with the size of the datasets. Consequently, one gets an iteration complexity with an explicit $\mathcal{O}(N^2)$ dependence on the dataset size $N$. In this paper, we explore a theoretical middle ground between mean-field variational families and full-rank families: structured variational families. We rigorously prove that certain scale matrix structures can achieve a better iteration complexity of $\mathcal{O}\left(N\right)$, implying better scaling with respect to $N$. We empirically verify our theoretical results on large-scale hierarchical models.
APA
Ko, J., Kim, K., Kim, W.C. & Gardner, J.R.. (2024). Provably Scalable Black-Box Variational Inference with Structured Variational Families. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:24896-24931 Available from https://proceedings.mlr.press/v235/ko24d.html.

Related Material