Provably Scalable Black-Box Variational Inference with Structured Variational Families

TL;DR

This paper introduces structured variational families for black-box variational inference, achieving better scalability with dataset size compared to full-rank families, and provides both theoretical proofs and empirical validation.

Contribution

It proposes a new class of structured variational families that improve scalability in BBVI, with rigorous theoretical analysis and empirical verification.

Findings

Structured variational families achieve $ ext{O}(N)$ complexity.

Full-rank families scale poorly with dataset size.

Empirical results confirm theoretical scalability improvements.

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.