Q6: Discussion of the granularity of the discretization

This question involves the sensitivity analysis of the proposed method to hyperparameters. In the original paper, we conducted experiments by setting different numbers of codewords, with results reported in $\underline{\text{Figures 2c and 3c of Sec.4}}$. Additionally, in $\underline{\text{Appendix D.5 of the revised paper}}$, we tested the codeword requirements for proteins of varying sizes. When the codeword count is low, the energy landscape becomes coarsely discretized, making modeling capacity a performance bottleneck. Distinct observational states may become indistinguishable due to an insufficient number of codewords, leading to a decline in prediction quality. As the preset number of codewords increases, the model has sufficient codewords to represent the states of interest, resulting in improved prediction performance. Further increasing the codeword count, however, does not lead to unlimited improvements; rather, the prediction performance gradually converges as the number of codewords grows.

Q7: Discussion of the noisy data

We tested the model's prediction performance under different noise levels in $\underline{\text{Appendix D.3 of the revised paper (Figure 9)}}$. Using the amplitude of the observed state values as a standard, the model remained stable and did not collapse when noise intensity was below 0.5, demonstrating sufficient robustness. Such robustness to noise can be attributed to its adaptive codebook learning model, which incorporates a reduced-order approach. By identifying a low-dimensional, compact representation of the original state space, PESLA inherently possesses the ability to filter out uncertainties such as noise-related errors.

Q3: Analysis of interpretablity

Thank you very much for your insightful suggestions. We have added analysis of the interpretability of the proposed method in the revised paper:

• We experimentally validated the impact of energy estimation on evolution prediction in $\underline{\text{Appendix D.1 of the revised paper}}$ (Table1 below). We replaced the predicted energy with dummy energy to train the evolution prediction module. As the correlation between the dummy energy and true energy decreases, the prediction error for evolution increases. When the energy correlation drops below 0.5, the evolution prediction error of the graph neural Fokker-Planck equation module gradually exceeds that of the optimal baseline, highlighting the importance of accurate energy estimation for evolution prediction.
• We experimentally validated the impact of evolution prediction on energy estimation in $\underline{\text{Appendix E of the revised paper}}$ (Table2 below). Specifically, we conducted ablation experiments by disabling the auxiliary loss term $L_{latent }$ that guides evolution prediction. The results show that without the guidance for evolution prediction, the correlation of the estimated energy also declines. This suggests that accurate evolution prediction plays a significant role in energy estimation within our collaborative learning framework.

Table1: Evolution prediction accuracy as a function of the correlation coefficient $\rho$ between the dummy energy and the true energy on the 2D Prinz potential. The last two columns report the results of the proposed model (PESLA) and the best baseline (NeuralMJP) as presented in the main text, serving as a reference.

Table2: Ablation study on the loss function and submodule for 2D Prinz Potential and Ecological Evolution. w/o * indicates the absence of the loss function * or module *. All experiments are run 10 times to obtain statistical values.

Q4: Analysis of computational complexity

We have added an analysis and experiments on computational complexity in $\underline{\text{Appendix F of the revised paper}}$ to address this question. During the training phase, the main computational bottleneck of the proposed method comes from codeword matching, with overall complexity growing linearly with the sample size and the number of predefined codewords. Once the adaptive codebook module is trained, the model retains only the activated codewords, significantly reducing inference costs. Overall, denoting the sample size as $N$, the predefined number of codewords as $K$, and the token activation rate as $r$, the time complexities for training and testing are $\mathcal{O}(NK)$ and $\mathcal{O}(rNK)$, respectively. Additionally, we compared the training time of all models under the same settings. The experimental results ($\underline{\text{Figure 11 of the revised paper}}$) are consistent with the analytical conclusions, showing that the proposed method’s training time is comparable to that of the baseline.

Q5: Analysis of consistency

We conducted a consistency analysis of the estimated energy landscapes under different hyperparameter settings and random seeds in $\underline{\text{Appendix D.2 of the revised paper}}$. The results show that for the same protein, models trained with varying codebook sizes or random seeds achieve an average correlation of over 0.9 in their predictions. This demonstrates that the proposed method is robust to hyperparameter selection, producing consistent results.

We would like to sincerely thank Reviewer iFFC for providing a detailed review and insightful suggestions. According to the reviewers' valuable suggestions, we have revised the manuscript and uploaded the latest version of the PDF. We endeavored to address all the comments and our reflections are now given below point by point.

Q1: Discussion of data bias and sparsity

A major challenge in data-driven modeling of dynamical systems is that observational trajectories may be biased or overly sparse. In the proposed method, the adaptive codebook learning module coarsens the encoding space through codeword mapping. This approach reduces the impact of limited observational coverage by consolidating system states with similar semantic features into a unified representation. A codeword-controlled region is essentially a continuous local area in the encoding space, within which there are many (infinite) unseen sample points during training. Directly using these unseen points for evolution prediction would result in significant errors. However, when encountering a new observational sample, the encoder projects it into the encoding space, and the adaptive codebook module can match it to the most similar codeword, allowing evolution prediction to be based on the features of this known codeword, thereby improving accuracy.

Furthermore, we have conducted experimental analysis in $\underline{\text{Appendix D.4 of the revised paper}}$ (Table1 below). Specifically, we performed cross-protein experiments on five protein datasets, using the adaptive codebook learning module trained on other proteins to test on unseen proteins. The results show that, while performance degradation is inevitable due to unobserved regions, the average transfer performance $\rho_t$ for energy prediction across all proteins reaches over 80% of the performance of specifically trained version.

Table1: Comparison of mean $\rho_t$, MJS, and TJS metrics for encoder $\Xi$, decoder $\Omega$, and codebook $C$ trained on specific protein data versus cross-protein data for Homeodomain, BBL, BBA, NTL9, and A3D. All experiments are run 10 times to obtain mean values.

Q2: Discussion of limitations

Thanks for your valuable suggestions. This work focuses on estimating the energy landscape of a class of energy-driven evolutionary systems. However, when a system is driven by non-conservative forces, an energy landscape does not exist, as in the case of motion in viscous fluids. Additionally, inferring energy landscapes becomes more challenging when the landscape is time-varying, such as in cases where climate change alters species fitness. For time-varying landscapes, where $E( * )$ need to adapt to $E( * ,t)$, further model design exploration will be necessary to accommodate these dynamics in the future.

Overall, we have added limitations in the $\underline{\text{Conclusion of the revised paper}}$.

Dear reviewer,

The discussion period is almost ending. Could you please confirm whether our responses have alleviated your concerns?

For your interest, we (1) discuss the data bias and sparsity; (2) discuss the limitations of the method; (3) add analysis of interpretablity; (4) add analysis of computational complexity; (5) discuss the granularity of the discretization; (6) discuss the noisy data. Please refer to the details in the responses for you.

If you have further comments, we are also very happy to have a discussion. Thank you very much!

Thank you for addressing my concerns and answering my questions in details. However, regarding the overal limitations of the work, I am unable to confidently raise the score beyond marginally above the acceptance threshold.

Q2: Explanation of “physics-inspired”

The loss term $L_{phy}$ is formally represented as a KL divergence but essentially serves to introduce a physical prior. Specifically, in the absence of direct supervisory signals to guide energy estimation, the joint optimization of energy estimation and trajectory prediction becomes challenging. Therefore, we aim to incorporate physical prior knowledge to partially constrain the optimization direction. This constraint should be sufficiently general to apply across various disciplinary scenarios, rather than being limited to a specific system.

In statistical physics, a system's energy is used to quantify the probability distribution of states, known as the Boltzmann distribution. This implies a direct relationship between the frequency of a state being observed and its energy level, which can help guide the optimization direction for energy estimation. Furthermore, although different fields may define this concept differently—such as 'fitness' in genetics or 'potential energy' in proteins—prior studies have shown consistency between these concepts and the notion of energy in statistical physics [1,2,3]. Thus, we incorporate the Boltzmann distribution from statistical physics as a guiding principle in the unsupervised task of energy estimation, promoting the joint optimization of energy and trajectory prediction.

Q3: Clarification of Terminology

Thanks for your scientific rigor. Here, we clarify the notation used in each loss term:

• $L _ {reconstruct}=-\log\mathbf{q} _ {\Xi,\Omega,C}(x)$, where $\mathbf{q} _ {\Xi,\Omega,C}$ represents the probability distribution generated by the encoder $\Xi$, decoder $\Omega$, and codebook $C$. In the original text, $\mathbf{q} _ {\Xi,\Omega,C}$ was mistakenly written as $\mathbf{q} _ {\Xi,\Psi,C}$; this has been corrected.
• $L _ {latent}=||\Phi(\mathbf{p}(c _ {t+\Delta t})) - \Psi(H(t+\Delta t))||$, where $\mathbf{p} _ {c _ {t+\Delta t}}$ represents the target distribution of states at the time $t+\Delta t$ in the evolution prediction task.
• $L _ {code}=-\mathbf{p}(c _ {t+\Delta t})\log{\mathbf{q}(c _ {t+\Delta t})}$, where $\mathbf{p} _ {c _ {t+\Delta t}}$ represents the target distribution of states at the time $t+\Delta t$. $\mathbf{q} _ {c _ {t+\Delta t}}$ represents the model-predicted distribution of states at time $t+\Delta t$.
• $L _ {phy} = D _ {\text{KL}}(\mathbf{p} \| \mathbf{q}) = \sum _ {i=0}^{K} \mathbf{p}(c _ i) \log \left(\frac{\mathbf{p}(c _ i)}{\mathbf{q}(c _ i)}\right)$, where $\mathbf{p}(c _ i)$ represents the empirical distribution of the codeword $c _ i$ in the training of evolution prediction. $\mathbf{q}(c _ i)$ represents the Boltzmann distribution derived from the estimated energy of the codeword $c _ i$.

Ref:

[1] Sella, Guy, and Aaron E. Hirsh. "The application of statistical physics to evolutionary biology." Proceedings of the National Academy of Sciences 102.27 (2005): 9541-9546.

[2] Noé, Frank, et al. "Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning." Science 365.6457 (2019): eaaw1147.

[3] Thirumalaiswamy, Amruthesh, Robert A. Riggleman, and John C. Crocker. "Exploring canyons in glassy energy landscapes using metadynamics." Proceedings of the National Academy of Sciences 119.43 (2022): e2210535119.

[4] Mao, Chengzhi, et al. "Discrete Representations Strengthen Vision Transformer Robustness." International Conference on Learning Representations.

[5] Yu, Lijun, et al. "Language Model Beats Diffusion-Tokenizer is key to visual generation." The Twelfth International Conference on Learning Representations.

Many thanks to Reviewer jDjA for providing thorough detailed comments. According to the reviewers' valuable suggestions, we have revised the manuscript and uploaded the latest version of the PDF. We endeavored to address all the comments and our reflections are now given below point by point.

Q1: Analysis of ablation study

Thank you very much for your insightful suggestions on this work. We have added experiments in $\underline{\text{Appendix E of the revised paper}}$ (Table2 below), where we analyze the role of each loss function term and validate the impact of the auxiliary loss function terms and submodules. Specifically, we:

• Disabled the loss term $L_{latent}$ guiding trajectory prediction and observed a decrease in prediction accuracy, along with a decline in energy estimation precision.
• Disabled the loss term $L_{phy}$ guiding energy estimation, which made the joint optimization of energy and trajectory prediction more challenging.
• Disabled the encoder $\Phi$ and decoder $\Psi$​ in the Graph Neural Fokker-Planck equation for the probability vector, resulting in a significant increase in trajectory prediction error and impaired energy estimation.

These results validate the contributions and necessity of the proposed loss function terms and modules.

Table2: Ablation study on the loss function and submodule for 2D Prinz Potential and Ecological Evolution. w/o * indicates the absence of the loss function * or module *. All experiments are run 10 times to obtain statistical values.

Additionally, we replaced discrete embeddings with continuous embeddings. Although modeling transition probabilities in a continuous state space is not directly compatible with the proposed Graph Neural Fokker-Planck Equation, we used RNN to model the transition probabilities in the encoded space for comparison between continuous and discrete embeddings. The test results on the Ecological Evolution dataset are as follows:

The predictive performance of continuous embeddings is significantly worse, primarily because they lose many of the advantages designed into our method. The discrete codebook learning aligns with the low-dimensional landscape where long-term evolution of the dynamical system unfolds (see Sec 3.1 of the paper). Continuous embedding methods face generalization issues when encountering samples not present in the training set [4,5]. In contrast, our discrete embedding approach mitigates the impact of limited observational coverage by coarsening system states with similar semantic features into a unified representation. Additionally, the dimensionality reduction achieved through discrete mapping filters out local uncertainties introduced by random noise in the raw observational data (see $\underline{\text{Appendix D.3 of the revised paper}}$), ensuring that the predictor focuses on the essential shape of the dynamics. This is crucial for modeling stochastic dynamical systems.

Thank you for addressing my questions thoroughly. Because most of my concerns have been resolved, I will raise my score. However, as I am not a complete expert in this field, I am unable to confidently champion the paper beyond a weak accept.

Q3: Explanation of Sec. 3.2

We greatly appreciate your valuable suggestions. In $\underline{\text{Sec. 3.2 of the revised paper}}$, we have expanded on the motivation behind employing the graph neural Fokker-Planck equation to model the energy-driven evolution of system states within the landscape space.

Specifically, in a one-dimensional probability distribution, the information is limited to the immediate state of the node. In contrast, by using a graph convolutional network (GCN) to encode the probability vector into a high-dimensional space, we can effectively capture richer relational information between nodes, leveraging the GCN’s ability to aggregate and integrate neighborhood features. This relational structure is crucial in the evolution modeled by graph neural Fokker-Planck equation, as the change in a node’s state often depends on its neighbors. A higher-dimensional representation allows the model to more accurately reflect these dependencies. Moreover, the high-dimensional space allows the model to leverage the nonlinear fitting capabilities of neural networks, enabling it to capture complex patterns in the energy landscape, thereby improving prediction accuracy.

Additionally, in $\underline{\text{Appendix E of the revised paper}}$, we conducted detailed ablation experiments to analyze the contributions of the probability encoding module and each loss function.

Q4: Comparison of experimental setup

In dynamical system modeling research, it is common and well-accepted in the community to use different trajectories generated by the same system for training and testing [2,3,4]. We comprehensively evaluate the proposed model's performance in evolution and energy prediction using both JS divergence and Pearson correlation coefficient, consistent with prior works [2,5,6]. Additionally, we apply the same evaluation settings and average results over multiple runs for all models to ensure fairness in the evaluation.

Ref:

[1] Wu, Tao, et al. "Predicting multiple observations in complex systems through low-dimensional embeddings." Nature Communications 15.1 (2024): 2242.

[2] Federici, Marco, et al. "Latent Representation and Simulation of Markov Processes via Time-Lagged Information Bottleneck." The Twelfth International Conference on Learning Representations.

[3] Seifner, Patrick, and Ramsés J. Sánchez. "Neural Markov jump processes." International Conference on Machine Learning. PMLR, 2023.

[4] Mardt, Andreas, et al. "VAMPnets for deep learning of molecular kinetics." Nature communications 9.1 (2018): 5.

[5] Faure, Andre J., et al. "Mapping the energetic and allosteric landscapes of protein binding domains." Nature 604.7904 (2022): 175-183.

[6] Yang, Huan, Zhaoping Xiong, and Francesco Zonta. "Construction of a deep neural network energy function for protein physics." Journal of Chemical Theory and Computation 18.9 (2022): 5649-5658.

Many thanks to Reviewer ciD9 for providing thorough insightful comments. According to the reviewers' valuable suggestions, we have revised the manuscript and uploaded the latest version of the PDF. We endeavored to address all the comments and our reflections are now given below point by point.

Q1: Discussion of generalization

This issue typically arises when observational data is biased or extremely sparse, which is an inherent challenge for data-driven modeling of dynamical systems. In practice, the available observation trajectories may not fully cover all the 'energy wells,' particularly localized wells with limited coverage. However, the proposed adaptive codebook learning module can match unseen states to the codeword whose features are most similar, grouping states with similar energy characteristics. For example, in a large energy well containing multiple sub-wells, even if certain sub-wells are unseen during training, the proposed model can at least capture the information of the overarching energy well. Furthermore, as you mentioned, each codeword in the proposed method does have a multivariate representation, sufficiently characterizing the local state features within its region. To validate this generalization capability, we conducted cross-protein dataset experiments in $\underline{\text{Appendix D.4 of the revised paper}}$ (Table1 below). We trained the adaptive codebook learning module on other proteins and tested it on unseen proteins. The average transfer performance $\rho_t$​ for energy prediction across all proteins reaches over 80% of the performance of specifically trained version, demonstrating the codebook learning module’s capability to generalize to out-of-distribution regions.

Table1: Comparison of mean $\rho_t$, MJS, and TJS metrics for encoder $\Xi$, decoder $\Omega$, and codebook $C$ trained on specific protein data versus cross-protein data for Homeodomain, BBL, BBA, NTL9, and A3D. All experiments are run 10 times to obtain mean values.

Q2: Discussion of observation function g

Thanks for this insightful question. In the study of complex system modeling, delay embedding theory demonstrates that manifolds of system evolution can be reconstructed from multi-step observation trajectories [1]. In the three experimental scenarios presented in the main text of our paper, the observation function retains all information about the state energy, allowing us to use single-step observations as input. We have added supplementary experiments on degraded observations in $\underline{\text{Appendix G of the revised paper}}$ (Table2 below). Specifically, we applied an observation function $g(x, y) = \begin{bmatrix} \cos(\frac{\pi}{4}) & \sin(\frac{\pi}{4}) \end{bmatrix} \begin{bmatrix} x \\ y \end{bmatrix}^T$ to the 2D Prinz Potential used in the main text, projecting the 2-dimensional system state coordinates onto the 1-dimensional diagonal of the energy landscape as a degraded observation state. The results show that when the observation state lacks complete information about the energy, our model can improve energy prediction by increasing the number of input observation steps. As the input lookback steps increase, the model gains access to the system's past evolution trajectories, improving prediction performance to over 85% of that under lossless observations.

Table2: Energy prediction as a function of lookback steps.

Dear reviewer,

The discussion period is almost ending. Could you please confirm whether our responses have alleviated your concerns?

For your interest, we (1) discuss the generalization of the method; (2) discuss the observation function g; (3) add explanation of Sec. 3.2; (4) add comparison of experimental setup. Please refer to the details in the responses for you.

If you have further comments, we are also very happy to have a discussion. Thank you very much!

As the rebuttal period is drawing to a close, we are concerned about the lack of communication with you. Could you please see whether our responses will influence your ratings? All the other reviewers gave our paper a rating above acceptance threshold. You are the sole reviewer who has expressed reservations, and we take your concerns very seriously. We have addressed all the issues you raised in our rebuttal, but have not yet received any feedback from you. We respectfully request your prompt response.

Q4: Analysis of ablation study

We have added experiments in $\underline{\text{Appendix E of the revised paper}}$ (Table2 below) to address this question. For each auxiliary loss function term or submodule, we:

• Disabled the loss term $L_{latent}$ guiding trajectory prediction and observed a decrease in prediction accuracy, along with a decline in energy estimation precision.
• Disabled the loss term $L_{phy}$ guiding energy estimation, which made the joint optimization of energy and trajectory prediction more challenging.
• Disabled the encoder $\Phi$ and decoder $\Psi$ in the Graph Neural Fokker-Planck equation for the probability vector, resulting in a significant increase in trajectory prediction error and impaired energy estimation.

These results validate the contributions and necessity of the proposed loss function terms and modules.

Table2: Ablation study on the loss function and submodule for 2D Prinz Potential and Ecological Evolution. w/o * indicates the absence of the loss function * or module *. All experiments are run 10 times to obtain statistical values.

Q5: Discussion of limitations of the method

Thanks for your valuable suggestions. This work focuses on estimating the energy landscape of a class of energy-driven evolutionary systems. However, when a system is driven by non-conservative forces, an energy landscape does not exist, as in the case of motion in viscous fluids. Additionally, inferring energy landscapes becomes more challenging when the landscape is time-varying, such as in cases where climate change alters species fitness. For time-varying landscapes, where $E( * )$ need to adapt to $E( * ,t)$, further model design exploration will be necessary to accommodate these dynamics in the future.

Overall, we have added limitations in the $\underline{\text{Conclusion of the revised paper}}$.

Q6: Explanation of Graph Neural Fokker-Planck Equation

Thank you very much for your insightful suggestions. In $\underline{\text{Sec. 3.2 of the revised paper}}$, we have provided a more detailed explanation of the motivation behind using the graph neural Fokker-Planck equation to model the energy-driven evolution of system states in the landscape space, along with further insights on the encoding of probability vectors.

Specifically, in a one-dimensional probability distribution, the information is limited to the immediate state of the node. In contrast, by using a graph convolutional network (GCN) to encode the probability vector into a high-dimensional space, we can effectively capture richer relational information between nodes, leveraging the GCN’s ability to aggregate and integrate neighborhood features. This relational structure is crucial in the evolution modeled by graph neural Fokker-Planck equation, as the change in a node’s state often depends on its neighbors. A higher-dimensional representation allows the model to more accurately reflect these dependencies. Moreover, the high-dimensional space allows the model to leverage the nonlinear fitting capabilities of neural networks, enabling it to capture complex patterns in the energy landscape, thereby improving prediction accuracy.

Additionally, in $\underline{\text{Appendix E of the revised paper}}$, we have included ablation experiments on encoding probability vectors to validate the necessity of this approach.

Many thanks to Reviewer c7fR for providing a detailed review and insightful questions. According to the reviewers' valuable suggestions, we have revised the manuscript and uploaded the latest version of the PDF. We endeavored to address all the comments and our reflections are now given below point by point.

Q1: Description and implementation of the overall model

Thanks for your valuable suggestions. We have added new paragraphs and an appendix section to summarize each module:

• Summarized the trainable parameters in the first paragraph of $\underline{\text{Section 3.3 of the revised paper}}$.
• Listed all modules and parameter shapes in $\underline{\text{Appendix A.2 of the revised paper}}$.
• Summarized the function of each loss term in $\underline{\text{Appendix E of the revised paper}}$.

While other reviewers, such as reviewer iFFC, have praised the writing of this paper, we sincerely appreciate your suggestions, which have helped us further enhance the clarity of each module’s presentation, making this work more accessible to a broader audience.

Q2: Analysis of transferability

In the original version of the paper, we trained separate models for each protein. We have added transfer experiments on five proteins in $\underline{\text{Appendix D.4 of the revised paper}}$ (Table1 below). We demonstrate that the encoder $\Xi$, decoder $\Omega$, and codebook $C$ can be reused across different proteins. By training on cross-protein data, these modules are expected to learn to capture general representations and project diverse protein states into a unified latent space. Notably, using the encoder, decoder, and codebook trained on other proteins, the average transfer performance $\rho_t$ for energy prediction across all proteins reaches over 80% of the performance of specifically trained version.

Table1: Comparison of mean $\rho_t$, MJS, and TJS metrics for encoder $\Xi$, decoder $\Omega$, and codebook $C$ trained on specific protein data versus cross-protein data for Homeodomain, BBL, BBA, NTL9, and A3D. All experiments are run 10 times to obtain mean values.

Q3: Analysis of scalability

We have added experiments in $\underline{\text{Appendix D.5 of the revised paper}}$ to address this question. In these experiments, the prediction performance for all proteins improves as the codebook size $K$ increases, approaching convergence before reaching $K=1000$. Using $K=1000$ as a benchmark, we found that the minimum codebook size required to achieve 90% performance increases with protein size ($\underline{\text{Figure 10}}$). However, for the proteins used in this study, a codebook size of $K=500$ is sufficient to meet all modeling requirements. Additionally, for unfamiliar systems, we recommend starting with a larger $K$ value, as the proposed method can automatically utilize only the necessary subset of codewords, remaining computationally efficient.

We also included a computational cost analysis in $\underline{\text{Appendix F of the revised paper}}$. This analysis shows that the time complexity of the proposed method scales linearly with the sequence length.

Dear reviewer,

The discussion period is almost ending. Could you please confirm whether our responses have alleviated your concerns?

For your interest, we (1) clarify the description and implementation of the overall model and loss terms; (2) add analysis on transferability; (3) add analysis on scalability; (4) add analysis on ablation study; (5) discuss the limitations of the method; (6) add explanation of graph neural Fokker-Planck equation. Please refer to the details in the responses for you.

If you have further comments, we are also very happy to have a discussion. Thank you very much!