SlowGuess/NeurIPS_2023_OCR huggingface dataset:
This algorithm produces correct filtered messages but due to its associativity can be computed in parallel. To complete the filtering algorithm, we outline how to evaluate  $(f_i,f_j)\otimes (g_i,g_j)$ . The integrals in Eqs. (119), (120) can be evaluated as simple functions of the parameters of  $f_{i},f_{j},g_{i},g_{j}$ :
Proposition D.3. Let
$$
f _ {i} (a, b) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \Phi_ {1 1, i} \\ - \phi_ {1, i} \\ \Phi_ {1 2, i} \\ \phi_ {2, i} \\ - \frac {1}{2} \Phi_ {2 2, i} \end{array} \right], \left[ \begin{array}{l} a a ^ {T} \\ a \\ a b ^ {T} \\ b \\ b b ^ {T} \end{array} \right] \right\rangle \right\}, \quad f _ {j} (b, c) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \Phi_ {1 1, j} \\ - \phi_ {1, j} \\ \Phi_ {1 2, j} \\ \phi_ {2, j} \\ - \frac {1}{2} \Phi_ {2 2, j} \end{array} \right], \left[ \begin{array}{l} b b ^ {T} \\ b \\ b c ^ {T} \\ c \\ c c ^ {T} \end{array} \right] \right\rangle \right\} \tag {140}
$$
$$
g _ {i} (a) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \Gamma_ {i} \\ \gamma_ {i} \end{array} \right], \left[ \begin{array}{l} a a ^ {T} \\ a \end{array} \right] \right\rangle \right\}, \quad g _ {j} (b) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \Gamma_ {j} \\ \gamma_ {j} \end{array} \right], \left[ \begin{array}{l} b b ^ {T} \\ b \end{array} \right] \right\rangle \right\} \tag {141}
$$
Then given the following definitions
$$
C _ {i j} = \Gamma_ {j} + \Phi_ {2 2, i} \tag {142}
$$
$$
\Gamma_ {i j} = \left(\Phi_ {1 1, i} - \Phi_ {1 2, i} C _ {i j} ^ {- 1} \Phi_ {1 2, i} ^ {T}\right) \tag {143}
$$
$$
\gamma_ {i j} = - \phi_ {1, i} + \Phi_ {1 2, i} C _ {i j} ^ {- 1} \left(\gamma_ {j} + \phi_ {2, i}\right) \tag {144}
$$
$$
P _ {i j} = \Gamma_ {j} + \Phi_ {2 2, i} + \Phi_ {1 1, i} \tag {145}
$$
We can compute  $(f_{ij},g_{ij}) = (f_i,g_i)\otimes (f_j,g_j)$  as:
$$
f _ {i j} (a, c) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \left(\Phi_ {1 1, i} - \Phi_ {1 2, i} P _ {i j} ^ {- 1} \Phi_ {1 2, i} ^ {T} - \Gamma_ {i j}\right) \\ - \left(\phi_ {1, i} - \Phi_ {1 2, i} P _ {i j} ^ {- 1} \left(\phi_ {2, i} - \phi_ {1, j} + \gamma_ {j}\right) + \gamma_ {i j}\right) \\ \Phi_ {1 2, i} P _ {i j} ^ {- 1} \Phi_ {1 2, j} \\ \phi_ {2, j} + \Phi_ {1 2, j} ^ {T} P _ {i j} ^ {- 1} \left(\phi_ {2, i} - \phi_ {1, j} + \gamma_ {j}\right) \\ - \frac {1}{2} \left(\Phi_ {2 2, j} - \Phi_ {1 2, j} ^ {T} P _ {i j} ^ {- 1} \Phi_ {1 2, j}\right) \end{array} \right], \left[ \begin{array}{c} a a ^ {T} \\ a \\ a c ^ {T} \\ c \\ c c ^ {T} \end{array} \right] \right\rangle \right\} \tag {146}
$$
$$
g _ {i j} (c) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \left(\Gamma_ {i} + \Gamma_ {i j}\right) \\ \gamma_ {i} + \gamma_ {i j} \end{array} \right], \left[ \begin{array}{c} c c ^ {T} \\ c \end{array} \right] \right\rangle \right\} \tag {147}
$$
The proof of this involves integrating the kernels of Gaussians, which is done many times in this supplement; it is left as an exercise for the reader.
The last task is to define the parallel smoother. The construction is very similar to the filter. Let
$$
\begin{array}{l} e _ {t} \left(z _ {t}, z _ {t + 1}\right) = \frac {m _ {z _ {t} , \omega_ {t + 1}} \left(z _ {t}\right)}{m _ {\omega_ {t + 1} , z _ {t + 1}} \left(z _ {t + 1}\right)} \cdot \exp \left\{\left\langle \omega_ {t + 1}, t \left(z _ {t}, z _ {t + 1}\right) \right\rangle \right\} (148) \\ = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \left(J _ {1 1, t} + F _ {t | t}\right) \\ - \left(h _ {1, t} - f _ {t | t}\right) \\ J _ {1 2, t} \\ h _ {2, t} + r _ {t + 1} - h _ {t + 1 | t + 1} \\ - \frac {1}{2} \left(J _ {2 2, t} + R _ {t + 1} - J _ {t + 1 | t + 1}\right) \end{array} \right], \left[ \begin{array}{c} z _ {t} z _ {t} ^ {T} \\ z _ {t} \\ z _ {t} z _ {t + 1} ^ {T} \\ z _ {t + 1} \\ z _ {t + 1} z _ {t + 1} ^ {T} \end{array} \right] \right\rangle \right\} (149) \\ = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} E _ {1 1, t} \\ - \epsilon_ {1, t} \\ E _ {1 2, t} \\ \epsilon_ {2, t} \\ - \frac {1}{2} E _ {2 2, t} \end{array} \right], \left[ \begin{array}{c} z _ {t} z _ {t} ^ {T} \\ z _ {t} \\ z _ {t} z _ {t + 1} ^ {T} \\ z _ {t + 1} \\ z _ {t + 1} z _ {t + 1} ^ {T} \end{array} \right] \right\rangle \right\} (150) \\ \end{array}
$$
Where we let  $e_{T}(z_{T}, z_{T + 1}) = e_{T}(z_{T})$  with  $E_{11,T} = J_{T|T}, \epsilon_{1,T} = -h_{T|T}$  and remaining parameters equal to 0 (as there is no  $z_{T + 1}$ ). Then the associative operator  $\oplus$  is given by
$$
e _ {i j} (a, c) = e _ {i} (a, b) \oplus e _ {j} (b, c) = \ker \left(\int_ {b} e _ {i} (a, b) e _ {j} (b, c) d b\right) \tag {151}
$$
Which can be computed (derivations omitted) as
$$
D _ {i j} = E _ {2 2, i} + E _ {1 1, j} \tag {152}
$$
$$
e _ {i j} (a, c) = \exp \left\{\left\langle \left[ \begin{array}{c} - \frac {1}{2} \left(E _ {1 1, i} - E _ {1 2, i} D _ {i j} ^ {- 1} E _ {1 2, i} ^ {T}\right) \\ - \left(\epsilon_ {1, i} - E _ {1 2, i} D _ {i j} ^ {- 1} \left(\epsilon_ {2, i} - \epsilon_ {1, j}\right)\right) \\ E _ {1 2, i} D _ {i j} ^ {- 1} E _ {1 2, j} \\ \epsilon_ {2, j} + E _ {1 2, j} ^ {T} D _ {i j} ^ {- 1} \left(\epsilon_ {2, i} - \epsilon_ {1, j}\right) \\ - \frac {1}{2} \left(E _ {2 2, j} - E _ {1 2, j} ^ {T} D _ {i j} ^ {- 1} E _ {1 2, j}\right) \end{array} \right], \left[ \begin{array}{c} a a ^ {T} \\ a \\ a c ^ {T} \\ c \\ c c ^ {T} \end{array} \right] \right\rangle \right\} \tag {153}
$$
Särkkä & García-Fernández [56] prove this operation is associative when operating on full distributions (i.e. without taking the kernel). The proof that we retain associativity when removing multiplicative constants mirrors the proof of Prop. D.1 and is omitted for brevity. Finally, we show that when these operations are done sequentially,
$$
e _ {t} \oplus \left(e _ {t + 1} \oplus \left(e _ {t + 1} \oplus \dots \left(e _ {T - 1} \oplus e _ {T}\right)\right)\right)) = q \left(z _ {t}\right) \tag {154}
$$
The base  $e_{T}(z_{t}) = q(z_{T})$  is true by construction. Then, plugging the inductive hypothesis and Eq. (148) into the definition of  $\oplus$  yields Eq. (103), proving that  $e_{t}(z_{t}) = q(z_{t})$  implies  $e_{t - 1}(z_{t - 1}) = q(z_{t - 1})$ .
D.3 HMM Inference
The discrete Markov model is time homogenous and requires only expected probabilities for the initial state  $k_{2}$  and expected probabilities for the transition matrix  $\pi_{ij}$ . The "observations" of the HMM are (in the case of the SLDS) messages passed from the LDS chain, which we denote  $\omega_{t}(k)$ . The setup is fundamentally the same as the LDS:
Input:  $\mathbb{E}_{q(\theta)}[t(\theta)]$ $\eta_{\mathrm{recog}}$
Forward Pass:
$$
\alpha_ {1} (k) = \mathbb {E} \left[ \log \pi_ {k} \right] + \omega_ {1} (k) \tag {155}
$$
$$
\text {f o r} t = 2, \dots , T \alpha_ {t} (k) = \omega_ {t} (k) + \log \operatorname {s u m e x p} _ {i = 1} ^ {K} \left[ \mathbb {E} \left[ \log \pi_ {i k} \right] + \alpha_ {t - 1} (i) \right] \tag {156}
$$
$$
\log Z = \log \operatorname {s u m e x p} _ {k = 1} ^ {K} \alpha_ {T} (k) \tag {157}
$$
Backward Pass:
$$
\beta_ {T} (k) = 0 \tag {158}
$$
$$
\text {f o r} t = T - 1, \dots , 1 \quad \beta_ {t} (k) = \operatorname {l o g s u m e x p} _ {j = 1} ^ {K} \left[ \omega_ {t + 1} (j) + \mathbb {E} [ \log \pi_ {k j} ] + \beta_ {t + 1} (j) \right] \tag {159}
$$
Return: log marginals  $\alpha_{t}(k) + \log \beta_{t}(k)$ ; log  $Z$
E Evaluating the SVAE loss
Johnson et al. [29] decompose the loss into three terms
$$
- \mathcal {L} [ q (\theta) q (z), \gamma ] = \overbrace {\operatorname {K L} (q (\theta) \| p (\theta)))} ^ {\text {P r i o r K L}} + \overbrace {\mathbb {E} _ {q (\theta)} \operatorname {K L} (q (z) \| p (z | \theta))} ^ {\text {L o c a l K L}} - \overbrace {\mathbb {E} _ {q (z)} [ p _ {\gamma} (x | z) ]} ^ {\text {R e c o n s t r u c t i o n}} \tag {160}
$$
Because  $p(\theta), p(z|\theta), q(\theta; \eta), q_{\phi}(z|x; \eta)$  are all exponential family distributions, the expectations here can be evaluated in closed form, with the exception of the reconstruction loss  $\mathbb{E}_{q(z)}[p_{\gamma}(x|z)]$ . The reconstruction loss can be estimated without bias using Monte Carlo samples and the reparameterization trick [32]. For this to work, we must be able to reparameterize the sample of  $q(z)$  such that we can back propagate through the sampling routine, which is not possible for discrete random variables. As a result, the latent variables which need to be sampled to reconstruct  $x$  (i.e. all variables which are in the same mean field component as a variable being fed into the encoder) must be continuous.
The first term of the loss, the prior KL, is a regularizer on the graphical model parameters and can be easily evaluated via Eq. (25). The rest of this section is devoted to computing the KL divergence between the prior and variational factor of the local latent variables  $z$ . We will start with the fully structured case where no mean field separation occurs.
Recalling that the optimal  $q(z)$  has parameters  $\omega$  given by Eq. (47), we can again invoke Eq. (25) to get
$$
\begin{array}{l} \mathbb {E} _ {q (\theta)} \mathrm {K L} (q (z) \| p (z | \theta)) = \mathbb {E} _ {q (\theta)} [ \langle \omega - t (\theta), E _ {q (z)} t (z) \rangle - \log Z (\omega) + \log Z (t (\theta)) ] (161) \\ = \left\langle \omega - \mathbb {E} _ {q (\theta)} [ t (\theta) ], E _ {q (z)} t (z) \right\rangle - \log Z (\omega) + \log \mathbb {E} _ {q (\theta)} [ Z (t (\theta)) ] (162) \\ = \left\langle \lambda_ {\phi} (x), E _ {q (z)} t (z) \right\rangle - \log Z (\omega) + \log \mathbb {E} _ {q (\theta)} [ Z (t (\theta)) ] (163) \\ \end{array}
$$
Namely, because  $\log q(z)$  is defined as the expected log prior  $p(z|\theta)$  plus the recognition potentials, these two distributions only differ in log space by those recognition potentials. The log partition function of the graphical model can be obtained by belief propagation, and the expected log partition function from the prior can be generally known in closed form as the sum of individual partition functions at each vertex or factor in the corresponding graphical model.
Example: LDS The LDS SVAE only outputs recognition potentials on a subset of the statics of  $q(z)$ . In particular, the encoder takes the observations  $x$  and outputs normal likelihood functions on each  $z$  independently. Thus the cross-temporal covariance statistics  $z_{t}z_{t + 1}^{T}$  have no corresponding recognition potentials. Further, we restrict the potentials on  $z_{t}z_{t}^{T}$  to be diagonal, meaning our observational uncertainty factorizes across dimensions. Thus the only statistics in  $t(z)$  for which  $\lambda_{\phi}(x)$  are non-zero are  $z_{t}$  and  $z_{ti}^{2}$  for each component  $i$  of  $t(z)$ . As a result, the inner product  $\langle \lambda_{\phi}(x), \mathbb{E}_{q(z)}t(z)\rangle$  sums over  $2TD$  terms for a sequence with  $T$  timesteps and  $D$ -dimensional latent variable at each step.
Local KL with structured mean field When there is mean field separation, the prior and variational factor differ in more than just the recognition potentials, as  $q(z)$  replaces some factors of  $p(z|\theta)$  with separable factors. We can rewrite
$$
\operatorname {K L} (q (z) \| p (z | \theta)) = \sum_ {m = 1} ^ {M} \operatorname {K L} (q \left(z _ {m}\right) \| p \left(z _ {m} \mid z _ {<   m}, \theta\right)) \tag {164}
$$
The natural parameters of  $q(z_{m})$  differ from the natural parameters of  $p(z_{m}|z_{<m},\theta)$  not only in the recognition potentials, but also for any factors which include a variable in "later" clusters  $z_{>m}$ .
Example: SLDS The SLDS SVAE has 2 mean field clusters,  $q(k)q(z)$  for the discrete variables  $k$  and continuous variables  $z$ . If we write our prior distribution  $p(\theta)p(k|\theta)p(z|k,\theta)$  then  $\mathrm{KL}(q(z)\parallel p(z|k,\theta))$  is computed exactly the same as in the LDS example above, as these two distributions differ only in the recognition potentials. The new challenge is computing  $\mathrm{KL}(q(k)\parallel p(k|\theta))$  which contains no recognition potentials but differs based on mean field messages from  $q(z)$  which are present in  $q(k)$  but not in  $p(k|\theta)$ .
In particular, the prior  $p(k|\theta)$  is a time series with no observations, while  $q(k)$  includes parameters at every time step which encode the log probability of being in that state at that time, given by Eq. (60):  $\mathbb{E}_{q(\theta)}[t(\theta_{\mathrm{trans}})] \cdot \mathbb{E}_{q(z)}[t(z_{t-1}, z_t)]$ . The corresponding sufficient statistics are the marginal probabilities of being in state  $k$  at time  $t$ . Note that we never need to compute the sufficient statistics corresponding to HMM transitions because we don't need them for the local KL nor do we need them to sample from the HMM (which we never do).
Surrogate Loss The surrogate loss differs from the true loss only in the reconstruction term, so it can similarly be written
$$
- \hat {\mathcal {L}} [ q (\theta) q (z), \phi ] = \overbrace {\operatorname {K L} (q (\theta) \| p (\theta))} ^ {\text {P r i o r K L}} + \overbrace {\mathbb {E} _ {q (\theta)} \operatorname {K L} (q (z) \| p (z | \theta))} ^ {\text {L o c a l K L}} - \langle \lambda_ {\phi} (x), \mathbb {E} _ {q (z)} t (z) \rangle \tag {165}
$$
This is trivial to evaluate, as the local KL contains the negation of the last term (see Eq. (163) for the structured case. The mean field case contains these terms as well in addition to ones corresponding to "pulled apart" factors) and thus the surrogate loss can be computed by omitting the recognition-potential-times-expected-statistics part of the local KL.
F Natural gradients
Here we dissect two claims in the main paper more closely. First, we will show how natural gradients with respect to  $\eta$  can be easily computed. Defining, as in the main paper,  $\mu_{\eta} = \mathbb{E}_{q(\theta ;\eta)}[t(\theta)]$ , we can use the chain rule to write:
$$
\frac {\partial \mathcal {L}}{\partial \eta} = \frac {\partial \mathcal {L}}{\partial \mu_ {\eta}} \frac {\partial \mu_ {\eta}}{\partial \eta} + \frac {\partial \mathcal {L}}{\partial \eta} _ {\text {d e t a c h} (\mu_ {\eta})} \tag {166}
$$
Where  $\frac{\partial\mathcal{L}}{\partial\eta}_{\mathrm{detach}(\mu_{\eta})}$  refers to the gradient of the loss with respect to  $\eta$ , disconnecting the computational graph at  $\mu_{\eta}$  (i.e. the dependence of the loss on  $\eta$  which is not accounted for by its dependence on  $\mu_{\eta}$ ).
Proposition F.1.
$$
\frac {\partial \mathcal {L}}{\partial \eta} _ {\text {d e t a c h} (\mu_ {\eta})} = 0 \tag {167}
$$
Proof.  $q(z)$  depends on  $\eta$  only through the expected statistics  $\mu_{\eta}$  which define the parameters of the graphical model factors. Further, the local KL, outlined in Eq. (163), only depends on  $q(\theta)$  through its expected statistics (noting that the log normalizers of  $p(z|\theta)$  are also statistics of  $q(\theta)$ ). Thus the only place in the loss that  $\eta$  appears on its own is the prior KL. Let  $\eta_0$  be the natural parameters of the prior  $p(\theta)$ . Then the prior KL, as per Eq. (25), is given by:
$$
\operatorname {K L} (q (\theta) \parallel p (\theta)) = \left\langle \eta - \eta_ {0}, \mu_ {\eta} \right\rangle - \log Z (\eta) + \log Z (\eta_ {0}) \tag {168}
$$
Then we see:
$$
\begin{array}{l} \frac {\partial}{\partial \eta} \mathrm {K L} (q (\theta) \| p (\theta)) = \frac {\partial \mathrm {K L} (q (\theta) \| p (\theta))}{\partial \mu_ {\eta}} \frac {\partial \mu_ {\eta}}{\partial \eta} + \mu_ {\eta} - \frac {\partial}{\partial \eta} \log Z (\eta) (169) \\ = \frac {\partial \mathrm {K L} (q (\theta) \| p (\theta))}{\partial \mu_ {\eta}} \frac {\partial \mu_ {\eta}}{\partial \eta} = (\eta - \eta_ {0}) \frac {\partial \mu_ {\eta}}{\partial \eta} (170) \\ \end{array}
$$
Where we use Prop. C.1 to cancel terms. Thus we conclude that the entire gradient of the loss with respect to  $\eta$  passes through  $\mu_{\eta}$
The main consequence of this theorem is that in all cases the gradient can be written
$$
\frac {\partial \mathcal {L}}{\partial \eta} = \frac {\partial \mathcal {L}}{\partial \mu_ {\eta}} \frac {\partial \mu_ {\eta}}{\partial \eta} \tag {171}
$$
And thus the natural gradient, which multiples the gradient by the inverse of  $\frac{\partial\mu_{\eta}}{\partial\eta}$ , can be easily computed without matrix arithmetic by skipping the portion of gradient computation corresponding to the  $\eta \rightarrow \mu_{\eta}$  mapping. Here is some sample Jax code which takes a function  $\mathsf{f}$  and returns a function which computes  $\mathsf{f}$  in the forward pass but in the backward pass treats the gradient as identity:
def straight_through(f):
def straightthrough_f(x):
    zero = x - jax.lax.stop_gradients(x)
    return zero + jax.lax.stop_gradients(f(x))
return straightthrough_f
Fisher information of transformed variables. Given a distribution  $q(\theta; \eta)$  with natural parameters  $\eta$ , the Fisher information matrix of  $\eta$  is given by:
$$
F _ {\eta} = \mathbb {E} _ {q (\theta)} \left[ \left(\frac {\partial \log q (\theta)}{\partial \eta}\right) ^ {T} \cdot \left(\frac {\partial \log q (\theta)}{\partial \eta}\right) \right] \tag {172}
$$
Note that for an exponential family distribution
$$
\frac {\partial \log q (\theta)}{\partial \eta} = t (\theta) - \mathbb {E} _ {q (\theta)} [ t (\theta) ] \tag {173}
$$
which clearly has expectation 0 under  $q(\theta)$ . Thus
$$
F _ {\eta} = \mathbb {E} _ {q (\theta)} \left[ (t (\theta) - \mathbb {E} _ {q (\theta)} [ t (\theta) ]) ^ {T} \cdot (t (\theta) - \mathbb {E} _ {q (\theta)} [ t (\theta) ]) \right] \tag {174}
$$
$$
= \operatorname {V a r} _ {q (\theta)} [ t (\theta) ] \tag {175}
$$
Exponential family theory gives us further equalities:
$$
\operatorname {V a r} _ {q (\theta)} [ t (\theta) ] = \frac {\partial^ {2}}{\partial \eta^ {2}} \log Z (\eta) = \frac {\partial}{\partial \eta} \mathbb {E} _ {q (\theta)} [ t (\theta) ] \tag {176}
$$
Thus deriving our desired property that  $F_{\eta} = \frac{\partial\mu_{\eta}}{\partial\eta}$
Now we consider a more general case, where  $\eta = f(\tilde{\eta})$  for some unconstrained, non-natural parameters  $\tilde{\eta}$  and continuous invertible function  $f$ . Then we see:
$$
\begin{array}{l} F _ {\tilde {\eta}} = \mathbb {E} _ {q (\theta)} \left[ \left(\frac {\partial \log q (\theta)}{\partial \tilde {\eta}}\right) ^ {T} \cdot \left(\frac {\partial \log q (\theta)}{\partial \tilde {\eta}}\right) \right] (177) \\ = \mathbb {E} _ {q (\theta)} \left[ \left(\frac {\partial \log q (\theta)}{\partial \eta} \cdot \frac {\partial \eta}{\partial \tilde {\eta}}\right) ^ {T} \cdot \left(\frac {\partial \log q (\theta)}{\partial \eta} \cdot \frac {\partial \eta}{\partial \tilde {\eta}}\right) \right] (178) \\ = \frac {\partial \eta}{\partial \tilde {\eta}} ^ {T} \cdot \mathbb {E} _ {q (\theta)} \left[ \left(\frac {\partial \log q (\theta)}{\partial \eta}\right) ^ {T} \cdot \left(\frac {\partial \log q (\theta)}{\partial \eta}\right) \right] \cdot \frac {\partial \eta}{\partial \tilde {\eta}} (179) \\ = \frac {\partial \eta}{\partial \tilde {\eta}} ^ {T} \cdot F _ {\eta} \cdot \frac {\partial \eta}{\partial \tilde {\eta}} (180) \\ \end{array}
$$
We can use this to compute the natural gradient with respect to the unconstrained parameters:
$$
\frac {\partial \mathcal {L}}{\partial \tilde {\eta}} F _ {\tilde {\eta}} ^ {- 1} = \left(\frac {\partial \mathcal {L}}{\partial \mu} \cdot \frac {\partial \mu}{\partial \eta} \cdot \frac {\partial \eta}{\partial \tilde {\eta}}\right) \cdot \left(\frac {\partial \eta}{\partial \tilde {\eta}} ^ {- 1} \cdot F _ {\eta} ^ {- 1} \cdot \frac {\partial \eta}{\partial \tilde {\eta}} ^ {- T}\right) = \frac {\partial \mathcal {L}}{\partial \mu} \cdot \frac {\partial \eta}{\partial \tilde {\eta}} ^ {- T} = \left(\frac {\partial \tilde {\eta}}{\partial \eta} \cdot \nabla_ {\mu} \mathcal {L}\right) ^ {T} \tag {181}
$$
Implementing natural gradients via automatic differentiation We argued in the main paper that to implement natural gradients we replace the reverse-mode backpropagation through the  $\tilde{\eta} \rightarrow \eta$  map with a forward-mode differentiation through the inverse operation. Assuming we have a forward map  $f$  and inverse function  $f\_inv$ , sample code is shown below:
f_natgrad =jax(custom_vjp(f)
def f_natgrad_fwd(input):
return f(input), f(input)
def f_natgrad_bwd(resids, grads):
return (jax.jvp(f_inv, (resids,)，(grads,)）[1]，)
f_natgrad.defvjp(f_natgrad_fwd，f_natgrad_bwd)
G Matrix-normal exponential family distributions
Recall that exponential family distributions are given by
$$
p (x; \eta) = \exp \left\{\langle \eta , t (x) \rangle - \log Z (\eta) \right\} \tag {182}
$$
The natural parameters often have coupled constraints, and so can be expressed as a simple function of canonical parameters which carry easily-interpretable information about the distribution (e.g. the natural parameters of a normal distribution can be expressed in terms of the mean and precision matrix of that distribution). Thus to define each distribution, we provide (a) the natural parameters  $\eta$  as a function of canonical parameters, (b) the sufficient statistics  $t(x)$ , (c) the expected sufficient statistics  $\mathbb{E}_{p(x;\eta)}[t(x)]$ , and (d) the log partition function  $\log Z(\eta)$ .
Note that sometimes we have symmetric positive definite (SPD) matrix-valued parameters. In every case, the corresponding statistics for those parameters are also symmetric. For example, given a  $n \times n$  SPD matrix, a distribution might have parameters and statistics:
$$
\eta = \left[ \begin{array}{c} S _ {i i} | i = 1, \dots , n \\ 2 \cdot S _ {i j} | i <   j \end{array} \right] \quad t (x) = \left[ \begin{array}{c} x _ {i} ^ {2} | i = 1, \dots , n \\ x _ {i} x _ {j} | i <   j \end{array} \right] \tag {183}
$$
We note that  $\langle \eta, t(x) \rangle = \sum_{i=1}^{n} \sum_{j=1}^{n} S_{ij} x_i x_j$  or equivalently the sum of elements of  $S \otimes xx^T$  (for column vector  $x$  containing each  $x_i$ ). In the following sections, we will abuse notation for compactness by simply writing:
$$
\eta = [ S ] \quad t (x) = \left[ x x ^ {T} \right] \tag {184}
$$
Constraints. To optimize these parameters via gradient descent, we must define constrained parameters as a function of unconstrained parameters as discussed in Sec.[natgrads]. Further, to implement natural gradients we must implement the inverse operation. The table below lists all such transformations.
Notation	Size	Constraint	η→η
Rn	n-length vector	n/a	n/a
Rn×m	n×m matrix	n/a	n/a
Rn+	n-length vector	All values positive	η = Softplus(η)
Rn> m	n-length vector	All values greater than m	η = Softplus(η)+m
Δn-1	n-length vector	Non-negative, sums to 1	η = Softmax(η)
Sn++	n×n matrix	symmetric positive definite	See below
Where Softmax is made invertible by appending a 0 to the end of  $\tilde{\eta}$ . For the SPD matrix parameters, we define  $S = \mathrm{diag}(\sigma) \cdot R \cdot \mathrm{diag}(\sigma)$  where  $\sigma \in \mathbb{R}_+^n$  is a vector of standard deviations and  $R$  is a correlation matrix, transformed to and from unconstrained space by the tensorflow Correlation Cholesky bijector.
G.1 Multivariate Normal distribution
If  $x \sim \mathcal{N}(\mu, \Sigma)$  for (i)  $x, \mu \in \mathbb{R}^n$ , and (ii)  $\Sigma \in S_{++}^{n}$ , then  $p(x; \mu, \Sigma)$  is given by:
$$
\eta = \left[ \begin{array}{c} \Sigma^ {- 1} \mu \\ - \frac {1}{2} \Sigma^ {- 1} \end{array} \right] \quad t (x) = \left[ \begin{array}{c} x \\ x x ^ {T} \end{array} \right] \quad \mathbb {E} _ {p (x; \mu , \Sigma)} [ t (x) ] = \left[ \begin{array}{c} \mu \\ \mu \mu^ {T} + \Sigma \end{array} \right] \tag {185}
$$
The log partition function is given by:
$$
\log Z (\eta) = \frac {1}{2} \mu^ {T} \Sigma^ {- 1} \mu + \frac {1}{2} \log | \Sigma | + \frac {n}{2} \log (2 \pi) \tag {186}
$$
Here,  $\pi \approx 3.14$  is a constant.
G.2 Normal inverse Wishart (NIW) distribution