SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Notation	Meaning
BP(.)	belief propagation function, which takes (natural) parameters of the graphical model and returns expected sufficient statistics.
γ	weights of the generative network (decoder)
η	natural parameters of q(θ;η)
η	unconstrained parameterization of η
Fη	Fisher information matrix of exponential family distribution with natural parameters η
g(ω;η,φ,x)	equation which block updating attempts to solve for the root of discrete local latent variable (subset of z)
k	approximate likelihood function outputted by recognition network (encoder)
l̂φ(z|x)	the loss function, the Evidence Lower Bound Objective (ELBO).
L̂	Surrogate ELBO, obtained by replacing true likelihood function with a conjugate surrogate.
MF(.)	mean field inference update, where messages are passed along residual edges (those edges removed by the structured mean field approximation).
μ	expected sufficient statistics of q(z), i.e. Eq(z)[t(z)]
μm	expected sufficient statistics of a particular mean field component q(zm). μm is defined analogously to ωm.
μη	expected sufficient statistics of q(θ), i.e. Eq(θ)[t(θ)]
ω	natural parameter of q(z) in any of its forms: q(z;η) (local VI), qφ(z|x;η) (SVAE VI) or ∏m qφ(zm|x;η), (Structured Mean Field SVAE VI)
ωm	natural parameter of each mean field cluster ∏m qφ(zm|x;η) in Structured Mean Field SVAE s.t. ω = ∪mωm. ωm is defined analogously to μm.
ω-m	natural parameter of all mean field clusters but m, ωm = ∪n≠mωn.
ω*(x,η,φ)	function which runs block updating to return the optimal natural parameters of q(z) with respect to surrogate loss funciton L̂
p0(z;η)	p0(z;η) ∝ exp{Eq(θ;η)[log p(z|θ)]}
φ	weights of recognition network (encoder)
q(θ;η),q(z;ω)	variational factors approximating the posteriors of global (θ) and local (z) latent variables.
qφ(z|x),qφ(z|x;η)	variational factor for z produced by amortized VI and SVAE VI respectively.
qφ(zm|x;η)	variational factor for subset zm⊆z
t(.)	sufficient statistics of an exponential family distribution.
θ	global latent variables defining graphical model factors
V	vertex set of the graphical model
x	observation
z	local (per-observation) latent variable
B.2 Notation in this appendix
Notation	Meaning
Sec. C	Exponential families and mean field objectives
KL	Kullback-Leibler divergence, defined in Eq. (25)
F	set of factors in the graphical model representation of q(z), introduced in Eq. (32).
λφ(x)	natural parameter of ℓφ(z|x), defined in Eq. (44).
rt,Rt	for Gaussian temporal models such as the LDS and SLDS, λφ(x) = {rt, -1/2Rt} as defined in Eq. (50).
θtrans	Matrix of LDS transition parameters for each of the K switching states in the SLDS, defined in Eq. (58).
Sec. D	Belief propagation in time series
h0,J0	natural parameters of p(z1|θ), defined in Eq. (64).
ωt	natural parameters of the LDS (or the continuous half of the SLDS) transition at time step t, defined in Eq. (65). Also used in the HMM as the mean field message from the LDS at step t.
h1,t,J11,t,J12,t,J22,t,h2,t	components of ωt, defined in Eq. (65).
mzi,ωi(zi)	belief propagation messages from variable zi to factor with parameter ωi, introduced in Eq. (69)
mωi,zi(zi)	belief propagation messages from factor with parameter ωi to variable zi, introduced in Eq. (70)
ft|t,Ft|t	parameters of mzt,ωt+1(zt), i.e. the filtered messages for the Gaussian inference at step t which has accrued evidence from the first t “observations”
ft+1|t,Ft+1|t	parameters of mωt+1,zt+1(zt+1), i.e. the predicted messages for the Gaussian inference at step t + 1 which has accrued evidence from the first t “observations”
ft|T,Ft|T	parameters of q(zt), i.e. the smoothed messages for the Gaussian inference at step t which has accrued evidence from all “observations”
ker(.)	returns an un-normalized distribution with no log partition function, defined in Eq. (111)
ft(zt-1,zt),φ.,t,Φ..,t	sequence of functions (and their parameters) used by the filtering step of the parallel Kalman filter, defined in Eq. (115)
gt(zt-1),γt,Γt	sequence of functions (and their parameters) used by the filtering step of the parallel Kalman filter, defined in Eq. (112)
et(zt,zt+1),εt,Et	sequence of functions (and their parameters) used by the smoothing step of the parallel Kalman smoother, defined in Eq. (148)
⊗	associative operator for parallel Kalman filtering, defined in Eq. (119) and (120)
⊕	associative operator for parallel Kalman smoothing, defined in Eq. (151)
C Exponential families and mean field objective
To rigorously define the SVAE inference routine, in Sec. C.1 we (re-)introduce exponential family and factor graph notation for defining distributions. Then, in Sec. C.2 we derive the optimizer of  $\hat{\mathcal{L}}[\dots]$  and the mean field message passing update MF(.). Sec. D then contains the necessary equations for belief propagation in the LDS and SLDS models, including proofs and derivations of our novel inference algorithm which parallelizes belief propagation across time steps.
C.1 Exponential families of distributions
Given an observation  $z \in \mathcal{Z}$ , an exponential family is a family of distributions characterized by a statistic function  $t(z): \mathcal{Z} \to \mathbb{R}^n$  with probabilities given by
$$
p (z; \eta) = \exp \left\{\langle \eta , t (z) \rangle - \log Z (\eta) \right\} \tag {19}
$$
Where  $\eta \in \mathbb{R}^n$  is a parameter vector known as the natural parameters, and the log partition function  $\log Z(\eta)$  is given by
$$
\log Z (\eta) = \log \int_ {z} \exp \left\{\langle \eta , t (z) \rangle \right\} d x \tag {20}
$$
and ensures that the distribution properly normalizes (to simplify notation, here we focus on continuous  $z$  but this definition can be easily extended to discrete observations by replacing the integral with a sum and restricting the output of  $t(z)$  to some discrete set). A single assignment of values to  $\eta$  characterizes a member of this exponential family, and so searching over a family means searching over valid values of  $\eta \in H$  where  $H$  is the set of parameters for which the integral in Eq. (20) is finite. We begin with a couple simple properties of exponential families:
Proposition C.1. The gradient of a distribution's log partition function with respect to its parameters is its expected statistics:
$$
\nabla_ {\eta} \log Z (\eta)) = \mathbb {E} _ {p (z; \eta)} [ t (z) ] \tag {21}
$$
$$
\begin{array}{l} \nabla_ {\eta} \log Z (\eta) = \frac {1}{\int_ {z} \exp \{\langle \eta , t (z) \rangle \} d x} \int_ {z} t (z) \exp \{\langle \eta , t (z) \rangle \} d x (22) \\ = \int_ {z} t (z) \exp \left\{\langle \eta , t (z) \rangle - \log Z (\eta) \right\} = \mathbb {E} _ {p (z; \eta)} [ t (z) ] (23) \\ \end{array}
$$
Proposition C.2. Given two distributions  $p(z; \eta_1), p(z; \eta_2)$  in the same exponential family with parameters  $\eta_1$  and  $\eta_2$ , the KL divergence between the two distributions is given by
$$
\begin{array}{l} \operatorname {K L} \left(p (z; \eta_ {1}) \| p (z; \eta_ {2}))\right) := \mathbb {E} _ {p (z; \eta_ {1})} \left[ \log \frac {p (z ; \eta_ {1})}{p (z ; \eta_ {2})} \right] (24) \\ = \left\langle \eta_ {1} - \eta_ {2}, \mathbb {E} _ {p (z; \eta_ {1})} [ t (z) ] \right\rangle - \log Z (\eta_ {1}) + \log Z (\eta_ {2}) (25) \\ \end{array}
$$
Proof. The log probabilities are linear in the statistics, so the expectation can be pulled in.
Definition C.3. Given an exponential family distribution  $p(\theta ;\eta)$
$$
p (\theta ; \eta) = \exp \left\{\langle \eta , t (\theta) \rangle - \log Z (\eta) \right\} \tag {26}
$$
we say that  $p(z|\theta)$  is conjugate to  $p(\theta; \eta)$  if one distribution's parameters is the other distribution's statistics, namely
$$
p (z | \theta) = \exp \left\{\langle t (\theta), (t (z), 1) \rangle - c \right\} = \exp \left\{\langle \hat {t} (\theta), t (z) \rangle - \log Z (\hat {t} (\theta)) \right\} \tag {27}
$$
Where  $\hat{t} (\theta)$  is the first  $\dim (t(z))$  entries of  $t(\theta)$  and  $c$  is a constant.
Note that  $t(\theta)$  and  $t(z)$  are not the same function with different inputs, but rather are unique to each distribution. As it is generally clear which function is needed, we share this notation, just as  $p(\theta), p(z)$  is understood to refer to different densities.
Further note that the sufficient statistics of  $p(\theta)$  either become parameters of  $p(z|\theta)$  or its normalizers. This can be notated two ways: either by appending 1s to the end of  $t(z)$  to align with the normalizers, or simply fold them into the log partition function. An example is provided below.
Example: NIw distribution. A normal distribution  $p(z;\mu ,\Sigma)$  has natural parameters and sufficient statistics 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] \tag {28}
$$
We bend notation here by concatenating vectors to matrices. Assume all matrices are appropriately flattened to make the inner product a scalar, for more info see Sec. G. The normal distribution's 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 {29}
$$
Meanwhile, the Normal Inverse Wishart (NIW) distribution has sufficient statistics:
$$
t (\mu , \Sigma) = \left[ \begin{array}{c} - \frac {1}{2} \Sigma^ {- 1} \\ \Sigma^ {- 1} \mu \\ - \frac {1}{2} \mu^ {T} \Sigma^ {- 1} \mu \\ - \frac {1}{2} \log | \Sigma | \end{array} \right] \tag {30}
$$
Note that the first two are parameters of  $p(z)$ , and the last two appear in the log partition function. A desirable property of conjugacy is that the posterior has a simple form, namely:
$$
p (\theta | z) \propto \exp \{\langle t (\theta), \eta + (t (z), 1) \rangle \} \tag {31}
$$
This follows easily from the fact that  $p(\theta | z) \propto p(\theta) p(z|\theta)$  which are both log-linear in  $t(\theta)$ .
Factor Graphs. The exponential family framework can be applied to complicated distributions over many variables, in which case it can be helpful to rephrase the distribution as a collection of factors on subsets of the variables. We represent the variable  $z$  as a graphical model with vertex set  $\mathcal{V}$  such that the vertices partition the components of  $z$ . Let the factor set  $\mathcal{F}$  be a set of subsets of  $\mathcal{V}$  such that  $z_{f} \subseteq z, f \in \mathcal{F}$ . Then we write our distribution
$$
p (z; \eta) = \exp \left\{\langle \eta , t (z) \rangle - \log Z (\eta) \right\} = \exp \left\{\sum_ {f \in \mathcal {F}} \langle \eta_ {f}, t \left(z _ {f}\right) \rangle - \log Z (\eta) \right\} \tag {32}
$$
These two ways of writing the distribution are trivially equal by defining  $\eta$  as the union of all  $\eta_{f}$  and combining  $t(z_{f})$  appropriately. We demonstrate where this notation can be useful in two examples.
C.2 Mean field Optima
Proposition C.4. Given a fixed distribution  $q(a)$  and a loss function of the form
$$
l [ q (a) q (b) ] = \mathbb {E} _ {q (a) q (b)} \left[ \log \frac {f (a , b)}{q (a) q (b)} \right] \tag {33}
$$
for some non-negative, continuous function  $f(a,b)$  with finite integral (such as an un-normalized density), the partially optimal  $q(b)$  across all possible distributions is given by
$$
\underset {q (b)} {\arg \max } l [ q (a) q (b) ] \propto \exp \left\{\mathbb {E} _ {q (a)} [ \log f (a, b) ] \right\} \tag {34}
$$
Proof. This is a simple extension of a proof in Johnson et al. [29]. Let  $\tilde{p}(b)$  be the distribution which satisfies  $\tilde{p}(b) \propto \exp\{\mathbb{E}_{q(a)}[\log f(a, b)]\}$ . We can drop terms in the loss which are constant with respect to  $q(b)$  and write:
$$
\begin{array}{l} \underset {q (b)} {\arg \max } \mathbb {E} _ {q (a) q (b)} \left[ \log \frac {f (a , b)}{q (b)} \right] = \underset {q (b)} {\arg \max } \mathbb {E} _ {q (b)} \left[ \mathbb {E} _ {q (a)} [ \log f (a, b) ] - \log q (b) \right] (35) \\ = \underset {q (b)} {\arg \max } \mathbb {E} _ {q (b)} [ \log \exp \mathbb {E} _ {q (a)} [ \log f (a, b) ] - \log q (b) ] (36) \\ = \underset {q (b)} {\arg \max } \mathbb {E} _ {q (b)} \left[ \log \frac {\exp \left\{\mathbb {E} _ {q (a)} [ \log f (a , b) ] \right\}}{q (b)} \right] (37) \\ = \arg \max  _ {q (b)} \mathbb {E} _ {q (b)} \left[ \log \frac {\tilde {p} (b)}{q (b)} + \operatorname {c o n s t} \right] (38) \\ = \underset {q (b)} {\arg \max } - \mathrm {K L} (q (b) \| \tilde {p} (b)) = \tilde {p} (b) (39) \\ \end{array}
$$
This is applicable to ELBO maximization problems, as well as the surrogate maximization problem arising in the SVAE:
Example: Exact Inference in the SVAE Applying proposition C.4 to the SVAE's surrogate objective, we see:
$$
\underset {q (z)} {\arg \max } \hat {\mathcal {L}} [ q (\theta ; \eta) q (z); \phi ] = \mathbb {E} _ {q (\theta ; \eta) q (z)} \left[ \log \frac {p (\theta) p (z | \theta) \hat {\ell} _ {\phi} (z | x)}{q (\theta ; \eta) q (z)} \right] \tag {40}
$$
$$
\propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta)} [ \log p (\theta) \cdot p (z | \theta) \cdot \hat {\ell} _ {\phi} (z | x) ] \right\} \tag {41}
$$
$$
\propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta)} [ \log p (\theta) ] + \mathbb {E} _ {q (\theta ; \eta)} [ \log p (z | \theta) ] + \log \hat {\ell} _ {\phi} (z | x) \right\} \tag {42}
$$
$$
\propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta)} [ \log p (z | \theta) ] + \log \hat {\ell} _ {\phi} (z | x) \right\} \tag {43}
$$
Recall that we have chosen  $\hat{\ell}_{\phi}(z|x)$  to be a conjugate likelihood function to  $p(z|\theta)$ . Let  $\lambda_{\phi}(x)$  be the outputs of the recognition network such that
$$
\hat {\ell} _ {\phi} (z | x) = \exp \left\{\left\langle \lambda_ {\phi} (x), t (z) \right\rangle \right\} \tag {44}
$$
$$
\underset {q (z)} {\arg \max } \hat {\mathcal {L}} [ q (\theta ; \eta) q (z); \phi ] \propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta)} [ \log p (z | \theta) ] + \log \hat {\ell} _ {\phi} (z | x) \right\} \tag {45}
$$
$$
\propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta)} [ \langle t (\theta), t (z) \rangle ] + \langle \lambda_ {\phi} (x), t (z) \rangle \right\} \tag {46}
$$
$$
= \exp \left\{\langle \mathbb {E} _ {q (\theta ; \eta)} [ t (\theta) ] + \lambda_ {\phi} (x), t (z) \rangle \right\} \tag {47}
$$
Demonstrating that the optimal  $q(z)$  is in the same exponential family as  $p(z|\theta)$ , with parameters  $\omega = \mathbb{E}_{q(\theta ;\eta)}[t(\theta)] + \lambda_{\phi}(x)$ . These parameters are the sum of expected parameters from the prior and the parameters of the fake observations.
Example: Linear Dynamical System. Consider the linear dynamical system, where  $p(z_{1}) = \mathcal{N}(\mu_{0},\Sigma_{0})$  and  $p(z_{t}|z_{t - 1}) = \mathcal{N}(A_{t}z_{t - 1} + b_{t},\Sigma_{t})$  (in practice, we use a time-homogenous model in which  $A,b,\Sigma$  do not depend on the time step  $t$ ). We can write this distribution
$$
p (z | \theta) = p \left(z _ {1}\right) \prod_ {t = 2} ^ {T} p \left(z _ {t} \mid z _ {t - 1}\right) = \exp \left\{\langle t \left(\theta_ {1}\right), t \left(z _ {1}\right) \rangle + \sum_ {t = 2} ^ {T} \langle t \left(\theta_ {t}\right), t \left(z _ {t - 1}, z _ {t}\right) \rangle - \log Z (t (\theta)) \right\} \tag {48}
$$
Where  $t(z_1), t(\theta_1)$  are given by the NIw example in Sec. C.1, and for  $t = 2,\ldots,T$ ,
$$
t \left(z _ {t - 1}, z _ {t}\right) = \left[ \begin{array}{c} z _ {t - 1} \\ z _ {t - 1} z _ {t - 1} ^ {T} \\ z _ {t - 1} z _ {t} ^ {T} \\ z _ {t} z _ {t} ^ {T} \\ z _ {t} \end{array} \right], \quad t \left(\theta_ {t}\right) = \left[ \begin{array}{c} - A _ {t} ^ {T} Q _ {t} ^ {- 1} b _ {t} \\ - \frac {1}{2} A _ {t} ^ {T} Q _ {t} ^ {- 1} A _ {t} \\ A _ {t} ^ {T} Q _ {t} ^ {- 1} \\ - \frac {1}{2} Q _ {t} ^ {- 1} \\ Q _ {t} ^ {- 1} b _ {t} \end{array} \right] \tag {49}
$$
Note that both  $t(\theta_t)$  and  $t(\theta_{t + 1})$  introduce parameters corresponding to the statistics  $z_{t}$ . In the context of the SVAE, we add observations to each time step,
$$
\hat {\ell} _ {\phi} (z | x) = \exp \left\{\sum_ {t} \left\langle \left[ \begin{array}{l} r _ {t} \\ - \frac {1}{2} R _ {t} \end{array} \right], \left[ \begin{array}{l} z _ {t} \\ z _ {t} z _ {t} ^ {T} \end{array} \right] \right\rangle \right\} \tag {50}
$$
which is conjugate to the likelihood  $p(z|\theta)$  so that the mean field optimum stays in the same exponential family.
Structured mean field. The main paper introduces the concept of dividing the latent variables into separate structured mean field components  $z_{m} \subset z, m \in \mathcal{M}$ . We outline the math in this case by noting that conjugacy allows us to break down statistics on groups of vertices into statistics on individual vertices.
$$
\langle t (\theta), t \left(z _ {t - 1}, z _ {t}\right) \rangle = \langle t \left(\theta , z _ {t - 1}\right), t \left(z _ {t}\right) \rangle = \langle t \left(\theta , z _ {t}\right), t \left(z _ {t - 1}\right) \rangle \tag {51}
$$
i.e. this expression is linear in the statistics of each argument. This proves important in deriving the optimal factor  $q(z_{m})$  given the other factors  $q(z_{-m})$ . By applying prop C.4 to the mean field surrogate objective with  $q(a) = q(\theta)\prod_{-m}q(z_m)$  and  $q(b) = q(z_{m})$ , we see:
$$
\underset {q \left(z _ {m}\right)} {\arg \max } \hat {\mathcal {L}} [ q (\theta ; \eta) \prod_ {m} q \left(z _ {m}\right); \phi ] \propto \exp \left\{\mathbb {E} _ {q (\theta ; \eta) q \left(z _ {- m}\right)} \left[ \log p (z | \theta) \right] + \log \hat {\ell} _ {\phi} (z | x) \right\} \tag {52}
$$
The expected-prior part of this expression can be rewritten (where  $z_{f,m} = z_f \cap z_m$  and  $z_{f, -m} = z_f \cap z_m^C$ ):
$$
\begin{array}{l} \mathbb {E} _ {q (\theta ; \eta) q (z - m)} [ \log p (z | \theta) ] \propto \mathbb {E} _ {q (\theta ; \eta) q (z - m)} \left[ \sum_ {f \in \mathcal {F}} \langle t \left(\theta_ {f}\right), t \left(z _ {f}\right) \rangle \right] (53) \\ = \mathbb {E} _ {q (\theta ; \eta) q (z _ {- m})} \left[ \sum_ {f \in \mathcal {F}} \langle t \left(\theta_ {f}, z _ {f, - m}\right), t \left(z _ {f, m}\right) \rangle \right] (54) \\ = \sum_ {f \in \mathcal {F}} \left\langle \mathbb {E} _ {q (\theta ; \eta) q (z _ {- m})} \left[ t \left(\theta_ {f}, z _ {f, - m}\right) \right], t \left(z _ {f, m}\right) \right\rangle (55) \\ \end{array}
$$
Note that because the joint sufficient statistics of  $z_{f, -m}$ ,  $\theta_f$  are linear in each input, we can distribute the expectations through the statistic function and write
$$
\mathbb {E} _ {q (\theta ; \eta) q (z _ {- m})} [ t (\theta_ {f}, z _ {f, - m}) ] = \mathrm {M F} \left(\mathbb {E} _ {q (\theta ; \eta)} [ t (\theta_ {f}) ], \mathbb {E} _ {q (z _ {- m})} [ t (z _ {f, - m}) ]\right) \tag {56}
$$
Where  $\mathsf{MF}(.)$  is a mean field function which takes the statistics of  $\theta_f, z_{f, -m}$  and produces the joint statistics  $t(\theta_f, z_{f, -m})$ . This function is linear in each of its inputs and generally cheap to compute.
Example: Switching Linear Dynamical System. The SLDS consists of the following generative model:  $p(k_{2}) = \mathrm{Cat}(\pi_{0})$ ,  $p(k_{t}) = \mathrm{Cat}(\pi_{k_{t - 1}})$ ,  $p(z_{1}) = \mathcal{N}(\mu_{0},\Sigma_{0})$ ,  $p(k_{2}) = \mathrm{Cat}(\pi_{0})$  and  $p(z_{t}|z_{t - 1}) = \mathcal{N}(A_{k_{t}}z_{t - 1} + b_{k_{t}},\Sigma_{k_{t}})$ .
In an exponential family form, the model consists of the following factors: (i) an initial-state factor on  $z_{1}$  matching the NIw example in Eqs. (10-12), (ii) an initial categorical factor on  $k_{2}$  (see Sec. G for the form), (iii) a transition factor for  $p(k_{t}|k_{t - 1})$  which has parameters  $\pi_{ij}$  and sufficient statistics  $1\{(k_{t} = i)\& (k_{t + 1} = j)\}$  for  $i,j = 1,\dots ,K$ , and (iv) transition factors corresponding to  $p(z_{t}|z_{t - 1},k_{t})$ , which we will examine further.
Recall from Eq. (49) the expression for the parameters of  $p(z_{t}|z_{t - 1})$  in the LDS. Then we can write:
$$
p \left(z _ {t} \mid z _ {t - 1}, k _ {t}\right) = \exp \left\{\langle t \left(\theta_ {k _ {t}}\right), t \left(z _ {t - 1}, z _ {t}\right) \rangle - \log Z \left(t \left(\theta_ {k _ {t}}\right)\right) \right\} \tag {57}
$$
We need to rewrite this to be an exponential family distribution of  $k_{t}$  as well (instead of only appearing as an index). Note that the sufficient statistics of a categorical variable is simply its one-hot encoding. If we define  $t(\theta_{\mathrm{trans}})$  to be a matrix whose  $k$ th row is given by  $t(\theta_k)$ , and  $\log Z(\theta_{\mathrm{trans}})$  to be a vector with components  $\log Z(\theta_k)$ , we write:
$$
p \left(z _ {t} \mid z _ {t - 1}, k _ {t}\right) = \exp \left\{\left\langle t \left(\theta_ {\text {t r a n s}}\right), t \left(k _ {t}\right) \cdot t \left(z _ {t - 1}, z _ {t}\right) ^ {T} \right\rangle + \left\langle - \log Z \left(t \left(\theta_ {\text {t r a n s}}\right)\right), t \left(k _ {t}\right) \right\rangle \right\} \tag {58}
$$
Thus the conditional distribution can be described by two factors, one which takes the outer product of the one-hot encoding  $t(k_{t})$  with the continuous observations (thus embedding them in the correct dimension for inner product with  $t(\theta)$ ), and one which simply takes the correct log normalizer based on what state we're in (again, by multiplying a one-hot encoding with a vector of log probabilities).
Given observations, exact inference in the SLDS is intractable. Thus we divide the graphical model via structured mean field inference into  $q(z)q(k)$ . This separates out the joint factors  $\langle t(\theta_{\mathrm{trans}}), t(k_t) \cdot t(z_{t-1}, z_t)^T \rangle$ . However, as noted in Eq. (51) we can reframe this as a factor on each mean field cluster:
$$
\begin{array}{l} \exp \left\{\langle t \left(\theta_ {\text {t r a n s}}\right), t \left(k _ {t}\right) \cdot t \left(z _ {t - 1}, z _ {t}\right) ^ {T} \rangle \right\} = \exp \left\{\langle t \left(z _ {t - 1}, z _ {t}\right), t \left(\theta_ {\text {t r a n s}}\right) ^ {T} \cdot t \left(k _ {t}\right) \rangle \right\} (59) \\ = \exp \left\{\langle t \left(k _ {t}\right), t \left(\theta_ {\text {t r a n s}}\right) \cdot t \left(z _ {t - 1}, z _ {t}\right) \rangle \right\} (60) \\ \end{array}
$$
Thus, applying Eq. (55), the optimal  $q(k)$  inherits the (expected) initial-state factors from  $p(k_2)$  and transition factors from  $p(k_t|k_{t-1})$ , as well as the normalizer factors  $\langle -\log Z(t(\theta_{\mathrm{trans}})), t(k_t) \rangle$  from Eq. (58), but replaces the joint factors of Eq. (58) which include  $z_{t-1}, z_t$  with
$$
\langle \mathbb {E} _ {q (\theta ; \eta)} [ t (\theta_ {\text {t r a n s}}) ] \cdot \mathbb {E} _ {q (z)} [ t (z _ {t - 1}, z _ {t}) ], t (k _ {t}) \rangle \tag {61}
$$
Which (when combined with the normalizer factor) is simply a vector of expected log probabilities that the transition at time  $t$  is caused by each state. Similarly, the continuous chain keeps its initial state factor corresponding to  $p(z_{1})$  but all of its transition factors are replaced by
$$
\left\langle \mathbb {E} _ {q (\theta ; \eta)} \left[ t \left(\theta_ {\text {t r a n s}}\right) \right] ^ {T} \cdot \mathbb {E} _ {q \left(k _ {t}\right)} \left[ t \left(k _ {t}\right) \right], t \left(z _ {t - 1}, z _ {t}\right) \right\rangle \tag {62}
$$
Because  $t(k_{t})$  is a one-hot encoding of  $k_{t}$ ,  $\mathbb{E}_{q(k_{t})}[t(k_{t})]$  is simply a probability vector of what value  $k_{t}$  takes. Thus the new parameter for  $t(z_{t-1}, z_{t})$  is simply a convex combination of the rows of  $t(\theta_{\mathrm{trans}})$ , yielding a time-inhomogeneous LDS with each transition computed as an expectation over discrete states. Our mean field optimization of the SLDS becomes:
- Run a belief propagation algorithm to get  $\mathbb{E}_{q(z)}[t(z)]$
- For each  $t$ , compute  $\mathbb{E}_{q(\theta)}[t(\theta_{\mathrm{trans}})] \cdot \mathbb{E}_{q(z)}[t(z_{t - 1},z_t)]$  as the new parameters in  $q(k)$  corresponding to  $t(k_{t})$
- Run a belief propagation algorithm to get  $\mathbb{E}_{q(k)}[t(k)]$
- For each  $t$ , compute  $\mathbb{E}_{q(\theta)}[t(\theta_{\mathrm{trans}})]^T \cdot \mathbb{E}_{q(k)}[t(k_t)]$  as the new parameters in  $q(z)$  corresponding to  $t(z_{t-1}, z_t)$
... and repeat until convergence.
D Belief propagation in time series
D.1 Sequential LDS Inference