A Framework for Improving the Reliability of Black-box Variational Inference

1A core strength of the Bayesian approach is that it is conceptually straightforward to carry out inference in arbitrary models, which enables the user to employ whatever model is most appropriate for the problem at hand. The flexibility and uncertainty quantification provided by Bayesian inference have led to its widespread use in statistics and machine learning, including in deep learning (for example,;;;;;; ). Using Bayesian methods in practice, however, typically requires using approximate inference algorithms to estimate quantities of interest such as posterior functionals (for example, means, covariances, predictive distributions, and tail probabilities) and measures of model fit (for example, marginal likelihoods and cross-validated predictive accuracy). Therefore, a core challenge in modern Bayesian statistics is the development of general-purpose (or black-box) algorithms that can accurately approximate these quantities for whatever model the user chooses. In machine learning, rather than using Markov chain Monte Carlo (MCMC), black-box variational inference (BBVI) has become the method of choice because of its scalability and wide-applicability. BBVI is implemented in many software packages for general-purpose inference such as Stan, Pyro, PyMC3, and TensorFlow Probability, which have seen widespread adoption by applied data analysts, statisticians, and data scientists.

2Variational inference methods aim to minimize a measure of discrepancy between a parameterized family of distributions and the posterior distribution, with the Kullback–Leibler divergence being the canonical choice of discrepancy. Conventional approaches to variational inference leverage conditional conjugacy and other model-specific structure to derive optimization algorithms. BBVI, on the other hand, uses stochastic optimization to avoid the need for model-specific derivations, thereby significantly broadening the applicability of variational methods. Ensuring the efficiency and reliability of the BBVI optimization requires careful selection of optimization method and the stochastic estimator of the discrepancy gradient. For example, using adaptive optimization procedures like Adam, RMSProp, and Adagrad can ensure stability. Due to its relatively small variance, the most common gradient estimator is the reparameterization gradient. However, sometimes alternatives such as the score function gradient are employed. No matter the choice of gradient estimator, various variance reduction strategies like control variates have been introduced to stabilize and speed up optimization. Recent research has furthered our understanding of convergence behavior, tackling theoretical challenges in stochastic optimization and providing new convergence guarantees.

3While some recent progress has been made in developing tools for assessing the accuracy of variational approximations, stochastic optimization methods for BBVI remain unreliable and require substantial hand-tuning of the number of iterations and optimizer tuning parameters. Moreover, there are few tools available for determining whether the variational parameters estimated by these frameworks are close to optimal in any meaningful sense and, if not, how to address the problem; More iterations? A different learning rate schedule? A smaller initial or final learning rate? demonstrate the absence of a reliable and coherent optimization methodology for BBVI. The authors synthesize and compare recent advances such as normalizing flows and gradient estimators using 30 benchmarked models. Despite the fact that these models vary greatly in the complexity and dimensionality of the posteriors, run each optimization for a fixed number of iterations (30,000) and for 5 different step sizes because the existing literature does not provide any compelling guidance for how to automate the choice of step size and reliably determine when the optimization has converged.

4In view of the significant limitations of existing BBVI optimization methodology, in this paper we aim to provide a practical, cohesive, and theoretically well-grounded optimization framework for BBVI. To ensure reliability and wide applicability, we develop a framework that is (1) automated, (2) intuitively adjustable by the user, and (3) robust to failure and tuning parameter selection. Our approach builds on a recent line of work inspired by, which uses a fixed learning rate γ that is adaptively decreased by a multiplicative factor ρ once the optimization iterates, which form a homogenous Markov chain, have converged. A benefit of this approach is that, for a given learning rate, a dramatically more accurate estimate of the optimal variational parameter can be obtained by using iterate averaging. However, as we have shown in previous work, existing convergence checks can be unreliable and stop too early. Since the learning rate is decreased by a constant multiplicative factor, decreasing it too early can slow down the optimization by an order of magnitude or more. Hence, it is crucial to develop methods that do not prematurely declare convergence. On the other hand, an optimization framework must also provide a termination criterion that triggers when it is no longer worthwhile to decrease the learning rate further, either because the current variational approximation is sufficiently accurate or because further optimization would be too time-consuming.

5The key idea that informs our solutions to these challenges is that we want, the target variational approximation for learning rate, to be close to the optimal variational approximation. We measure closeness in terms of symmetrized Kullback–Leibler (KL) divergence and show that closeness in symmetrized KL divergence can be translated into bounds on other widely used accuracy metrics like Wasserstein distance. Figure 1 summarizes our proposed framework, which we call robust and automated black-box VI (RABVI). The primary contributions of this paper are in steps 2, 3, and 4. In step 2, to determine convergence at a fixed step size, we build upon our approach in, where we establish that the scale-reduction factor, which is widely used to determine convergence of Markov chain Monte Carlo algorithms, can be combined with a Monte Carlo standard error (MCSE) cutoff to construct a convergence criteria. We improve upon our previous proposal by:

6adaptively finding the size of the convergence window, which may need to be large for challenging or high-dimensional distributions over the model parameters, and

7developing a new rigorous MCSE cutoff condition that guarantees the symmetrized KL divergence between and the estimate of obtained via iterate averaging will be small.

8In step 3, we leverage recent results that characterize the bias of the stationary distribution of stochastic gradient descent (SGD) with a fixed learning rate to estimate and without access to. In step 4, these estimates enable the use of a termination criterion that compares (i) the predicted relative decrease in the KL divergence if the smaller learning were used and (ii) the predicted computation required to converge with the learning rate. By trading off between (i) and (ii), the criterion enables the user to balance desired accuracy against computational cost. Figure 2 provides an example of the faster convergence, higher accuracy, and greater reliability achievable using RABVI compared to alternative optimization algorithms and demonstrates how the user can trade off accuracy and computation by adjusting the accuracy threshold.

9In summary, by drawing on recent developments in theory and methods for fixed–learning-rate stochastic optimization, tools from MCMC methodology and results from functional analysis, RABVI delivers a number of benefits:

10it relies on rigorously justified automation techniques, including automatic learning rate adaptation;

11it has an interpretable, user-adjustable accuracy parameter along with a small number of additional intuitive tuning parameters;

12it detects inaccurate estimates of the optimal variational approximation; and

13it can flexibly incorporate additional or future methodological improvements related to variational inference and stochastic optimization.

14We demonstrate through synthetic comparisons and real-world model and data examples that RABVI provides high-quality black-box approximate inferences. We make RABVI available as part of the open source Python package VIABEL.1

15In this section, we briefly review relevant background about Bayesian and variational inference.

16Let denote a parameter vector of interest, and let x denote observed data. A Bayesian model consists of a prior density and a likelihood. Together, the prior and likelihood define a joint distribution over the data and parameters. The Bayesian posterior distribution π is the conditional distribution of given fixed data x, with x suppressed in the notation since it is always fixed throughout this work. To write this conditional, we define the unnormalized posterior density and the marginal likelihood, or evidence,. Then the posterior is. Typically, practitioners report posterior summaries, such as point estimates and uncertainties, rather than the full posterior. For, typical summaries include the mean, the covariance, and [a, b] interval probability, where is equal to one when C is true and zero otherwise.

17In most Bayesian models, it is infeasible to efficiently compute quantities of interest such as posterior means, variances, and quantiles. Therefore, one must use an approximate inference method that produces an approximation q to the posterior. The summaries of q may in turn be used as approximations to the summaries of. One approach, variational inference, is widely used in machine learning. Variational inference aims to minimize some measure of discrepancy over a tractable family of approximating distributions:

18The variational family is chosen to be tractable in the sense that, for any, we are able to efficiently calculate relevant summaries either analytically or using independent and identically distributed samples from q.

19In variational inference, the standard choice for the discrepancy is the Kullback–Leibler (KL) divergence. Note that the KL divergence is asymmetric in its arguments. The direction is most typical in variational inference, largely out of convenience; the unknown marginal likelihood Z appears in an additive constant that does not influence the optimization and computing gradients require estimating expectations only with respect to, which is chosen to be tractable. Minimizing is equivalent to maximizing the evidence lower bound (ELBO; ):

20While numerous other divergences have been used in the literature (for example,;;;;; ), we focus on since it is the most common choice; the default or only choice in widely used frameworks such as Stan, Pyro, and PyMC3; and easiest to estimate when using simple Monte Carlo sampling to approximate the gradient.

21Black-box variational inference (BBVI) methods have greatly extended the applicability of variational inference by removing the need for model-specific derivations and enabling the use of more flexible approximation families. This flexibility is a result of using simple Monte Carlo (and automatic differentiation) to approximate the (gradients of the) expectations that define common choices of the discrepancy objective. To estimate the optimal variational parameter λ*, BBVI methods commonly use stochastic optimization schemes which at iteration k are of the form

22where is the descent direction and is the learning rate (also called the step size). Standard stochastic gradient descent corresponds to taking, a (usually unbiased) stochastic estimate of the gradient. We are particularly interested in adaptive stochastic optimization methods (for example,;; ) that use a smoothed and/or rescaled version of as the descent direction. For example, RMSProp tracks an exponential moving average of the squared gradient,, which is used to rescale the current stochastic gradient:. Or, Adam tracks an exponential moving average of the gradient as well as the squared gradient and uses both to rescale the current stochastic gradient:. The benefits of adaptive algorithms include that they tend to be more stable and are scale invariant, so the learning rate can be set in a problem-independent manner.

23If the learning rate is fixed so that, then we can view the iterates produced by Eq. (1) as a homogenous Markov chain, which under certain conditions will have a stationary distribution. Stochastic optimization with a fixed learning rate exhibits two distinct phases: a transient (a.k.a. warm-up) phase during which iterates make rapid progress toward the optimum, followed by a stationary (a.k.a. mixing) phase during which iterates oscillate around the mean of the stationary distribution,.

24The mean is a natural target because even if the variance of the individual iterates means they are far from, can be a much more accurate approximation to. For example, the following result quantifies the bias of standard fixed–learning-rate SGD ( provide similar results for momentum-based SGD algorithms):

25Theorem 1 (, Theorem 4)) Under regularity conditions on the objective function and the unbiased stochastic gradients, there exist constant vectors such that2

26and a matrix such that

27Remark 2 The regularity conditions required by Theorem 1 are mostly mild. For example, the stochastic gradients must be unbiased and have finite variance that does not grow too quickly away from the optimum. However, it does require the stronger assumptions that the objective function is smooth and strongly convex. While these conditions do not hold globally, we do not view it be a significant problem in practice because near the optimum we expect the objective function to be locally smooth and strongly convex.

28Theorem 1 shows that, at stationarity, a single iterate will satisfy (with high probability) while its expectation will satisfy. Therefore, when the learning rate is small, is a substantially better estimator for than. In practice the iterate average (that is, the sample mean)

29provides an estimate of, where is the iteration at which the optimization has reached the stationary phase and is the number of iterations used to compute the average. Using as an estimate of is known as Polyak–Ruppert averaging.

30When using iterate averaging, it is crucial to ensure the iterate average accurately approximates. Considering a stationary Markov chain, we can compute the Monte Carlo estimate of the mean of the Markov Chain at stationarity. Then the notion of effective sample size (ESS) aids in quantifying the accuracy of this Monte Carlo estimate. Further, the effective sample size can also be used to define the Monte Carlo standard error (MCSE) when the Markov chain satisfies a central limit theorem. We can efficiently estimate the ESS and also approximate the MCSE (see Appendix B.1 for details). We denote the estimates of ESS and MCSE for the ith component using iterates as and respectively. use the conditions and to determine when to stop iterate averaging. However, no rigorous justification is given for the MCSE threshold.

31A benefit of using a fixed learning rate is that it can be adaptively and automatically decreased once the iterates reach stationarity. For example, if the initial learning rate is, after reaching stationarity the learning rate can decrease to, where is a user-specified adaptation factor. The process can be repeated: when stationarity is reached at learning rate, the learning rate can decrease to. In this way the learning rate is not decreased too early (when the iterates are still making fast progress toward the optimum) or too late (when the accuracy of the iterates is no longer improving). Compare this adaptive approach to the canonical one of setting a schedule such as, which requires the choice of three tuning parameters. These tuning parameters can have a dramatic effect on the speed of convergence, particularly when is far from.

32The question of how to determine when the stationary phase has been reached has a long history with recent renewed attention. One line of work is based on finding an invariant function that has expectation zero under the stationary distribution of the iterates, then using a test for whether the empirical mean of the invariant function is sufficiently close to zero. An alternative approach developed in makes use of the potential scale reduction factor, perhaps the most widely used MCMC diagnostic for detecting stationarity. The standard approach to computing is to use multiple Markov chains. If we have chains and iterates in each chain, then, where and are estimates of, respectively, the between-chain and within-chain variances. In the split- version, each chain is split into two before carrying out the computation above, which helps with detecting non-stationarity and allows for use even when J = 1. Let denote the split- value computed from, the ith dimension of the last W iterates. uses the stationarity condition,

33We now summarize our criteria when designing a robust and automatic optimization framework for BBVI.

34A robust method should not be too sensitive to the choice of tuning parameters. It should also work well on a wide range of “typical” problems. To achieve this we design an adaptive methods for setting parameters (such as the window size for detecting convergence) that are problem-dependent.

35An automatic method should require minimal input from the user. Any inputs that are required should be clearly necessary (for example, the model and the data) or be intuitive to an applied user who is not an expert in variational inference and optimization. Therefore, we require the parameters of any adaptation scheme to either be intuitive or not require adjustment by the user. Examples of intuitive parameters include the maximum number of iterations, maximum runtime, and, when defined appropriately, accuracy.

36We ensure these criteria are satisfied when designing the two core components of a BBVI stochastic optimization framework with automated learning rate scheduling:

37A termination rule for stopping the optimization once the final approximation is close to the optimal approximation (Section 4).

38A learning rate scheduler, which must detect stationarity and determine how many iterates to average before decreasing the learning rate (Section 5).

39The development of our termination rule will proceed in three steps. First, we will select an appropriate discrepancy measure between distributions. Next, we will design an idealized termination rule based on this discrepancy measure. Finally, we will develop an implementable version of the idealized termination rule that satisfies the criteria from Section 3.

40To develop a termination rule, we must specify a measure of how close a variational approximation returned by the optimization algorithm,, is to the optimal variational approximation. But the answer to this question depends upon choosing an appropriate measure of the discrepancy between and the posterior. Based on the discussion in Section 2.1, the goal should be for quantities such as,, and to be close to, respectively,,, and. The interval probabilities are already on an interpretable scale, so ensuring that is much less than 1 is an intuitive notion of accuracy. Since establishes the relevant scale of the problem for means and standard deviations, so it is appropriate to ensure that and are much less than 1.

41While we want to choose a discrepancy measure that guarantees the accuracy of mean, covariance, and interval probabilities, ideally it would also guarantee other plausible expectations of interest (for example, predictive densities) are accurately approximated. The Wasserstein distance provides one convenient metric for accomplishing this goal, and is widely used in the analysis of MCMC algorithms and in large-scale data asymptotics (for example,;;;;;; ). For and a positive-definite matrix, we define the (p, )-Wasserstein distance between distributions and as

42where the infimum is over the set of couplings between and; that is, Borel measures on such that and (, Defs. 6.1 & 1.1). Small (p, )-Wasserstein distance implies many functionals of the two distributions are close relative to the scale determined by.

43Specifically, we have the following result, which is an immediate corollary of, Theorem 3.4).

44Proposition 3 If for any, then

45If, then, for,

46More generally, small (p, )-Wasserstein distance for any guarantees the accuracy of expectations for any function f with small Lipschitz constant with respect the metric; that is, when is small.

47While the Wasserstein distance controls the error in mean and covariance estimates, it does not provide strong control on interval probability estimates. The KL divergence, however, does, since for distributions and, for all a < b and i (see Appendix B.2). As we show next, in many scenarios we can bound the Wasserstein distance by the KL divergence and therefore enjoy the benefits of both. Our result is based on the following definition, which makes the notion of the scale of a distribution precise:

48Definition 4 For and a positive-definite matrix, the distribution is said to be (p, )-exponentially controlled if

49Specially, establishes the appropriate scale for uncertainty with respect to. For example, if, then it is a straightforward exercise to confirm that is (2, 1.782V)-exponentially controlled.

50The following result establishes the relevant link between the KL divergence and the Wasserstein distance via Definition 4.

51Proposition 5 If η is (p, )-exponentially controlled, then for all absolutely continuous with respect to,

52where.

53Proof The result follows from, Corollary 2.3) after the change-of-variable and using the fact that the KL divergence is invariant under diffeomorphisms, then applying Eq. (4).∎

54If and could operate over different scales, then we can use the symmetrized KL divergence. Indeed, it follows from Proposition 5 that if is small, then the (p, )-Wasserstein distance is small whenever either or is (p, )-exponentially controlled.

55Based on our developments in Section 4.1, we will define our termination rule in terms of the symmetrized KL divergence. Recall that denotes the optimal variational approximation to and denotes an estimate of. Since the total variation and (1, )-Wasserstein distances are controlled by the square root of the KL divergence, we focused on the square root of the symmetrized KL divergence. For the current learning rate, let denote the target –learning-rate variational approximation. The termination rule we propose is based on the trade-off between the improved accuracy of the approximation if the learning rate were reduced to and the time required to reach that improved accuracy. To quantify the improved accuracy, we introduce a user-chosen target accuracy target for. If the user expects to be large, then setting to a moderate value such as 1 or 10 could give acceptable performance. If the user expects to be small, then setting to a value such as 0.1 or 0.01 might be more appropriate. Using, define the relative SKL improvement

56where the first term measures the relative improvement of the approximation if the learning rate were reduced and the second term measures the current accuracy relative to the desired accuracy. To quantify the time to obtain the relative accuracy improvement, we use the number of iterations to reach convergence for the fixed learning rate. Letting denote the number of iterations required to reach the target –learning-rate variational approximation, we define the relative iteration increase

57where denotes the number of iterations the user would consider “small”. Combining RSKL and RI, we obtain the inefficiency index, the relative improvement in accuracy times the relative increase in runtime. Thus, we can interpret as quantifying how much greater the increase in runtime cost (above a baseline of iterations) will be compared to the reduction in error (down to a target error of ). For example, means the increase in runtime cost is twice as large as the reduction in error. Our idealized SKL inefficiency termination rule triggers when, where is a user-specified inefficiency threshold that allows the user to trade off accuracy with computation, but only up to the point where.

58The idealized SKL inefficiency termination rule cannot be directly implemented since is unknown; and if it were known, it would be unnecessary to run the optimization algorithm. However, we will show that it is possible to obtain a good estimate of the symmetrized KL divergence between the approximation obtained with a given learning rate and the optimal approximation without access to.

59Recall that denotes the target –learning-rate variational approximation. With a slight abuse of notation, we let the optimal zero–learning-rate approximation refer to the optimal approximation:. Our approach is motivated by Theorem 1 and in particular the form of the bias in Eq. (2). We first consider the still-common setting when is the family of mean-field Gaussian distributions, where the parameter corresponds to the distribution.

60Proposition 6 Let be the family of mean-field Gaussian distributions. If Eq. (2) holds and, then there is a constant depending only on A and such that

61See Appendix A.1 for the proof. Assuming that the current learning rate is, then the previous learning rate was. Let denote the symmetrized KL divergence between the optimal variational approximations obtained at each of these learning rates. In principle we can use Proposition 6 to estimate C by

62and then estimate that

63There are, however, two problems with the tentative approach just outlined. The first problem is that Theorem 1 only holds for standard SGD; however, adaptive SGD algorithms are widely used in practice. Indeed, we observe empirically that ) with for RMSProp (Fig. C.6) and (with a point estimate at 1.2) for Adam Fig. C.7). Hence RMSProp and Adam both appear to have larger errors than SGD when the step size is small. To get the accuracy of SGD but also adaptivity, we modify the adaptive gradient methods to behave asymptotically (in the number of iterations) like SGD. In the cases of RMSProp and Adam, we propose averaged RMSProp (avgRMSProp) and averaged Adam (avgAdam), which use the squared gradient update

64for. Hence, is the averaged squared gradient over all iterations (, §4). As long as the SGD Markov chain is ergodic and at stationarity, converges almost surely to a constant and hence the SGD bias analysis also applies to avgRMSProp and avgAdam.

65The second problem is that Proposition 6 only holds for the mean-field Gaussian variational family. However, other variational families such as normalizing flows are of substantial practical interest. Therefore, we consider the weaker assumption that either Eq. (2) holds or there exist constant vectors, and a constant such that

66Adding the latter assumption, we have the following generalization of Proposition 6, which holds for any sufficiently regular variational family:

67Proposition 7 Let be the variational approximation family. If (i) Eq. (5) holds for some or Eq. (2) holds (in which case let κ = 1), (ii), and (iii) for all, is three-times continuously differentiable with respect to, then for some,

68Moreover, C depends on only and either (if ) or A (if ).

69See Appendix A.2 for the proof. Using Proposition 7 and omitting terms, we have. To improve the reliability of the estimates based on Proposition 7, we propose to use the symmetrized KL estimates between the variational approximations obtained at successive fixed learning rates. Let denote the initial learning rate, so that after t learning rate decreases, the learning rate is. Let and assume the current learning rate is.

70Depending on the optimization algorithm, we can estimate κ (or set κ = 1 if using modified adaptive SGD algorithm with a mean-field Gaussian variation family) and C using a regression model of the form

71where. Given the estimate for C and the estimate for κ (or ), we obtain the estimated relative SKL,

72Because we use the regression model in Eq. (6) in a low-data setting, we place (weak) priors on log C and σ:

73where Cauchy+ is the Cauchy distribution truncated to nonnegative values. If we use an adaptive stochastic optimization algorithm then we also place a prior on κ:

74Also, because we expect early SKL estimates to be less informative about C (and κ) due to the influence of terms, we use a weighted regression with the likelihood term for having weight

75The weight formula enables the amplification of the significance of the most recent observations, with down-weighting becomes more significant after there are about 3 additional observations. On the other hand, the power of 1/4 ensures a gradual reduction in weight, preventing a steep drop-off in importance.

76We use the posterior mean(s) to estimate (and ). Figure 3 validates that, in the case of avgAdam, the log of the learning rate and symmetrized KL divergence have approximately a linear relationship and that our regression approach to estimating C leads to reasonable estimates of. See Fig. C.2 for similar results for other target distributions with avgAdam.

77To estimate the relative iteration increase RI, we need to estimate the number of iterations to reach convergence at the next learning rate. It is reasonable to assume that there is exponential growth in the number of iterations to reach convergence as the learning rate decreases since stochastic gradient algorithms to converge at a polynomial rate. Recall that is the number of iterations to reach convergence at the current learning rate. We fit a weighted least square regression model of the form

78where. We then use the coefficient estimates and to predict the number of iterations required for convergence at the next learning rate to be. We use the same weights given in Eq. (7) for observations of the regression model due to the non-linear behavior of the earlier convergence iterate estimates. Figure 4 demonstrates that linear relationship in Eq. (8) does in fact hold and that our weighted least square regression model predicts the number of convergence iterations quite accurately. The estimated relative iterations is then.

79Using the estimates and we obtain the termination rule. Figure 5 shows that when the user chosen target accuracy, the termination rule triggers when the square root of the symmetrized KL divergence is approximately equal to. Figures C.3 to C.5 shows similar results of other Gaussian targets and posteriordb models and data sets (see Section 7.2 for details).

80For a fixed learning rate, computing the iterate average defined in Eq. (3) requires determining the iteration at which stationarity is reached and the number of iterations to use for computing the average. We address each of these in turn.

81We investigate two approaches to detecting stationarity: the SASA+ algorithm of and the -based criterion from. We make several adjustments to both approaches to reduce the number of tuning parameters and to make the remaining ones more intuitive. In our empirical findings, we have observed that the criterion outperforms the SASA+ criterion. Therefore, we describe the former here and the latter in Appendix B.3.

82Let denotes the split- of the ith component of the last W iterates and define

83An value close to 1 indicates the last W iterates are close to stationarity. In MCMC applications having is desirable. uses the weaker condition since iterate averaging does not require the same level of precision as MCMC. take the window size W = 100, but in more challenging and high-dimensional problems a fixed smaller W is insufficient. Therefore, we instead search over window sizes between a minimum window size Wmin and 0.95k to find the one that minimizes. The minimum window size is necessary to ensure the values are reliable. We use the upper bound 0.95k to always allow a small amount of “warm-up” without sacrificing more than 5% efficiency. Therefore, we estimate using a grid search over 5 equally spaced values ranging from to 0.95k and require as the stationarity condition.

84Figure 6 compares our adaptive SASA+ and adaptive criteria to the criterion used in with a fixed window size of W = 800 and ΔELBO rule from, which is used Stan’s ADVI implementation (cf. the results of ). We do not use W = 100 as is done by because it was too small to detect convergence. Additionally, Fig. 6 compares to another convergence detection approach proposed by (described in Appendix B.4), where they use a distance-based statistic to detect convergence. While adaptive SASA+, ΔELBO, fixed window size, and the distance-based statistic approach sometimes trigger too early or too late or SASA+ use iterations before it reaches the convergence, adaptive consistently triggers when the full window suggests convergence has been reached. See Fig. C.1 for additional Gaussian target examples.

85After detecting convergence to stationarity, we need to find large enough to ensure the iterative average is sufficiently close to the mean. But what is close enough? Building on our discussion in Section 3, we aim to ensure the error in the variational parameter estimates are small relative to the scale of uncertainty. For mean-field Gaussian distributions, the following result allows us to make such a guarantee precise.

86Proposition 8 Let be the family of mean-field Gaussian distributions. Let denote an approximation to. Define and. If there exists such that and, then

87See Appendix A.4 for the proof. Based on Proposition 8, for mean-field Gaussian variational families we use the iterate average once the mean MCSEs and are less than ε.For other variational families we rely on the less rigorous condition that is less than. We also require the effective sample sizes of all parameters to be at least 50 to ensure the MCSE estimates are reliable.

88Because the MCSE check requires computing d ESS values, it can be computationally expensive, especially for high-dimensional models. Therefore, it is important to optimize when conducting the checks.A well-known approach in such situations is the “doubling trick.” Let denote the window size when convergence is detected, and Let denote the window size when convergence is detected, and let denote the minimal window size that satisfies the MCSE check. The doubling trick would suggest checking at iteration numbers for in which case the total computational cost is within a factor of of the optimal scenario in in which the check is only done at and. However, we can potentially do substantially better by accounting for the different computational cost of the optimization versus the MCSE check.

89Proposition 9 Assume that the cost of the MCSE check using K iterates is and the cost of K iterations of optimization is. Let,, and. If the MCSE check is done on iteration numbers for then the total computational cost will be within a factor of of optimal.

90See Appendix A.5 for the proof. Since and is monotonically decreasing in r, when; that is, is negligible compared to, we recover the doubling rule since. However, as long as r is significantly greater than zero, the worst-case additional cost factor can be substantially less than 4. Therefore we carry out the MCSE check on iteration numbers with r estimated based on the actual runtimes of the optimization so far and the first MCSE check.

91Combining our innovations from Sections 4 and 5 leads to our complete framework. When is fixed, our proposal from Section 5 is summarized in Algorithm 1, which we call fixed–learning-rate automated stochastic optimization (FASO). Combining the termination rule from Section 4 with FASO, we get our complete framework, robust and automated black-box variational inference (RABVI), which we summarize in Algorithm 2. We will verify the robustness of RABVI through numerical experiments. RABVI is automatic since the user is only required to provide a target distribution and the only tuning parameters we recommend changing from their defaults are defined on interpretable, intuitive scales:

92accuracy threshold: The symmetrized KL divergence accuracy threshold can be set based on the expected accuracy of the variational approximation. If the user expects to be large, then we recommend choosing. If the user expects to be fairly small, then we recommend choosing. Our experiments suggest is a good default value.

93inefficiency threshold τ: We recommend setting the inefficiency threshold, as this weights accuracy and computation equally. A larger value (for example, 2) could be chosen if accuracy is more important while a smaller value (for example, 1/2) would be appropriate if computation is more of a concern.

94maximum number of iterations: The maximum number of iterations can be set by the user based on their computational budget. RABVI will warn the user if the maximum number of iterations is reached without convergence, so the user can either increase or accept the estimated level of accuracy that has been reached.

95We expect the remaining tuning parameters will typically not be adjusted by the user. We summarize our recommendations:

96initial learning rate: When using adaptive methods such as RMSProp or Adam, the initial learning rate can essentially be set in a problem-independent manner. We use in all of our experiments. If using non-adaptive methods, a line search rule such as the one proposed in could be used to find a good initial learning rate.

97minimum window size: We recommend taking so that that each of the split- values are based on at least 100 samples.

98small iteration number: The value of should represent a number of iterations the user considers to be fairly small (that is, not requiring too much computational effort). We use for our experiments, but it could be adjusted by the user.

99initial iterate average relative error threshold: We recommend scaling with since more accurate iterate averages are required for sufficiently accurate symmetrized KL estimates. Therefore, we take by default.

100adaptation factor ρ: We recommend taking ρ = 0.5 because using a smaller ρ value could lead to too few values for the estimation of C (and κ) and using a larger ρ value would make the algorithm too slow.

101Monte Carlo samples M: We find that M = 10 provides a good balance between gradient accuracy and computational burden but the performance is fairly robust to the choice of M as long as it is not too small.

102Unless stated otherwise, all experiments use avgAdam to compute the descent direction, mean-field Gaussian distributions as the variational family, and the tuning parameters values recommended in Section 6. We fit the regression model for C (and k) in Stan, which result in extremely small computational overhead of less than 0.5%. We compare RABVI to FASO, Stan’s ADVI implementation, SGD using an exponential decay of the learning rate, and fixed-learning rate versions of RMSProp, Adam, and a windowed version of Adagrad (WAdagrad), which is the default optimizer in PyMC3. Moreover, we compare RABVI with exponential decay and cosine learning rate schedules using Adam and RMSProp optimization methods. We run all the algorithms that do not have a termination criterion for iterations and for the fixed–learning-rate algorithms we use learning rate in an effort to balance speed with accuracy. For exponential decay, we use a learning rate of, where denotes the initial learning rate, denotes the decay rate, k denotes the iteration, and s denotes the decay step. We choose,, and so that the final learning rate is approximately 0.0001. For cosine schedule, we use a learning rate of, where and denote the minimum and maximum values of learning rate, k denotes the current iteration, and K denotes the maximum number of iterations. We choose, to make it comparable with other methods.

103We use symmetrized KL divergence as the accuracy measure when we can compute the ground-truth optimal variational approximation. Otherwise, we use the following metrics (where and are, respectively, the posterior mean and standard deviation vectors):

104Relative mean error, where is the variational approximation to.

105Relative standard deviation error, where is the variational approximation to.

106Under coverage error of the variational approximation to the 95% credible intervals, where and is the variational estimate of the central 95% credible interval for parameter.

107Maximum mean discrepancy (MMD), where and are independent and is the squared exponential kernel.

108First, to explore optimization accuracy relative to the optimal variational approximation, we consider Gaussian targets of the form. In such cases, we can compute the ground-truth optimal variational approximation either analytically (because the distribution belongs to the mean-field variational family and hence ) or numerically using deterministic optimization (since the KL divergence between Gaussians is available in closed form). Specifically, we consider the following covariances that aid in assessing our framework across a range of condition numbers from 1 to around 9000:

109Identity covariance:

110Diagonal non-identity covariance:

111Uniform covariance with correlation

112Banded covariance with maximum correlation

113Diagonal non-identity banded covariance with maximum correlation

114Diagonal identity (except first entry) uniform covariance with maximum correlation

115Diagonal identity (except first entry) banded covariance with maximum correlation

116In our selection, we specifically included diagonal identity matrices (with the exception of the first entry) combined with either uniform or banded covariance structures, showcasing a maximum covariance of 0.8. This setting results in weaker correlation between the first component and the others. This choice was strategic to achieve higher condition numbers (around 5000 and 9000 respectively), given that correlations set at 0.8 or yield condition numbers around 400 and 80, respectively.

117Figures 7, C.8 and C.9 show the comparison of RABVI to FASO, Stan’s ADVI implementation, SGD with exponential decay learning rate (SGD-ED), Adam with exponential decay and cosine learning rates (Adam-ED, Adam-C), RMSProp with exponential decay and cosine learning rates (RMSProp-ED, RMSProp-C), and fixed-learning rate versions of RMSProp, Adam, Windowed Adagrad (WAdagrad). The findings demonstrate that RABVI consistently outperforms ADVI, SGD-ED, both adaptive learning rate versions of RMSProp, and the fixed-learning rate methods in a majority of the cases. While Adam and SGD-ED occasionally reach performance levels similar to RABVI, they tend to converge more slowly and with less reliability. Additionally, despite Adam-ED and Adam-C closely matching RABVI’s performance on most problems, they lack a dependable mechanism for determining when to terminate the optimization at a desired accuracy level. Although a validation data set can be used, this requires the availability of such a data set and would allow for control of the approximation accuracy. On the other hand, by varying the accuracy threshold, the quality of the final RABVI approximation also varies such that.

118To demonstrate the flexibility of our framework, we used RABVI with a variety of optimization methods: RMSProp, avgRMSProp, avgAdam, natural gradient descent (NGD), and stochastic quasi-Newton (SQN). See Appendices B.5 and B.6 for details. Figure C.10 shows that avgAdam and avgRMSProp optimization methods have a similar improvement in symmetrized KL divergence between optimal and estimated variational approximation for all cases. NGD is not stable for the diagonal non-identity covariance structure and SQN does not perform well with uniform covariance structure. Even though RMSProp shows an improvement in accuracy for large step sizes, accuracy does not improve as the step size decreases.

119To validate the robustness and reliability of RABVI across realistic use cases, we consider 18 diverse data set/model pairs found in the posteriordb package3 (see Appendix C for details). The posteriordb package contains a wide range of real-world data and models and is specifically designed to provide realistic performance evaluations of approximate posterior inference algorithms. The accuracy was computed based on ground-truth estimates obtained using the posterior draws included in posteriordb package if available. Otherwise, we ran Stan’s dynamic HMC algorithm to obtain the ground truth (4 chains for 50,000 iterations each). To stabilize the optimization, we initialize the variational parameter estimates using RMSProp for the initial learning rate only. A comparison across optimization methods validates our choice of avgAdam over alternatives (Fig. C.11).

120To evaluate RABVI’s effectiveness in real-world applications, we compared it against alternative optimization methods with both fixed and adaptive learning rate schedules. Based on the results described in Section 7.1, we opt to compare to Adam using either a fixed, exponential decay, or cosine learning rate schedule since they perform best overall in the Gaussian target experiments. Additionally, we include FASO, which used avgAdam, as another benchmark. Figure 8 shows RABVI is more consistent than all the alternative methods. While these methods sometimes matched RABVI’s performance, RABVI’s ability to identify an appropriate stopping point contributes to its overall efficiency, setting it apart from the competition.

121First, we investigate the accuracy and algorithmic robustness of RABVI. In terms of robustness, Figures 9, C.12 and C.13 validate our termination criteria since after reaching the termination point there is no considerable improvement in the accuracy for most of the models and data sets. While in many cases the mean estimates are quite accurate, the standard deviation estimates tended to be poor, which is consistent with typical behavior of mean-field approximations. To examine whether RABVI can achieve more accurate results with more flexible variational families, we conduct the same experiment using multivariate Gaussian approximation family and normalizing flow family using real-NVP flow with 2 hidden layers, 8 hidden units, and 3 coupling layers. We employ FASO in our real-NVP experiments because the complexity of the approximation family prevents us from obtaining a closed form for the SKL divergence, which is necessary for computing the termination rule in RABVI. In some cases the accuracy of the mean and/or standard deviations estimates improve (bball_0, dogs_log, 8schools_c, hudson_lynx, hmm_example, nes2000, and sblrc). However, the results are inconsistent and sometimes worse due to the higher-dimensional, more challenging optimization problem. Our findings underscore the necessity of supplementing an improved optimization framework like RABVI with diagnostics for assessing the accuracy of the posterior approximation.

122We additionally benchmarked the runtime and accuracy of RABVI to Stan’s dynamic HMC algorithm, for which we ran 1 chain for 25,000 iterations including 5,000 warmup iterations. We measure runtime in terms of the number of likelihood evaluations and compared the relative error between the methods at the RABVI termination rule trigger point or final likelihood evaluation of HMC (whichever comes first). Figures 10, C.15 and C.16 show that RABVI tends to provide similar or better posterior mean estimates (the exceptions are gp_pois_regr, hudson_lynx, and sblrc). However, the RABVI standard deviation estimates tend to be less accurate even when using the full-rank Gaussian variational family. This could be because the optimization of full-rank Gaussian is more challenging having more variational parameters to estimate. Figure C.17 shows that, in terms of the MMD, the HMC approximation is closer to the target than BBVI as one would expect. Overall, the MMD values for RABVI are reasonably small.

123We also compared the 95% quantiles posterior under coverage error between RABVI and FASO methods using different approximation families and MCMC. Figure C.18 shows that HMC and real-NVP flows do not undercover the posterior. However, the mean-field and full-rank Gaussian families do.

124To validate the robustness of RABVI to different choices of algorithm tuning parameters, we consider the Gaussian targets and two posteriordb data set/model pairs: dogs (logistic mixed-effects model) and arK (AR(5) time-series model). We vary one tuning parameter while keeping the recommended default values for all others. We consider the following values for each parameter (default in bold):

125initial learning rate

126minimum window size

127initial iterate average relative error threshold

128inefficiency threshold

129Monte Carlo samples.

130We repeat each experimental condition 10 times to verify the robustness of different initializations of the variational parameters. Figures 11 and C.19 to C.25 suggest that overall the accuracy and runtime of RABVI is not too sensitive to the choice of the tuning parameters. However, extreme tuning parameter choices (for example, or ) can lead to longer runtimes.

131As we have shown through both theory and experiments, RABVI, our stochastic optimization framework for black-box variational inference, delivers a number of benefits compared to existing approaches:

132The user only needs to, at most, adjust a small number of tuning parameters which empowers the user to intuitively control and trade off computational cost and accuracy. Moreover, RABVI is robust, both in terms of accuracy and computational cost, to small changes in all tuning parameters.

133Our framework can easily incorporate different stochastic optimization methods such as adaptive versions, natural gradient descent, and stochastic quasi-Newton methods. In practice, we found that the averaged versions of RMSProp and Adam we propose perform particularly well. However the performance of RABVI will benefit from future innovations in stochastic optimization methodology.

134RABVI allows for any choice of tractable variational family and stochastic gradient estimator. For example, in many cases we find accuracy improves when using the full-rank Gaussian variational family rather than the mean-field one.

135Our empirical results also highlight some of the limitations of BBVI, which sometimes is less accurate than dynamic HMC when given equal computational budgets. However, BBVI can be further sped up using, for example, data subsampling when the data set size is large (which was not the case for the posteriodb data sets from our experiments).