SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} p (\mathbf {x} _ {1}, \mathbf {x} _ {2}, \mathbf {x} _ {3} \mid \mathbf {x} _ {0}) = \sum_ {c _ {1}, c _ {2}, c _ {3}} p (c _ {1}) p (\mathbf {x} _ {1} \mid \mathbf {x} _ {0}, c _ {1}) \cdot A _ {c _ {1}, c _ {2}} p (\mathbf {x} _ {2} \mid \mathbf {x} _ {1}, c _ {2}) \cdot A _ {c _ {2}, c _ {3}} p (\mathbf {x} _ {3} \mid \mathbf {x} _ {2}, c _ {3}) \\ = \sum_ {c _ {1}, c _ {2}} p (c _ {1}) A _ {c _ {1}, c _ {2}} p (\mathbf {x} _ {1} \mid \mathbf {x} _ {0}, c _ {1}) \cdot p (\mathbf {x} _ {2} \mid \mathbf {x} _ {1}, c _ {2}) \cdot \left(\sum_ {c _ {3}} A _ {c _ {2}, c _ {3}} p (\mathbf {x} _ {3} \mid \mathbf {x} _ {2}, c _ {3})\right) \\ = \sum_ {c _ {2}} \left(\sum_ {c _ {1}} p \left(c _ {1}\right) A _ {c _ {1}, c _ {2}} p \left(\mathbf {x} _ {1} \mid \mathbf {x} _ {0}, c _ {1}\right)\right) \cdot p \left(\mathbf {x} _ {2} \mid \mathbf {x} _ {1}, c _ {2}\right) \cdot \left(\sum_ {c _ {3}} A _ {c _ {2}, c _ {3}} p \left(\mathbf {x} _ {3} \mid \mathbf {x} _ {2}, c _ {3}\right)\right) \\ = \sum_ {c _ {2}} \pi_ {c _ {2}} \underbrace {\left(\sum_ {c _ {1}} \frac {\pi_ {c _ {1}} A _ {c _ {1} , c _ {2}}}{\pi_ {c _ {2}}} \mu_ {c _ {1}}\right)} _ {\bar {\mu} _ {c _ {2}}} \cdot \mu_ {c _ {2}} \cdot \underbrace {\left(\sum_ {c _ {3}} A _ {c _ {2} , c _ {3}} \mu_ {c _ {3}}\right)} _ {\dot {\mu} _ {c _ {2}}} \tag {5} \\ \end{array}
$$
where  $\pi_{c_i} = p(c_i)$ . Since  $\mathbf{A}$  has full rank and the probability measures  $\mu_1, \ldots, \mu_C$  are linearly independent, the probability measures  $\{\bar{\mu}_{c_2} = \sum_{c_1} \frac{\pi_{c_1} A_{c_1, c_2}}{\pi_{c_2}} \mu_{c_1} | c_2 = 1, \ldots, C\}$  are linearly independent, and the probability measures  $\{\dot{\mu}_{c_2} = \sum_{c_3} A_{c_2, c_3} \mu_{c_3} | c_2 = 1, \ldots, C\}$  are also linearly independent. Thus, applying Theorem 9 of [46], there exists a permutation  $\sigma$  of  $\{1, \ldots, C\}$  such that,  $\forall i \in \{1, \ldots, C\}$ :
$$
\tilde {\mu} _ {i} = \mu_ {\sigma (i)}
$$
$$
\sum_ {j} \tilde {A} _ {i, j} \tilde {\mu} _ {j} = \sum_ {j} A _ {\sigma (i), j} \mu_ {j}
$$
This gives easily  $\forall i\in \{1,\ldots ,C\}$  ..
$$
\sum_ {j} \tilde {A} _ {i, j} \mu_ {\sigma (j)} = \sum_ {j} A _ {\sigma (i), \sigma (j)} \mu_ {\sigma (j)}.
$$
Since the conditional distributions  $\mu_{i}$  are linearly independent, we can establish the equivalence between  $\tilde{\mathbf{A}}$  and  $\mathbf{A}$  via permutation  $\sigma$ ,
$$
\tilde {A} _ {j, i} = A _ {\sigma (j), \sigma (i)}, \tag {6}
$$
then the theorem is proved.
For notational simplicity, and without loss of generality, we assume the components are ordered such that  $c = \sigma(c)$ . That leads us to the identifiability of the nonstationarity in the system i.e. up to label swapping of the hidden states, the conditional emission distributions  $p(\mathbf{x}_t | \mathbf{x}_{t-1}, c_t)$  and transition matrix  $\mathbf{A}$  are identifiable, hence providing us a bridge to further leverage the temporal independence condition in the latent space to establish the identifiability result for demixing function or in other words the latent variables  $\mathbf{z}_t$ .
A.2 Identifiability of Latent Causal Processes
$$
p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t - 1}, c _ {t}\right) = \prod_ {k = 1} ^ {n} p \left(z _ {k t} \mid \mathbf {z} _ {t - 1}, c _ {t}\right). \tag {7}
$$
Second, the observed data  $\mathbf{x}_t\in \mathbb{R}^n$  is generated by a nonlinear mixing function as in Eq. (1) which is assumed to be bijective with inverse given by  $\mathbf{z}_t = \mathbf{g}^{-1}(\mathbf{x}_t)$ . Let  $\eta_{kt}(c_t)\triangleq \log p(z_{kt}|\mathbf{z}_{t - 1},c_t)$ , and assume that  $\eta_{kt}(c_t)$  is twice differentiable in  $z_{kt}$  and is differentiable in  $z_{l,t - 1}$ ,  $l = 1,2,\dots,n$ . Note that the parents of  $z_{kt}$  may be only  $c_t$  and a subset of  $\mathbf{z}_{t - 1}$ ; if  $z_{l,t - 1}$  is not a parent of  $z_{kt}$ , then  $\frac{\partial\eta_k}{\partial z_{l,t - 1}} = 0$ .
Theorem 2. Suppose there exists an invertible function  $\hat{\mathbf{g}}^{-1}$ , which is the estimated demixing function that maps  $\mathbf{x}_t$  to  $\hat{\mathbf{z}}_t$ , i.e.,
$$
\hat {\mathbf {z}} _ {t} = \hat {\mathbf {g}} ^ {- 1} \left(\mathbf {x} _ {t}\right) \tag {8}
$$
$$
\mathbf {\hat {v}} _ {k, t} (c) \triangleq \left(\frac {\partial^ {3} \eta_ {k t} (c)}{\partial z _ {k , t} ^ {2} \partial z _ {1 , t - 1}}, \frac {\partial^ {3} \eta_ {k t} (c)}{\partial z _ {k , t} ^ {2} \partial z _ {2 , t - 1}},..., \frac {\partial^ {3} \eta_ {k t} (c)}{\partial z _ {k , t} ^ {2} \partial z _ {n , t - 1}}\right) ^ {\intercal}. \tag {9}
$$
$$
\mathbf {s} _ {k t} \triangleq \left(\mathbf {v} _ {k t} (1) ^ {\intercal}, \dots , \mathbf {v} _ {k t} (C) ^ {\intercal}, \frac {\partial^ {2} \eta_ {k t} (2)}{\partial z _ {k t} ^ {2}} - \frac {\partial^ {2} \eta_ {k t} (1)}{\partial z _ {k t} ^ {2}}, \dots , \frac {\partial^ {2} \eta_ {k t} (C)}{\partial z _ {k t} ^ {2}} - \frac {\partial^ {2} \eta_ {k t} (C - 1)}{\partial z _ {k t} ^ {2}}\right) ^ {\intercal}, \tag {10}
$$
Proof. Combining (1) and (6) gives  $\mathbf{z}_t = (\mathbf{g}^{-1}\circ \hat{\mathbf{g}})(\hat{\mathbf{z}}_t) = \mathbf{h}(\hat{\mathbf{z}}_t)$ , where  $\mathbf{h}\triangleq \mathbf{g}^{-1}\circ \hat{\mathbf{g}}$ . Since both  $\mathbf{g}$  and  $\hat{\mathbf{g}}$  are invertible,  $\mathbf{h}$  is invertible. Let  $\mathbf{H}_t$  be the Jacobian matrix of the transformation  $\mathbf{h}(\hat{\mathbf{z}}_t)$ , and denote by  $\mathbf{H}_{kit}$  its  $(k,i)$ th entry.
First, it is straightforward to see that if the components of  $\hat{\mathbf{z}}_t$  are mutually independent conditional on previous  $\hat{\mathbf{z}}_{t - 1}$  and current  $c_{t}$ , then for any  $i\neq j$ ,  $\hat{z}_{it}$  and  $\hat{z}_{jt}$  are conditionally independent given  $\hat{\mathbf{z}}_{t - 1}\cup (\hat{\mathbf{z}}_t\setminus \{\hat{z}_{it},\hat{z}_{jt}\})\cup \{c_t\}$ . Mutual independence of the components of  $\hat{\mathbf{z}}_t$  conditional on  $\hat{\mathbf{z}}_{t - 1}$  implies that  $\hat{z}_{it}$  is independent from  $\hat{\mathbf{z}}_t\setminus \{\hat{z}_{it},\hat{z}_{jt}\}$  conditional on  $\hat{\mathbf{z}}_{t - 1}$  and  $c_{t}$ , i.e.,
$$
p \left(\hat {z} _ {i t} \mid \hat {\mathbf {z}} _ {t - 1}, c _ {t}\right) = p \left(\hat {z} _ {i t} \mid \hat {\mathbf {z}} _ {t - 1} \cup \left(\hat {\mathbf {z}} _ {t} \setminus \left\{\hat {z} _ {i t}, \hat {z} _ {j t} \right\}\right), c _ {t}\right).
$$
At the same time, it also implies  $\hat{z}_{it}$  is independent from  $\hat{\mathbf{z}}_t\setminus \{\hat{z}_{it}\}$  conditional on  $\hat{\mathbf{z}}_{t - 1}$  and  $c_{t}$ , i.e.,
$$
p (\hat {z} _ {i t} | \hat {\mathbf {z}} _ {t - 1}, c _ {t}) = p (\hat {z} _ {i t} | \hat {\mathbf {z}} _ {t - 1} \cup (\hat {\mathbf {z}} _ {t} \setminus \{\hat {z} _ {i t} \}), c _ {t}).
$$
Combining the above two equations gives
$$
p (\hat {z} _ {i t} | \hat {\mathbf {z}} _ {t - 1} \cup (\hat {\mathbf {z}} _ {t} \setminus \{\hat {z} _ {i t} \}), c _ {t}) = p (\hat {z} _ {i t} | \hat {\mathbf {z}} _ {t - 1} \cup (\hat {\mathbf {z}} _ {t} \setminus \{\hat {z} _ {i t}, \hat {z} _ {j t} \}), c _ {t}),
$$
i.e., for  $i \neq j$ ,  $\hat{z}_{it}$  and  $\hat{z}_{jt}$  are conditionally independent given  $\hat{\mathbf{z}}_{t-1} \cup (\hat{\mathbf{z}}_t \setminus \{\hat{z}_{it}, \hat{z}_{jt}\}) \cup \{c_t\}$ .
We then make use of the fact that if  $\hat{z}_{it}$  and  $\hat{z}_{jt}$  are conditionally independent given  $\hat{\mathbf{z}}_{t-1} \cup (\hat{\mathbf{z}}_t \setminus \{\hat{z}_{it}, \hat{z}_{jt}\}) \cup \{c_t\}$ , then
$$
\frac {\partial^ {2} \log p (\hat {\mathbf {z}} _ {t} , \hat {\mathbf {z}} _ {t - 1} , c _ {t})}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}} = 0,
$$
assuming the cross second-order derivative exists [47]. Since  $p(\hat{\mathbf{z}}_t, \hat{\mathbf{z}}_{t-1}, c_t) = p(\hat{\mathbf{z}}_t | \hat{\mathbf{z}}_{t-1}, c_t)p(\hat{\mathbf{z}}_{t-1}, c_t)$  while  $p(\hat{\mathbf{z}}_{t-1}, c_t)$  does not involve  $\hat{z}_{it}$  or  $\hat{z}_{jt}$ , the above equality is equivalent to
$$
\frac {\partial^ {2} \log p (\hat {\mathbf {z}} _ {t} \mid \hat {\mathbf {z}} _ {t - 1} , c _ {t})}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}} = 0. \tag {11}
$$
Then for any  $c_t$ , the Jacobian matrix of the mapping from  $(\mathbf{x}_{t-1}, \hat{\mathbf{z}}_t)$  to  $(\mathbf{x}_{t-1}, \mathbf{z}_t)$  is  $\begin{bmatrix} \mathbf{I} & \mathbf{0} \\ * & \mathbf{H}_t \end{bmatrix}$ , where  $*$  stands for a matrix, and the (absolute value of the) determinant of this Jacobian matrix is  $|\mathbf{H}_t|$ . Therefore  $p(\hat{\mathbf{z}}_t, \mathbf{x}_{t-1}|c_t) = p(\mathbf{z}_t, \mathbf{x}_{t-1}|c_t) \cdot |\mathbf{H}_t|$ . Dividing both sides of this equation by  $p(\mathbf{x}_{t-1}|c_t)$  gives
$$
p \left(\hat {\mathbf {z}} _ {t} \mid \mathbf {x} _ {t - 1}, c _ {t}\right) = p \left(\mathbf {z} _ {t} \mid \mathbf {x} _ {t - 1}, c _ {t}\right) \cdot \left| \mathbf {H} _ {t} \right|. \tag {12}
$$
Since  $p(\mathbf{z}_t \mid \mathbf{z}_{t-1}, c_t) = p(\mathbf{z}_t \mid \mathbf{g}(\mathbf{z}_{t-1}), c_t) = p(\mathbf{z}_t \mid \mathbf{x}_{t-1}, c_t)$  and similarly  $p(\hat{\mathbf{z}}_t \mid \hat{\mathbf{z}}_{t-1}, c_t) = p(\hat{\mathbf{z}}_t \mid \mathbf{x}_{t-1}, c_t)$ , Eq. 12 tells us
$$
\log p \left(\hat {\mathbf {z}} _ {t} \mid \hat {\mathbf {z}} _ {t - 1}, c _ {t}\right) = \log p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t - 1}, c _ {t}\right) + \log | \mathbf {H} _ {t} | = \sum_ {k = 1} ^ {n} \eta_ {k t} \left(c _ {t}\right) + \log | \mathbf {H} _ {t} |. \tag {13}
$$
Its partial derivative w.r.t.  $\hat{z}_{it}$  is
$$
\begin{array}{l} \frac {\partial \log p (\hat {\mathbf {z}} _ {t} \mid \hat {\mathbf {z}} _ {t - 1} , c _ {t})}{\partial \hat {z} _ {i t}} = \sum_ {k = 1} ^ {n} \frac {\partial \eta_ {k t} (c _ {t})}{\partial z _ {k t}} \cdot \frac {\partial z _ {k t}}{\partial \hat {z} _ {i t}} - \frac {\partial \log | \mathbf {H} _ {t} |}{\partial \hat {z} _ {i t}} \\ = \sum_ {k = 1} ^ {n} \frac {\partial \eta_ {k t} (c _ {t})}{\partial z _ {k t}} \cdot \mathbf {H} _ {k i t} - \frac {\partial \log | \mathbf {H} _ {t} |}{\partial \hat {z} _ {i t}}. \\ \end{array}
$$
Its second-order cross-derivative is
$$
\frac {\partial^ {2} \log p (\hat {\mathbf {z}} _ {t} \mid \hat {\mathbf {z}} _ {t - 1} , c _ {t})}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}} = \sum_ {k = 1} ^ {n} \left(\frac {\partial^ {2} \eta_ {k t} (c _ {t})}{\partial z _ {k t} ^ {2}} \cdot \mathbf {H} _ {k i t} \mathbf {H} _ {k j t} + \frac {\partial \eta_ {k t} (c _ {t})}{\partial z _ {k t}} \cdot \frac {\partial \mathbf {H} _ {k i t}}{\partial \hat {z} _ {j t}}\right) - \frac {\partial^ {2} \log | \mathbf {H} _ {t} |}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}}. \tag {14}
$$
The above quantity is always 0 according to Eq. (11). Therefore, for each  $l = 1,2,\dots,n$  and each value  $z_{l,t - 1}$ , its partial derivative w.r.t.  $z_{l,t - 1}$  is always 0. That is,
$$
\frac {\partial^ {3} \log p \left(\hat {\mathbf {z}} _ {t} \mid \hat {\mathbf {z}} _ {t - 1} , c _ {t}\right)}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t} \partial z _ {l , t - 1}} = \sum_ {k = 1} ^ {n} \left(\frac {\partial^ {3} \eta_ {k t} (c _ {t})}{\partial z _ {k t} ^ {2} \partial z _ {l , t - 1}} \cdot \mathbf {H} _ {k i t} \mathbf {H} _ {k j t} + \frac {\partial^ {2} \eta_ {k t} (c _ {t})}{\partial z _ {k t} \partial z _ {l , t - 1}} \cdot \frac {\partial \mathbf {H} _ {k i t}}{\partial \hat {z} _ {j t}}\right) \equiv 0, \tag {15}
$$
where we have made use of the fact that entries of  $\mathbf{H}_t$  do not depend on  $z_{l,t - 1}$ . Using different values  $r$  for  $c_{t}$  in Eq. (14) take the difference of this equation across them gives
$$
\begin{array}{l} \frac {\partial^ {2} \log p (\hat {\mathbf {z}} _ {t} | \hat {\mathbf {z}} _ {t - 1} ; r + 1)}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}} - \frac {\partial^ {2} \log p (\hat {\mathbf {z}} _ {t} | \hat {\mathbf {z}} _ {t - 1} ; r)}{\partial \hat {z} _ {i t} \partial \hat {z} _ {j t}} \\ = \sum_ {k = 1} ^ {n} \left[ \left(\frac {\partial^ {2} \eta_ {k t} (r + 1)}{\partial z _ {k t} ^ {2}} - \frac {\partial^ {2} \eta_ {k t} (r)}{\partial z _ {k t} ^ {2}}\right) \cdot \mathbf {H} _ {k i t} \mathbf {H} _ {k j t} + \left(\frac {\partial \eta_ {k t} (r + 1)}{\partial z _ {k t}} - \frac {\partial \eta_ {k t} (r)}{\partial z _ {k t}}\right) \cdot \frac {\partial \mathbf {H} _ {k i t}}{\partial \hat {z} _ {j t}} \right] \equiv 0. \tag {16} \\ \end{array}
$$
If for any value of  $\mathbf{z}_t$ ,  $\mathbf{s}_{1t}, \mathbf{\dot{s}}_{1t}, \mathbf{s}_{2t}, \mathbf{\dot{s}}_{2t}, \dots, \mathbf{s}_{nt}, \mathbf{\dot{s}}_{nt}$  are linearly independent, to make the above equation hold true, one has to set  $\mathbf{H}_{kit} \mathbf{H}_{kjt} = 0$  or  $i \neq j$ . That is, in each row of  $\mathbf{H}_t$  there is only one non-zero entry. Since  $h$  is invertible, then  $\mathbf{z}_t$  must be an invertible, component-wise transformation of a permuted version of  $\hat{\mathbf{z}}_t$ .
So far, the identifiability result has been established without observing the nonstationarity indicators such as domain indices.
A.3 Discussion on Assumptions in Theorem 2
This condition was initially introduced in GCL [11], namely, "sufficient variability", to extend the modulated exponential families [9] to general modulated distributions. Essentially, the condition says that the nonstationary domains  $c$  must have a sufficiently complex and diverse effect on the transition distributions. In other words, if the underlying distributions are composed of relatively many domains of data, the condition generally holds true. Loosely speaking, the sufficient variability holds if the modulation of by  $c$  on the conditional distribution  $q(z_{it}|\mathbf{z}_{\mathrm{Hx}},c)$  is not too simple in the following sense:
1. Higher order of  $k$  ( $k > 1$ ) is required. If  $k = 1$ , the sufficient variability cannot hold;
2. The modulation impacts  $\lambda_{ij}$  by  $\mathbf{u}$  must be linearly independent across domains  $c$ . The sufficient statistics functions  $q_{ij}$  cannot be all linear, i.e., we require higher-order statistics.
Further details of this example can be found in Appendix B of [11] and Appendix S1.4.1 of [18]. In summary, we need the domains denoted by  $c$  to have diverse (i.e., distinct influences) and complex impacts on the underlying data generation process.
All experiments are done in a GPU workstation with CPU: Intel i7-13700K, GPU: NVIDIA RTX 4090, Memory: 128 GB. The code can be found via https://github.com/xiangchensong/nctrl.
B.2 Prior Likelihood Derivation
Let us start with an illustrative example of stationary latent causal processes consisting of two time-delayed latent variables, i.e.,  $\mathbf{z}_t = [z_{1,t}, z_{2,t}]$  with maximum time lag  $L = 1$ , i.e.,  $z_{i,t} = f_i(\mathbf{z}_{t-1}, \epsilon_{i,t})$  with mutually independent noises. Let us write this latent process as a transformation map  $\mathbf{f}$  (note that we overload the notation  $f$  for transition functions and for the transformation map):
$$
\left[ \begin{array}{l} z _ {1, t - 1} \\ z _ {2, t - 1} \\ z _ {1, t} \\ z _ {2, t} \end{array} \right] = \mathbf {f} \left(\left[ \begin{array}{l} z _ {1, t - 1} \\ z _ {2, t - 1} \\ \epsilon_ {1, t} \\ \epsilon_ {2, t} \end{array} \right]\right). \tag {17}
$$
By applying the change of variables formula to the map  $\mathbf{f}$ , we can evaluate the joint distribution of the latent variables  $p(z_{1,t-1}, z_{2,t-1}, z_{1,t}, z_{2,t})$  as:
$$
p \left(z _ {1, t - 1}, z _ {2, t - 1}, z _ {1, t}, z _ {2, t}\right) = p \left(z _ {1, t - 1}, z _ {2, t - 1}, \epsilon_ {1, t}, \epsilon_ {2, t}\right) / \left| \det  \mathbf {J} _ {\mathbf {f}} \right|, \tag {18}
$$
where  $\mathbf{J}_{\mathbf{f}}$  is the Jacobian matrix of the map  $\mathbf{f}$ , which is naturally a low-triangular matrix:
$$
\mathbf {J _ {f}} = \left[ \begin{array}{c c c c} 1 & 0 & 0 & 0 \\ 0 & 1 & 0 & 0 \\ \frac {\partial z _ {1 , t}}{\partial z _ {1 , t - 1}} & \frac {\partial z _ {1 , t}}{\partial z _ {2 , t - 1}} & \frac {\partial z _ {1 , t}}{\partial \epsilon_ {1 , t}} & 0 \\ \frac {\partial z _ {2 , t}}{\partial z _ {1 , t - 1}} & \frac {\partial z _ {2 , t}}{\partial z _ {2 , t - 1}} & 0 & \frac {\partial z _ {2 , t}}{\partial \epsilon_ {2 , t}} \end{array} \right].
$$
Given that this Jacobian is triangular, we can efficiently compute its determinant as  $\prod_{i}\frac{\partial z_{i,t}}{\partial\epsilon_{i,t}}$ . Furthermore, because the noise terms are mutually independent, and hence  $\epsilon_{i,t} \perp \epsilon_{j,t}$  for  $j \neq i$  and  $\epsilon_t \perp \mathbf{z}_{t-1}$ , we can write the RHS of Eq. 18 as:
$$
\begin{array}{l} p \left(z _ {1, t - 1}, z _ {2, t - 1}, z _ {1, t}, z _ {2, t}\right) = p \left(z _ {1, t - 1}, z _ {2, t - 1}\right) \times p \left(\epsilon_ {1, t}, \epsilon_ {2, t}\right) / | \det  \mathbf {J} _ {\mathbf {f}} | \quad (\text {b e c a u s e} \epsilon_ {t} \perp \mathbf {z} _ {t - 1}) \\ = p \left(z _ {1, t - 1}, z _ {2, t - 1}\right) \times \prod_ {i} p \left(\epsilon_ {i, t}\right) / \left| \det \mathbf {J} _ {\mathbf {f}} \right| \quad (\text {b e c a u s e} \epsilon_ {1, t} \perp \epsilon_ {2, t}) \tag {19} \\ \end{array}
$$
Finally, by canceling out the marginals of the lagged latent variables  $p(z_{1,t-1}, z_{2,t-1})$  on both sides, we can evaluate the transition prior likelihood as:
$$
p \left(z _ {1, t}, z _ {2, t} \mid z _ {1, t - 1}, z _ {2, t - 1}\right) = \prod_ {i} p \left(\epsilon_ {i, t}\right) / \left| \det  \mathbf {J} _ {\mathbf {f}} \right| = \prod_ {i} p \left(\epsilon_ {i, t}\right) \times \left| \det  \mathbf {J} _ {\mathbf {f}} ^ {- 1} \right|. \tag {20}
$$
Now we generalize this example and derive the prior likelihood below.
Let  $\{f_i^{-1}\}_{i = 1,2,3\dots}$  be a set of learned inverse transition functions that take the estimated latent causal variables, and output the noise terms, i.e.,  $\hat{\epsilon}_{i,t} = f_{i}^{-1}\left(\hat{z}_{i,t},\{\hat{\mathbf{z}}_{t - \tau},c_{t}\}\right)$ .
Design transformation  $\mathbf{A} \rightarrow \mathbf{B}$  with low-triangular Jacobian as follows:
$$
\underbrace {\left[ \hat {\mathbf {z}} _ {t - L} , \dots , \hat {\mathbf {z}} _ {t - 1} , \hat {\mathbf {z}} _ {t} \right] ^ {\top}} _ {\mathbf {A}} \text {m a p p e d t o} \underbrace {\left[ \hat {\mathbf {z}} _ {t - L} , \dots , \hat {\mathbf {z}} _ {t - 1} , \hat {\epsilon} _ {i , t} \right] ^ {\top}} _ {\mathbf {B}}, w i t h \mathbf {J} _ {\mathbf {A} \rightarrow \mathbf {B}} = \left(\begin{array}{c c}\mathbb {I} _ {n L}&0\\*&\operatorname {d i a g} \left(\frac {\partial f _ {i , j} ^ {- 1}}{\partial \bar {z} _ {j t}}\right)\end{array}\right). \tag {21}
$$
Similar to Eq. 20, we can obtain the joint distribution of the estimated dynamics subspace as:
$$
\log p (\mathbf {A}) = \underbrace {\log p \left(\hat {\mathbf {z}} _ {t - L} , \dots , \hat {\mathbf {z}} _ {t - 1}\right) + \sum_ {j = 1} ^ {n} \log p \left(\hat {\epsilon} _ {i , t}\right)} _ {\text {B e c a u s e o f m u t u a l l y i n d e p e n d e n t n o i s e a s s u m p t i o n}} + \log \left(| \det \left(\mathbf {J} _ {\mathbf {A} \rightarrow \mathbf {B}}\right) |\right). \tag {22}
$$
$$
\log p \left(\hat {\mathbf {z}} _ {t} \mid \{\hat {\mathbf {z}} _ {t - \tau} \} _ {\tau = 1} ^ {L}, c _ {t}\right) = \sum_ {j = 1} ^ {n} \log p \left(\hat {\epsilon} _ {i, t} \mid c _ {t}\right) + \sum_ {i = j} ^ {n} \log \left| \frac {\partial f _ {i} ^ {- 1}}{\partial \hat {z} _ {i , t}} \right| \tag {23}
$$
B.3 Derivation of ELBO
Then the second part is to maximize the Evidence Lower BOund (ELBO) for the VAE framework, which can be written as:
$$
\begin{array}{l} \mathrm {E L B O} \triangleq \log p _ {\mathrm {d a t a}} (\mathbf {X}) - D _ {K L} \left(q _ {\phi} (\mathbf {Z} | \mathbf {X}) | | p _ {\mathrm {d a t a}} (\mathbf {Z} | \mathbf {X})\right) \\ = \mathbb {E} _ {\mathbf {Z} \sim q _ {\phi} (\mathbf {Z} | \mathbf {X})} \log p _ {\text {d a t a}} (\mathbf {X} | \mathbf {Z}) - D _ {K L} \left(q _ {\phi} (\mathbf {Z} | \mathbf {X}) \mid \mid p _ {\text {d a t a}} (\mathbf {Z} | \mathbf {X})\right) \\ = \mathbb {E} _ {\mathbf {Z} \sim q _ {\phi} (\mathbf {Z} | \mathbf {X})} \log p _ {\text {d a t a}} (\mathbf {X} | \mathbf {Z}) - \mathbb {E} _ {\mathbf {Z} \sim q _ {\phi} (\mathbf {Z} | \mathbf {X})} \left[ \log q _ {\phi} (\mathbf {Z} | \mathbf {X}) - \log p _ {\text {d a t a}} (\mathbf {Z}) \right] \\ = \mathbb {E} _ {\mathbf {Z} \sim q _ {\phi} (\mathbf {Z} | \mathbf {X})} \left[ \log p _ {\text {d a t a}} (\mathbf {X} | \mathbf {Z}) + \underbrace {\log p _ {\text {d a t a}} (\mathbf {Z})} _ {\mathbb {E} _ {\mathbf {c}} \left[ \sum_ {t = 1} ^ {T} \log p \left(\mathbf {z} _ {t} \mid \mathbf {z} _ {t - 1} , c _ {t}\right) \right]} - \log q _ {\phi} (\mathbf {Z} | \mathbf {X}) \right] \tag {24} \\ = \mathbb {E} _ {\mathbf {z} _ {t}} \left[ \underbrace {\sum_ {t = 1} ^ {T} \log p _ {\mathrm {d a t a}} (\mathbf {x} _ {t} | \mathbf {z} _ {t})} _ {- \mathcal {L} _ {\mathrm {R e c o n}}} + \underbrace {\mathbb {E} _ {\mathbf {c}} \left[ \sum_ {t = 1} ^ {T} \log p _ {\mathrm {d a t a}} (\mathbf {z} _ {t} | \mathbf {z} _ {\mathrm {H x}} , c _ {t}) \right] - \sum_ {t = 1} ^ {T} \log q _ {\phi} (\mathbf {z} _ {t} | \mathbf {x} _ {t})} _ {- \mathcal {L} _ {\mathrm {K L D}}} \right] \\ \end{array}
$$
B.4 Synthetic Dataset Generation
We generated two synthetic datasets (A and B) with different nonlinear mixing functions. In this section we will introduce the detailed implementation of the generation. The generation can be split into steps (1) sample  $c_{t}$  from a Markov chain, (2) generate  $\mathbf{z}_{t}$  with different transition functions  $f_{c_t}$  with respect to  $c_{t}$ , and (3) generate observation  $\mathbf{x}_t$  via mixing function  $\mathbf{g}$ .
B.4.1 Sample  $c_t$  from Markov chain
We first randomly initialized a Markov chain with transition matrix  $\mathbf{A}$  and sample 20,000 steps.
B.4.2 Generation of latent variables  $\mathbf{z}_t$
We first randomly initialized  $|C| = 5$  different transition functions  $\{f_1, f_2, \dots, f_{|C|}\}$  with different MLPs, and generate  $\mathbf{z}_t = f_{c_t}(\mathbf{z}_{\mathrm{Hx}})$ . The dimensions are set to 8 for fair comparison.
B.4.3 Generation of observations  $\mathbf{x}_t$
The difference between datasets A and B is the mixing function. We use a two-layer randomly initialized MLP for dataset A and a three-layer MLP for dataset B. For each linear layer in the MLP, we use condition number of the weight matrix to filter out ones that are not "invertible".
B.5 Modified CartPole Dataset Generation
Similar to the synthetic datasets, we also sample from a Markov chain and get  $c_{t}$ . For the modified CartPole, we initialized 5 different environments which have different combinations of hyperparameters such as gravity, pole mass, etc. A detailed comparison is listed in Table 1.
Table 1: Different config for different Modified CartPole environments.
Environment ID	Gravity	Pole Mass	Noise Scale
0	9.8	0.2	0.01
1	24.79	0.5	0.01
2	3.7	1.0	0.01
3	11.15	1.5	0.01
4	0.62	2.0	0.01
At each time step  $t$  the environment will load the corresponding hyperparameters for given  $c_{t}$  and update the states  $\mathbf{z}_t$  according to the configuration given  $c_{t}$ . The nonlinear mixing function from states to observations  $\mathbf{x}_t$  is fixed by a rendering method in the gym package.
B.6 MoSeq Dataset
In the MoSeq dataset, the observations  $\mathbf{x}_t$  are taken to be the first 10 principal components of depth camera video data of mice exploring an open field. The dataset consists of 20-minute depth camera recordings of 24 mice. In preprocessing, the videos are cropped and centered around the mouse centroid and then filtered to remove recording artifacts. Finally, the preprocessed video is projected onto the top principal components to obtain a 10-dimensional time series.
B.7 Mean Correlation Coefficient