Scaling up Continuous-Time Markov Chains Helps Resolve Underspecification

1 Introduction . Modeling the time evolution of physical processes with discrete states is an important machine learning problem that spans a wide range of application domains , including molecular dynamics ( Crommelin & Vanden-Eijnden , 2006 ) , phylogenetics ( Suchard et al. , 2001 ) , and computational medicine ( Liu et al. , 2015 ) . Continuous-time markov chains have found success in such modeling tasks , supported by extensive research into learning and inference in such models ( Opper & Sanguinetti , 2008 ; Perkins , 2009 ; Archambeau & Opper , 2011 ; Rao & Teh , 2012 ) . While the majority of previous work has focused on modeling a relatively small number of discrete states , many interesting applications involve the interaction between sets of items , which results in exponentially large state spaces . For example , when modeling the time progression of cancer , we want to consider a number of genetic alterations ( e.g. , mutations , copy number variations , etc . ) that exhibit complex time-related dependencies ( Beerenwinkel et al. , 2014 ) . To model how sets of n such alterations evolve over time , we need to consider a state space of size 2n . What makes the problem even more challenging is the fact that the available data sets are typically cross-sectional ; that is , they consist of ( unordered ) sets of items observed at unknown time points without any further information about the history of the underlying process . In this paper , we focus on the problem of learning a particular parametric family of continuous-time Markov chains from such cross-sectional data , and show that the resulting problem is in general underspecified . The issue of underspecification is many-faceted and has been shown to permeate a 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . large number of machine learning systems , often leading to poor generalization , lack of robustness , and spurious relationships when interpreting the resulting models ( D ’ Amour et al. , 2020 ) . In our setting , we explore a perhaps surprising remedy : including a number of additional ( approximately ) independent items can help determine the time order of process events , and hence resolve underspecification . We theoretically show that these extra items act as a “ background clock ” , since counting the number of their occurrences in a data sample can help us estimate the time at which that sample was observed . In sharp contrast to common practice , which limits the analysis to a small subset of relevant , highly-interacting items , our insight suggests that scaling up the learning procedure can be crucial for the robustness of the inferred models . We thus show that items deemed a priori unimportant to the application at hand may in fact be particularly valuable for recovering the time properties of the underlying physical process . Existing learning approaches are , unfortunately , not well-suited for practically applying this insight ; for example , the state-of-the-art method by Schill et al . ( 2019 ) for learning continuous-time Markov chains to model cancer progression , scales exponentially in the number of items considered , thus limiting the analysis to around 20 items . To alleviate this issue , we propose an approximate likelihood maximization method that relies on a fast gradient approximation . On problems of 20 items our approach runs almost 1000 times faster that the state of the art , while it can also scale to problems involving hundreds of items . In experiments on real cancer data , we demonstrate how some previous results may have been artifacts of underspecification , and how scaling up the analysis can result in more robust models of cancer progression . 2 Problem setup . Given a ground set V = { 1 , . . . , n } , we consider a continuous-time Markov chain ( CTMC ) { X ( t ) } t≥0 on state space 2V ( Grimmett & Stirzaker , 2001 ) . Thus , the states of the Markov chain are subsets of V , and can equivalently be identified as binary vectors of size n. CTMCs are commonly represented by a generator matrix Q ∈ R2n×2n . Since the rows and columns of this matrix correspond to states in 2V , we will use qS R to denote the entry of Q at row indexed by S and column indexed by R. For S ̸= R , the entry qS R represents the infinitesimal rate of transitioning from state S to state R , that is , qS R = lim δt→0+ P ( Xt+δt = R |Xt = S ) δt , for S ̸= R , with qS R ≥ 0 . The quantity q̃S : = −qS S represents the total infinitesimal rate of leaving state S , with q̃S ≥ 0 . Since the total rate of leaving S is equal to the sum of the rates of transitioning to any other state R , we have q̃S = ∑ R∈2V qS R , for all S ∈ 2V . We are interested in modeling physical processes that induce an increasing accumulation of items over time , for example , the accumulation of genetic alterations when modeling cancer progression . To this end , we impose the following two constraints on the CTMC . First , we assume that the chain always starts from the empty state , that is , X ( 0 ) = ∅ . Second , we only allow transitions that add a single item to the current set , that is , for S ̸= R , qS R ̸= 0 only if R = S ∪ { i } for some i ∈ V . Since items are never removed , this implies that limt→∞Xt = V . In this paper , we focus on learning the CTMC from a given data set . Since learning directly the exponentially sized generator Q is a hopeless endeavor , we impose a particular parametric form on this matrix , recently introduced by Schill et al . ( 2019 ) . We define a parameter matrix Θ ∈ Rn×n ( also denoted by θ as a vector ) , and parameterize the off-diagonal entries of the generator matrix as qS S∪ { j } ( θ ) = exp ( θjj + ∑ i∈S θij ) , for all S ⊆ V , and j ∈ V \S . The off-diagonal elements of Θ encode positive ( attractive ) or negative ( repulsive ) pairwise interactions between items in V . The presence of item i increases ( θij > 0 ) or decreases ( θij < 0 ) the rate of adding item j by a multiplicative factor of wij : = eθij . The diagonal elements θjj encode the intrinsic rate of adding element j when ignoring all other interactions . Obtaining a sample from a CTMC under all aforementioned constraints can be done using the following sequential procedure . At each step , given the current state S , we draw a “ holding time ” h Algorithm 1 : Sampling a set from the marginal CTMC Input : Parameters θ Draw tobs ∼ Exp ( 1 ) t← 0 , S ← ∅ while t < tobs do Draw h ∼ Exp ( q̃S ( θ ) ) Compute pi = qS S∪ { i } ( θ ) /q̃S ( θ ) , for i ∈ V \ S Draw x ∼ Cat ( V \ S , ( pi ) i∈V \S ) t← t+ h , S ← S ∪ { x } return S from an exponential distribution with parameter q̃S . This represents how much time will pass until we add the next item . Then , we draw the next item x ∈ V \S according to a categorical distribution with probabilities ∝ qS S∪ { x } . The result is a sequence of items ( σ1 , . . . , σn ) , σi ∈ V , together with a sequence of holding times ( h1 , . . . , hn ) , hi ∈ R. Ideally , our data set would consist of a number of such item sequences and holding times ; we would then proceed to learn the parameters θ by maximizing the likelihood of our CTMC model . Unfortunately , such detailed data is typically not available in practice . In particular , it is rather common in biomedical applications to only have access to cross-sectional data , i.e. , to only observe the current state ( in the form of an unordered subset of V ) at a particular point in time without knowing how the process reached that state . Furthermore , it is often the case that the time point of each observation can only be roughly estimated or , worse , is completely unknown . Similarly to previous work , we assume that we are given a data setD = { S ( 1 ) , . . . , S ( N ) } , S ( i ) ⊆ V , with each S ( i ) representing a state observed at a potentially different unknown time Tobs . More specifically , we assume that Tobs is a random variable following an exponential distribution with density p ( t ) = e−t ( Gerstung et al. , 2009 ; Schill et al. , 2019 ) , and investigate the problem of maximizing the marginal likelihood , p ( D ; θ ) = N∏ i=1 p ( S ( i ) ; θ ) = N∏ i=1 ∫ ∞ 0 p ( S ( i ) | t ; θ ) p ( t ) dt . ( 1 ) Algorithm 1 shows how to sample a set from p ( S ; θ ) . The difference from our previous description of sampling from a CTMC lies in the fact that the procedure is only run until we reach the observation time tobs , and the result is an unordered set without any information about the holding times . 3 The value of “ unimportant ” items . Our ultimate goal is to retrieve the ordered interaction structure present in the data by learning the parameter vector θ . Before attempting to maximize the marginal likelihood of eq . ( 1 ) , though , it is natural to first ask whether the limited information available in the data is sufficient to infer such ordered interactions in the first place . We start with a simple example to showcase this issue . 3.1 Warmup example . Consider a ground set V = { 1 , 2 } and true model parameters θ∗11 = θ∗21 = 0 , and θ∗12 = −θ∗22 = α , for some α > 0 . When α is sufficiently large , the probability of item 2 occurring before item 1 goes to zero , therefore the model encodes the fact that item 1 is a prerequisite ( and therefore , always appears before ) item 2 . As a consequence , we will also have p ( { 2 } ; θ∗ ) ≈ 0 for the resulting marginal distribution of sets . One may intuitively think that this could be enough information to infer the true time order ; the following proposition shows that this is not the case . Proposition 1 . There is a one-dimensional family of parameters θ = θ ( s ) , and s1 , s2 ∈ R , such that p ( · ; θ ( s ) ) = p ( · ; θ∗ ) , for all s1 < s < s2 . The proof follows by simple algebra , and can be found in Appendix A together with an illustration of the parameter family . Interestingly , for s → s1 , the model θ ( s ) encodes the fact that , when we observe both items , item 2 always appears before item 1 , which is the opposite of what is encoded in our true model .

Summary: The authors model the time evolution of discrete states (e.g. genetic mutations) using a continuous-time Markov chains. They show that the resulting learning task is generally underspecified given cross-sectional data. The authors suggest including additional independent items that can help determine time order, and resolve this underspecification.  The authors implement an approximate likelihood maximization method for learning continuous-time Markov chains, which shows good performance on high-dimensional data, is faster than competing methods, and shows promising results on synthetic and real cancer data.