Amortized Causal Discovery: Learning to Infer Causal Graphs from Time-Series Data

1 INTRODUCTION . Inferring causal relations in observational time-series is central to many fields of scientific inquiry ( Berzuini et al. , 2012 ; Spirtes et al. , 2000 ) . Suppose you want to analyze fMRI data , which measures the activity of different brain regions over time — how can you infer the ( causal ) influence of one brain region on another ? This question is addressed by the field of causal discovery ( Glymour et al. , 2019 ) . Methods within this field allow us to infer causal relations from observational data - when interventions ( e.g . randomized trials ) are infeasible , unethical or too expensive . In time-series , the assumption that causes temporally precede their effects enables us to discover causal relations in observational data ( Peters et al. , 2017 ) ; with approaches relying on conditional independence tests ( Entner and Hoyer , 2010 ) , scoring functions ( Chickering , 2002 ) , or deep learning ( Tank et al. , 2018 ) . All of these methods assume that samples share a single underlying causal graph and refit a new model whenever this assumption does not hold . However , samples with different underlying causal graphs may share relevant information such as the dynamics describing the effects of causal relations . fMRI test subjects may have varying brain connectivity but the same underlying neurochemistry ; social networks may have differing structure but comparable interpersonal relationships ; different stocks may relate differently to one another but obey similar market forces . Despite a range of relevant applications , inferring causal relations across samples with different underlying causal graphs is as of yet largely unexplored . In this paper , we propose a novel causal discovery framework for time-series that embraces this aspect : Amortized Causal Discovery ( Fig . 1 ) . In this framework , we learn to infer causal relations across samples with different underlying causal graphs but shared dynamics . We achieve this by separating the causal relation prediction from the modeling of their dynamics : an amortized encoder predicts the edges in the causal graph , and a decoder models the dynamics of the system under the predicted causal relations . This setup allows us to pool statistical strength across samples and to achieve significant improvements in performance with additional training data . It also allows us to infer causal relations in previously unseen samples without refitting our model . Additionally , we show that Amortized Causal Discovery allows us to improve robustness under hidden confounding by modeling the unobserved variables with the amortized encoder . Our contributions are as follows : • We formalize Amortized Causal Discovery ( ACD ) , a novel framework for causal discovery in time-series , in which we learn to infer causal relations from samples with different underlying causal graphs but shared dynamics . • We propose a variational model for ACD , applicable to multi-variate , non-linear data . • We present experiments demonstrating the effectiveness of this model on a range of causal discovery datasets , both in the fully observed setting and under hidden confounding . 2 BACKGROUND : GRANGER CAUSALITY . Granger causality ( Granger , 1969 ) is one of the most commonly used approaches to infer causal relations from observational time-series data . Its central assumption is that causes precede their effects : if the prediction of the future of time-series Y can be improved by knowing past elements of time-series X , then X “ Granger causes ” Y . Originally , Granger causality was defined for linear relations ; we follow the more recent definition of Tank et al . ( 2018 ) for non-linear Granger causality : Definition 2.1 . Non-Linear Granger Causality : Given N stationary time-series x = { x1 , ... xN } across time-steps t = { 1 , ... , T } and a non-linear autoregressive function gj , such that xt+1j = gj ( x ≤t 1 , ... , x ≤t N ) + ε t j , ( 1 ) where x≤tj = ( ... , x t−1 j , x t j ) denotes the present and past of series j and ε t j represents independent noise . In this setup , time-series i Granger causes j , if gj is not invariant to x ≤t i , i.e . if ∃ x′≤ti 6= x ≤t i : gj ( x ≤t 1 , ... , x ′≤t i , ... , x ≤t N ) 6= gj ( x ≤t 1 , ... , x ≤t i , ... x ≤t N ) . Granger causal relations are equivalent to causal relations in the underlying directed acyclic graph if all relevant variables are observed and no instantaneous1 connections exist ( Peters et al. , 2013 ; 2017 , Theorem 10.1 ) . Many methods for Granger causal discovery , including vector autoregressive ( Hyvärinen et al. , 2010 ) and more recent deep learning-based approaches ( Khanna and Tan , 2020 ; Tank et al. , 2018 ; Wu et al. , 2020 ) , can be encapsulated by a particular framework : 1 . Define a function fθ ( an MLP in Tank et al . ( 2018 ) , a linear model in Hyvärinen et al . ( 2010 ) ) , which learns to predict the next time-step of the test sequence x . 2 . Fit fθ to x by minimizing some loss L : θ ? = argminθ L ( x , fθ ) . 3 . Apply some fixed function h ( e.g . thresholding ) to the learned parameters to produce the Granger causal graph estimate for x : Ĝx = h ( θ ? ) . For instance , Tank et al . ( 2018 ) infer the Granger causal relations through examination of the weights θ ? : if all outgoing weights wij between time-series i and j are zero , then i does not Granger-cause j . 1connections between two variables at the same time step The shortcoming of this approach is that , when we have S samples x1 , . . . , xS with different underlying causal graphs , the parameters θ must be optimized separately for each of them . As a result , methods within this framework can not take advantage of the information that might be shared between samples . This motivates us to question : can we amortize this process ? 3 AMORTIZED CAUSAL DISCOVERY . We propose Amortized Causal Discovery ( ACD ) , a framework in which we learn to infer causal relations across samples with different underlying causal graphs but shared dynamics . To illustrate : Suppose you want to infer synaptic connections ( i.e . causal relations ) between neurons based on their spiking behaviour . You are given a set of S recordings ( i.e . samples ) , each containing N time-series representing the firing of N individual neurons . Even though you might record across different populations of neurons with different wiring , the dynamics of how neurons connected by synapses influence one another stays the same . ACD takes advantage of such shared dynamics to improve the prediction of causal relations . It can be summarized as follows : 1 . Define an encoding function fφ which learns to infer Granger causal relations of any sample xi in the training setXtrain . Define a decoding function fθ which learns to predict the next time-step of the samples under the inferred causal relations . 2 . Fit fφ and fθ toXtrain by minimizing some loss L : fφ ? , fθ ? = argminfφ , fθ L ( Xtrain , fφ , fθ ) . 3 . For any given test sequence xtest , simply output the Granger causal graph estimate Ĝxtest : Ĝxtest = fφ ? ( xtest ) . By dividing the model into two parts , an encoder and a decoder , ACD can use the activations of fφ ? to infer causal structure . This increases the flexibility of our approach greatly compared to methods that use the learned weights θ ? such as the prior Granger causal discovery methods described in Section 2 . In this section , we describe our framework in more detail , and provide a probabilistic implementation thereof . We also extend our approach to model hidden confounders . Preliminaries We begin with a datasetX = { xs } Ss=1 of S samples , where each sample xs consists of N stationary time-series xs = { xs,1 , . . . , xs , N } across timesteps t = { 1 , ... , T } . We denote the t-th time-step of the i-th time-series of xs as xts , i.We assume there is a directed acyclic graph G1 : Ts = { V1 : Ts , E1 : Ts } underlying the generative process of each sample . This is a structural causal model ( SCM ) ( Pearl , 2009 ) . Its endogenous ( observed ) variables are vertices vts , i ∈ V1 : Ts for each time-series i and each time-step t. Every set of incoming edges to an endogenous variable defines inputs to a deterministic function gts , i which determines that variable ’ s value 2 . The edges are defined by ordered pairs of vertices E1 : Ts = { ( vts , i , vt ′ s , j ) } , which we make two assumptions about : 1 . No edges are instantaneous ( t = t′ ) or go back in time . Thus , t < t′ for all edges . 2 . Edges are invariant to time . Thus , if ( vts , i , v t+k s , j ) ∈ E1 : Ts , then ∀1 ≤ t′ ≤ T − k : ( vt ′ s , i , v t′+k s , j ) ∈ E1 : Ts . The associated structural equations gts , i are invariant to time as well , i.e . gts , i = gt ′ s , i ∀t , t′ . The first assumption states that causes temporally precede their effects and makes causal relations identifiable from observational data , when no hidden confounders are present ( Peters et al. , 2013 ; 2017 , Theorem 10.1 ) . The second simplifies modeling : it is a fairly general assumption which allows us to define dynamics that govern all time-steps ( Eq . ( 2 ) ) . Throughout this paper , we are interested in discovering the summary graph Gs = { Vs , Es } ( Peters et al. , 2017 ) . It consists of vertices vs , i ∈ Vs for each time-series i in sample s , and has directed edges whenever they exist in E1 : Ts at any time-step , i.e . Es = { ( vs , i , vs , j ) | ∃t , t′ : ( vts , i , vt ′ s , j ) ∈ E1 : Ts } . Note that while G1 : Ts is acyclic ( due to the first assumption above ) , the summary graph Gs may contain ( self- ) cycles . Amortized Causal Discovery The key assumption for Amortized Causal Discovery is that there exists some fixed function g that describes the dynamics of all samples xs ∈ X given their past 2The SCM also includes an exogenous ( unobserved ) , independently-sampled error variable v as a parent of each vertex v , which we do not model and thus leave out for brevity . observations x≤ts and their underlying causal graph Gs : xt+1s = g ( x ≤t s , Gs ) + εts . ( 2 ) There are two variables in this data-generating process that we would like to model : the causal graph Gs that is specific to sample xs , and the dynamics g that are shared across all samples . This separation between the causal graph and the dynamics allows us to divide our model accordingly : we introduce an amortized causal discovery encoder fφ which learns to infer a causal graph Gs given the sample xs , and a dynamics decoder fθ that learns to approximate g : xt+1s ≈ fθ ( x≤ts , fφ ( xs ) ) . ( 3 ) We formalize Amortized Causal Discovery ( ACD ) as follows . Let G be the domain of all possible summary graphs on xs : Gs ∈ G. Let X be the domain of any single step , partial or full , observed sequence : xts , x ≤t s , xs ∈ X . The model consists of two components : a causal discovery encoder fφ : X → G which infers a causal graph for each input sample , and a decoder fθ : X × G → X which models the dynamics . This model is optimized with a sample-wise loss ` : X×X→ R which scores how well the decoder models the true dynamics of xs , and a regularization term r : G→ R on the inferred graphs . For example , this function r may enforce sparsity by penalizing graphs with more edges . Note , that our formulation of the graph prediction problem is unsupervised : we do not have access to the true underlying graph Gs . Then , given some dataset Xtrain with S samples , we optimize : fφ ? , fθ ? = argminfφ , fθ L ( Xtrain , fφ , fθ ) ( 4 ) where L ( Xtrain , φ , θ ) = S∑ s=1 T−1∑ t=1 ` ( xt+1s , fθ ( x ≤t s , fφ ( xs ) ) ) + r ( fφ ( xs ) ) . ( 5 ) See Appendix B for a proof of the consistency of the loss ` and a discussion on regularization r. Once we have completed optimization , we can perform causal graph prediction on any new input test sample xtest in two ways – we can feed xtest into the amortized encoder and take its output as the predicted edges ( Eq . 6 ) ; or we can instantiate our estimate Ĝtest ∈ G which will be our edge predictions , and find the edges which best explain the observed sequence xtest by minimizing the ( learned ) decoding loss with respect to Ĝtest , which we term Test-Time Adaptation ( TTA ) ( Eq . 7 ) : ĜEnc = fφ ? ( xtest ) ; ( 6 ) ĜTTA = argminĜtest∈G L ( xtest , Ĝtest , fθ ? ) . ( 7 ) By separating the prediction of causal relations from the modeling of their dynamics , ACD yields a number of benefits . ACD can learn to infer causal relations across samples with different underlying causal graphs , and it can infer causal relations in previously unseen test samples without refitting ( Eq . ( 6 ) ) . By generalizing across samples , it can improve causal discovery performance with increasing training data size . We can replace either fφ or fθ with ground truth annotations , or simulate the outcome of counterfactual causal relations . Additionally , ACD can be applied in the standard causal discovery setting , where only a single causal graph underlies all samples , by replacing the amortized encoder fφ with an estimated graph Ĝ ( or distribution over G ) in Eq . ( 4 ) .

Summary: The authors proposed a framework called Amortized Causal Discovery (ACD) for recovering causal relationships in time series where samples are generated from models with different underlying causal graphs but shared dynamics. This framework is applicable in settings such as modeling neural spiking trains where the dynamics of how neurons react to the activities of other neurons remain the same. In the proposed framework, they considered a causal discovery encoder $f_{\theta}$, which tries to extract causal relationships and map it to a latent space. Moreover, there is a dynamic decoder $f_{\phi}$, which provides one-step predictions. The proposed architecture is similar to a variational auto-encoder and the encoder part is based on graph neural networks.