Predicting the Energy Landscape of Stochastic Dynamical System via  Physics-informed Self-supervised Learning

Authors: Ruikun Li, Huandong Wang, Qingmin Liao, Yong Li
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: dynamical system, energy landscape, deep learning

Abstract:
Energy landscapes play a crucial role in shaping dynamics of many real-world complex systems. System evolution is often modeled as particles moving on a landscape under the combined effect of energy-driven drift and noise-induced diffusion, where the energy governs the long-term motion of the particles.
Estimating the energy landscape of a system has been a longstanding interdisciplinary challenge, hindered by the high operational costs or the difficulty of obtaining supervisory signals. Therefore, the question of how to infer the energy landscape in the absence of true energy values is critical. In this paper, we propose a physics-informed self-supervised learning method to learn the energy landscape from the evolution trajectories of the system. It first maps the system state from the observation space to a discrete landscape space by an adaptive codebook, and then explicitly integrates energy into the graph neural Fokker-Planck equation, enabling the joint learning of energy estimation and evolution prediction. Experimental results across interdisciplinary systems demonstrate that our estimated energy has a correlation coefficient above 0.9 with the ground truth, and evolution prediction accuracy exceeds the baseline by an average of 17.65\%. The code is available at https://github.com/tsinghua-fib-lab/PESLA.

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730215786038

thread_type:
official_review

<thread_object>
creator: reviewer_iffc

<content>
summary:
This paper proposed a new framework PESLA to estimate the energy landscapes from the evolution trajectories of the system. The framework includes two stage, first it uses an adaptive codebook to obtain a discrete landscape space from the observation trajectories. Second, it uses a graph neural ODE based on the Fokker-Planck equation to predict system evolution.

strengths:
- The paper is clearly written. 

- The method is trained in a self-supervised manner, hence it can conduct estimation without knowing true energy landscape.

- The methods were evaluated on three interdisciplinary systems and all demonstrated an improvement compared with SOTA baselines.

weaknesses:
- The author mentioned that “observable evolutionary trajectories typically cover only a limited portion of the vat state space”, however, the paper did not explain how the purposed method perform in the situation when the available data is biased or very sparse.

- The method relies on the assumption that the system is driven by energy and a well-defined low dimensionality energy landscape exists. Therefore, the assumption might limit the applicability, which was not discussed in the paper.

- The method is lack of interpretablity. The training fully relies on the final trajectory prediction errors to avoid knowing true energy landscape, however, it is not discussed how reliable or meaningful of the predicted trajectory and the estimated energy landscape. For example, dose the correlation of the estimated energy landscape and true energy landscape influence the final trajectory prediction?

questions:
- What is the computational complexity of the proposed methods and its comparison to existing methods? 

- Are the estimate energy landscape consistent under different hyperparameters?

- Does the granularity of the discretization of the energy landscape influence the accuracy of system evolution estimation?

- Considering that real-world trajectory data is usually noisy and sparse, how robust is the method to such data?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730266260336

thread_type:
official_review

<thread_object>
creator: reviewer_jdja

<content>
summary:
Estimating the energy landscape is challenging because obtaining direct energy measurements is costly. To bypass this process, the paper proposes a physics-informed self-supervised learning approach, where the energy landscape is estimated from historical evolution trajectories instead of direct energy signals. This method uses discrete codebook embeddings, assuming that energy landscapes generally have low intrinsic dimensionality. Additionally, a physics-informed graph neural Fokker-Planck architecture and physics-inspired regularization are employed to predict system states more accurately.

strengths:
The paper begins with a strong motivation, highlighting the challenges in obtaining direct energy values. Additionally, it provides a solid rationale for incorporating discrete codebook embeddings into their method, based on the inherently low dimensionality of energy landscapes.

weaknesses:
1. The ablation studies in this paper seem too limited. Since multiple modules are introduced in their method, additional experiments are needed to clarify the individual contributions of each component to the overall performance. For instance, it would be helpful to examine the effect of using discrete versus continuous embeddings. Similarly, given that five types of losses are included, an analysis of each loss’s specific impact would provide valuable insights. If these results were included, I would increase my rating.

2. Regarding  $L_{phy}$, could you explain in more detail why this loss is referred to as “physics-inspired”? How it differs from standard machine learning loss terms?

3. There are five types of losses ( L_{reconstruct}, L_{vq}, L_{latent}, L_{code}, L_{phy}) in total. Could you clarify the exact meaning of the terms $\mathbf{p}$ and $\mathbf{q}$ in each loss?

questions:
See Weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730544267967

thread_type:
official_review

<thread_object>
creator: reviewer_cid9

<content>
summary:
The authors propose an algorithm for discovering the energy function that may explain data trajectories assumed to be driven by the gradient of the energy plus noise. Their model learns a partitioning of the space with one discrete symbol per region and an associated energy value, along with a decoder (and associated encoder) mapping to the observed continuous space. They assume that the dynamics is driven by the energy difference between nearby partitions and Markovian transitions between regions. The model is evaluated on three tasks and compared with several baselines.

strengths:
This is an interesting learning problem, also requiring to learn a latent explanatory space for sampled continuous time trajectories. The authors compare against multiple baselines across everal environments. The results appear encouraging.

weaknesses:
(1) I do not see how such an approach can generalize outside of the visited regions, i.e., the cluster centers corresponding to the codebook entries; in general, the most interesting energy functions can have a very large (if not exponential) number of energy wells, and one needs to generalize from the visited wells to new ones. The reason is that many factors may be composed that give rise to stable solutions. Hence a multivariate representation of the discrete identities must be constructed (think about language, which is discrete but allows generating a huge number of legal combinations of words).

(2) It is not clear to me that g in general will preserve all information about the energy of the state, so it may be necessary for the encoder (from observations to states) to look at the whole past trajectory in order to make a probabilistic guess about the state proxy.

(3) the main model (sec 3.2) is not sufficiently understandable and could use more motivation for its different parts. Why these equations in particular (4 and 5).

(4) The experimental setup seems to lack comparisons with *published results* from earlier work, which makes it difficult to know if the comparisons are fair.

(5) The test set includes only trajectories from the same system for which a training trajectory is given, i.e., there is no form of generalization to new systems (e.g., new molecules).

questions:
* to address (3) above, please provide better justifications for the main model and the associated training losses

* to address (4) and (5), please compare with published results on recovering the energy and fitting new trajectories of new molecules or systems
</content>
</thread_object>

</thread>
<thread>
creation date:
1730671179742

thread_type:
official_review

<thread_object>
creator: reviewer_c7fr

<content>
summary:
This paper proposes to infer the energy landscape of complex, stochastic dynamical systems. A distinguishing feature is the adoption of vector quantization techniques for dimension reduction. Dynamical information was inferred via a graph neural Fokker-Planck equation method. A regularization term was introduced to constrain the long-term prediction behavior between the empirical distribution p and the Boltzmann distribution. The novelty of bringing together these ideas is appreciated. Experiments were performed on 3 tests: 2d Prinz potential, ecological evolution, and fast-folding peptides. The method part of the paper is not very clear and somewhat difficult to follow. There is no ablation tests, so it's hard to see what the various loss terms or model designs offer. There is no discussion of the limitations of the method.

strengths:
The idea of bringing together several novel methods for modeling stochastic dynamical systems is very interesting. Three different examples cover different application areas and add to the credibility of the method. The capability to learn unknown energy landscape from stochastic dynamical systems is a potentially high-impact contribution.

weaknesses:
* The description of the overall model is not clear to me. At least not easy to follow. For example, the encoder, decoder, Phi, Psi, Xi,... were defined in different places and some were not clearly defined. What space and dimension does Phi map from and to? Please make it easier to follow. And the dimension of latents such as H is not clear. Please collect the equations in a coherent way, e.g. all the different losses, and what are the losses sampled over?

* Overall, the writing of the paper is too hand-waving. As a referee I don't like guesswork.

questions:
* Please discuss the issue of transferrability. Can the trained model be transferred to unseen structures? Was a separate model trained for each protein? Which of the sub-modules, encoder, decoder, etc, can be reused?

* Scalability. To model a bigger protein, how large does the codebook need to be? What about training and inference cost for longer sequences?

* Spell out KNN (k-nearest neighbor?)

* How are encoder Ξ, decoder Ω implemented? GCNN was mentioned but I still have no idea what's in the model.

* There is no ablation tests, so it's hard to see what the various loss terms or model designs offer. 

* There is no discussion of the limitations of the method.

* Consider adding some explanation of the graph FPE method. FPE models the evolution of the probability p, not its encoding Phi(p) as in eq (4). I suppose it is explained in the cited reference. But since it is a central piece of this paper, please explain how this works.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733118055267

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comments by Authors

comment:
# General Comments by Authors

We sincerely thank all the reviewers for their insightful reviews and valuable comments, which are instructive for us to improve our paper further.

The reviewers generally believe that this paper addresses an important problem within the community, acknowledging that the proposed method is "**novel**", "**a potentially high-impact contribution**", and "**well-supported by solid rationale**",  the paper "**has  strong motivation**", "**addresses an interesting learning problem**", and "**is  clearly written**", and the experiments "**cover a wide range of application domains**", "**are cross-disciplinary**", and "**enhance the credibility of the method**".

The reviewers also raised insightful and constructive concerns. We made every effort to address each of the reviewers’ points by providing sufficient evidence and the required results. Below are the updates and responses to all of the reviewers' suggestions and concerns.

1. **Robustness of experiments**

   We provided results demonstrating the consistency of the proposed method under different hyperparameters. Additionally, we included experiments with varying noise strengths to support the claim of the method's robustness. These experiments address Reviewer **iFFC**'s concerns.

2. **Analysis of transferability**

   We added cross-system transfer experiments to assess the model's generalization performance when available data is sparse or biased, addressing the questions raised by Reviewer **C7fR** and **iFFC**.

3. **Analysis of scalability**

   We provided an analysis of the model's computational cost as it varies with data size and codeword count, and added experiments to examine how the required codeword count increases with system scale. These experiments support the scalability of the proposed method and address the concerns raised by Reviewer **iFFC** and **C7fR**.

4. **Analysis of interpretablity**

   We validated through experiments how energy estimation and evolutionary prediction influence each other’s performance, which further strengthens the motivation of this paper. This was acknowledged by Reviewer **iFFC**.

5. **Additional ablation study**

   We added ablation experiments on the loss function and key modules, addressing the concerns raised by Reviewer **C7fR** and **jDjA**.

6. **Clarification of experiment protocols**

   Following Reviewer **ciD9**'s suggestion, we clarified the consistency and fairness of the experimental setup and evaluation metrics with respect to prior work. Additionally, we added supplementary experiments on the degradation of the observation function g to ensure the completeness of the experimental setup.

7. **Polished writings**

   We refined the introduction of the proposed method and model components, and added analyses of the model architecture, loss function, and applicable scenarios to make the work more accessible to the audience. This was acknowledged by Reviewers **iFFC**, **jDjA**, and **C7fR**.

All updates are highlighted in red. Compared with the first submissions, $\underline{\text{the revised paper}}$ has **additional 7 pages**.

We sincerely appreciate the active and constructive feedback from the reviewers. All reviewers who have participated in the discussion so far have acknowledged that we have addressed their concerns and provided positive rating.  Although Reviewer **ciD9** has not provided further feedback or participated in the discussion, our responses to similar concerns (generalization and method motivation) have been acknowledged by the other reviewers (**C7fR** and **iFFC**). Additionally, we have provided the additional experiments and clarifications in response to his questions regarding the experimental setup. We believe that the concerns have been thoroughly addressed in our responses and revisions.

The valuable suggestions from reviewers are very helpful for us to revise the paper to a better shape. We'd be very happy to answer any further questions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper is on estimating the gradient dynamics of a scalar energy field. For dimensionality reduction, vector quantization techniques are proposed wherein the latent space is discretized, and the energy field is approximated using the energy of a codeword. To respect locality of codeword regions and model the time-dependent evolution of the probability distribution over codewords, a graph-based Fokker-Planck equation is formulated. Experiments are reported on three problems: 2d Prinz potential, ecological evolution, and fast-folding peptides. 

Strengths: Good combination of ideas: vector-quantization, self-supervised formulation etc.  Encouraging experimental results.

Weaknesses:  Presentation clarity is inconsistent.  Generalizability of the approach outside of visited regions. Limited ablations and lack of comparisons with published results.