Riemannian Black Box Variational Inference

Authors: mykola  lukashchuk
wouter w l nuijten
dmitry  bagaev
i̇smail s ¸enöz
bert  de vries
Venue: NeurIPS_cc_2024_Workshop_BDU
Published: 2024-09-05 14:19:50

Abstract: We introduce Riemannian Black Box Variational Inference (RBBVI) for scenarios lacking gradient information of the model with respect to its parameters. Our method constrains posterior marginals to exponential families, optimizing variational free energy using Riemannian geometry and gradients of the log-partition function. It excels with black-box or nondifferentiable models, where popular methods fail. We demonstrate efficacy by inferring parameters from the SIR model and tuning neural network learning rates. The results show competitive performance with gradient-based (NUTS) and gradient-free (Latent Slice Sampling) methods, achieving better coverage and matching Bayesian optimization with fewer evaluations. RBBVI extends variational inference to settings where model gradients are unavailable, improving efficiency and flexibility for real-world applications.

Main paper: 
 introduction: 
 Bayesian Inference focuses on updating beliefs about hypotheses based on newly available evidence. Previously too costly, it is now more feasible thanks to Markov-Chain Monte Carlo (MCMC) algorithms that can sample from difficult posterior distributions and Variational Inference methods that can directly optimize an approximate posterior distribution. However, widely used MCMC methods, such as the No-U-Turn Sampler (NUTS) [12] and well-known variational inference methods such as Automatic Differentiation Variational Inference (ADVI) [17] or Stochastic Variational Inference (SVI) [11] require that the gradient of the log-probability density function with respect to the parameters be computable.Gradient-free methods, such as Ensemble Slice Sampling (ESS) [14] , propose a way to sample from the posterior distribution over parameters even when the gradient is unavailable. Although these methods are gradient-free and nonparametric, they require significant evaluations of the logprobability density function, which may be computationally expensive. This paper introduces Riemannian Black Box Variational Inference (RBBVI) for approximating posterior marginal distributions, addressing computational challenges in two ways: by constraining the posterior to a predefined exponential family [13] , we balance approximation accuracy with computational cost; optimizing variational free energy allows for early stopping, potentially reducing function evaluations compared to sampling methods. RBBVI is particularly well-suited for approximate inference in probabilistic models with blackbox components. This makes the proposed method ideal for scenarios where nondifferentiable and computationally expensive operations are frequent or for hyperparameter tuning, where gradients might not be easily accessible. To demonstrate this, we address the problem of inferring the optimal learning rate of gradient descent while training a neural network. In this situation, the gradient of the loss function with respect to the learning rate is not directly available. Since training a neural network is costly, it is crucial to minimize the number of executions during hyperparameter tuning. 
riemannian variational inference: 
 We are concerned with Bayesian inference in a generative model p(θ, y) with θ unobserved parameters and y observed data. We are interested in the inference task of determining p(θ|y). In this paper, we assume a factorizable generative model with independent priors over (θ 1 , . . . , θ n ) p(y, θ 1 , . . . , θ n ) = p(y|θ 1 , . . . , θ n )EQUATIONwhere p(y, |θ 1 , . . . , θ n ) is a generative process that is possibly not differentiable with respect to θ.Variational inference addresses this by introducing an approximate posterior distribution q on which we want to maximize the evidence lower bound (ELBO) or equivalently minimize the Free EnergyEQUATIONwhere the goal is to computeEQUATIONfor a given observation y.The problem (2b) is a nonparametric one. To cast it into a parametric optimization problem, we assume the following functional form of q(θ 1 , . . . , θ n ) to be a factorized posterior i q i (θ i ), where each constrained q i (θ i ) to be a known exponential familyEQUATIONwith natural parameters λ i , sufficient statistics T i , a base measure κ i , and known log-partition function A i . This turns the problem (2b) into the following parametric problem λ * = argmin λ=(λi,...,λn)EQUATIONThe fruitful idea is to apply gradient-based methods to the Free Energy objective F with respect to the parameters λ. This approach was first proposed by Amari [2] with the following gradient updateEQUATIONwhere F is the Fisher information matrix. The update rule (5) has gained further popularity, and Khan [15] has developed extensions to the rule. The typical challenge in applying the rule ( 5) is that λ usually lies within an open constraint set, so the update may not always meet the constraints. Lin [20] addresses this issue for positive-definite constraints (e.g., Gamma, Gaussian) using Riemannian gradient descent. It is crucial to note that even when the gradient of p(y, θ 1 , . . . , θ n ) with respect to (θ 1 , . . . , θ n ) does not exist or is intractable, the gradient ∂ λ F can still exist and be tractable, depending on the specific choice of q.Similarly to Lin's approach, we assume that each λ i belongs to a constraint set Λ i that forms a Riemannian manifold with F i as its metric. For readers interested in a comprehensive treatment of these concepts, we recommend the book by Absil et al. [1] . For readers new to manifold theory, a manifold M can be thought of as a set that, at any point x ∈ M can be approximated by a vector space. This vector space is called the tangent space T x M. An important detail is that the dimensionality of T x M remains constant for each point x in the manifold. A Riemannian manifold is equipped with a Riemannian metric, which defines a smooth inner product on each tangent space, allowing us to measure distances and angles between tangent vectors. The Riemannian gradient ∂f (x) at a point x on M is the unique tangent vector in T x M that satisfies ⟨∂f (x), ξ⟩ x = Df (x)[ξ] for all ξ ∈ T x M, where ⟨•, •⟩ x denotes the Riemannian metric in x. To move along the manifold in the direction of a tangent vector, we use a retraction R x : T x M → M. Specifically, for a tangent vector ξ ∈ T x M, the curve γ(t) = R x (tξ) satisfies γ(0) = x and γ(0) = ξ.We equip each Λ i with a Riemannian metric F i and a known closed-form retraction R i . We treat Λ = i Λ i as a product manifold with the product retraction R λ (ξ) = (R 1 λ1 (ξ i ), . . . , R n λn (ξ n )), leading to the following update ruleEQUATIONwhereEQUATION)The problem in implementing the scheme (6a) then boils down to a fast computation of ∂ λ F . 
gradient computation: 
 To compute ∂ λ F , we employ the integration by parts as described in [21, Appendix: The Gradient of the ELBO] for the Euclidean case, the proof can be trivially extended to the Riemannian setting via [7, Proposition 8.59 ]. The integration by parts yields the following expression for the gradient of the Free EnergyEQUATIONThe form of ∂ λ F given in the equation (7) leads to an important observation: the differentiability of F (both in the Euclidean and Riemannian senses) does not depend on the properties of p(y, θ 1 , . . . , θ n ) with respect to (θ 1 , . . . , θ n ). Instead, it depends solely on the existence of ∂ λi A(λ i ). When p(y, θ 1 , . . . , θ n ) is differentiable in terms of (θ 1 , . . . , θ n ), a Monte Carlo method with favorable convergence properties is used to estimate the gradient (6a) with the reparameterization trick [16] . This approach can be extended to nondifferentiable functions by approximating the nondifferentiable function with a differentiable surrogate (smooth approximation), as in the RELAX method [10] and then applying the reparameterization trick.The key observation of our study is as follows: we can construct a first-order approximation to ∂ λ F using the form (7) without an explicit smooth approximation of p(y, θ 1 , . . . , θ n ) via∂ λ F ≈      E q λ 1 (θ1) log q λ 1 (θ1) p(y,θ1,µ2,...,µn) (T 1 (θ 1 ) -∂ λ1 A 1 (λ 1 )). . .EQUATIONwhereµ j = E q λ j (θj ) [θ j ].The error term for the approximation (8) can be obtained from Theorem 1 provided in Appendix E applied to each component of the gradient. For each i, we define a function h i that depends on all variables except θ iEQUATIONApplying Theorem 1 to each h θi , we getEQUATIONwhere the approximation error is the limit expression from Theorem 1.To estimate the right-hand side of equation ( 8) we employ REINFORCE estimator [25] corrected on the exponential family sufficient statistics (for details see Appendix C)EQUATIONAlgorithm 1 in Appendix D presents the complete implementation of the procedure described in Equation (6a). This algorithm integrates all key components of our method, offering a comprehensive description of RBBVI. 
experimental results: 
 Experiments presented in this section were implemented in Python [23] and Julia [6] . The code is publicly available at https://github.com/biaslab/GradientFreeVI. For the SIR model experiments in Subsection 4.1, we utilized the model implementation from Frost's SIR repository 1 , integrating their Turing.jl [9] model specification. This integration allowed us to compare with the NUTS [12] and LSS (Latent Slice Sampling) [19] samplers using the Turing.jl framework.To evaluate our method, we use the Susceptible-Infected-Recovered (SIR) epidemiological model [4] . This provides a benchmark in an "ideal" scenario where gradients are readily available, allowing comparison against gradient-based and gradient-free approaches. The SIR model dynamics are governed by the following system of ordinary differential equationsEQUATIONwhere S, I, and R represent the susceptible, infected, and recovered populations respectively, C tracks cumulative cases, N = S + I + R is the total population, β is the infection rate, c = 10 is the contact rate, and γ = 0.25 is the recovery rate. We simulate daily infections by solving this system numerically with a population size of 1000 and observe infections through a Poisson distribution.Our task is to estimate the infection rate β and the initial proportion of infected i 0 = I(0) N given only observed data, with uniform priors on both parameters.We simulate daily infections using known parameters and model observed infections with a Poisson distribution. Our task is to estimate the infection rate β and the initial proportion of infected i 0 given only observed data, with uniform priors. Our generative model isi 0 ∼ Uniform(0, 1)β ∼ Uniform(0, 1)EQUATIONFor our approximate posterior, we use Beta distributions for i 0 and β, as they represent proportions in [0, 1].We compared RBBVI with NUTS [12] and LSS samplers [19] . Table 1 (the table is provided in Appendix A) reveals distinct trade-offs between the methods. Under limited computational resources, our approach demonstrates strong performance, achieving superior coverage compared to nonparametric alternatives. The result highlights the key advantage of our parametric approach: efficient uncertainty quantification with restricted computational resources. However, as computational budgets increase, gradient-based methods like NUTS show their strengths through superior MSE values, while LSS also achieves excellent coverage. These results position our method as particularly valuable when computational efficiency is crucial or gradient information is unavailable. Meanwhile, we acknowledge that gradient-based methods like NUTS should be preferred when gradient information and substantial computational resources are available.In this experiment, we tune the learning rate for training a simple neural network on the MNIST dataset [18] . We cast this as an inverse problem, defining our forward generative model as training the network with a given learning rate εp(y, X, ε) = p(y|X, ε)p(ε)p(X)with p(X) uniform andEQUATIONwhereEQUATIONHere, L is the loss function, X v , X t , y v , y t are validation and training splits, and f w is the neural network with parameters w. The function f ε represents the trained neural network, where the optimal parameters w(ε) depend on the learning rate ε used during training. Crucially, each forward model evaluation requires training a network with the given learning rate.We compare our method against several hyperparameter optimization techniques: Bayesian Optimization with Gaussian Processes (GP) [22] using different kernels: Radial Basis Function (RBF) and Matérn kernel [24] with ν = 1.0 and ν = 2.5, Tree-structured Parzen Estimator (TPE) [5] ; and running inference over the same probabilistic model with LSS sampler [19] .The results in Table 2 (the table is provided in Appendix A) show that our method achieves comparable or better performance than advanced Bayesian optimization techniques with significantly fewer tuning steps. It provides direct uncertainty quantification for the optimal learning rate, enabling effective ensemble training. Our approach balances exploration and exploitation better than other methods, achieving high accuracy for individual models and ensembles. The detailed analysis of the experiment is provided in Appendix B. 
conclusion: 
 In this work, we introduced RBBVI and demonstrated its application in two key areas: obtaining a posterior distribution over the unknown parameters of an ordinary differential equation (ODE) and tuning the learning rate of a neural network. Future research could investigate the assumption of a joint prior distribution over the parameters as mentioned in Equation 1. Additionally, our method could be utilized as a subroutine within larger probabilistic models. Notably, our algorithm was used as an inference subroutine in RxInfer.jl [3] in our experiments. •, and ⋆ represent low, medium, and high computational budgets, respectively. These budgets are defined by the maximum number of calls to an ODE solver: up to 50,000 for low, 300,000 for medium, and 2,500,000 for high. Coverage is calculated as the proportion of times the true parameter value falls within the 95% Bayesian credible interval of the estimated posterior distribution. Higher coverage indicates that the method's uncertainty quantification is more reliable. Our RBBVI demonstrates particularly strong performance under low computational budgets, achieving superior coverage for both parameters compared to nonparametric alternatives and comparable MSE. This highlights RBBVI's ability to provide reliable uncertainty quantification even with limited computational resources. While NUTS shows better MSE due to its gradient-based nature, and LSS performs well with higher budgets, RBBVI's parametric approach offers a compelling advantage in resource-constrained scenarios. As computational budgets increase (• and ⋆), all methods show improved performance, with RBBVI maintaining competitive coverage and MSE values. This comparison demonstrates that our parametric method offers a robust alternative to nonparametric approaches, particularly excelling in scenarios where computational resources are limited while maintaining competitive performance at higher budgets. The table shows the mean and mode of the posterior distribution over the learning rate, the test accuracy using the mode (Acc), the test accuracy of an ensemble of 50 models trained with learning rates sampled from the posterior (Acc @ 50), and the number of tuning steps. Our method (Gamma, Inverse Gamma) provides uncertainty estimates and competitive performance with fewer tuning steps than Bayesian optimization methods (TPE, GP variants). GP methods provide uncertainty estimates of validation loss but not of optimal learning rate. Further analysis is provided in Appendix B. 
 
 Appendix: This section will elaborate on the number of tuning iterations needed in the experiments of subsection 4.2. To tune the learning rate, we train neural networks on a subset of the MNIST training dataset of size 3000. In addition to our method, we tuned the learning rate using Bayesian optimization with a Gaussian Process (GP) maximizer [22] and a tree-structured Parzen Estimator (TPE) [5] . Bayesian optimization methods do not have access to gradients. They are, therefore, prone to exploring low-probability regions of the search space, resulting in more iterations of training the neural network with candidate learning rates, which is very costly [8] . On the other hand, our method will search for a high-probability region of the search space and explore this region. The advantage is that we need fewer neural network training iterations to find a candidate learning rate, but the disadvantage is that we do not explore the entire search space. For every method, performance is evaluated on the validation dataset, a split of 1000 samples from the test dataset. The proposed learning rates are used in the Adam optimizer [16] . All experiments are run on a MacBook Pro 2021 M1 CPU, and none of the individual experiments take longer than 2 hours of CPU time.With our method, we have to choose the functional form of the posterior distribution. To this extent, we choose three different distributions from the Exponential family: Exponential, Gamma, and Inverse Gamma Distributions. We run 100 iterations of variational inference for each distribution, each taking 30 samples. The total number of neural networks trained will be 3000 for every distribution. However, since we expect the Variational Free Energy to decrease and converge gradually, we employ an early stopping criterion that stops the tuning procedure when the difference between the Variational Free Energy of two iterations is less than 0.029. The choice for this stopping criterion was based on empirical results, and different values can be considered.Since we know that the learning rate is typically a small number, we set the prior for every method to a distribution with the expected value at 1 300 . For the shape parameter of the Gamma distributions, we choose 1 as this corresponds to an Exponential distribution.We tune the learning rate with Bayesian optimization using a GP maximization procedure and a TPE. For the GP, we have to choose the kernel that is used. We run the tuning procedure with three different kernels: The Radial basis function (RBF) kernel and two variants of the Matern kernel with hyperparameters 1 and 2.5 [24] . The loss in the validation data set is optimized for these Bayesian optimization methods.However, unlike our method, Bayesian optimization methods do not incrementally improve their estimate of the best learning rate. Instead, they explore a predefined search space to find the optimal learning rate. Therefore, the best learning rate seen so far in terms of the number of neural networks trained is a piecewise linear function with no guarantee of finding a better learning rate at a later stage. For this reason, employing a stopping criterion is impossible, as the best learning rate estimate is not adjusted incrementally. Figure 1 : Convergence plots for the different methods. On the left-hand side, we see the convergence of the Variational Free Energy using our variational inference method. Here, the crosses indicate when our stopping criterion hit and which corresponding learning rate is used in subsection 4.2. On the right, we see the validation loss of the optimal learning rate found so far plotted over the number of neural networks trained. The optimal value so far is visualized because both the GP optimizers and the TPE explore the entire search space instead of incrementally improving their suggestions. This emphasizes that we cannot employ a stopping criterion for the convergence of these methods since some training curves remain constant for hundreds of iterations before finding a new optimal value. We use the optimal found learning rate for the experiments of subsection 4.2.The GP maximization procedure explores the space of log-learning rates between -7 and 1. This means a learning rate between 0.0009 and 1 is considered.Maximizing the validation loss with a GP brings a significant disadvantage: fitting a GP has a computational cost of O(n 3 ), with n being the number of samples. To iteratively fit a GP every time we train a new neural network and determine the optimal learning rate, therefore, it has a computational cost of O(n 4 + nc) with n the number of tuning iterations and c is the cost of training a neural network. This contrasts with the O(nc) run-time cost of variational inference. Because of this, we were able to run only 1000 iterations of the GP.For the TPE, we describe the configuration space with a log-normal density with parameters µ = -5 and σ = 3 .Since our method and the Bayesian Optimization methods optimize a different performance metric, we report them separately. In Figure 1 , we see the curves of both methods. We see the Variational Free Energy over the number of neural networks trained, and we see the validation loss of the optimal learning rate found after n neural networks trained. In addition, in Figure 1a , we can see the point at which our early stopping criterion triggers.As we can see, our method gradually improves its suggestions for the optimal learning rate before converging. Although we cannot directly do gradient steps to improve our performance metric, this suggests that our method does not explore low-probability regions of the Bayesian posterior since the Variational Free Energy never spikes up to the value of the initial guess. In addition, because of this behavior, we can employ a stopping criterion in our method since we do not need to explore the entire search space.We are interested in the estimation of the ∂ λi F (8), rephrasing we are interested in the expectation of the following functionEQUATIONfor that we employ the REINFORCE estimator [25] with a control variate . Specifically, we use the score function of the variational approximation as a control variate, which for exponential family distributions has the following formEQUATIONThis choice maintains the generic nature of our algorithm while allowing us to easily compute the expectation and the covariance matrix of the control variateE q λ i (θi) [f i (λ i , θ i )] = 0 (18a) E q λ i (θi) [f i (λ i , θ i ) T f i (λ i , θ i )] = F i (λ i ) (18b)which results in the following estimatorEQUATIONOur RBBVI approximates posterior distributions without requiring gradients of the log-probability density function. The algorithm implements an iterative optimization procedure that leverages Riemannian geometry and a gradient-free estimator based on the REINFORCE estimator [25] with a control variate. This approach enables optimization even for nondifferentiable forward models.Algorithm 1 provides a detailed implementation, synthesizing all key components into a comprehensive approach. For conciseness, we denote the Riemannian metrics on each Λ i as ⟨v,EQUATIONBefore formally stating the main result, let us introduce a technique for smoothing continuous functions, which can be particularly useful for nondifferentiable processes. Consider a continuous function h : R n → R. We can define a smoothed version of this function as followsEQUATIONwhere σ > 0 is a smoothing parameter. This new function h σ is differentiable even if the original function h is not. Moreover, as σ → 0, h σ converges to h.To illustrate this technique, let us consider the Wiener process (also known as Brownian motion) as an example of a continuous but nowhere differentiable function. Let {W (t)} t≥0 be a standard Wiener process defined in a probability space (Ω, F, P). Although we will not explore its full properties, it is worth noting that W (t) is almost surely continuous but not differentiable at any point.Now, let us apply our smoothing technique to a sample path of the Wiener process. Figure 2 shows a realization of W (t) in the interval [-1, 1], together with three smoothed approximations W σ (t) for different values of σ. In this figure, we observe the original Wiener process W (t) (black line), which is highly irregular and nondifferentiable, alongside three smoothed approximations W σ(t) for decreasing values of σ. The green line (σ = 0.707) shows a very smooth approximation that captures the general trend but misses fine details, the red line (σ = 0.224) presents a moderately smooth approximation that captures more detail while remaining differentiable, and the blue line (σ = 0.095) offers a closer approximation that nearly overlaps with the original process while maintaining differentiability. As σ decreases, we observe that W σ (t) approaches W (t) more closely, illustrating how our smoothing technique can approximate nondifferentiable functions with differentiable ones. This example provides intuition for the more general result we are about to present, which extends this idea to expectations over arbitrary continuous functions.Let f : R n → R be a smooth function and q be a continuous probability distribution with support on an open subset U ⊆ R n . We are interested in computing E q(x) [f ( for i = 1 to n do 4:µ i ← E q t-1 (θi) [θ i ] 5:end for 6:for i = 1 to n do 7:λ 0 ← parameters of q t-1 (θ i )8:for k = 0, 1, 2, . . . do 9:for s = 1 to S do 10:Θ[s] ∼ q λ k (θ) 11: h i [s] = T i (Θ[s]) -∂ λi A i (λ i ) 12: f i [s] = (log q λi (Θ[s]) -log p i (y, µ 1 , . . . , Θ[s], . . . , µ n )) h i [s] 13:end for Update: end for 21:EQUATIONq t (θ i ) = h i (θ i ) exp(T i (θ i ) ⊤ λ k+1 i -A(λ k+1 i )) 22:end forfor i = 1 to n do 24:µ i ← E q t (θi) [θ i ] end if 30: end for 31: return q N (θ i ) for i = 1, . . . , n with Residual Term, we know that for any point x * in which one f is analytic, exists such x ∈ R n that the following property identity holdsEQUATION)Using the fact (22) assuming that f is analytic in µ = E q(x) [x], we obtain the following identityEQUATIONit means that first order approximation to the expectation E q(x) [f (x)] is f (µ). Then for µ = E q [x] there exists a x ∈ U such that the following identity holds First, note that h σ (x) is infinitely differentiable for all σ > 0.EQUATIONBy the definition of expectation and the properties of h σ (x), we haveE q [h(x)] = lim σ→0 E q [ h σ (x)]For any fixed σ > 0, we can apply Taylor's theorem to h σ (x) around µ. There exists a point xσ on the line segment between x and µ such that h σ (x) = h σ (µ) + ∇ h σ (µ) T (x -µ) + 1 2 (x -µ) T ∇ 2 h σ (x σ )(x -µ)Taking the expectation with respect to q(x) on both sidesE q h σ (x ] = h σ (µ) + E q ∇ h σ (µ) T (x -µ)] + 1 2 E q [(x -µ) T ∇ 2 h σ (x σ )(x -µ)Note that E q [(x -µ)] = 0, so the second term on the right-hand side vanishes. Now, let σ → 0. We know that lim σ→0 h σ (x) = h(x) pointwise, solim σ→0 h σ (µ) = h(µ)By the continuity of h and the boundedness of the Gaussian kernel, we can apply the dominated convergence theorem to swap the limit and the expectationlim σ→0 E q h σ (x) = E q [h(x)].For the Hessian term, there exists a subsequence σ k → 0 such that xσ k converges to some x ∈ U . This is because U is open and contains the line segment between x and µ.Combining all these results, we obtain the statement of the theoremE q [h(x)] = h(µ) + lim σ→0 1 2 E q [(x -µ) T ∇ 2 h σ (x σ )(x -µ)].Workshop on Bayesian Decision-making and Uncertainty, 38th Conference on Neural Information Processing Systems (NeurIPS 2024).https://github.com/epirecipes/sir-juliahttps://github.com/orgs/ReactiveBayes/discussions/56 



Reviews: review: title: title: Good paper, clearly written
main_review: main_review: The authors propose a new method for approximating posterior marginal distributions in Bayesian Inference for non-differentiable models.

The paper shows:
1. Strong technical explanation of the underlying concept. 
2. The objective is clearly stated and justified. 
3. The examples are appropriate and of high relevance for real-world applications.
title: title: The paper proposes a scheme to approximate a joint density $p(y,\theta_1,\dots,\theta_i,\dots,\theta_n)$ with $p(y,\mu_1,\dots,\theta_i,\dots,\mu_n)$, where $\mu_j$ is the expectation of $\theta_j$ under its marginal. This is then employed in the computation of gradients of the ELBO w.r.t. the natural parameters of a mean-field variational family belonging to an exponential family of distributions. The advantages of this approach are unclear and a comparison with the exact computation of the likelihood is not presented. The Riemann part is introduced to deal with parameters that belong to constrained sets.
main_review: main_review: The authors propose a gradient-free computational algorithm for approximate Bayesian inference. Gradient-free shall be intended in the sense that we do not need to compute gradients of the likelihood. A mean-field variational distribution is considered, where each marginal is an exponential family. The natural parameters are then learned iteratively by computing natural gradient steps, and the authors aim at addressing the issue that some of these parameters might belong in a constrained set (e.g. positive definite covariance matrix for multivariate normal). They do so along the lines of [1], by performing natural gradient descent in a Riemannian manifold. 

Despite the clear steps involved in the process, I wish the authors clarified some of the aspects specific to working with Riemann manifolds to non-technical audiences. Why are retractions available in closed form? Can examples be provided? The authors proceed by focusing on the computation of the gradient of the ELBO w.r.t. the natural parameters. Despite being available in closed form, computation of such gradients still requires evaluation of the joint density which can be computationally burdensome.  The novelty in the proposed approach is in using an approximation to the joint density given by substituting each $\theta_j$ with its expectation when updating $\theta_i$, where $i\neq j$. I am not convinced by the application of Theorem 1, as it is unclear to me how much we are losing by employing a first order approximation. Also, what happened to $\prod_{j\neq i} q_{\lambda_j}(\theta_j)$ at the numerator in (7), and why does it disappear in (8)? I would like to see a proper derivation for that approximation. I also wish the authors would provide a discussion of this with a clear comparison when we compute the actual joint density, as to me applying such a crude approximation does not appear to be standard practice. Moreover, I suggest changing the notation between $h_{\theta_i}$ and $h_i$ as it can get confusing. To make matters worse, $h_i$ and $f_i$ are flipped in Algorithm 1. It is also unclear to me how substituting the joint density $p(y,\theta_1,\dots,\theta_i,\dots,\theta_n)$ with $p(y,\mu_1,\dots,\theta_i,\dots,\mu_n)$ provides an actual advantage, if the bottleneck is evaluating the likelihood. Maybe there are cases where this is the case, but this should be discussed as it is central to the claims that the paper puts forward.

The results lack a clear description of the specific choices for each baseline. I invite the authors to consider visualizations or other metrics, as coverage and MSE alone can paint a misleading picture. A coverage of 1.00 is not desirable if MSE is consistently higher than other baselines, as the method could simply be predicting inaccurate estimates with a very high variance. Performing on par with GPs on Bayesian optimization tasks would be very impressive. However, the task (MNIST classification) does not seem to be sensitive to the choice of the learning rate as different choices do not lead to dramatic differences in Accuracy, hampering the very need for a Bayesian optimization procedure to be performed in the first place. Moreover, only relatively old or not-so-well-known methods are compared. I would also like to see visualizations of how the method behaves since this is a simple 1D scenario, and the quality of the uncertainty estimates it provides. I think the paper has potential to improve if good performance is shown on Bayesian optimization tasks, but there needs to be a more challenging problem and a more thorough exploration of the results.

I invite the authors to correct the numerous typos and syntax errors in the paper, as they make it quite hard to read.

### **Strengths & Weaknesses**
Strengths:
- The paper addresses the interesting question of performing VI with no gradients of the likelihood.
- The paper highlights how natural gradient updates are a viable option when dealing with such problems, and the possibility of operating in Riemann manifolds to address constraints on the parameters.

Weaknesses:
- The title is misleading since using Riemann gradient descent has been proposed by previous work cited in the paper ([1]). I believe it should be focuse on the specific problem addressed by the paper (i.e. when gradients of the likelihood aren't available)
- Lack of novelty: besides highlighting how the problem of gradient-free VI can be addressed by considering variational distributions belonging to an exponential family and using natural gradients, the novelty of this paper lies in an approximation technique for such gradient. However, the extent to which such approximation is good enough is not discussed. Most importantly, the extent to which such approximation helps in addressing the problem of evaluating the joint density within the closed form gradient update is also not discussed.
- Experimental results are not convincing.

### References
[1] Lin, Wu, Mark Schmidt, and Mohammad Emtiyaz Khan. "Handling the positive-definite constraint in the Bayesian learning rule." International conference on machine learning. PMLR, 2020