Deep Generative Markov State Models
We propose a deep generative Markov State Model (DeepGenMSM) learning framework for inference of metastable dynamical systems and prediction of trajectories. After unsupervised training on time series data, the model contains (i) a probabilistic encoder that maps from high-dimensional configuration space to a small-sized vector indicating the membership to metastable (long-lived) states, (ii) a Markov chain that governs the transitions between metastable states and facilitates analysis of the long-time dynamics, and (iii) a generative part that samples the conditional distribution of configurations in the next time step. The model can be operated in a recursive fashion to generate trajectories to predict the system evolution from a defined starting state and propose new configurations. The DeepGenMSM is demonstrated to provide accurate estimates of the long-time kinetics and generate valid distributions for molecular dynamics (MD) benchmark systems. Remarkably, we show that DeepGenMSMs are able to make long timesteps in molecular configuration space and generate physically realistic structures in regions that were not seen in training data.
The manuscript introduces a new technique for learning Markov State Models from molecular trajectory data. The model consists of several novel ideas, including a generative component that makes it possible to reconstruct molecular structures from the latent space.  Quality: The methodological contributions are convincing, and backed up by both source code and detailed derivations in supporting material (and references to relevant literature). Although the experiments focus on systems of modest size, and are thus mostly proof-of-concept, the modelling strategy itself has several non-standard components, which in my view makes the manuscript a significant Machine Learning contribution. The authors are honest about the strenths and weaknesses of their work.  Clarity: The manuscript is of high quality. It is well written, and the text is accompanied with helpful figures. My only minor concern regarding clarity is that some of the subfigures in Fig 2 - 4 are difficult to read due to their small size.  Originality: The presented methods are new, and represent a clear deviation from earlier work in this area.  Significance: The authors present a convincing case that their approach has fundamental advantages over current techniques for estimating Markov State Models. It is thus likely to have considerable impact on the MSM field. The modelling strategy itself should be of interest more broadly in the Machine Learning community.   Suggestions for corrections:  In the text following eq (1), perhaps you should more explicitly introduce the hyperparameter m as the number of metastable states, to make it clearer to readers not familiar with the MSM literature.  The discussion regarding the rewiring trick is a little brief. Perhaps you could explicitly write out the approximation either in the text or in supporting material. It's also not quite clear how many samples you draw in this approximation.  Eq 7 is not quite clear to me. You write that \rho(y) is the empirical distribution, but how should I think of an empirical distribution over a continuous space y? Or are we talking about a distribution over the fixed set of input structures (in which case I would assume they are all equally likely - i.e. the exact same structure is never visited twice). Please elaborate.  I found your Energy Distance procedure interesting, but could you perhaps briefly state why you would choose this over e.g. a Kullback-Leibler based approach? (since you use the Kullback-Leibler divergence later for measuring the difference between distributions (fig 2d).  The dialanine system seems to be of rather modest size compared to the systems typically studied with MD. It would be nice if the paper concluded with some some statements on the extent to which the DeepResampleMSM and DeepGenMSM scale to such systems.   Minor corrections: Line 42: "is conserves" -> "conserves"  Line 130: "an generated trajectory" -> "any generated trajectory"?  Update after rebuttal period:  I only raised minor issues in my review, which the authors have addressed in their rebuttal. I therefore stand by my original (high) rating of the manuscript.