SlowGuess/NeurIPS_2023_OCR huggingface dataset:
which can easily be computed by the Forward-Backward algorithm and luckily all of it is differentiable to the HMM training parameters  $\theta_{\mathrm{HMM}}$  (transition matrix  $\mathbf{A}$  and transition function parameters  $\theta_f$ ).
Then the second part is to maximize the Evidence Lower BOund (ELBO) for the VAE framework, which can be written as (complete derivation steps are in Appendix B.3):
$$
\begin{array}{l} \operatorname {E L B O} \triangleq \log p _ {\text {d a t a}} (\mathbf {X}) - D _ {K L} \left(q _ {\phi} (\mathbf {Z} | \mathbf {X}) | | p _ {\text {d a t a}} (\mathbf {Z} | \mathbf {X})\right) \\ = \underbrace {\mathbb {E} _ {\mathbf {z} _ {t}} \sum_ {t = 1} ^ {T} \log p _ {\text {d a t a}} (\mathbf {x} _ {t} | \mathbf {z} _ {t})} _ {- \mathcal {L} _ {\text {R e c o n}}} + \underbrace {\mathbb {E} _ {\mathbf {c}} \left[ \sum_ {t = 1} ^ {T} \log p _ {\text {d a t a}} (\mathbf {z} _ {t} | \mathbf {z} _ {\mathrm {H x}} , c _ {t}) - \sum_ {t = 1} ^ {T} \log q _ {\phi} (\mathbf {z} _ {t} | \mathbf {x} _ {t}) \right]} _ {- \mathcal {L} _ {\mathrm {K L D}}} \tag {13} \\ \end{array}
$$
We use mean-squared error (MSE) for the reconstruction likelihood loss  $\mathcal{L}_{\mathrm{Recon}}$ . The KL divergence  $\mathcal{L}_{\mathrm{KLD}}$  is estimated via a sampling approach since with a learned nonparametric transition prior, the distribution does not have an explicit form. Specifically, we obtain the log-likelihood of the posterior, evaluate the prior  $\log p(\hat{\mathbf{z}}_t|\hat{\mathbf{z}}_{\mathrm{Hx}},c_t)$  in Eq. (9), and compute their mean difference in the dataset as the KL loss:  $\mathcal{L}_{\mathrm{KLD}} = \mathbb{E}_{\hat{\mathbf{z}}_t\sim q(\hat{\mathbf{z}}_t|\mathbf{x}_t)}\log q(\hat{\mathbf{z}}_t|\mathbf{x}_t) - \log p(\hat{\mathbf{z}}_t|\hat{\mathbf{z}}_{\mathrm{Hx}},c_t)$ .
We evaluate the identifiability results of NCTRL on a number of simulated and real-world temporal datasets. We first introduce the evaluation metrics and baselines and then discuss the datasets we used in our experiments. Lastly, we show the experiment results discuss the performance, and make comparisons.
Mean Correlation Coefficient (MCC) To evaluate the identifiability of the latent variables, we compute the Mean Correlation Coefficient (MCC) on the test dataset. MCC is a standard metric in the ICA literature for continuous variables which measure the identifiability of the learned latent causal processes. MCC is close to 1.0 when latent variables are identifiable up to permutation and component-wise invertible transformation in the noiseless case.
Mean Square Error (MSE) for estimating A As introduced in the theory, the A is identifiable in our setting, which means that our proposed method can provide accurate estimation for the transition matrix A, to valid such a claim, we use mean square error (MSE) to capture the distance between the estimated  $\hat{\mathbf{A}}$  and ground truth A.
Accuracy for estimating  $c_t$  We also test the accuracy for estimating the discrete domain indices  $c_t$  supplementary to the MSE for  $\mathbf{A}$  since in theory, the  $\mathbf{A}$  is identifiable but the  $c_t$  is generally not identifiable, which is relatively easy to understand as an analogy in Hidden Markov Models, the transition matrix is identifiable but we can only "infer" the best possible discrete variables but cannot establish identifiability for it.
It is also worth mentioning that the MSE and Accuracy are influenced by the permutation, which is also true in clustering evaluation problems. Here we explored all permutations and selected the best possible assignment for evaluation.
The following identifiable nonlinear ICA methods are used: (1) BetaVAE [22] which ignores both history and nonstationarity information. (2) i-VAE [12] and TCL [9] which leverage nonstationarity to establish identifiability but assume independent factors. (3) SlowVAE [16], and PCL [10] which exploit temporal constraints but assume independent sources and stationary processes. (4) TDRL [18] which assumes nonstationary, causal processes but with observed domain indices. (5) HMNLICA [14] which considers the unobserved nonstationary part in the data generation process but doesn't allow any causally related time-delayed relations.
5.3 Simulated Results
We generate two synthetic datasets corresponding to different complexity of the nonlinear mixing function g. Both synthetic datasets satisfy our identifiability conditions in the theorems following the procedures in Appendix B.4. As in Table 1, NCTRL can recover the latent processes under unknown nonstationary distribution shifts with high MCCs  $(>0.95)$ . The baselines that do not exploit history (i.e., BetaVAE, i-VAE, TCL), with independent source assumptions (SlowVAE, PCL), consider limited nonstationary cases (TDRL) distort the identifiability results. The only baseline that considers the unknown nonstationarity in the domain indices (HMN
Table 1: Experiment results of two synthetic datasets on baselines and proposed NCTRL, we run the experiments with five different random seeds and calculate the average with standard derivation. The best results are shown in bold.
Method	Mean Correlation Coefficien (MCC)
Dataset A	Dataset B	Ave.
BetaVAE	44.02 ± 3.11	47.48 ± 10.58	45.75
i-VAE	89.74 ± 3.38	44.50 ± 0.25	67.12
TCL	37.12 ± 0.60	56.33 ± 3.77	46.73
SlowVAE	33.84 ± 0.60	53.92 ± 3.56	43.88
PCL	42.41 ± 2.87	63.66 ± 2.77	53.04
HMNLICA	59.82 ± 4.94	57.25 ± 1.45	58.54
TDRL	83.99 ± 1.92	72.02 ± 2.76	78.01
NCTRL	98.85 ± 0.30	99.01 ± 0.24	98.93
LICA) explored the Markov Assumption but doesn't allow a time-delayed causal process and hence suffers a poor result (MCC 0.58).
On the other hand, the difference between dataset A and dataset B is the nonlinearity in the mixing function, dataset A has a relatively simple nonlinear mixing function, on the contrary, dataset B has more complex nonlinearity. Some variability has been observed among the relative performance ranks of different baselines. For example, i-VAE showed a great discrepancy between the two datasets, which revived the weakness of capturing complex nonlinearity in the unknown nonstationary distribution shift environments. Again we also observed that our proposed method can constantly recover the latent independent components with high MCC which indicates on both datasets the model is identifiable and the estimation algorithm is highly effective. To further validate if NCTRL successfully recovered the Markov transition matrix  $\mathbf{A}$  and inferred the domain indices  $c_{t}$  with high accuracy. We further examine the accuracy for estimating nonstationary domain indices  $c_{t}$  and the mean square error estimating the transition matrix  $\mathbf{A}$ . As shown in Table 2 the result is consistent with our theory in which the transition matrix  $\mathbf{A}$  is identifiable and we can estimate it with
Table 2: Supplementary experiment results of two synthetic datasets on estimating domain indices  $c_{t}$  and transition matrix  $\mathbf{A}$  in NCTRL, we run the experiments with five different random seeds and calculate the average with standard derivation.
Unknown Nonstationary Metrics
Dataset	Accuracy estimating ct	MSE estimating A
A	89.96 ± 0.24	1.01 ×10-3± 1.67 ×10-4
B	89.84 ± 0.29	1.08 ×10-3± 1.89 ×10-4
high accuracy. For the nonstationary domain indices  $c_{t}$  even though there is no identifiability result governing the estimation accuracy, it can still be inferred pretty well since it is nothing but a decoding problem in Hidden Markov Models.
5.4 Real-world Applications
Video data - Modified CartPole Environment We evaluate NCTRL on the modified CartPole [23] video dataset and compare the performances with the baselines. Modified Cartpole is a nonlinear dynamical system with cart positions  $x_{t}$  and pole angles  $\theta_{t}$  as the true state variables. The dataset descriptions are in Appendix B.5. Similar to the synthetic dataset, we randomly initialize a Markov chain and roll out a series of  $c_{t}$ , and configure the CartPole environment with respect to the  $c_{t}$ . Specifically, we use five domains with different configurations of cart mass, pole mass, gravity, and noise levels. Together with the two discrete actions (i.e., left and right). By doing so, the nonstationarity is enforced, and since we can control and access all intermediate states in the system, all metrics including MCC and  $c_{t}$  accuracy together with A MSE can be easily calculated. We fit NCTRL with two-dimensional causal factors. We set the latent size  $n = 2$  and the lag number  $L = 2$ . In Fig. 3, the latent causal processes are recovered, as seen from (a) high MCC for the latent causal processes; (b) the latent factors are estimated up to component-wise transformation; and (c) the latent traversals confirm the two recovered latent variables correspond to the position and pole angle.
(b)
Figure 3: Modified Cartpole results: (a) MCC for causally-related factors; (b) scatterplots between estimated and true factors; and (c) latent traversal on a fixed video frame
Table 3: Experiment results of CartPole dataset. The best results are shown in bold.
Mean Correlation Coefficient (MCC)
BetaVAE	i-VAE	TCL	SlowVAE	SKD	TDRL	NCTRL
57.54	60.14	65.07	63.16	73.24	85.26	96.06
Similar to Table 1 and 2, we compare our NCTRL with baseline methods. In addition, we also compare with SKD [24], a state-of-the-art sequential disentangle representation learning method without identifiability guarantee. In Table 3 and 4 we can see that compared with TDRL, our NCTRL can recover the latent processes under unknown nonstationary distribution shifts with high MCCs ( $>0.95$ ) with highly accurate estimated transition matrix  $\mathbf{A}$  and high quality inferred  $c_{t}$ . Specifically, by comparing the result of SKD, the MCC for SKD is better than a variety of baselines,
Table 4: Supplementary experiment results of CartPole datasets on estimating domain indices  $c_{t}$  and transition matrix A in NCTRL, we run the experiments with five different random seeds and calculate the average with standard derivation.
Unknown Nonstationary Metrics
Accuracy estimating ct	MSE estimating A
79.23 ± 5.33	5.01 ×10-2± 1.23 ×10-2
however, we can see the distinction between well-disentangled models and identifiable models, only the models with identifiability can find the ground truth latent variables with theoretical guarantee.
Video data - MoSeq Dataset We test NCTRL framework to analyze mouse behavior video data from Wiltschko et al. [19], which represents the original application to clustering mouse behavior<sup>3</sup>, details of this dataset are available in Appendix B.6. Since there are no ground truth independent components in this particular real-world dataset, we analyze it by several visualizations to see if different domains can be properly identified and if the patterns in the recovered independent components are consistent with the recovered domain indices. We analyze the first video clip of mouse behavior data and visualize the two phases we discovered and segmented in Fig 4. We can clearly see from Fig 4 that there are different phases with the upper one actively moving and the lower one inactive. The recovered independent components showed a consistent pattern with the recovered phase or domain indices as shown in Fig 4.
Figure 4: Result visualization of MoSeq dataset. (Active, Inactive) show two representative video frames for the active and inactive phases and (Independent Components) visualize the discovered independent components with corresponding phases tagged with different colors.
Causal Discovery from Time Series Understanding the causal structure in time-series data is pivotal in areas such as machine learning [1], econometrics [2], and neuroscience [3]. A bulk of the research in this realm emphasizes determining the temporal causal links among observed variables. The primary techniques employed are constraint-based methods [25], which use conditional independence tests to ascertain causal structures, and score-based methods [26, 27], where scores are utilized to oversee a search operation. Some researchers also proposed a combination of these two methods [28, 29]. Additionally, Granger causality [30] and its nonlinear adaptations [31, 32] have gained widespread acceptance in this context.
3Dataset can be accessed via https://dattalab.github.io/moseq2-website/index.html
Nonlinear ICA for Time Series Recently, the significance of temporal structures and nonstationarities has been recognized in achieving identifiability within nonlinear ICA. Time-contrastive learning (TCL [9]) utilizes the independent sources principle, focusing on data segments' variability. On the other hand, Permutation-based contrastive (PCL [10]) offers a learning approach that distinguishes true independent sources from shuffled ones under the uniformly dependent assumption. HMNLICA [14] integrates nonlinear ICA with an HMM to address non-stationarity without segmenting data manually. The i-VAE [12] approach employs VAEs to capture the actual joint distribution between observed and auxiliary non-stationary domains, assuming an exponential families conditional distribution. The recent advancements in nonlinear ICA for time series include LEAP [17], (i-)CITRIS [33, 34], and TDRL [18]. While LEAP introduces a novel condition emphasizing nonstationary noise, TDRL delves deeper into a non-parametric environment within a nonstationary context. In contrast, CITRIS recommends utilizing intervention target data for pinpointing latent causal aspects, avoiding certain constraints but necessitating active intervention access.
Sequential Disentanglement Majority of existing work about sequential disentanglement focuses on architecture based on dynamical variational autoencoder (VAE) [35]. Early works [36, 37] separate dynamic factors from static factors using probabilistic methods. Then auxiliary tasks with self-supervisory signals [38] were introduced. C-DSVAE [39] utilized contrastive penalty terms with data augmentation to introduce additional inductive biases. In R-WAE [40], Wasserstein distance was introduced to replace KL divergence. To deal with video disentanglement [41, 42] explored generative adversarial network (GAN) architectures and [43] introduced a recurrent model with adversarial loss. FAVAE, [44] proposed a factorizing VAE and [45] proposed to learn hierarchical features. Finally, SKD [24] introduced a spectral loss term that leads to structured Koopman matrices and disentanglement.
Conclusion. In this paper, we first established an identifiability theory for general sequential data with nonstationary causally-related processes under unknown distribution shifts. Then we presented NCTRL, a principled framework to recover the time-delayed latent causal variable identify their causal relations from measured data, and decode high-quality domain indices under Markov assumption. Experiment results on both synthetic datasets and real-world video datasets showed that our proposed method can recover the latent causal variables and their causal relations purely from measured data with the observation of auxiliary variables or domain indices.
**Limitation.** The basic limitation of this work is that the nonstationary domain indices are assumed to follow a Markov chain. Also, this work highly relies on the latent processes to have no instantaneous causal relations but only time-delayed influences. If the resolution of the time series is much lower, then it is usually violated and one has to find a way to deal with instantaneous causal relations. Extending our theories and framework to address the scenarios when more flexibility in the domain indices transition is allowed (i.e. beyond discrete variables following a Markov chain) and to address instantaneous dependency or instantaneous causal relations will be some of our future work.
Boarder Impacts. This work proposes a theoretical analysis and technical methods to learn the causal representation from time-series data, which facilitate the construction of more transparent and interpretable models to understand the causal effect in the real world. This could be beneficial in a variety of sectors, including healthcare, finance, and technology. In contrast, misinterpretations of causal relationships could also have significant negative implications in these fields, which must be carefully done to avoid unfair or biased predictions.
This project has been graciously funded by NGA HM04762010002, NSF IIS1955532, NSF CNS2008248, NIGMS R01GM140467, NSF IIS2123952, NSF BCS2040381, an Amazon Research Award, NSF IIS2311990, and DARPA ECOLE HR00112390063. This project is also partially supported by NSF Grant 2229881, the National Institutes of Health (NIH) under Contract R01HL159805, a grant from Apple Inc., a grant from KDDI Research Inc., and generous gifts from Salesforce Inc., Microsoft Research, and Amazon Research.
[35] Laurent Girin, Simon Leglaive, Xiaoyu Bie, Julien Diard, Thomas Hueber, and Xavier Alameda-Pineda. Dynamical Variational Autoencoders: A Comprehensive Review. 2021.
[36] Wei-Ning Hsu, Yu Zhang, and James Glass. Unsupervised learning of disentangled and interpretable representations from sequential data. Advances in neural information processing systems, 30, 2017.
[37] Yingzhen Li and Stephan Mandt. Disentangled sequential autoencoder. arXiv preprint arXiv:1803.02991, 2018.
[38] Yizhe Zhu, Martin Renqiang Min, Asim Kadav, and Hans Peter Graf. S3VAE: Self-supervised sequential VAE for representation disentanglement and data generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6538-6547, 2020.
[39] Junwen Bai, Weiran Wang, and Carla P Gomes. Contrastively disentangled sequential variational autoencoder. Advances in Neural Information Processing Systems, 34, 2021.
[40] Jun Han, Martin Renqiang Min, Ligong Han, Li Erran Li, and Xuan Zhang. Disentangled recurrent wasserstein autoencoder. In International Conference on Learning Representations, 2021. URL https://openreview.net/forum?id=07ms4LFdsX.
[41] Ruben Villegas, Jimei Yang, Seunghoon Hong, Xunyu Lin, and Honglak Lee. Decomposing motion and content for natural video sequence prediction. In International Conference on Learning Representations, 2017. URL https://openreview.net/forum?id=rkEFLFqee.
[42] Sergey Tulyakov, Ming-Yu Liu, Xiaodong Yang, and Jan Kautz. Mocogan: Decomposing motion and content for video generation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 1526-1535, 2018.
[43] Emily L Denton and Vighnesh Birodkar. Unsupervised learning of disentangled representations from video. Advances in neural information processing systems, 30, 2017.
[44] Masanori Yamada, Heechel Kim, Kosuke Miyoshi, Tomoharu Iwata, and Hiroshi Yamakawa. Disentangled representations for sequence data using information bottleneck principle. In Asian Conference on Machine Learning, pages 305-320. PMLR, 2020.
[45] Shengjia Zhao, Jiaming Song, and Stefano Ermon. Learning hierarchical features from deep generative models. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 4091-4099. PMLR, 06-11 Aug 2017. URL https://proceedings.mlr.press/v70/zhao17c.html.
[46] Elizabeth S. Allman, Catherine Matias, and John A. Rhodes. Identifiability of parameters in latent structure models with many observed variables. Annals of Statistics, 37:3099-3132, 2008.
[47] Alessio Spantini, Daniele Bigoni, and Youssef Marzouk. Inference via low-dimensional couplings. The Journal of Machine Learning Research, 19(1):2639-2709, 2018.
"Temporally Disentangled Representation Learning under Unknown Nonstationarity"
Appendix organization:
A Identifiability 15
A.1 Identifiability of Nonstationary Hidden States 15
A.2 Identifiability of Latent Causal Processes 16
A.3 Discussion on Assumptions in Theorem 2 18
B Implementation Details 18
B.1 Reproducibility 18
B.2 Prior Likelihood Derivation 18
B.3 Derivation of ELBO 19
B.4 Synthetic Dataset Generation 20
B.4.1 Sample  $c_{t}$  from Markov chain 20
B.4.2 Generation of latent variables  $\mathbf{z}_t$  20
B.4.3 Generation of observations  $\mathbf{x}_t$  20
B.5 Modified CartPole Dataset Generation 20
B.6 MoSeq Dataset 21
B.7 Mean Correlation Coefficient 21
B.8 Network Architecture 21
A Identifiability
Assume we observe  $n$ -dimensional time-series data at discrete time steps,  $\mathbf{X} = \{\mathbf{x}_1,\mathbf{x}_2,\dots ,\mathbf{x}_T\}$ , where each  $\mathbf{x}_t$  is generated from time-delayed causally related hidden components  $\mathbf{z}_t\in \mathbb{R}^n$  by the invertible mixing function:
In addition to latent components  $\mathbf{z}_t$ , there is an extra hidden component  $c_t$  which is a discrete variable with cardinality  $|c_t| = C$ , it follows first-order Markov process controlled by a  $C \times C$  transition matrix  $\mathbf{A}$ , in which the  $i, j$ -th entry  $A_{i,j}$  is the probability to transit from state  $i$  to  $j$ .
$$
z _ {i t} = f _ {i} \left(\left\{z _ {j, t - 1} \mid z _ {j, t - \tau} \in \mathbf {P a} \left(z _ {i t}\right) \right\}, c _ {t}, \epsilon_ {i t}\right) \quad \text {w i t h} \quad \epsilon_ {i t} \sim p _ {\epsilon_ {i} | c _ {t}} \tag {3}
$$
A.1 Identifiability of Nonstationary Hidden States
Theorem 1. (identifiability of the nonstationarity with Markov Assumptions) Suppose the observed data is generated following the nonlinear ICA framework as defined in Eqs. (1), (2) and (3). And Suppose the following assumptions (Markov Assumptions) hold:
Use  $\mu_1, \ldots, \mu_C \in \mathbb{R}^n$  to denote nonparametric probability distributions of the  $C$  emission distributions  $\mu_c = p(\mathbf{x}_t | \mathbf{x}_{t-1}, c)$ . Then the parameters  $\mathbf{A}$  and  $M = (\mu_1, \ldots, \mu_C)$  are identifiable given the distribution,  $\mathbb{P}_{\mathbf{A}, M}^{(3)}$ , of at least 4 consecutive observations  $\mathbf{x}_t, \mathbf{x}_{t+1}, \mathbf{x}_{t+2}, \mathbf{x}_{t+3}$ , up to label swapping of the hidden states, that is:
If  $\widetilde{\mathbf{A}}$  is a  $C\times C$  transition matrix, if  $\widetilde{\pi} (c)$  is a stationary distribution of  $\widetilde{\mathbf{A}}$  with  $\widetilde{\pi} (c) > 0\forall c\in$ $\{1,\ldots ,C\}$ , and if  $\tilde{M} = (\tilde{\mu}_1,\dots ,\tilde{\mu}_C)$  are  $C$  probability distributions on  $\mathbb{R}^n$  that verify the equality of the distribution functions  $\mathbb{P}_{\widetilde{\mathbf{A}},\tilde{M}}^{(3)} = \mathbb{P}_{\mathbf{A},M}^{(3)}$ , then there exists a permutation  $\sigma$  of the set  $\{1,\dots ,C\}$  such that for all  $k,l = 1,\dots ,C$  we have  $\tilde{A}_{k,l} = A_{\sigma (k),\sigma (l)}$  and  $\tilde{\mu}_k = \mu_{\sigma (k)}$ .
Proof. Suppose we have:
$$
\tilde {p} \left(\mathbf {x} _ {1}, \dots , \mathbf {x} _ {T}\right) = p \left(\mathbf {x} _ {1}, \dots , \mathbf {x} _ {T}\right) \tag {4}
$$
where  $p(\mathbf{x}_1,\dots ,\mathbf{x}_T)$  has transition matrix  $\mathbf{A}$  and emission distributions  $(\mu_{1},\ldots ,\mu_{C})$  , similarly for  $\tilde{p} (\mathbf{x}_1,\dots ,\mathbf{x}_T)$
We consider four consecutive observations  $\mathbf{x}_0, \mathbf{x}_1, \mathbf{x}_2, \mathbf{x}_3$  and corresponding four discrete elements  $c_0, c_1, c_2, c_3$ .