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 tackle the problem of time evolution of discrete sets through the view of continuous-time markov chains. The authors provide a theoretical justification of why these problems in general underspecified in the setting of cross-sectional data. They go on to demonstrate that the underspecification can be overcome by use of additional items / features that can help determine time order by deriving a bound on the mean and variance of the observed time for these additional items. An approximate formulation for likelihood maximization is then presented without any assumptions of prior knowledge of the additional items or the structure of the additional items. The theoretical framework is then tested using simulation studies and demonstrate that accuracy of state sequence recovery can be significantly improved with use of the additional items. The framework is also evaluated using a real world cancer dataset from TCGA by ordering of pairs of mutation and copy number aberration events.