SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g} = \boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} _ {0} \boldsymbol {g} + \left[ (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g} \right] ^ {2} - \left(\boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g}\right) ^ {2}. \tag {30}
$$
As shorthand, let  $\rho = \pmb{g}^{\top}\pmb{\Sigma}\pmb{g}$ . Then from eq. (30) we see that  $\rho$  satisfies the quadratic equation
$$
\rho (1 + \rho) = \boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} _ {0} \boldsymbol {g} + \left[ (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g} \right] ^ {2}.
$$
Note that there are no unknowns on the right side of this equation. The correct solution is given by the positive root since  $\rho = \pmb{g}^{\top}\pmb{\Sigma}\pmb{g} > 0$ . Also note that  $\rho = (\pmb{\mu} - \pmb{\theta})^{\top}\pmb{g}$  from eq. (23).
It is useful to define one final intermediate quantity before solving for  $\pmb{\mu}$ . Let
$$
\varepsilon_ {0} = \Sigma_ {0} g - \mu_ {0} + \theta .
$$
Note that  $\varepsilon_0$  simply measures the degree to which the parameters of  $q_0(\pmb{\theta})$  violate the desired constraint  $\nabla_{\pmb{w}}\log q(\pmb{\theta}) = \nabla_{\pmb{w}}\log p(\pmb{\theta},\pmb{y})$ . Put another way, if  $\varepsilon_0 = 0$ , then we have the trivial solution  $\pmb{\mu} = \pmb{\mu}_0$  and  $\pmb{\Sigma} = \pmb{\Sigma}_0$ .
Now we have everything to express the solution for  $\pmb{\mu}$  in a highly intuitive form; in particular, it will be immediately evident that  $\pmb{\mu} \rightarrow \pmb{\mu}_0$  as  $\varepsilon_0 \rightarrow 0$ . Starting from eqs. (23) and (28), we find
$$
\begin{array}{l} \boldsymbol {\mu} - \boldsymbol {\mu} _ {0} = \boldsymbol {\theta} - \boldsymbol {\mu} _ {0} + \boldsymbol {\Sigma} \boldsymbol {g}, (31) \\ = \boldsymbol {\theta} - \boldsymbol {\mu} _ {0} + \left[ \boldsymbol {\Sigma} _ {0} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} - (\boldsymbol {\mu} - \boldsymbol {\theta}) (\boldsymbol {\mu} - \boldsymbol {\theta}) ^ {\top} \right] \mathbf {g}, (32) \\ = \varepsilon_ {0} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g} - (\boldsymbol {\mu} - \boldsymbol {\theta}) (\boldsymbol {\mu} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g}, (33) \\ = \varepsilon_ {0} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g} - (\boldsymbol {\mu} - \boldsymbol {\mu} _ {0} + \boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g}, (34) \\ = \varepsilon_ {0} - \rho (\boldsymbol {\mu} - \boldsymbol {\mu} _ {0}) + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) [ (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) - (\boldsymbol {\mu} - \boldsymbol {\theta}) ^ {\top} ] \boldsymbol {g}, (35) \\ = \varepsilon_ {0} - \rho (\boldsymbol {\mu} - \boldsymbol {\mu} _ {0}) + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\mu}) ^ {\top} \boldsymbol {g}, (36) \\ = \varepsilon_ {0} - (\rho \mathbf {I} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) \boldsymbol {g} ^ {\top}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\mu}). (37) \\ \end{array}
$$
Note what has happened here: eq. (32) is a system of nonlinear equations for  $\mu$ , but in eq. (37), all the nonlinearity has been expressed in terms of  $\rho$ . Since  $\rho$  can be determined via eq. (30), we arrive effectively at a system of linear equations for  $\mu$ . Collecting terms, we obtain
$$
\left[ (1 + \rho) \mathbf {I} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) \boldsymbol {g} ^ {\top} \right] (\boldsymbol {\mu} - \boldsymbol {\mu} _ {0}) = \varepsilon_ {0}. \tag {38}
$$
We thus arrive at the closed-form update
$$
\boldsymbol {\mu} = \boldsymbol {\mu} _ {0} + \left[ (1 + \rho) \mathbf {I} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) \boldsymbol {g} ^ {\top} \right] ^ {- 1} \varepsilon_ {0} \tag {39}
$$
It is evident from this update that  $\pmb{\mu} \rightarrow \pmb{\mu}_0$  as  $\varepsilon_0 \rightarrow 0$ . The matrix inverse in this update can also be computed efficiently from the Woodbury matrix identity.
In sum, the joint update for  $\mu$  and  $\Sigma$  can be efficiently computed as follows:
1. Set  $\pmb{g} = \nabla_{\pmb{w}}\log p(\pmb {\theta},\pmb {y})$  and  $\varepsilon_0 = \Sigma_0g - \mu_0 + \pmb{\theta}$
2. Solve  $\rho (1 + \rho) = \pmb{g}^{\top}\pmb{\Sigma}_{0}\pmb{g} + \left[(\pmb{\mu}_{0} - \pmb{\theta})^{\top}\pmb{g}\right]^{2}$  for  $\rho > 0$
3. Compute  $\pmb{\mu} = \pmb{\mu}_0 + [(1 + \rho)\mathbf{I} + (\pmb{\mu}_0 - \pmb{\theta})\pmb{g}^\top ]^{-1}\varepsilon_0$
4. Compute  $\pmb{\Sigma} = \pmb{\Sigma}_0 + (\pmb{\mu}_0 - \pmb{\theta})(\pmb{\mu}_0 - \pmb{\theta})^\top - (\pmb{\mu} - \pmb{\theta})(\pmb{\mu} - \pmb{\theta})^\top$
Note that it straightforward to solve for the value of  $\rho$  in step 2 of this procedure. In particular, from the quadratic formula, we obtain the positive root
$$
\rho = \frac {\sqrt {1 + 4 \left(\boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} _ {0} \boldsymbol {g} + \left[ (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g} \right] ^ {2}\right)} - 1}{2}.
$$
We can solve the linear equation for  $\mu$  in step 3 of this procedure most efficiently by using the Sherman-Morrison formula
$$
\left[ a \mathbf {I} + u g ^ {\top} \right] ^ {- 1} = \frac {1}{a} \left(\mathbf {I} - \frac {u g ^ {\top}}{a + u ^ {\top} g}\right), \quad \text {f o r e v e r y} u, g, a.
$$
Applying this formula in step 3, we find that
$$
\boldsymbol {\mu} = \boldsymbol {\mu} _ {0} + \frac {1}{1 + \rho} \left[ \mathbf {I} - \frac {(\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) \boldsymbol {g} ^ {\top}}{1 + \rho + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g}} \right] \varepsilon_ {0}. \tag {40}
$$
Next we relate this update to the difference in the scores of the target distribution and the current estimate of the variational approximation. Since  $\nabla_{\pmb{\theta}}\log q_0(\pmb {\theta}) = -\pmb{\Sigma}_0^{-1}(\pmb {\theta} - \pmb {\mu}_0)$ , we have that
$$
\varepsilon_ {\mathbf {0}} = \Sigma_ {0} \left(\boldsymbol {g} - \Sigma_ {0} ^ {- 1} (\boldsymbol {\mu} _ {0} + \boldsymbol {\theta})\right) = \Sigma_ {0} \left(\boldsymbol {g} - \nabla_ {\boldsymbol {\theta}} \log q _ {0} (\boldsymbol {\theta})\right).
$$
It follows, by substituting the above expression into Eq. 40, that
$$
\boldsymbol {\mu} = \boldsymbol {\mu} _ {0} + \frac {1}{1 + \rho} \left[ \mathbf {I} - \frac {(\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) \boldsymbol {g} ^ {\top}}{1 + \rho + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top} \boldsymbol {g}} \right] \boldsymbol {\Sigma} _ {0} (\nabla_ {\boldsymbol {\theta}} \log p (\boldsymbol {x}, \boldsymbol {\theta}) - \nabla_ {\boldsymbol {\theta}} \log q _ {0} (\boldsymbol {\theta}))
$$
which gives the update in the statement of the theorem.
Finally we prove that if  $\Sigma_0$  is positive definite, then so is  $\Sigma$ . To do so, we show explicitly that all of the eigenvalues of  $\Sigma$  are positive. We begin by rewriting our results for  $\Sigma$  in eq. (28) and  $\rho$  in eq. (30) in a more convenient form. As shorthand, let
$$
\mathbf {M} _ {0} = \boldsymbol {\Sigma} _ {0} + (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) (\boldsymbol {\mu} _ {0} - \boldsymbol {\theta}) ^ {\top}, \tag {41}
$$
so that  $\mathbf{M}_0$  captures the first two terms on the right side of eq. (28). Note that  $\mathbf{M}_0$  is positive-definite, a fact that we will exploit repeatedly in what follows. In addition, recall that  $\mu -\theta = \Sigma g$  from eq. (23). Thus with this notation we can rewrite eqs. (28) and (30) as
$$
\boldsymbol {\Sigma} = \mathbf {M} _ {0} - (\boldsymbol {\Sigma} \boldsymbol {g}) (\boldsymbol {\Sigma} \boldsymbol {g}) ^ {\top}, \tag {42}
$$
$$
\rho (1 + \rho) = \boldsymbol {g} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}. \tag {43}
$$
Now let  $e$  be any normalized eigenvector of  $\Sigma$ ; we want to show that its corresponding eigenvalue  $\lambda_e$  is positive. From eq. (42), it follows that
$$
\begin{array}{l} \lambda_ {e} = e ^ {\top} \boldsymbol {\Sigma} e (44) \\ = \boldsymbol {e} ^ {\top} [ \mathbf {M} _ {0} - (\Sigma \boldsymbol {g}) (\Sigma \boldsymbol {g}) ^ {\top} ] \boldsymbol {e} (45) \\ = e ^ {\top} \mathbf {M} _ {0} e - \lambda_ {e} ^ {2} \left(e ^ {\top} g\right) ^ {2}. (46) \\ \end{array}
$$
Note that if  $e^{\top}g = 0$ , then it follows trivially that  $\lambda_{e} = e^{\top}\mathbf{M}_{0}e > 0$ . So we only need to consider the non-trivial case  $e^{\top}g \neq 0$ . To proceed, we note that
$$
\left(\boldsymbol {e} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}\right) ^ {2} = \left(\boldsymbol {e} ^ {\top} \mathbf {M} _ {0} ^ {\frac {1}{2}} \mathbf {M} _ {0} ^ {\frac {1}{2}} \boldsymbol {g}\right) ^ {2} \leq \left(\boldsymbol {e} ^ {\top} \mathbf {M} _ {0} \boldsymbol {e}\right) \left(\boldsymbol {g} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}\right), \tag {47}
$$
where we have used the Cauchy-Schwartz inequality to bound  $(e^{\top}\mathbf{M}_{o}\pmb {g})$  in terms of  $(e^{\top}\mathbf{M}_0e)$ , the latter of which appears in eq. (46). Substituting this inequality into eq. (46), we find that
$$
\lambda_ {e} \geq \frac {\left(\boldsymbol {e} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}\right) ^ {2}}{\boldsymbol {g} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}} - \lambda_ {e} ^ {2} \left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) ^ {2}. \tag {48}
$$
To prove that  $\lambda_{e} > 0$  we need one more intermediate result. Focusing on the rightmost term in this equality, we note that
$$
\lambda_ {e} \left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) = \boldsymbol {e} ^ {\top} \Sigma \boldsymbol {g} = \boldsymbol {e} ^ {\top} \left[ \mathbf {M} _ {0} - (\Sigma \boldsymbol {g}) (\Sigma \boldsymbol {g}) ^ {\top} \right] \boldsymbol {g} = \boldsymbol {e} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g} - \lambda_ {e} \left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) \left(\boldsymbol {g} ^ {\top} \Sigma \boldsymbol {g}\right), \tag {49}
$$
and rearranging the terms in this equation, we find
$$
\boldsymbol {e} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g} = \lambda_ {e} (\boldsymbol {e} ^ {\top} \boldsymbol {g}) (1 + \boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g}). \tag {50}
$$
This intermediate result is useful because it relates the two terms on the right side of eq. (48). In particular, using eq. (50) to eliminate the term  $e^{\top}\mathbf{M}_0\pmb{g}$  in eq. (48), we find:
$$
\begin{array}{l} \lambda_ {e} \geq \frac {\left[ \lambda_ {e} (\boldsymbol {e} ^ {\top} \boldsymbol {g}) (1 + \boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g}) \right] ^ {2}}{\boldsymbol {g} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}} - \lambda_ {e} ^ {2} (\boldsymbol {e} ^ {\top} \boldsymbol {g}) ^ {2} \\ = \lambda_ {e} ^ {2} \left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) ^ {2} \left[ \frac {\left(1 + \boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g}\right) ^ {2}}{\boldsymbol {g} ^ {\top} \mathbf {M} _ {0} \boldsymbol {g}} - 1 \right] \\ = \lambda_ {e} ^ {2} (\boldsymbol {e} ^ {\top} \boldsymbol {g}) ^ {2} \left[ \frac {(1 + \rho) ^ {2}}{\rho (1 + \rho)} - 1 \right] \\ = \frac {\lambda_ {e} ^ {2} (\boldsymbol {e} ^ {\top} \boldsymbol {g}) ^ {2}}{\rho}, \\ > 0, \\ \end{array}
$$
where the final inequality follows because the individual terms  $\lambda_e^2$ ,  $(e^\top g)^2$ , and  $\rho$  are all strictly positive; note that  $\lambda_e$  cannot be equal to zero because this contradicts the equality in eq. (46). This completes the proof. Perhaps it is useful that this derivation also gives upper bounds on  $\lambda_e$ , namely
$$
\frac {1}{\lambda_ {e}} \geq \frac {\left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) ^ {2}}{\rho} \Rightarrow \lambda_ {e} \leq \frac {\rho}{\left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) ^ {2}} = \frac {\boldsymbol {g} ^ {\top} \boldsymbol {\Sigma} \boldsymbol {g}}{\left(\boldsymbol {e} ^ {\top} \boldsymbol {g}\right) ^ {2}}. \tag {51}
$$