Latent Representation and Simulation of Markov Processes via Time-Lagged Information Bottleneck

We express our gratitude to Reviewer Yi96 for their recognition of our work’s value and for the insightful questions and comments, that have prompted updates to both our main text and supplementary material.

1. Related work

   In response to the reviewer’s comment, we have restructured the related work section. This revision aims to more clearly highlight the distinctions between our contributions and existing literature, thereby providing a better context for our work.

2. Autoinformation and Data Processing Inequality

   The assertion in section 2.2 stems from the Data Processing Inequality (Theorem 2.8.1, page 32 in [1]). The underlying intuition is that processing the original data sequence cannot generate new temporal information. A detailed proof for this statement is included in Appendix B.3 and relies on assumption A.1 (appendix B.1), which is based on the graphical model illustrated in Figure 2a. We included a reference to the appendix following the statement in the main text to clarify this.

3. Runtime

   Addressing the reviewers’ comments, we decided to include a detailed analysis and comparison of the simulation runtimes by adding dedicated supplementary material (Appendix G), which is referenced in the experimental section. For further information, we direct the reviewer to our shared answer where these details are elaborated upon.

References

[1] Thomas M. Cover, Joy A. Thomas "Elements of information theory." (1991).

4. Comparison with VAMPnet

   The main purpose of including the VAMPnet visualization in Figures 3c and 4a is to provide intuitive insights regarding the effect of the training objective and encoder choice on the resulting representation. More specifically, comparing the TICA and VAMPnet representations showcases the effect of choosing a linear vs non-linear encoder. Furthermore, Since VAMPnet employs the same encoder architecture as the T-InfoMax model, this comparison also visually emphasizes the differences between maximizing linear and non-linear autoinformation. This aspect is particularly significant, as it helps in comprehending the subtle yet impactful distinctions brought about by different maximization strategies.

   It is important to note that, following previous research [5], the quantitative evaluations of the fidelity of the unfolded latent simulations for VAMPnet refer to higher-dimensional representations (32 dimensions).

References

[1] Oord, Aaron van den, Yazhe Li, and Oriol Vinyals. "Representation learning with contrastive predictive coding." arXiv preprint arXiv:1807.03748 (2018).

[2] Radford, Alec, et al. "Learning transferable visual models from natural language supervision." International conference on machine learning. PMLR, 2021.

[3] Chen, Wei, Hythem Sidky, and Andrew L. Ferguson. "Capabilities and limitations of time-lagged autoencoders for slow mode discovery in dynamical systems." The Journal of Chemical Physics 151.6 (2019).

[4] Kawaguchi, Kenji, et al. "How Does Information Bottleneck Help Deep Learning?." arXiv preprint arXiv:2305.18887 (2023).

[5] Sidky, Hythem, Wei Chen, and Andrew L. Ferguson. "Molecular latent space simulators." Chemical Science 11.35 (2020): 9459-9467.

We thank reviewer mH7G for their acknowledgment of the novelty of our work and for providing insightful comments, which we address as follows:

1. Lack of comprehensive discussion regarding the mutual information lower-bound and estimation method

   The decision to employ a contrastive bound on $I(z_t;z_{t+\tau})$ instead of $I(z_t;x_{t+\tau})$ was driven by several factors:

   • Computational efficiency: $z_{t+\tau}$ is usually much lower-dimensional than $x_{t+\tau}$. Using contrastive learning to maximize information between representations becomes computationally more efficient than employing a decoder to predict $x_{t+\tau}$.

   • Additional Independence Condition: the condition (i) $I(z_t;z_{t+\tau})=I(x_t;x_{t+\tau})$ is stronger than (ii) $I(z_t;x_{t+\tau})=I(x_t;x_{t+\tau})$ since it requires $I(x_t;x_{t+\tau}|z_{t+\tau})=0$, which is not directly implied by (ii). This additional independence is crucial in our Lemma 3 (refer to Appendix B.7), providing a stronger foundation for our theoretical framework. Although the results from Lemma 1 and 2 can be extended under condition (ii), the specific independence in (i) plays an important role in our methodology.

   • Successful examples: Numerous contemporary studies have successfully employed contrastive methods to maximize information among representations, especially in handling structured, high-dimensional data sequences [1,2]. This contrasts with the limitations observed in reconstruction-based mutual information maximization methods, particularly in the context of molecular dynamics, as highlighted by previous research [3].

2. Explicit superfluous information minimization vs Dimensionality Reduction

   We appreciate the reviewer's insights regarding dimensionality reduction and its relationship with information loss. In response, we'd like to elaborate on several pertinent aspects:

   • Non-Guaranteed Compression: Reducing the dimensionality of the data does not necessarily result in compression. In cases where high-dimensional data resides on a lower-dimensional manifold, non-linear encoding methods have the potential to map data into a reduced dimensionality without incurring information loss.

   • Generalization: Contemporary research on generalization bounds in representation-based prediction underscores the significance of superfluous information measures [4]. These bounds are more directly related to the minimization of irrelevant information rather than solely focusing on the dimensionality of the data. This perspective aligns with our approach, as we emphasize the explicit reduction of superfluous information.

   • Applicability: When dealing with low-dimensional data, such as the Prinz-2D dataset depicted in Figure 3, dimensionality reduction may not be a feasible or a necessary strategy. However, our approach to minimizing superfluous information retains its applicability and effectiveness even in these scenarios. By focusing on the reduction of non-essential information, our methods can still contribute to the compression and enhanced understanding of the data without necessarily reducing its dimensionality.

3. Quantitative Evaluation of Dimensionality Reduction

   We quantitatively assess representations by examining the autoinformation at various timescales, as depicted in Figures 4b, 12, and 13. Additionally, we measure the amount of information carried about relevant $y^{slow}$ and irrelevant $y^{fast}$ features (Figure 3b). Despite the subtlety of qualitative differences in Figure 4a, a consistent decrease in information at faster timescales is observed across Figures 4b and 12. Figure 11 offers a visual representation of how different $\tau$ settings impact the representations of the Villin dataset, further clarifying the effects of our information reduction approach.

   While our method indeed facilitates dimensionality reduction, it's important to emphasize that our primary objective is to leverage the representation for latent simulations. Consequently, we assert that evaluating the statistical fidelity of simulations unrolled from each representation provides a meaningful metric for assessing representation quality. This evaluation becomes particularly significant considering that the transition and prediction models, along with the training procedures, remain constant across all fixed representations.

Thank authors for the reply. Some more questions and comments:

1. For the difference between $I(z_t, z_{t+\tau})$ and $I(z_t,x_{t+\tau})$, I agree that the former one can be estimated more efficiently. But under the assumption $I(z_t,x_{t+\tau})=I(x_t,x_{t+\tau})$, we can also obtain the conclusion $I(z_t,x_{t+\tau'})=I(x_t,x_{t+\tau'})$ for $\tau'>\tau$, which is similar to Lemma 3.

2. In Figure 4b, it seems that there is a significant reduction in autoinformation when $\beta=0.1$ compared to $\beta=0$. However, the authors show the results for $\beta=0.1$ in Figure 4a. I am curious to understand the rationale behind the authors' choice of $\beta$. What criteria or methodology guided the selection of $\beta$?

3. I would like to reemphasize that VAMPnet is a method for estimating dominant spectral components. In other words, the more accurately the model estimates, the more precisely it forgets those fast components. Therefore, I do not consider that it is reasonable to compare the AIG of VAMPnet with other methods. Authors should highlight this point in the analysis of experiments.

5. How does the model scale with the amount of data available for training?

   In our experiments, due to data availability constraints, we focused on a single long simulation for each molecule, exploring all relevant states and approaching equilibrium distribution. We experimented with lag times up to 204.8 ns, yielding meaningful molecular representations. We believe that longer training trajectories could enable us to extend the time scale further. Additionally, using more training trajectories could improve the accuracy of our autoinformation estimates, as discussed in Appendix C.2. Enhanced autoinformation estimates would not only allow the model to more effectively capture features at the chosen time scale but also provide more precise quantitative estimations, such as those shown in Figure 4, aiding in diagnosis and evaluation.

References

[1] Chen, Wei, Hythem Sidky, and Andrew L. Ferguson. "Capabilities and limitations of time-lagged autoencoders for slow mode discovery in dynamical systems." The Journal of Chemical Physics 151.6 (2019).

We are grateful to the reviewer for their positive remarks about our presentation and for the insightful comments and questions raised.

To address the main points:

1. Architectural choices

   We acknowledge that the main text could have provided more comprehensive details regarding the selection of encoder, transition, prediction, and mutual information maximization models. Due to the constraints of space, our focus was primarily on delineating a general family of models, directing readers to relevant literature for specifics on architectural choices. Our architecture selection was primarily driven by two key principles:

   • Flexibility: Our primary criterion was to choose architectures known for their effectiveness in modeling similar problems. This led us to opt for contrastive learning methods over reconstruction-based methods for mutual information maximization [1], use TorchMD encoders for invariant molecular representations, and select advanced normalizing flows like Flow++ for transitions, instead of simpler mixture models.

   • Computational cost: Another important consideration was computational efficiency. While we acknowledge that more flexible transition models like diffusion models might yield more precise transitions, they could also slow down latent simulations. Exploring this trade-off is an exciting avenue for future research.

   We recognize that the performance of our models could potentially be enhanced with more nuanced choices for the various architectures. However, it is crucial to emphasize that the primary objectives of this submission are to establish a formal foundation for the T-InfoMax and T-IB losses and to demonstrate the efficacy of our proposed inference and representation learning method. We aim to show that, with reasonable architectural choices, our approach compares favorably against different training objectives using the same architectures.

2. How to choose $\tau$

   The selection of $\tau$ essentially serves as a task-specific hyper-parameter. For instance, in weather prediction, $\tau$ might be set to one day for daily forecasts, but could be reduced for predictions requiring higher temporal resolution, such as hourly rain forecasts. Similarly, in molecular dynamics, the time scale for $\tau$ can be aligned with the expected duration of folding/unfolding events or other significant state transitions.

   Figure 11 in Appendix F illustrates the impact of varying $\tau$ on the representation and the corresponding autoinformation values. A smaller $\tau$ tends to retain more informative details, capturing faster processes and a greater number of clusters. Conversely, a larger $\tau$ emphasizes slower processes, discarding fast dynamics and speeding up the latent simulation process due to the larger temporal jumps.

   An important consideration in choosing $\tau$ is the necessity of having multiple training examples with pairs spaced $\tau$ apart. Consequently, the selection of $\tau$ is often bounded by the feasibility of simulating trajectory lengths greater than $\tau$. In our experiments, we opted for a $\tau$ that is several orders of magnitude shorter than the total trajectory length. However, we propose that the total simulation time $T$ of available trajectories need not be significantly larger than $\tau$, especially if multiple parallel simulations are utilized instead of a single long trajectory. This approach is particularly relevant in the context of molecular dynamics since producing long ground truth trajectories is computationally intensive. Exploring this efficiency in leveraging parallel trajectories represents a promising direction for future research.

3. Does this parameter remain fixed?

   Once an encoder is trained with T-IB or T-InfoMax for a given $\tau$, it is possible to learn transition models $q_\phi(z_t|z_{t-\tau’})$ for any larger $\tau’\ge\tau$. This implies that we can retrain the transition model for a larger temporal scale $\tau’>\tau$ without needing to modify the encoder or the predictor $q_\psi(y_t|z_t)$.

   However, with T-IB, it's often advantageous to retrain or fine-tune the encoder for a new lag time $\tau’$. This adjustment allows the encoder to disregard processes slower than $\tau’$, thereby simplifying the representation. When updating the encoder, it's necessary to retrain both the transition and predictive models as well.

   Conversely, reducing the time-scale $\tau’<\tau$ necessarily requires re-training the encoder to incorporate faster processes. Utilizing previously trained encoder, transition, and prediction models as starting points should speed up the training.

4. Can you provide run times?

   We thank the reviewer for highlighting the need for runtime measurements. We have now included estimates for Molecular Dynamics runtime in Appendix G.1, Latent Simulation runtime in G.2, and training time in G.3. These estimations have been referenced in the main text.