Toward Temporal Causal Representation Learning with Tensor Decomposition

TL;DR

This paper introduces CaRTeD, a novel framework combining temporal causal representation learning with irregular tensor decomposition, providing theoretical convergence guarantees and improved performance on EHR data.

Contribution

It proposes a new causal formulation for latent clusters and a joint learning framework that enhances tensor decomposition with causal insights, filling theoretical gaps.

Findings

Outperforms state-of-the-art methods on synthetic and real EHR data

Provides convergence guarantees for the proposed algorithm

Improves interpretability of causal representations in complex data

Abstract

Temporal causal representation learning is a powerful tool for uncovering complex patterns in observational studies, which are often represented as low-dimensional time series. However, in many real-world applications, data are high-dimensional with varying input lengths and naturally take the form of irregular tensors. To analyze such data, irregular tensor decomposition is critical for extracting meaningful clusters that capture essential information. In this paper, we focus on modeling causal representation learning based on the transformed information. First, we present a novel causal formulation for a set of latent clusters. We then propose CaRTeD, a joint learning framework that integrates temporal causal representation learning with irregular tensor decomposition. Notably, our framework provides a blueprint for downstream tasks using the learned tensor factors, such as modeling latent structures and extracting causal information, and offers a more flexible regularization design to enhance tensor decomposition. Theoretically, we show that our algorithm converges to a stationary point. More importantly, our results fill the gap in theoretical guarantees for the convergence of state-of-the-art irregular tensor decomposition. Experimental results on synthetic and real-world electronic health record (EHR) datasets (MIMIC-III), with extensive benchmarks from both phenotyping and network recovery perspectives, demonstrate that our proposed method outperforms state-of-the-art techniques and enhances the explainability of causal representations.