Natural Gradient VI: Guarantees for Non-Conjugate Models

Thank you for taking the time to thoroughly review our manuscript and share constructive suggestions. We address each of your weaknesses and major comments in detail below.

Weaknesses

The fact that the ELBO is relatively smooth under compactness is unsurprising. Compactness is very strong, and most regularity assumptions start to hold under this assumption due to continuity. However, in addition to this, the relative smoothness result requires both the log-likelihood and its gradient to be Lipschitz. The need for both at the same time is quite strong.

We agree that compactness is a strong assumption, and that requiring both the log-likelihood and its gradient to be Lipschitz is restrictive. These conditions partly reflect the inherent complexity of NGVI under the relative smoothness framework. As illustrated in Appendix F, for univariate case, while it is possible to relax the compactness assumption on $\mu$, the necessary and sufficient conditions for relative smoothness involve higher-order (third and fourth order) derivatives of the log-probability, highlighting the technical difficulty of characterizing smoothness in this setting. We will explicitly address the limitations of these assumptions in the revision and emphasize that developing approaches to relax them is an important direction for future work.

It is unclear to me if Assumption 1 can be simply assumed away. In [1], Assumption 1 was explicitly established by studying special cases. However, that justification does not apply here since [1] did not study NGVI with Monte Carlo gradients. Given that the assumptions used in the paper are already strong enough, I would expect Assumption 1 to be verified explicitly.

Thank you for raising this point, and we agree that Assumption 1 requires careful verification. In Appendix I, we have explicitly checked that Assumption 1 holds for logistic regression model. We will make this clarification more explicit in the revision.

We did not find Assumption 1 in [1] and assume you meant to refer to some other paper. Could you please share the correct reference?

The experiments in Section 5 have a correctness issue.

Thank you very much for catching this issue. You are correct that there was an error in the gradient computation for Prox-SGD. We have corrected this mistake and rerun all the experiments. The updated results on MNIST dataset (see below) show that Prox-SGD converges as expected and achieves performance comparable to Proj-SGD. We will include these corrected results in the revised version. Thank you again for pointing this out.

We report the median negative ELBO after 1000 epochs (6000 iterations) across 5 runs with tuned step size; values in parentheses indicate the first and third quartiles. These results confirm that Prox-SGD converges with the corrected implementation and achieves performance comparable to Proj-SGD.

We also report the median number of iterations required for the loss to fall below 700 with tuned step size; values in parentheses indicate the first and third quartiles. The results below also show that Prox-SGD converges at a rate comparable to Proj-SGD. However, SNGD algorithms converge much faster than SGD algorithms.

It is unclear why Assumption 2 is needed and what it implies. Why has this to be assumed when projections to a compact set are already being used?

The key condition we need in the proof of Theorem 4.2 is $\omega^+=\omega_*$ (see Appendix H.3), where $\omega^+\coloneqq\arg\min_{\omega' \in\tilde{\Omega}}\\{\langle\nabla\ell(\omega),\omega'\rangle+2\rho D_{A^\*}(\omega',\omega)\\},\quad \omega_*\coloneqq \arg\min_{\omega'\in\Omega}\\{\langle\nabla\ell(\omega),\omega'\rangle+2\rho D_{A^*}(\omega',\omega)\\}.$When $\omega\in\partial \tilde\Omega$, this condition does not necessarily hold, so we impose Assumption 2 to avoid this scenario.

However, it is possible to replace Assumption 2 with another algorithm-independent assumption. The intuition is that if the gradients $\nabla \ell(\omega)$ point towards the interior of $\tilde \Omega$ when $\omega\in\partial \tilde\Omega$, then for sufficiently large $\rho$, we have $\omega_*\in\tilde\Omega$, and hence $\omega^+=\omega_*$.

Since $\nabla A(\eta)$ and $\nabla A^*(\omega)$ are inverse operators of one another (see Appendix C.1), (35) implies $\omega_\*=(\nabla A^\*)^{-1}\left(\nabla A^\*(\omega)-\frac{1}{2\rho}\nabla\ell(\omega)\right)=\omega-\frac{1}{2\rho}\nabla^2 A(\eta)\nabla \ell(\omega)+o(1/\rho)=\omega-\frac{1}{2\rho}(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega)+o(1/\rho)$as $\rho\to\infty$. Then we impose the following assumption.

Assumption 2'. For any $\omega\in\partial \tilde\Omega$, let $\mathbf{n}_\omega$ be the outward normal direction at $\omega$. Then it holds that$\left<-(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega),\mathbf{n}_\omega\right><0.$

Let $W\coloneqq\\{\omega_t:0\leq t\leq T\\}\cap \partial \tilde\Omega$ denote the set of iterates that lie on the boundary of $\tilde \Omega$. Note that $W$ is a finite set. If Assumption 2' is satisfied, there exists some $\varepsilon>0$ such that $\left<-(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega),\mathbf{n}_\omega\right><-\varepsilon$for all $\omega\in W$. Then we have $\left<\mathbf{n}_\omega, \omega_*-\omega\right><-\frac{\varepsilon}{2\rho}+o(1/\rho)<0$

for some sufficiently large $\rho$. Hence $\omega_*\in\tilde\Omega$.

In the following, we present an example where Assumption 2' holds. We consider a univariate Bayesian linear regression model with a single data point $(x, y)$, where $x,y,z\in \mathbb{R},p(z)=\mathcal{N}(0, 1)$ and $p(y|x,z)=\mathcal{N}(xz, 1)$. It can be shown that Assumption 2' is satisfied if the following set of inequalities holds:$-U(x^2+1)<xy<U(x^2+1),\text{ and }D^{-1}<x^2+1<D.$Therefore, by carefully choosing $U$ and $D$, Assumption 2' can be satisfied.

Major Comments

The title is too non-descriptive. It is not obvious that this paper is about a theoretical analysis of the convergence of NGVI.

We will revise our title to a more descriptive one reflecting the main theoretical contribution of the paper, e.g., ``Natural Gradient VI: Guarantees for Non-Conjugate Models''.

Discuss the assumptions in the abstract, especially since they are strong. The fact that the domain is enforced to be bounded and both the log-likelihood and its gradient need to be Lipschitz continuous should be stated. Citing the parametrization considered in the paper would also be useful.

Thank you for this suggestion. We will revisit the abstract mentioning the key assumptions and the parameterization considered in the paper.

What is also not clear to me is why the paper focuses on joint likelihoods with an explicit prior-likelihood pair. Since the paper considers the non-conjugate setting, it should suffice to just consider a generic unnormalized target, which should simplify the setup and the notation. Similarly, the regularity results like Theorem 3.2 and Proposition 2,3 assume regularity of the likelihood $\log p(\mathcal{D}|z)$. This doesn't seem to be necessary. Why not just assume the regularity of the joint? (and equivalently, the "target unnormalized density")

Thank you for this observation. We adopt the explicit prior–likelihood decomposition in order to leverage the structure of the KL term: by isolating $D_{\mathrm{KL}}(q(z)\|p(z))$ we can directly verify that this term is 1-1 relative smooth with respect to $A^*$ (see line 166). The remaining term is the expected negative log-likelihood $-\mathbb{E}_q[\log p(\mathcal{D}|z)]$, for which we establish separate regularity assumptions.

Line 211: Isn't this about "strong" convexity, not just convexity? If it is the case, this should be made more precise.

Yes, you are correct. In this case, $\mu_H>0$, which corresponds to strong convexity rather than mere convexity.

Proposition 2: I think the constant $\mu_C$ can easily be improved. If the diagonal of a triangular matrix is strictly positive, then the diagonal elements correspond to its eigenvalues. Also, the eigenvalues coincide with the singular values. Then I think, for all $i$, $\sigma_{\mathrm{max}}(G_i)\leq\max\left(1,2C_{11},\ldots,2C_{dd}\right)\leq1+2\max_{j}C_{jj}\leq 1+2\sqrt{D}$ holds?

Thank you for your thoughtful attempt to improve the bound. However, the argument does not apply in this case because $G_i$ is not symmetric, so its eigenvalues do not necessarily coincide with its singular values. The singular values must therefore be computed directly from the definition.

---

We also thank the reviewer for the detailed suggestions on citations, terminology and presentation in minor comments. We have already included additional related work in Appendix B, where more papers on NGVI and mirror descent are cited.

We hope our responses have addressed your concerns, and we will incorporate your constructive and thoughtful remarks in the revision, which we believe will greatly improve the overall quality of our work.

References

[1] David Wingate and Theophane Weber. Automated variational inference in probabilistic programming. arXiv preprint arXiv:1301.1299, 2013.

We sincerely appreciate the reviewer’s careful evaluation and insightful feedback. Below, we provide detailed responses to the identified weaknesses and questions.

Weaknesses

There are a few minor typos such as The realative smoothness -> the relative smoothness

Thank you for pointing this out. We will correct these typos.

Can you please clarify the steps in the equalities between lines 176 and 177. I do not understand how the right hand side is obtained from the expressions in the middle.

Since $\mu=\xi,\sigma^2=\Xi-\xi^2$, the chain rule yields $\frac{\partial \mu}{\partial \xi}=1,\frac{\partial \sigma^2}{\partial \xi}=-2\xi, \frac{\partial \mu}{\partial \Xi}=0$ and $\frac{\partial \sigma^2}{\partial \Xi}=1$. Applying Lemma 3.1 together with the results above gives the formulas shown between lines 176 and 177.

On line 126 please establish the PL inequality when first used and describe what PL abbreviates to make the manuscript more self-contained.

Thank you for the suggestion. The abbreviation "PL" first appeared in line 216. We introduced the full term “Polyak–Lojasiewicz” in line 152, and in the revision we will make it clear at this point that we will use the abbreviation “PL.”

Questions

Why are there no error bars on Figure 2 – the results appear somewhat anecdotal in the current presentation? - it would be important here to see how the convergence behavior is across many initializations.

In Figure 2, Poisson regression admits closed-form gradients (see line 1005 in Appendix E). Therefore, for a fixed initialization, the setup is deterministic, and error bars are not necessary.

To examine the convergence behavior under different initializations, we conduct an additional experiment, where $\mu_0\sim \mathrm{Unif}([-3, 0])$ is sampled 10 times, and we compare the performance of SNGD and Proj-SNGD with different step sizes. We report the median number of iterations required for the function value gap to fall below $2\times 10^{-4}$; values in parentheses indicate the first and third quartiles. The results indicate that projection still accelerates convergence under different initializations. We will include these new results in the revised version of the paper.

For $\mu_0>0$, this effect is less pronounced, as SNGD exhibits more stable behavior in this regime.

The experimentation is rather limited considering one synthetic and one real dataset. The paper would improve by establishing the convergence properties on more problems to highlight that the observed benefits especially of proj-SNGD when compared to SNGD can be observed in general. The gains are interesting but with the limited experimentation it could be established more convincingly by considering more problems than presently given in the main paper and small additional experimentation on MNIST given in the supplementary material.

We have conducted additional experiments on new datasets to further highlight the advantages of Proj-SNGD over SNGD.

1. Madelon dataset [1] ($n=2600,d=500$). For SNGD and Proj-SNGD algorithms with different choices of $U$ and $D$, we report the median number of iterations required for the loss to fall below 3000 (denoted as “NA” if not achieved in 2000 iterations); values in parentheses indicate the first and third quartiles (similar to Figure 4).

2. A subset of the CIFAR-10 dataset [2] containing images of cats and dogs ($n=10000,d=3072$). Similarly, we report the median number of iterations required for the loss to fall below 35000 (denoted as “NA” if not achieved in 15000 iterations); values in parentheses indicate the first and third quartiles.

The results above suggest that, with properly chosen hyperparameters $U$ and $D$, Proj-NGVI can outperform NGVI.

It would further be interesting to increase the problem domains systematic from very smooth to less smooth and probe the strengths and limitations of the derived theoretical results relying on smoothness and projections as well as provide a general guidance on tuning U and D in the projection procedure as briefly investigated in the supplementary.

Thank you for raising this practical point. An illustrative experiment on the effect of the domain for the MNIST dataset is presented in Figure 4, Appendix J.2. Based on this, a broad guideline for the choice of $U$ and $D$—intended as intuition rather than a strict prescription—is as follows:

(1) $U$ and $D$ should not be too small, otherwise the optimal solution may fall outside $\tilde \Omega$. For instance, in the left panel of Figure 4, Proj-SNGD with $U=2,D=10$ converges to a suboptimal solution.

(2) $U$ and $D$ should not be too large, otherwise the algorithm will not benefit from projection.

In our experiments, we run a small number of iterations with several candidate values of $U$ and $D$, and select the sets of parameters that lead to faster convergence.

---

We appreciate the reviewer for thoughtful remarks and suggestions, which we believe will improve the overall quality of the manuscript.

References

[1] Isabelle Guyon. Madelon. UCI Machine Learning Repository, 2004. DOI:https://doi.org/10.24432/C5602H.

[2] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.

We thank the reviewer for the thoughtful and helpful feedback. We will address each of the weaknesses and questions in detail.

Weaknesses

I found some incomplete sentences to be edited, e.g., page 2, line 72, "In the Euclidean set up, ..." and page 5 line 167, "the Hessian of the two functions coincide".

Thank you for pointing this out. We will revise these sentences in the revision.

Also, I suggest providing more references behind characterizations they are using. For instance, the duality of the natural parameter in the exponential family and the dual expectation parameter are stated without providing further references.

Thank you for your suggestion. Most of the content in Section 2.1 is based on [1], which provides detailed discussion of the duality between the natural and expectation parameters in exponential families. We will revise the paper to make this clearer.

Questions

While the finite error bounds in Theorem 4.1 and Theorem 4.2 are established, they are conditional upon the choice of the step size. The current theorem only demonstrates the error bound for some step size (not necessarily any step size schedule). Can you further elaborate on influence of the step size choice in error bounds and their rate? Will it hold regardless of how large the initial step size is chosen?

Our theorems use specific step size schemes, which are standard in non-convex optimization literature. For Theorem 4.1, the step-size is constant and is provided in line 1251, Appendix H.2: \gamma_t=\gamma=\min\left\\{\frac{1}{2L},\sqrt{\frac{\lambda_0}{V^2LT}}\right\\}. For Theorem 4.2, we follow the same step size scheme as Theorem 4.7 in [2], which we did not explicitly mention in the paper:$\gamma_t=\begin{cases} \frac{1}{2L},&\text{if }t\leq T/2\text{ and }T\leq\frac{6L}{\mu_B},\\\\ \frac{6}{\mu_B(t-\lceil T/2\rceil)+12L}, &\text{otherwise}. \end{cases}$ In particular, it is important in both cases that $\gamma_t \leq 1/(2L)$ to avoid divergence. We acknowledge that this could cause confusion and will revise the main text to clearly state the step size schemes.

In the projection step introduced in the Section 4.1., one needs to calibrate the projection radius. How should practitioners decide projection radiuses? Can you perhaps elaborate more on the influence of the choice of projection radiuses? I understand that the authors have stated moderate values of $U$ and $D$ work empirically, but it's not clear what they mean by moderate values.

By "moderate values", we mean that:

1. $U$ and $D$ should not be too small, otherwise the optimal solution may fall outside $\tilde \Omega$. For instance, in the left panel of Figure 4 on page 52, Proj-SNGD with $U=2,D=10$ converges to a suboptimal solution.

2. $U$ and $D$ should not be too large, otherwise the algorithm will not benefit from projection.

In our experiments (see Appendix J.2), we run a small number of iterations with several candidate values of $U$ and $D$, and select the sets of parameters that lead to faster convergence.

Perhaps more extensive numerical experiments would be beneficial to justify the projection step.

We have conducted additional experiments on new datasets to further justify the projection step.

1. Madelon dataset [3] ($n=2600,d=500$). For SNGD and Proj-SNGD algorithms with different choices of $U$ and $D$, we report the median number of iterations required for the loss to fall below 3000 (denoted as “NA” if not achieved in 2000 iterations); values in parentheses indicate the first and third quartiles (similar to Figure 4).

2. A subset of the CIFAR-10 dataset [4] containing images of cats and dogs ($n=10000,d=3072$). Similarly, we report the median number of iterations required for the loss to fall below 35000 (denoted as “NA” if not achieved in 15000 iterations); values in parentheses indicate the first and third quartiles.

The results above suggest that projection with appropriately chosen hyperparameters can accelerate the convergence of NGVI.

---

Thank you again for the insightful questions and helpful suggestions. We hope our responses have addressed your concerns, and we will revise the manuscript accordingly to improve the clarity and completeness of our manuscript.

References

[1] Martin J Wainwright, Michael I Jordan, et al. Graphical models, exponential families, and variational inference. Foundations and Trends in Machine Learning, 1(1–2):1-305, 2008.

[2] Ilyas Fatkhullin and Niao He. Taming nonconvex stochastic mirror descent with general bregman divergence. In International Conference on Artificial Intelligence and Statistics, pages 3493–3501. PMLR, 2024.

[3] Isabelle Guyon. Madelon. UCI Machine Learning Repository, 2004. DOI:https://doi.org/10.24432/C5602H.

[4] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.

We thank the reviewer for the careful reading and valuable feedback. Below, we address each of the weaknesses and questions in detail.

Weaknesses

The main limitation, as emphasized by the authors, is Assumption 2. Unlike Assumptions 1 and 3, it cannot be checked directly and is susceptible to not being satisfied even in the restricted setting considered in the paper.

We agree that Assumption 2 is primarily a technical condition and can be challenging to verify beforehand. The key condition we need in the proof of Theorem 4.2 is $\omega^+=\omega_*$ (see Appendix H.3), where $\omega^+\coloneqq\arg\min_{\omega' \in\tilde{\Omega}}\\{\langle\nabla\ell(\omega),\omega'\rangle+2\rho D_{A^\*}(\omega',\omega)\\},\quad \omega_*\coloneqq \arg\min_{\omega'\in\Omega}\\{\langle\nabla\ell(\omega),\omega'\rangle+2\rho D_{A^*}(\omega',\omega)\\}.$When $\omega\in\partial \tilde\Omega$, this condition does not necessarily hold, so we impose Assumption 2 to avoid this scenario.

However, it is possible to replace Assumption 2 with another algorithm-independent assumption. The intuition is that if the gradients $\nabla \ell(\omega)$ point towards the interior of $\tilde \Omega$ when $\omega\in\partial \tilde\Omega$, then for sufficiently large $\rho$, we have $\omega_*\in\tilde\Omega$, and hence $\omega^+=\omega_*$.

Since $\nabla A(\eta)$ and $\nabla A^*(\omega)$ are inverse operators of one another (see Appendix C.1), (35) implies $\omega_\*=(\nabla A^\*)^{-1}\left(\nabla A^\*(\omega)-\frac{1}{2\rho}\nabla\ell(\omega)\right)=\omega-\frac{1}{2\rho}\nabla^2 A(\eta)\nabla \ell(\omega)+o(1/\rho)=\omega-\frac{1}{2\rho}(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega)+o(1/\rho)$as $\rho\to\infty$. Then we impose the following assumption.

Assumption 2'. For any $\omega\in\partial \tilde\Omega$, let $\mathbf{n}_\omega$ be the outward normal direction at $\omega$. Assume that$\left<-(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega),\mathbf{n}_\omega\right><0.$

Let $W\coloneqq\\{\omega_t:0\leq t\leq T\\}\cap \partial \tilde\Omega$ denote the set of iterates that lie on the boundary of $\tilde \Omega$. Note that $W$ is a finite set. If Assumption 2' is satisfied, there exists some $\varepsilon>0$ such that $\left<-(\nabla^2 A^\*(\omega))^{-1}\nabla\ell(\omega),\mathbf{n}_\omega\right><-\varepsilon$for all $\omega\in W$. Then we have $\left<\mathbf{n}_\omega, \omega_*-\omega\right><-\frac{\varepsilon}{2\rho}+o(1/\rho)<0$

for some sufficiently large $\rho$. Hence $\omega_*\in\tilde\Omega$.

In the following, we present an example where Assumption 2' holds. We consider a univariate Bayesian linear regression model with a single data point $(x, y)$, where $x,y,z\in \mathbb{R},p(z)=\mathcal{N}(0, 1)$ and $p(y|x,z)=\mathcal{N}(xz, 1)$. It can be shown that Assumption 2' is satisfied if the following set of inequalities holds:$-U(x^2+1)<xy<U(x^2+1),\text{ and }D^{-1}<x^2+1<D.$Therefore, by carefully choosing $U$ and $D$, Assumption 2' can be satisfied. We will include the derivation of this fact in the revision.

While the paper is well-structured, its clarity could be improved. This is partly because each result relies on its own set of technical assumptions. A dedicated section outlining all the assumptions would significantly improve readability and be a valuable addition. It would be helpful to have more than just the two (three, counting the appendices) assumptions explicitly stated in the paper.

Thank you for your suggestion. We will add a list of assumptions and a summary of key external results used in the proofs, especially when the proof mainly involves verifying these conditions. This will make the paper easier to follow and more readable.

The results are limited to specific problems in a purely Gaussian setting. They do not seem to be easily generalizable to other settings, as a significant part of the arguments rely on the properties of Gaussian distributions.

We acknowledge that our analysis relies on properties of the Gaussian family. We choose this setting due to its analytical simplicity and its widespread use in practice. This restriction is standard in the literature (e.g., [1,2]), as relaxing it would require fundamentally different techniques. We view extending the results beyond the Gaussian setting as an important direction for future work.

There is no explicit construction of the step-size sequence.

Thank you for the comment. Our theorems use specific step size schemes, which are standard in non-convex optimization literature. For Theorem 4.1, the step-size is constant and is provided in line 1251, Appendix H.2: \gamma_t=\gamma=\min\left\\{\frac{1}{2L},\sqrt{\frac{\lambda_0}{V^2LT}}\right\\}. For Theorem 4.2, we follow the same step size scheme as Theorem 4.7 in [3], which we did not explicitly mention in the paper:$\gamma_t=\begin{cases} \frac{1}{2L},&\text{if }t\leq T/2\text{ and }T\leq\frac{6L}{\mu_B},\\\\ \frac{6}{\mu_B(t-\lceil T/2\rceil)+12L}, &\text{otherwise}. \end{cases}$ In particular, it is important in both cases that $\gamma_t \leq 1/(2L)$ to avoid divergence. We acknowledge that this could cause confusion and will revise the main text to clearly state the step size schemes.

Questions

Lines 284-286: “empirical results indicate that […] even if Assumption 2 does not hold, the performance of Proj-SNGD is no worse and sometimes better than SGD.” Can you expand on this?

In our experiments (Figures 2 and 3), Assumption 2 is not satisfied. This can be inferred from the different loss behaviors of SNGD and Proj-SNGD: such differences indicate that projection occurred during Proj-SNGD, and when projection occurs, the iterate must lie on the boundary of $\tilde\Omega$. We will make this explicit in the revision to avoid confusion. Despite this, Proj-SNGD converges faster and demonstrates greater robustness to step size tuning. These observations suggest that Proj-SNGD can outperform SGD algorithms even when Assumption 2 does not hold.

Is the big O characterization of $L$ given by Theorem 3.2 practically relevant?

The Big-O characterization of $L$ in Theorem 3.2 is primarily of theoretical interest, as it provides a conservative estimate of the real smoothness parameter. It offers insight into how $L$ scales with problem parameters, including $d,U,D$ and the properties of the log likelihood function, which can guide step size selection in practice. In particular, the dependence on dimension $d$ is at most linear, and the dependence on $U$ and $D$ is polynomial of degree at most 3. This characterization is practically relevant because $L$ directly influences the convergence rate of Proj-SNGD, and implies the polynomial iteration complexity in these parameters.

I don’t really understand how (35) is obtained. What are the steps that lead to this equation?

(35) is obtained by explicitly solving the optimization problem in line 1323. By definition, the Bregman divergence equals $D_{A^\*}(\omega',\omega)=A^\*(\omega')-A^\*(\omega)-\left<\nabla A^\*(\omega),\omega'-\omega\right>.$ Substituting it into the problem in line 1323 and applying the first-order optimality condition yields $\nabla\ell(\omega)+2\rho(\nabla A^\*(\omega')-\nabla A^\*(\omega))=0$, which coincides with (35). The second derivative $2\rho \nabla^2 A^\*(\omega')\succ 0$ because $A^\*$ is strictly convex.

In the Poisson regression experiment, what is the function value gap?

Function value gap at iteration $t$ is the difference between the function value at iteration $t$ and the optimal value, i.e., $\ell(\omega_t)-\ell^*$.

What is the difference between Figures 1 and 2? Does projecting achieve the same effect as changing the initialization in this example?

Figures 1 and 2 correspond to the same experimental setting. The left panel of Figure 1 is identical to the left panel of Figure 2. The right panel of Figure 1 illustrates the performance of SNGD under a different initialization, and the right panel of Figure 2 shows the performance of our proposed Proj-SNGD with the original initialization. Both projection and changing the initialization in this example can stabilize and accelerate the training process, but the training trajectories of the two are not identical. The main insight from this experiment is that the objective is less well-behaved in the nonconvex setting (Figure 1), and that our Proj-SNGD converges even with initializations for which original SNGD initially diverges (Figure 2).

---

We also thank the reviewer for the helpful suggestions on clarity and presentation in the additional remarks, and we will incorporate these suggestions in the revision to improve readability. We believe that the reviewer’s feedback will help us improve the quality of the manuscript.

References

[1] Kaiwen Wu and Jacob R. Gardner. Understanding Stochastic Natural Gradient Variational Inference. In International Conference on Machine Learning, pages 53398-53421. PMLR, 2024.

[2] Navish Kumar, Thomas Möllenhoff, Mohammad Emtiyaz Khan, and Aurelien Lucchi. Optimization guarantees for square-root natural-gradient variational inference. Transactions on Machine Learning Research, 2025.

[3] Ilyas Fatkhullin and Niao He. Taming nonconvex stochastic mirror descent with general bregman divergence. In International Conference on Artificial Intelligence and Statistics, pages 3493–3501. PMLR, 2024.