Paper: Latent Representation and Simulation of Markov Processes via Time-Lagged Information Bottleneck (ICLR-2024)

Reviewer comment (Writing):
The related works section currently reads like a list of related works. It will be far better to re-write this section as a discussion of how current literature relates to this proposed work and how this proposed work builds on top of this existing literature, and how it is better than its closest related works.

Author response:
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.

Reviewer comment (Theory):
Can the authors comment on "Since $z_t$ is derived from $x_t$, the autoinformation for $x_t$ sets an upper-bound for the autoinformation for $z_t$. Given that the encoder is time-independent, how can we expect the autoinformation between $x_t$ to be related in any way to the autoinformation in $z_t$?

Author response:
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.

Reviewer comment (Experiments):
The paper sets up the reader to expect a solution that is faster than current available methods for simulation. I was expecting experiments measuring how long it takes for the simulations to be run compared to existing methods. Can the authors provide that?

Author response:
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.

Reviewer comment (Theory):
Comparison with Past-Future IB Method: The paper builds upon the past-future IB method proposed by Tiwary et al., but it introduces an optimization objective for I(z_t, z_t+tau) instead of I(z_t, x_t+tau). While theoretically, optimizing I(z_t, z_t+tau) is expected to yield a lower bound on I(z_t, x_t+tau), the paper lacks a comprehensive discussion of the differences between the two approaches. A deeper exploration of the distinctions and the implications of this choice is needed to provide a clearer understanding.

Author response:
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].

Reviewer comment (Experiments):
Choice of Estimation Method: In the machine learning domain, there are various methods for estimating lower bounds on mutual information. The paper selects a specific method, but it would benefit from a more thorough justification for this choice. Providing a comparative analysis of the selected method against other available options would enhance the paper's robustness.

Author response:
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].

Reviewer comment (Theory):
Complexity of Additional Terms in Formula (8): The introduction of two additional terms in Formula (8) may appear somewhat perplexing. From an information bottleneck theory perspective, the goal is to discard as much irrelevant information as possible. However, dimensionality reduction inherently leads to information loss. The rationale and impact of introducing these terms need further clarification. Moreover, the experimental results suggest that the inclusion of these terms does not significantly affect the dimensionality reduction outcome, requiring a more detailed explanation.

Author response:
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.

Reviewer comment (Evaluation):
Quantitative Evaluation of Dimensionality Reduction: The paper lacks a robust quantitative evaluation of the dimensionality reduction results. While the field may lack a general method for quantitatively assessing dimensionality reduction outcomes, the authors could consider exploring quantitative evaluation approaches, such as measuring mutual information under different tau conditions, to provide a more concrete assessment of the method's performance.

Author response:
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.

Reviewer comment (Experiments):
Comparison with VAMPnet: VAMPnet, as mentioned in the weaknesses, primarily serves the purpose of extracting dominant or slow components and estimating kinematic properties. It may not be a direct dimensionality reduction method. Comparing VAMPnet with the proposed approach for dimensionality reduction may not be entirely fair or informative.

Author response:
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).

Reviewer comment (Experiments):
The method has many moving parts and some choices are not well motivated. For example, different choices of architectures for different distributions, training with contrastive estimation. A study that sees which part is contributing to the performance would be nice.

Author response:
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.

Reviewer comment (Evaluation):
As far I can see the authors do not provide actual run times. What should be compared is the total time to obtain the training data, train the model and perform the simulation.

Author response:
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.

Reviewer comment (Experiments):
How to choose $\tau$ and once chosen, does this parameter remain fixed and the model has to be retrained for different choice?

Author response:
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.

Reviewer comment (Evaluation):
Can you provide run times?

Author response:
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.

Reviewer comment (Experiments):
How does the model scale with the amount of data available for training?

Author response:
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.