SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Corollary 26. Take a logistic regression model with inputs  $a_{n}\in \mathbb{R}^{d}$  and outputs  $x_{n}\in \{-1, + 1\}$ . Let  $p(z) = \mathcal{N}(z|0,\Sigma)$  and  $p(x_{n} = 1|z) = \mathrm{Sigmoid}(x_{n}z^{\top}a_{n}) = 1 / (1 + \exp (-x_{n}z^{\top}a_{n}))$ . Then  $-\log p(z,x)$  is  $\mu$ -strongly convex for  $\mu = \lambda_{\min}\left(\Sigma^{-1}\right)$  and  $M$ -smooth for  $M = \lambda_{\max}(\Sigma^{-1} + \frac{1}{4} AA^{\top})$ .
$$
\begin{array}{l} \phi \left(z ^ {\top} a _ {n}, x _ {n}\right) = - \log p \left(x _ {n} \mid z\right) \\ = \log \left(1 + \exp \left(- x _ {n} z ^ {\top} a _ {n}\right)\right). \\ \end{array}
$$
$$
\phi (\alpha , \beta) = \log (1 + \exp (- \alpha \beta))
$$
$$
\phi^ {\prime} (\alpha , \beta) = - \frac {\beta}{1 + \exp (\alpha \beta)}
$$
$$
\phi^ {\prime \prime} (\alpha , \beta) = \frac {\beta^ {2} \exp (\alpha \beta)}{(1 + \exp (\alpha \beta)) ^ {2}}
$$
The second derivative is non-negative and goes to zero if  $\alpha \to -\infty$  or  $\alpha \to +\infty$ . It is maximized at  $\alpha = 0$  in which case  $\phi'' = \frac{1}{4}\beta^2$ . But of course, in this model, the second input is  $\beta = x_{n} \in \{-1, +1\}$ , meaning  $\beta^2 = \text{always}$ . Thus we have that  $\theta_{\min} = 0$  and  $\theta_{\max} = 1$ . It follows from Theorem 24 that  $-\log p(z,x)$  is  $\mu$ -strongly convex for  $\mu = \lambda_{\min}(\Sigma^{-1})$  and  $M$ -smooth for  $M = \lambda_{\max}(\Sigma^{-1} + \frac{1}{4} AA^\top)$ .
In the particular case where  $\Sigma = I$ , we get that  $\mu = 1$  and  $M = \lambda_{\max}(1 + \frac{1}{4} AA^{\top}) = 1 + \frac{1}{4}\sigma_{\max}(A)^2$ .
Theorem 27. Take a hierarchical model of the form  $p(\theta, z, x) = p(\theta) \prod_{i} p(z_i | \theta) p(x_i | \theta, z_i)$ , where all log-densities are twice-differentiable. Then,  $\log p(\theta, z, x)$  is smooth jointly over  $(\theta, z)$  if and only if the spectral norm of the second derivative matrices  $\nabla_{\theta \theta^{\top}} \log p(\theta, z, x)$ ,  $\nabla_{\theta z_i^{\top}} \log p(\theta, z_i, x_i)$ , and  $\nabla_{z_i z_i^{\top}} \log p(\theta, z_i, x_i)$  are all bounded uniformly over  $(\theta, z)$ .
Proof. The following conditions are all equivalent:
1.  $\log p(\theta ,z,x)$  is smooth as a function of  $(\theta ,z)$
2.  $\left\| \nabla^2 \log p(\theta, z, x) \right\|_2$  is bounded above (uniformly over  $(\theta, z)$ , where  $\nabla^2$  denotes the Hessian with respect to  $(\theta, z)$  and  $\| \cdot \|_2$  denotes the spectral norm.)
3.  $\left\| \nabla^2 \log p(\theta, z, x) \right\|$  is bounded above (uniformly over  $(\theta, z)$ , where  $\nabla^2$  denotes the Hessian with respect to  $(\theta, z)$  and  $\| \cdot \|$  denotes a block infinity norm on top of spectral norm inside each block.
4.  $\left\| \nabla_{\theta \theta^{\top}}^{2}\log p(\theta ,z,x)\right\|_{2},\left\| \nabla_{\theta z_{i}^{\top}}^{2}\log p(\theta ,z,x)\right\|_{2}$  and  $\left\| \nabla_{z_i z_j^\top}\log p(\theta ,z,x)\right\| _2$  are all bounded above (uniformly over  $(\theta ,z)$ )
5.  $\left\| \nabla_{\theta \theta^{\top}}^{2}\log p(\theta ,z,x)\right\|_{2},\left\| \nabla_{\theta z_{i}^{\top}}^{2}\log p(\theta ,z_{i},x_{i})\right\|_{2},$  and  $\left\| \nabla_{z_i z_j^\top}\log p(\theta ,z,x)\right\| _2$  are all bounded above (uniformly over  $(\theta ,z)$ )
6.  $\left\| \nabla_{\theta \theta^{\top}}^{2}\log p(\theta ,z,x)\right\|_{2},\left\| \nabla_{\theta z_{i}^{\top}}^{2}\log p(\theta ,z_{i},x_{i})\right\|_{2}$  and  $\left\| \nabla_{z_i z_i^\top}\log p(\theta ,z,x)\right\| _2$  are all bounded above (uniformly over  $(\theta ,z)$ )
7.  $\left\| \nabla_{\theta \theta^{\top}}^{2}\log p(\theta ,z,x)\right\|_{2},\left\| \nabla_{\theta z_{i}^{\top}}^{2}\log p(\theta ,z_{i},x_{i})\right\|_{2},$  and  $\left\| \nabla_{z_i z_i^\top}\log p(\theta ,z_i,x_i)\right\| _2$  are all bounded above (uniformly over  $(\theta ,z))$
Corollary 28. Take a hierarchical logistic regression model defined by  $p(\theta, z, x) = p(\theta) \prod_{i} p(z_i | \theta) \prod_{j} p(x_{ij} | z_i)$  where  $\{\theta, z_i\} \in \mathbb{R}^d$ ,  $x_{ij} \in \{-1, +1\}$ , and
$$
\theta \sim \mathcal {N} (0, \Sigma)
$$
$$
z _ {i} \sim \mathcal {N} (\theta , \Delta)
$$
$$
\begin{array}{c c c} x _ {i j} & \sim & p (x _ {i j} | z _ {i}) \end{array}
$$
$$
p (x _ {i j} = 1 | z _ {i}) = \operatorname {S i g m o i d} \left(x _ {i j} z _ {i} ^ {\top} a _ {i j}\right),
$$
where  $\operatorname{Sigmoid}(b) = \frac{1}{1 + \exp(-b)}$ . Then,  $-\log p(\theta, z, x)$  is convex and smooth with respect to  $(\theta, z)$ .
Proof. It's immediate that  $-\log p$  is convex with respect to  $(\theta, z)$ . We can show that this function is smooth, by verifying the three conditions of Theorem 27.
-  $\nabla_{\theta \theta^{\top}}^{2}\log p(\theta ,z,x) = -\Sigma^{-1} - \Delta^{-1}$  is constant.
-  $\nabla_{\theta z_i}^2\log p(\theta ,z_i,x_i) = \Delta^{-1}$  is also constant.
-  $p(z_i|\theta)$  is a Gaussian with constant covariance, so is smooth with a constant independent of  $\theta$ . And we can write that  $\log p(x_{ij}|z_i) = \log \left(1 + \exp \left(-x_{ij}z_i^\top a_{ij}\right)\right)$ . This function is also smooth over  $z_i$  with a constant that doesn't depend on  $\theta$ . Thus  $\log p(z_i|\theta) + \sum_j\log p(x_{ij}|z_i)$  is smooth over  $z_i$  with a smoothness constant that is independent of  $\theta$ .
8 Estimators and variance bounds
We remember that unless specified otherwise, we consider that  $w = (m,C) \in \mathbb{R}^d \times V$ , where  $V$  is a vector subspace of  $\mathcal{M}^d$ .
8.1 Variance bounds for the estimators
Theorem 1. Let  $T_w(u) = Cu + m$  for  $w = (m, C)$ . Let  $\phi : \mathbb{R}^d \to \mathbb{R}$  be  $M$ -smooth, suppose that  $\phi$  is stationary at  $\bar{m}$ , and define  $\bar{w} = (\bar{m}, 0)$ . Then
$$
\underset {\mathsf {u} \sim \mathcal {N} (0, I)} {\mathbb {E}} \| \nabla_ {w} \phi (T _ {w} (\mathsf {u})) \| _ {2} ^ {2} \leq (d + 1) M ^ {2} \| m - \bar {m} \| _ {2} ^ {2} + (d + 3) M ^ {2} \| C \| _ {F} ^ {2} \leq (d + 3) M ^ {2} \| w - \bar {w} \| _ {2} ^ {2}.
$$
Proof. This bound is proven by [20, Thm. 3] in the case that  $V = \mathcal{M}^d$ . We conclude that this bound holds for every subspace  $V$  of  $\mathcal{M}^d$  by using the inequality in Lemma 17.2.
Theorem 2. Suppose that  $\log p(\cdot ,x)$  is  $M$  -smooth and has a maximum (or stationary point) at  $\bar{m}$  and define  $\bar{w} = (\bar{m},0)$ . Then, for every  $w$  and every solution  $w^{*}$  of the VI problem,
$$
\begin{array}{l} \mathbb {E} \left\| \mathbf {g} _ {\text {e n e r g y}} (\mathfrak {u}) \right\| _ {2} ^ {2} \leq (d + 3) M ^ {2} \| w - \bar {w} \| _ {2} ^ {2} \tag {6} \\ \leq 2 (d + 3) M ^ {2} \| w - w ^ {*} \| _ {2} ^ {2} + 2 (d + 3) M ^ {2} \| w ^ {*} - \bar {w} \| _ {2} ^ {2}. \\ \end{array}
$$
Proof. It is a direct consequence of Theorem 1 and Young's inequality
$$
\begin{array}{l} \mathbb {E} \left\| \mathbf {g} _ {\text {e n e r g y}} \right\| _ {2} ^ {2} \leq (d + 3) M ^ {2} \| w - \bar {w} \| _ {2} ^ {2} \\ = (d + 3) M ^ {2} \| w - w ^ {*} + w ^ {*} - \bar {w} \| _ {2} ^ {2} \\ \leq 2 (d + 3) M ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} + 2 (d + 3) M ^ {2} \left\| w ^ {*} - \bar {w} \right\| _ {2} ^ {2}. \\ \end{array}
$$
Proof. The difference between  $\mathbf{g}_{\mathrm{ent}}$  and  $\mathbf{g}_{\mathrm{energy}}$  is the addition of the constant vector  $\nabla h(w)$ . Thus, we have that
$$
\begin{array}{l} \mathbb {E} \left\| \mathbf {g} _ {\mathrm {e n t}} \right\| _ {2} ^ {2} = \mathbb {E} \left\| \mathbf {g} _ {\mathrm {e n e r g y}} + \nabla h (w) \right\| _ {2} ^ {2} \\ \leq 2 \mathbb {E} \| \mathrm {g} _ {\text {e n e r g y}} \| _ {2} ^ {2} + 2 \mathbb {E} \| \nabla h (w) \| _ {2} ^ {2} \\ \stackrel {(6)} {\leq} 2 (d + 3) M ^ {2} \| w - \bar {w} \| _ {2} ^ {2} + 2 \| \nabla h (w) \| _ {2} ^ {2}. \\ \end{array}
$$
It remains to bound the final term. We can do this using the closed-form expression for the entropy gradient (see Lemma 18.2), plus the assumption that  $w \in \mathcal{W}_L \subset \mathcal{W}$ , to see that
$$
\begin{array}{l} \left\| \nabla h (w) \right\| _ {2} ^ {2} = \left\| \operatorname {p r o j} _ {V} \left(C ^ {- \top}\right) \right\| _ {F} ^ {2} \\ \leq \left\| C ^ {- \top} \right\| _ {F} ^ {2} \\ = \sum_ {i} \sigma_ {i} \left(C ^ {- \top}\right) ^ {2} \\ = \sum_ {i} \sigma_ {i} (C) ^ {- 2} \\ \leq \sum_ {i} \left(\frac {1}{\sqrt {L}}\right) ^ {- 2} \\ = d L. \\ \end{array}
$$
We conclude with Young's inequality
$$
\begin{array}{l} \mathbb {E} \| \mathbf {g} _ {\text {e n t}} \| _ {2} ^ {2} \leq 2 (d + 3) M ^ {2} \| w - \bar {w} \| _ {2} ^ {2} + d L \\ = 2 (d + 3) M ^ {2} \left\| w - w ^ {*} + w ^ {*} - \bar {w} \right\| _ {2} ^ {2} + d L \\ \leq 4 (d + 3) M ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} + 4 (d + 3) M ^ {2} \left\| w ^ {*} - \bar {w} \right\| _ {2} ^ {2} + d L. \\ \end{array}
$$
For the next Theorem, we will need the following technical Lemma :
Lemma 29. Let  $\mathfrak{u} \sim \mathcal{N}(0, I)$ ,  $A \in \mathcal{M}^d$  and  $b \in \mathbb{R}^d$ . Then
$$
\mathbb {E} \| A u + b \| ^ {2} (1 + \| u \| ^ {2}) = (d + 1) \| b \| ^ {2} + (d + 3) \| A \| _ {F} ^ {2}.
$$
Proof. Develop the squares to write
$$
\mathbb {E} \| A u + b \| ^ {2} (1 + \| u \| ^ {2}) = \mathbb {E} _ {u} \| A u \| ^ {2} + 2 \langle A ^ {\top} b, u \rangle + \| b \| ^ {2} + \| A u \| ^ {2} \| u \| ^ {2} + 2 \langle A ^ {\top} b, u \| u \| ^ {2} \rangle + \| b \| ^ {2} \| u \| ^ {2}.
$$
Since  $\mathbb{E}\mathfrak{u} = 0$ , we immediately see that  $\mathbb{E}\langle A^{\top}b,\mathfrak{u}\rangle = 0$ . Because of symmetry, the third-order moment  $\mathbb{E}\mathfrak{u}\| \mathfrak{u}\| ^2$  is also zero, so we obtain  $\mathbb{E}\langle A^{\top}b,\mathfrak{u}\| \mathfrak{u}\| ^2\rangle = 0$ . We also have the second-order moment  $\mathbb{E}\| \mathfrak{u}\| ^2 = d$ , which means that  $\mathbb{E}\| b\| ^2\| \mathfrak{u}\| ^2 = d\| b\| ^2$ . For the other terms, we use the identities
$$
\mathbb {E} \mathbf {u u} ^ {\top} = I, \quad \mathbb {E} \mathbf {u u} ^ {\top} \mathbf {u u} ^ {\top} = (d + 2) I,
$$
from [20, Lemma 9] It allows us to write
$$
\mathbb {E} \| A u \| ^ {2} = \mathbb {E} \operatorname {t r} \left(A ^ {\top} A u u ^ {\top}\right) = \operatorname {t r} \left(A ^ {\top} A\right) = \| A \| _ {F} ^ {2},
$$
$$
\mathbb {E} \| A u \| ^ {2} \| u \| ^ {2} = \operatorname {t r} \left(A ^ {\top} A u u ^ {\top} u u ^ {\top}\right) = (d + 2) \| A \| _ {F} ^ {2}.
$$
The conclusion follows after gathering all those identities.
$$
\begin{array}{l} \mathbb {E} \| \mathrm {g} _ {\mathrm {S T L}} \| _ {2} ^ {2} \leq 8 (d + 3) M ^ {2} \| w - w ^ {*} \| _ {2} ^ {2} + 2 (d + 3) K ^ {2} \| w - \hat {w} \| _ {2} ^ {2} \tag {10} \\ \leq 4 (d + 3) \left(K ^ {2} + 2 M ^ {2}\right) \| w - w ^ {*} \| _ {2} ^ {2} + 4 (d + 3) K ^ {2} \| w ^ {*} - \hat {w} \| _ {2} ^ {2}. \\ \end{array}
$$
Proof. Let  $w \in \mathcal{W}_M$  be fixed. For the duration of the proof, we take  $v$  to be a second copy of  $w$  (held constant under differentiation with respect to  $w$ ). Now, we can rearrange the estimator in Eq. 9 by making appear  $q_{w^*}(T_w(\mathfrak{u}))$ , namely
$$
\mathbf {g} _ {\mathrm {S T L}} (\mathfrak {u}) = \nabla_ {w} \log \frac {q _ {v} (T _ {w} (\mathfrak {u}))}{p (T _ {w} (\mathfrak {u}) , x)} = \nabla_ {w} \log \frac {q _ {v} (T _ {w} (\mathfrak {u}))}{q _ {w ^ {*}} (T _ {w} (\mathfrak {u}))} + \nabla_ {w} \log \frac {q _ {w ^ {*}} (T _ {w} (\mathfrak {u}))}{p (T _ {w} (\mathfrak {u}) , x)}.
$$
Introducing  $\phi(z) \coloneqq \log q_v(z) - \log q_{w^*}(z)$ , and recalling  $r(z) = \log p(z, x) - \log q_{w^*}(z)$ , we have
$$
\mathbf {g} _ {\mathrm {S T L}} = \nabla_ {w} \phi (T _ {w} (\mathfrak {u})) - \nabla_ {w} r (T _ {w} (\mathfrak {u})).
$$
We assume that  $\log p(\cdot, x)$  is  $M$ -smooth, and that  $V = \mathcal{M}^d$ ,  $\mathcal{T}^d$  or  $\mathcal{S}^d$ , which implies that  $w^* \in \mathcal{W}_M$  (see Lemma 20). This in turn guarantees that  $\log q_{w^*}$  is also  $M$ -smooth (see Lemma 21). Thus,  $r$  is the sum of two  $M$ -smooth functions, so it is  $K$ -smooth with  $K \in [0, 2M]$ . Using Young's inequality, together with Theorem 1 applied to  $r$ , we obtain
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {u}} \left\| \mathbf {g} _ {\text {S T L}} (\mathbf {u}) \right\| ^ {2} \leq 2 \mathbb {E} _ {\mathbf {u}} \| \nabla_ {w} r (T _ {w} (\mathbf {u})) \| ^ {2} + 2 \mathbb {E} _ {\mathbf {u}} \| \nabla_ {w} \phi (T _ {w} (\mathbf {u})) \| ^ {2} \\ \leq 2 (d + 3) K ^ {2} \| w - \hat {w} \| ^ {2} + 2 \mathbb {E} _ {\mathbf {u}} \| \nabla_ {w} \phi \left(T _ {w} (\mathfrak {u})\right) \| ^ {2}, \tag {19} \\ \end{array}
$$
where  $\hat{w} = (\hat{m},0)$ , with  $\hat{m}$  being a stationary point of  $r$ . Now it remains to control the last term of inequality (19). Apply the chain rule onto  $\phi(T_w(u))$  to write (with the help of Lemma 37)
$$
\nabla_ {w} \phi (T _ {w} (\mathfrak {u})) = \left(\nabla_ {z} \phi (T _ {w} (\mathfrak {u})), \operatorname {p r o j} _ {V} (\nabla_ {z} \phi (T _ {w} (\mathfrak {u})) \mathfrak {u} ^ {\top})\right).
$$
This gives us directly that
$$
\begin{array}{l} \left\| \nabla_ {w} \phi \left(T _ {w} (\mathfrak {u})\right) \right\| ^ {2} = \left\| \nabla_ {z} \phi \left(T _ {w} (\mathfrak {u})\right) \right\| ^ {2} + \left\| \operatorname {p r o j} _ {V} \left(\nabla_ {z} \phi \left(T _ {w} (\mathfrak {u})\right) \mathbf {u} ^ {\top}\right) \right\| _ {F} ^ {2} \\ \leq \| \nabla_ {z} \phi (T _ {w} (\mathfrak {u})) \| ^ {2} + \| \nabla_ {z} \phi (T _ {w} (\mathfrak {u})) \mathfrak {u} ^ {\top} \| _ {F} ^ {2} \\ = \| \nabla_ {z} \phi \left(T _ {w} (\mathfrak {u})\right) \| ^ {2} (1 + \| \mathfrak {u} \| ^ {2}), \tag {20} \\ \end{array}
$$
where we used that  $\| vu^{\top}\|_{F}^{2} = \| v\|^{2}\| u\|^{2}$  for any vectors  $v$  and  $u$ . Computing  $\nabla_z\phi$  amounts to computing the gradient of the Gaussian density  $q_w$ . From its definition (see Eq. 17) we have:
$$
\nabla_ {z} \log q _ {w} (z) = - \left(C C ^ {\top}\right) ^ {- 1} (z - m) = - \Sigma^ {- 1} (z - m).
$$
Recalling the definition of  $\phi = \log q_v - \log q_{w^*}$ , and writing  $v = (m,C)$ ,  $w^{*} = (m_{*},C_{*})$ ,  $\Sigma = CC^{\top}$ ,  $\Sigma_{*} = C_{*}C_{*}^{\top}$ , we obtain
$$
\nabla_ {z} \phi (z) = \left(\Sigma_ {*} ^ {- 1} - \Sigma^ {- 1}\right) z + \Sigma^ {- 1} m - \Sigma_ {*} ^ {- 1} m _ {*}.
$$
$$
\nabla_ {z} \phi \left(T _ {w} (\mathfrak {u})\right) = \left(\Sigma_ {*} ^ {- 1} - \Sigma^ {- 1}\right) (C \mathfrak {u} + m) + \Sigma^ {- 1} m - \Sigma_ {*} ^ {- 1} m _ {*} = A \mathfrak {u} + b, \tag {21}
$$
where  $A \coloneqq (\Sigma_{*}^{-1} - \Sigma^{-1})C$  and  $b \coloneqq \Sigma_{*}^{-1}m - \Sigma_{*}^{-1}m_{*}$ . We can now combine (20) and (21) and use Lemma 29 to write
$$
\begin{array}{l} \mathbb {E} _ {\mathbf {u}} \| \nabla_ {w} \phi (T _ {w} (\mathbf {u})) \| ^ {2} \leq \mathbb {E} _ {\mathbf {u}} \| A \mathbf {u} + b \| ^ {2} (1 + \| \mathbf {u} \| ^ {2}) \\ = (d + 1) \| b \| ^ {2} + (d + 3) \| A \| _ {F} ^ {2} \\ \leq (d + 3) \left(\| b \| ^ {2} + \| A \| _ {F} ^ {2}\right). \tag {22} \\ \end{array}
$$
Let us now introduce the function  $\kappa : \mathbb{R}^d \times \mathcal{M}^d \to \mathbb{R} \cup \{+\infty\}$ , defined by  $\kappa(w) := \mathbb{E}_{z \sim q_w} \log q_w(z) - \log q_{w^*}(z)$  if  $C$  is invertible,  $+\infty$  otherwise. This function is nothing but the Kullback-Liebler divergence between the two Gaussian distributions  $q_w$  and  $q_{w^*}$ , and can be computed in closed-form on its domain:
$$
\kappa (w) = \frac {1}{2} \left(\log \det  (\Sigma_ {*}) - \log \det  (\Sigma) - d + \operatorname {t r} (\Sigma_ {*} ^ {- 1} \Sigma) + \langle \Sigma_ {*} ^ {- 1} (m _ {*} - m), (m _ {*} - m) \rangle\right),
$$
where we note as before  $w = (m, C)$ ,  $w^{*} = (m_{*}, C_{*})$ ,  $\Sigma = CC^{\top}$ ,  $\Sigma_{*} = C_{*}C_{*}^{\top}$ . This allows us to compute its gradient at  $w \in \mathcal{W}_M$ :
$$
\begin{array}{l} \nabla_ {m} \kappa (w) = \Sigma_ {*} ^ {- 1} (m _ {*} - m) \\ \nabla_ {C} \kappa (w) = \frac {- 1}{2} \nabla_ {C} \log \det  (C C ^ {\top}) + \frac {1}{2} \nabla_ {C} \operatorname {t r} (\Sigma_ {*} ^ {- 1} C C ^ {\top}) \\ = \left(\Sigma_ {*} ^ {- 1} - \Sigma^ {- 1}\right) C. \\ \end{array}
$$
We see that  $\nabla_{m}\kappa (w) = b$ , and  $\nabla_{C}\kappa (w) = A$ , which means that  $\| b\| ^2 +\| A\| _F^2 = \| \nabla_w\kappa (w)\| ^2$ . From (22) and the fact that  $\nabla_w\kappa (w^*) = 0$  (because  $w^{*}$  is a minimizer of  $\kappa$ ), we deduce that
$$
\mathbb {E} _ {\mathbf {u}} \| \nabla_ {w} \phi (T _ {w} (\mathbf {u})) \| ^ {2} \leq (d + 3) \| \nabla_ {w} \kappa (w) - \nabla_ {w} \kappa (w ^ {*}) \| ^ {2}.
$$
Now we remind that  $w^{*} \in \mathcal{W}_{M}$  which implies that  $\log q_{w^{*}}$  is  $M$ -smooth. This means that  $\mathbb{E}_{z \sim q_w} \log q_{w^*}$  is  $M$ -smooth as well, according to Domke [13, Theorem 1]. On the other hand, we know (see the proof of Lemma 16) that  $\mathbb{E}_{z \sim q_w} \log q_{w^*} = h(w)$  which is  $M$ -smooth on  $\mathcal{W}_M$  (see Lemma 18.3). All this implies that  $\kappa$  is  $2M$ -smooth on  $\mathcal{W}_M$ , from which we conclude
$$
\mathbb {E} _ {\mathfrak {u}} \| \nabla_ {w} \phi (T _ {w} (\mathfrak {u})) \| ^ {2} \leq (d + 3) \| \nabla_ {w} \kappa (w) - \nabla_ {w} \kappa (w ^ {*}) \| ^ {2} \leq (d + 3) 4 M ^ {2} \| w - w ^ {*} \| ^ {2}.
$$
The Theorem's main inequality (10) follows after plugging the above inequality into (19). The second inequality follows after using Young's inequality
$$
\begin{array}{l} \mathbb {E} \left\| \mathbf {g} _ {\mathrm {S T L}} \right\| _ {2} ^ {2} \leq 2 (d + 3) K ^ {2} \left\| w - \hat {w} \right\| _ {2} ^ {2} + 8 (d + 3) M ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} \\ = 2 (d + 3) K ^ {2} \left\| w - w ^ {*} + w ^ {*} - \hat {w} \right\| _ {2} ^ {2} + 8 (d + 3) M ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} \\ \leq 4 (d + 3) K ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} + 4 (d + 3) K ^ {2} \left\| w ^ {*} - \hat {w} \right\| _ {2} ^ {2} + 8 (d + 3) M ^ {2} \left\| w - w ^ {*} \right\| _ {2} ^ {2} \\ = 4 (d + 3) \left(K ^ {2} + 2 M ^ {2}\right) \| w - w ^ {*} \| _ {2} ^ {2} + 4 (d + 3) K ^ {2} \| w ^ {*} - \hat {w} \| _ {2} ^ {2}. \\ \end{array}
$$
To conclude the proof, it remains to check that  $K = 0$  whenever  $p(\cdot |x)$  is Gaussian. In that case, because our main objective function  $f(w)$  is the divergence between  $q_{w}$  and  $p(\cdot |x)$ , it is clear that  $f$  is minimized whenever  $q_{w} = p(\cdot |x)$ . Therefore, without loss of generality, we can assume that  $p(\cdot |x) = q_{w^*}$ . This implies that the residual  $r$  is constant  $(r(z) = \log p(x))$ , and so is 0-smooth.
8.2 Quadratically bounded estimators
Theorem 30. The estimators obey  $\mathbb{E}\left\| \mathbf{g}\right\| _2^2\leq a\left\| w - w^*\right\| ^2 +b$  with the following constants:
Estimator	a	b	valid w
genergy	2(d+3)M2	2(d+3)M2||w* - \bar{w}||2	any
gent	4(d+3)M2	4(d+3)M2||w* - \bar{w}||2 + dL	w ∈ W_L
gSTL	4(d+3)(K2+2M2)	4(d+3)K2||w* - \hat{w}||2	w ∈ W_M
Proof. Combine the results of Theorems 2, 3 and 4.
9 Optimization proofs
This section contains all the proofs for the Theorems stated in Section 3.
9.1 Anytime Convergence Theorem for Strongly Convex
Theorem 31 ( Theorem 3.2 in [53]). Suppose we are given a sequence of  $w^t$  iterates such that for a step size  $\gamma_t \leq \frac{1}{2\mathcal{L}}$  and given constants  $\mu > 0$  and  $\sigma^2 > 0$  we have that
$$
\mathbb {E} \left[ \| w ^ {t + 1} - w ^ {*} \| ^ {2} \right] \leq (1 - \gamma_ {t} \mu) \mathbb {E} \left[ \| w ^ {t} - w ^ {*} \| ^ {2} \right] + 2 \sigma^ {2} \gamma_ {t} ^ {2}. \tag {23}
$$
By switching to a decaying stepsize according to
$$
\gamma_ {t} = \left\{ \begin{array}{l l} \frac {1}{2 \mathcal {L}} & \text {f o r t <   t ^ {*}} \\ \frac {1}{\mu} \frac {2 t + 1}{(t + 1) ^ {2}} & \text {f o r t \geq t ^ {*}} \end{array} \right.
$$
where  $t^* = 4\lfloor \mathcal{L} / \mu \rfloor$  we have that
$$
\mathbb {E} \left[ \| w ^ {T + 1} - w ^ {*} \| ^ {2} \right] \leq \frac {1 6 \lfloor \mathcal {L} / \mu \rfloor^ {2}}{(T + 1) ^ {2}} \| w ^ {0} - w ^ {*} \| ^ {2} + \frac {\sigma^ {2}}{\mu^ {2}} \frac {8}{T + 1}. \tag {24}
$$
Proof. We divide the proof for steps  $t$  that are great or smaller than  $t^*$ . For  $t \leq t^*$ , we have by unrolling (23) starting from  $t^*, \ldots, 1$ , with  $\gamma_t \equiv \gamma$  gives
$$
\mathbb {E} \left[ \| w ^ {t ^ {*}} - w ^ {*} \| ^ {2} \right] \leq (1 - \gamma \mu) ^ {t ^ {*}} \| w ^ {0} - w ^ {*} \| ^ {2} + \frac {2 \sigma^ {2} \gamma}{\mu}. \tag {25}
$$
Now consider  $t \geq t^*$  for which we have that  $\gamma_t = \frac{1}{\mu} \frac{2t + 1}{(t + 1)^2}$  which when inserted into (23) gives
$$
\mathbb {E} \left[ \| w ^ {t + 1} - w ^ {*} \| ^ {2} \right] \leq \frac {t ^ {2}}{(t + 1) ^ {2}} \mathbb {E} \left[ \| w ^ {t} - w ^ {*} \| ^ {2} \right] + \sigma^ {2} \frac {2}{\mu^ {2}} \frac {(2 t + 1) ^ {2}}{(t + 1) ^ {4}}.
$$
Multiplying through by  $(t + 1)^2$  and re-arranging gives
$$
\begin{array}{l} (t + 1) ^ {2} \mathbb {E} \left[ \| w ^ {t + 1} - w ^ {*} \| ^ {2} \right] \leq t ^ {2} \mathbb {E} \left[ \| w ^ {t} - w ^ {*} \| ^ {2} \right] + \sigma^ {2} \frac {2}{\mu^ {2}} \frac {(2 t + 1) ^ {2}}{(t + 1) ^ {2}} \\ \leq t ^ {2} \mathbb {E} \left[ \| w ^ {t} - w ^ {*} \| ^ {2} \right] + \frac {8 \sigma^ {2}}{\mu^ {2}}, \\ \end{array}
$$
where we used that  $\frac{2t + 1}{t + 1} \leq 2$ . Summing up over  $t = t^*, \ldots, T$  and using telescopic cancellation gives
$$
(T + 1) ^ {2} \mathbb {E} \left[ \| w ^ {T + 1} - w ^ {*} \| ^ {2} \right] \leq (t ^ {*}) ^ {2} \mathbb {E} \left[ \| w ^ {t ^ {*}} - w ^ {*} \| ^ {2} \right] + (T - t ^ {*}) \frac {8 \sigma^ {2}}{\mu^ {2}}.
$$
Now using that for  $t \leq t^*$  we have that  $\gamma = \frac{1}{\mathcal{L}}$  and (25) holds, thus
$$
\begin{array}{l} (T + 1) ^ {2} \mathbb {E} \left[ \| w ^ {T + 1} - w ^ {*} \| ^ {2} \right] \leq (t ^ {*}) ^ {2} \left((1 - \gamma \mu) ^ {t ^ {*}} \| w ^ {0} - w ^ {*} \| ^ {2} + 2 \sigma^ {2} \frac {\gamma}{\mu}\right) + (T - t ^ {*}) \frac {8 \sigma^ {2}}{\mu^ {2}} \\ \leq \left(t ^ {*}\right) ^ {2} \left(\| w ^ {0} - w ^ {*} \| ^ {2} + 2 \sigma^ {2} \frac {\gamma}{\mu}\right) + (T - t ^ {*}) \frac {8 \sigma^ {2}}{\mu^ {2}}, \\ \end{array}
$$
where we used that  $(1 - \gamma \mu) \leq 1$ . Substituting  $\gamma = \frac{1}{2\mathcal{L}}$ ,  $t^* = 4\lfloor \mathcal{L} / \mu \rfloor$  and multiplying by  $(T + 1)^2$  gives
$$
\begin{array}{l} \mathbb {E} \left[ \| w ^ {T + 1} - w ^ {*} \| ^ {2} \right] \leq \frac {1 6 \lfloor \mathcal {L} / \mu \rfloor^ {2}}{(T + 1) ^ {2}} \| w ^ {0} - w ^ {*} \| ^ {2} + \frac {\sigma^ {2}}{\mu^ {2} (T + 1) ^ {2}} \left(8 (T - t ^ {*}) + (t ^ {*}) ^ {2} \frac {\mu}{\mathcal {L}}\right) \\ \leq \frac {1 6 \left\lfloor \mathcal {L} / \mu \right\rfloor^ {2}}{(T + 1) ^ {2}} \| w ^ {0} - w ^ {*} \| ^ {2} + \frac {\sigma^ {2}}{\mu^ {2}} \frac {8 (T - 2 t ^ {*})}{(T + 1) ^ {2}} \\ \leq \frac {1 6 \lfloor \mathcal {L} / \mu \rfloor^ {2}}{(T + 1) ^ {2}} \| w ^ {0} - w ^ {*} \| ^ {2} + \frac {\sigma^ {2}}{\mu^ {2}} \frac {8}{T + 1} \\ \end{array}
$$
which concludes the proof of (24).
9.2 Complexity Lemma for Convex
Lemma 32. Suppose we are given a sequence  $\bar{w}^T$  and constants  $a, b, B > 0$  such that
$$
\mathbb {E} \left[ f (\bar {w} ^ {T}) \right] - \inf  f \leq \gamma \left(\frac {a \| w ^ {0} - w ^ {*} \| ^ {2}}{1 - \theta^ {T}} + b\right)
$$
holds for
$$
\theta := \frac {1}{1 + B \gamma^ {2}}.
$$
If  $\gamma = \frac{A}{\sqrt{T}}$  for  $A \geq \sqrt{\frac{2}{B}}$  and  $T \geq 2$ , then
$$
\mathbb {E} \left[ f (\bar {w} ^ {T}) \right] - \inf  f \leq \frac {A}{\sqrt {T}} \left(2 a \| w ^ {0} - w ^ {*} \| ^ {2} + b\right).
$$
Thus,  $T = \mathcal{O}(\epsilon^{-2})$  iterations are sufficient to guarantee that  $\mathbb{E}\left[f(\bar{w}^T)\right] - \inf f\leq \epsilon$ .
Proof. First we show that  $\frac{1}{1 - \theta^T} \leq 2$  is equivalent to
$$
T \geq \frac {\log 2}{\log \left(1 + B A ^ {2} / T\right)}. \tag {26}
$$
Indeed this follows since
$$
\frac {1}{1 - \theta^ {T}} \leq 2 \quad \Leftrightarrow
$$
$$
\theta^ {T} \leq \frac {1}{2} \quad \Leftrightarrow
$$
$$
T \geq \frac {\log 2}{\log 1 / \theta} \quad \Leftrightarrow
$$
$$
T \geq \frac {\log 2}{\log (1 + B \gamma^ {2})} \Leftrightarrow
$$