Fast Black-box Variational Inference through Stochastic Trust-Region Optimization
We introduce TrustVI, a fast second-order algorithm for black-box variational inference based on trust-region optimization and the "reparameterization trick." At each iteration, TrustVI proposes and assesses a step based on minibatches of draws from the variational distribution. The algorithm provably converges to a stationary point. We implemented TrustVI in the Stan framework and compared it to two alternatives: Automatic Differentiation Variational Inference (ADVI) and Hessianfree Stochastic Gradient Variational Inference (HFSGVI). The former is based on stochastic first-order optimization. The latter uses second-order information, but lacks convergence guarantees. TrustVI typically converged at least one order of magnitude faster than ADVI, demonstrating the value of stochastic second-order information. TrustVI often found substantially better variational distributions than HFSGVI, demonstrating that our convergence theory can matter in practice.
SUMMARY OF THE PAPER:
The paper transfers concepts known in the optimization community and applies them to variational inference. It introduces a new method to optimize the stochastic objective function of black box variational inference. In contrast to the standard SGD, which optimizes the objective using only first order derivatives, the proposed methods takes estimates of the Hessian into account and approximates the objective function on a small trust region around the current iterate by a quadratic function. The algorithm automatically adapts the size of the trust region based on estimates of the quality of the quadratic model. The paper proves theoretically that the algorithm converges, and shows experimentally that convergence is typically faster than standard SGD.
GENERAL IMPRESSION / POSITIVE ASPECTS:
The paper is well written and the proposed algorithm is potentially highly relevant given its generality and the reported improvements in speed of convergence. The experimental part summarizes results from a very large number of experiments, albeit it is not clear to me whether this includes any experiments on large-scale datasets that require minibatch sampling.
WHAT COULD BE IMPROVED:
1. I appreciate the fact that the paper includes an unambiguous definition of the algorithm and a rigorous proof that the proposed algorithm converges. To make the method more accessible, it would however be helpful to also briefly explain the algorithm in more intuitive terms. For example, a geometric interpretation of \gamma, \lambda, \eta, and \delta_0 would help practitioners to choose good values for these hyperparameters.
2. An often cited advantage of VI is that it scales to very large models when minibatch sampling (stochastic VI) is used. Minibatch sampling introduces an additional source of stochasticity. The paper discusses only stochasticity due to the black box estimation of the expectation under q(z). How does the proposed algorithm scale to large models? Are the convergence guarantees still valid when minibatch sampling is used and do we still expect similar speedups?
3. It is not clear to me how $\sigma_k$, defined in Condition 4, can be obtained in practice.
MINOR REMARKS:
1. Algorithm 1: R^d should be R^D in the first line.
2. Lines 108-114 discuss a set of hyperparameters, whose motivation becomes clear to the reader only after reading the convergence proofs. It might be clearer to present the hyperparameters in a table with a row for each hyperparameter, and columns for the symbol, a short and descriptive name, and the interval of allowed values.
3. Eq. 7: which norm is used for the Hessian matrix here?
4. Line 154: It may be helpful to point out at this point that one draws *new* samples from p_0, i.e., one may not reuse samples that were used to generate g_k and H_k. Otherwise, $\ell'_{ki}$ are not i.i.d. The algorithm box states this clearly, but it was not immediately clear to me from the main text.