Predicting the Energy Landscape of Stochastic Dynamical System via  Physics-informed Self-supervised Learning
Authors: Ruikun Li
Huandong Wang
Qingmin Liao
Yong Li
Venue: ICLR 2025 Poster (2025)
Keywords: dynamical system
energy landscape
deep learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We propose a self-supervised learning framework to predict the energy landscape of stochastic dynamical systems without supervisory signals for energy, while collaboratively achieving accurate energy estimation and evolution prediction.

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.

Published as a conference paper at ICLR 2025
PREDICTING THE ENERGY LANDSCAPE OF STOCHAS- TIC DYNAMICAL SYSTEM VIA PHYSICS-INFORMED SELF-SUPERVISED LEARNING
Ruikun Li
Huandong Wang'
Shenzhen International Graduate School Tsinghua University
Department of Electronic Engineering
BNRist, Tsinghua University
1irk61 20gma .com
wanghuandong nighua edu
Qingmin Liao
Yong Li
Shenzhen International Graduate School Tsinghua University
Department of Electronic Engineering BNRist, Tsinghua University 1 Loygnn0@0tss (onnggg7@tsinghua edu. cn
1 lawgmmtsiinggeeee edu.
ABSTRACT
Energy landscapes play 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 dif- fusion, where the energy governs the long-term motion of the particles. Estimat- ing the energy landscape of system has been a longstanding interdisciplinary challenge, hindered by the high operational costs or the difficulty of obtaining su- pervisory signals. Therefore, the question of how infer the energy landscape the absence of true energy values is critical. In this paper, we propose physics- informed self-suppr learning method learn the energy landscape from the evolution trajectories of the system. It first maps the system state from the obser- vation space to a discrete landscape space by an adaptive codebook, and then ex- plicitly integrates energy into the graph neural Fokker- fokker-Planck equation, enabling the joint learning of energy estimation and evolution prediction. Experimental re- sults across interdisciplinary systems demonstrate that estimated energy has correlation coefficient above 0.9 with the ground truth, and evolution prediction accuracy exceeds the baseline by average of 17.65%. The code is available githubcom/ssinghu---l-llliiiiii
INTRODUCTION
Energy landscapes are inherent in many stochastic dynamical systems in nature, such as the potential energy surface of protein conformations (Norn et al.) 2021). the fitness landscape of species evolu- tion (Papkou et al. 2023 Poelwijk et al.) (007). and the fractal energy landscapes of soft glassy materials. The evolution of these systems can modeled as particles moving on the landscape under the combined effect of energy-driven drift and noise-induced diffusion. The structure of the energy landscape governs the long-term motion of particles, forming the deterministic aspect of the dynamics, while inherent random noise disrupts the movement along the energy gradient, driving exploration across energy barriers (Blount et al., 2018 Kryazhimskiy low-energy regions exist in the landscape, the combined effect of the energy gradient and noise in- duces high-frequency movement within individual regions and fow-frequency transitions between different regions (Lin et al. 2024). In this context, energy landscapes have been applied to guide the generation of stable molecular structures (Noé et al., (219)) and direct the evolution pro- teins (Packer & Liu 2015 Greenbury al.) 2022), and more recently, they have been incorporated as physical knowledge into deep learning for predicting system evolution (Guan et al. 2024 Wang
al., 2014).. When multiple
al. 2024b Ding al. 2024).
"Corresponding author

---

Published as a conference paper at ICLR 2025
Due to its fundamental role in governing the system dynamics, estimating the energy landscape of dynamical systems has become an essential research problem across various disciplines. [Couce ee al.] (2024) cultivate 50,000 generations of bacteria measure the fitness effects of mutations, while Sarkisyan et al. (2016) measure tens of thousands of luminescent protein genotypic sequences to construct the functional landscape. These manual experimental approaches are not only costly but also heavily reliant on expert knowledge. With the success of deep learning in numerous dis- ciplines (Jumper et al. 2021; Han et al. 2023; Wang et al.) 2023] Chen et al.) 2024), several deep learning models have been proposed to estimate energy or equivalent quantities based on molec- ular spatial structures (Zhang et al. 2018), species genotypes (Tonner et al.] (022). or population compositions (Skwara et al., 2023). These methods still require high-cost annotations to provide supervisory signals for energy, which limits their practicality. In real-world scenarios, it is typi- cally more accessible to obtain abundant low-cost evolutionary trajectories of the system, which inherently embeds information about energy-driven drift (Weinstein et al. (022). Therefore, an im- portant research question arises: can estimate the energy landscape only based on the system's evolution trajectories in a data-driven manner?
However, estimating the energy landscape from evolutionary trajectories remains complicated problem with the following challenges. First, observable evolutionary trajectories typically cover only a limited portion of the vast state space. For instance, there are approximately 10t1 potential triple mutants of a typical protein, while available high-throughhut measurement techniques can only handle around 10 10" distinct genotypes, covering just a small fraction of the mutational space surrounding the natural sequence (Tonner et al., 2022). Second, distilling energy information from evolutionary trajectories requires building a model incorporating the energy landscape and the dis- tribution of trajectory data, thereby establishing connections between them. Classical Markov state models (Noé et al., 2019) establish this connection by strictly assuming that sampled data follow a Boltzmann distribution derived from the energy, which unrealistically demands that trajectories are fully sampled from a thermodynamic equilibrium state. In contrast, existing self-supervised learn- ing methods (Kamyshanska & Memisevic, 2014) treat neural networks black-box models to fit data distributions, completely disregarding the guidance of physical knowledge in terms of energy and system evaluation. Currently, there is still effective model that organically integrates Al techniques and physical knowledge for energy estimation without supervisory signals.
In this paper, we propose Physics-informed Energy Self-supervised Landscape Analysis (PESLA) method to estimate the energy landscape from historical evolution trajectories in a self self-super manner. PESLA maps the system state from the observed space discretized latent space via vec- tor quantization techniques (Van Den Oord et al., 2017). Through adaptively learning a codebook to partition the vast state space, our model concentrates on the essential shapes of the energy land- scape in discrete domains, thus disregarding the negligible information of the energy landscape and overcoming the challenge posed by limited observations. Then, PESLA utilizes the self self-superv signal from the prediction error of the system state to guide energy estimation. In this process, a graph neural ODE inspired by the Fokker-Planck equation is utilized to model the time evolution of probability distributions across different discretized states, and a physics-inspired regularization constraint is employed to integrate the prior knowledge of Boltzmann distribution of long-term dy- namies (Saoo & Nakagawa 2014), without relying on the assumption of thermodynamic equilib- rium sampling. These physics-inspired architectures serve as the bridge to distill information of the energy landscape from the system dynamies, thereby enabling the self-supervised learning of the system's 's energy landscape.
Our contribution can be summarized as follows:
We develop a novel framework to estimate the energy landscape of the system only utilizing the archi-
self-supervision signal from predicting the system state, where the tecture of graph neural Fokker-Plancl equation and physics-inspired regularization serves as the bridge to distill information of energy landscape from the system dynamics.
We develop a discrete encoding method of the system state to coarsen the continuous energy landscape using codebook obtained through vector quantization techniques. It allows our model to concentrate on the essential shapes of the energy landscape, effectively disregarding its negligible information and enhancing the sample efficiency of limited observational data. Experimental results across interdisciplinary systems demonstrate that PESLA reliably estimates system energy with absolute correlation coefficients above 0.9 and achieves 17.65% higher evo- lution prediction accuracy compared to of-e--oftthe-ar baselines.
2

---

Published as a conference paper at ICLR 2025
2 BACKGROUND AND PROBLEM SETUP
Let us consider stochastic dynamical process which can be described by the following differential equation:
ds; = f(s))dt + (ss,
It g(s))
(2)
Specifically, represents a system with latent state variable St E S whose evolution is driven by deterministic drag force f(sz) and random force described by white noise (st, t)dW (t). While the state variable &, is hidden and cannot be observed directly, the observable measurement E X' of the system is derived through a transformation, or nonlinear and can even represent a mapping from continuous space to discrete space, thereby describing systems with discrete observable metrics, such as ecological evolution.
S X, which can be either linear
More specifically, we focus on systems where the force [(+) conservative. This implies the existence of energy function E(st), also referred to the energy landscape, such that f(s)) -v)(s). Then, the dynamic equation can be be rewritten as:
dsc = V ( dt + (se, (t),
The energy landscapes measure the thermodynamic stability of given state. Low-energy regions induce drift that draws the system state into them with greater probability and duration, manifesting e B(a)/KTT where k is Boltzmann constant and 7 represents temperature. For evolution starting from any initial state distribution, the system's long-term dynamies will eventually drift toward the Boltzmann distribution defined by the energy landscape. Examples of such energy landscapes in different disciplines include fitness landscapes in ecology Papkou et al. (2023), potential energy in molecular dynamics Chmiela et al. (217)) and free energy in glassy materials Charbonneau et al. 2014)).
thermodynamically as the Boltzmann distribution,
Learning problem In this paper, our primary objective is to estimate the energy landscape of stochastic dynamical system based on its evolution trajectories, without the true energy as a super- visory signal. More formally, the input of this learning problem is set of the N-step evolution trajectory XN {Ie, Ji=0 of the stochastic dynamical system in the D-dimensional observation space X. Then, for an arbitrary observable state z, the objective of this learning problem is twofold: (1) building transformation & map the observable measurement to latent feature E(x) that determines the energy of the system; (2) estimating the energy É(E(x)) as an approximation of the true energy E(g (z)). Since the true energy E(g (xx) is unavailable as supervisory signal in the learning process, the estimated energy (EE()) only required to be linear transformation of the true energy.
3 METHOD
In this section, we introduce a Physics-informed Energy Self-supervised Landscape Analysis (PESLA) method, which learns predict the energy landscape through a self-supervised evolu- tion prediction task, as shown in Figure I First, we develop an adaptive codebook learning module to instantiate the mapping & from the observed space to the energy landscape. This approach inte- grates concepts from reduced-order modeling of complex systems to mitigate uncertainties caused by limited sample coverage. Next, we explicitly incorporate the energy function into a graph neural fokker-Plande equation to model the system's evolution the energy landscape. Additionally, we introduce physics-insping regularization constraints into the optimization objective to eliminate the assumption of thermodynamic equilibrium sampling.
3.1 ADAPTIVE CODEBOOK LEARNING
Constructing the energy landscape involves learning the transformation È from the observed space X' to the latent space S where the energy landscape resides. Previous studies have shown that, despite the high dimensionality of the state space, the long-term dynamics of systems unfold very low-lomeent dimensional manifold in the form of reduced-orde model (Vlachas et al.] 2022 Thibcault
et al.
2024 Li et al. 2024). This suggests that the energy landscape, which shapes the system's
3

---

Published as a conference paper at ICLR 2025
Energy Landscape
(b) Codeword Topology
Creph Tucker-Haenck
Ercnde
Dccode 4
Trojectory
X
Adaptive Cudebook
Framework of PESLA. (a) The energy landscape with evolution trajectories; (b) Parti- tioning the state space with an adaptive codebook to form the codewords with graph topology and modeling the time evolution of probability across neighhoring regions by graph neural Fokker- Planck equation.
Figure 1:
long-term evolution, has inherently low dimensionality, Similar phenomena have been observed in natural language processing and image representation, where of discrete codewords is sufficient to capture the essential representation of the original data (Van Den Oord et al. 2017 Razavi et al., 2019). Therefore, modeling the energy landscape as discrete reduced-ore -rrdrr model in the latent space S offers promising approach to addressing the challenge of the vast state space (Noé et al., 2009).
To implement such reduced-order approach and identify the energy landscape in the latent space S, we enhance the autoencoder with a learnable codebook = {or E 22 | 1, ..., K} to discretize the latent space of the encoded data. The transformations & and g between the observed space and the latent space are parameterized by and Ω, respectively. Specifically, after a sample x is encoded into a latent vector s, it is mapped to the most similar codeword c,, which then serves the input to the probabilistic decoder for reconstructing x. This k-nearest neighbor (KNN) style discrete aggregation partitions the latent space into multiple local regions (as shown in Figure [p), each uniquely represented by the energy of a codeword, thereby forming the low-dimensional landscape space. The encoding function maps the original space to the landscape space, capturing semantic features to ensure similar states fall into the same codeword region, thereby reducing reconstruction error. We emphasize that this design allows for optimal scaling of the state space partitioning from the limited coverage of observed data, rather than simple equidistant grid binning, as shown in Figure Baa (center). This adaptive scaling ensures the maximal utilization of codewords, enhancing the robustness to the preset number of codewords. Through the adaptive codebook encoding, the observed trajectories are mapped onto the energy landscape in the form of codewords, i.e., {Ct, N
3.2 GRAPH NEURAL FOKKER-PLANE
EQUATION
In the latent space S, the time evolution of the system state influenced by the combined effect of energy-driven drift and diffusion caused by inherent random noise, theoretically modeled by the Fokker-Plande equation (Risken| 1996). On the discretized energy landscape, we extend the traditional Fokker-Planck equation into graph neural differential equation, enabling joint learning of energy estimation and evolution prediction.
We construct the codeword topology A (aıj)x<K based the adjacency relationships of the codeword regions (as shown in Figure and estimate the energy of each codeword as E(c,) as the energy landscape
{A, C, E{+)} system evolution. At this point, have projected the origi- nal observed trajectory onto a now-dimensional energy landscape, obtaining the transition trajectory of the system state on the codeword topology. Predicting the temporal evolution of the system means modeling the time-dependent evolution of the probability distribution over the codewords. The ef- fects of energy and noise on this evolution are modeled by the Graph Fokker- Planck equation Chow
et al. 20122) as:
dp:
Σ B,,
((Ess + Blog Pi Jps+
Σ B
((Ess + Blog
Pj Pi
Jp:
Σ
3(p; p.),
dt
E,
(4)
where p; denotes the probability of node and E,g = E, E, E is positive constant which governs the noise strength. Denoting p(ct,) K-dimensional probability distribution at time lo.
4

---

Published as a conference paper at ICLR 2025
naively obtain the conditional probability distribution p(Ct++Ai (ta) of öt by performing a time integration of the Pokker-Planck equation on the initial condition p(c,.)-
after diffusion time
However, considering that the evolution of node's s state often depends its neighbors, projecting the onc-dimensional probability vector into a highherdimmnsionall space with stronger representa- tional capacity helps capture this rich relational structure. We employ graph convolutional neural networks (GCN) based probability encoder, (tt) = ((p(c,)) )), introducing neighborhood infor- mation through positional encoding (Chamberlain al., 2021b Yuan et al.) 2024a) to lift the one- dimensional probability vector into a high-dimensional representation. Thus, the time evolution of conditional state probabilities p ( [t.)) on the energy landscape modeled as a graph neural
diffusion process (Chamberlain et al. 2021a| Yuan et al. 2024b), formalized as
tottr OH(!) dt,
H(oo + At) H(to)
Øt
(5)
p(Ct)-ar(ct)) V (H + At)).
We extend Chow et al. (2012)'s 'ss theory by designing a graph neural Fokker-Plande equation to ex- plicitly model state diffusion driven by energy differences between neighhoring codewords
0 H Σ Way(B (log Ho, log (o + + o(k,)())))) HeJ, (6) 01
where Wiy is calculated by neighborhood attention. The learnable coefficient 3g represents the strength of noise acting between neighboring codewords, while is the scaling factor for the sigmoid activation function. The ablation study can be found in Appendix EE] where we demonstrate that modeling in the encoded probability space performs significantly better than directly modeling the probability vector.
3.3 TRAINING
The trainable parameters include the encoder E, decoder Ω codebook C, probability encoder Φ probability decoder V neighborhood attention weights W, coefficient vector ßg. and energy func- tion E(*). The detailed model architecture is provided in Appendix A.2 In the following, introduce the training procedure for the model.
Adaptive codebook learning and evolution prediction form joint learning task. The optimization objective for the former is to minimize the negative log-likelihood of the reconstructed distribution, i.c., Lreconstraat log (=).,c(x). In our experiments, we use a Gaussian prior distribution de- coder with negative log-likelihood loss for continuous systems, and cross-entropy loss for discrete systems. Additionally, the loss function Lug for updating codeword is consistent with the one pro- posed by Van Den Oord et al. (017)). Similarly, we minimize the negative log- ligeliiid in both the latent space and the landscape space for the evolution prediction task. In the latent space, we minimize the 1.2 error Trtaeet = [(((((++))) V (T(((L + At))],, while in the landscape space, we use cross-entropy Looute = - (p(++a)) log ((++At))
With the mapping of adaptive codebook, we can estimate the distribution p(c:) of observed sam- ples within the landscape space and employ the corresponding negative log-probability as reference energies guide energy estimation. However, this approach fails when evolution trajectories not sampled from a thermodynamic equilibrium state. Proven by statistical mechanics, the state probability distribution evolving in the form of Pokker-Planck equation converges the Boltzmann distribution. Although we cannot expect all sample data to be drawn from a thermodynamic equi- librium state, the long-term evolution of states will eventually converge the Boltzmann distribu- tion. This suggests incorporating regularization term into the long-term prediction task, expressed as the KL divergence between the empirical distribution p and the Boltzmann distribution q, i.c.,
Overall, conduct the training process by opti- mizing the aforemmenioned objectives Deeconstruct + Log + Liatent + Leade + Lpky- Detailed training strategies are provided in Section
Lphy (pl q) E p(c.) log
p(c:)
q(c,)
4. and ablation studies can be found in Appendix
5

---

Published as a conference paper at ICLR 2025
(a)
(b)
2D Prinz Potential
Ecological Evolution
Ours
0.8
APE
Ours APE
305
100
10%
20%
30%
40%
50%
60%
70%
80%
907
100%
10%
20%
30%
40%
50%
60%
70%
90%
80%
100%
Data Size
PT
PF
MSM
0.7762 = 0.0089
0.4354 + 0.0118
MSM
-0.8104 + 0.0012
-0.4787 + 0.0018
APE
0.5755 + 0.0261
0.4449 + 0.0738
APE
0.7063 0.0217
0..711 1 0.0465
PESLA
0.9290 1 0 0.0342
0.7419 0.0318
PESLA
0.9067 0.0100
-0.7582 0.0241
Figure 2: as a function of data size (top), and comparison across different methods (bottom); (b): trajectory energy correlation pr function of data size (top), and comparison across different methods (bottom).
Visualization of the results the energy estimation. (a): full-space energy correlation PF
4 EXPERIMENTS
We conduct experiments on three classic dynamical systems from different disciplines evaluate the accuracy of PESLA in (I) energy estimation and (2) evolution prediction. For fairness, use the same data preprocessing and apply grid search to fine-tune the learning rates and hyperparameters for all models. We perform 10 independent training and testing runs for each model to calculate the mean and standard deviation all evaluation metrics in each experiment.
4.1 SETUP
Baselines For the energy estimation task, we employ the Markov state model (MSM) (Majewski et al. 2023) and autoencoder potential energy (APE) (Kamyshanska & Memisevic 2014) base- lines. For the evolution prediction task, we compare PESLA with NeuralMJP (Seifner & Sánchez
2023 T-IB (Federici et al. 2024), VAMPNets (Mardt et al. 2018), and SDE-Net (Kong et al, 2020 Details on the implementation and hyperparameter searching of these baseline algorithms
be found in Appendix
Evaluation Metrics We evaluate the accuracy of energy estimation from two perspectives. The trajectory energy correlation pT represents the Pearson correlation coefficient between the predicted and true energies for all samples along a new trajectory, assessing predictive performance within the regions covered by training data. The full-space energy correlation AF measures the correlation coefficient for the energy of system states uniformly across the entire state space, accounting for unseen areas during training. For the evolution prediction task, all metrics are measured from M reference trajectories XT, M unfolding from randomly initialized system states, where 7 denotes the lag time of each step. All models are tasked with predicting evolution trajectories starting from these initial states, covering the same time span as the reference trajectories. We evaluate the accuracy of the predicted distributions by calculating the Jensen-Shannon divergence between the marginal (SJ and transition (S probability distributions of the predicted and reference trajectories across all states. For systems with continuous state space, discretize it into evenly spaced grid partitions, following previous work Federici et al. (2024): Arts al. (2023). Further details can be found in Appendix A.5
Training strategy We first train encoder E, decoder Ω2 and the feature vectors of the codewords C to construct the landscape topology. Then, we freeze them and train the parameters of the graph neural Fokker-Planck equation and energy function E(*) on the landscape. For all models, we use the Adam optimizer, with the learning rate decaying exponentially by factor of 0.99 each epoch.

---

Published as a conference paper at ICLR 2025
Robustness
(a)
Potential Surface
Codebook
Energy
1.5
720
Energy (Pred)
0.15
0.10
250
400
600
800
1000
Code Nar
-16
X
0
Potential (CT)
0.6
(b)
0.6
Taaeetary
Full Space
MJS 4
TJS +
250
4.00
632
800
1036
NeuralMJP
0.1463 0.0214
0.2282 1 0.0264
Code Num
T-IB
0.1668 0.0104
0.3385 = 0.0115
Coousted Coart Crunt
Occuparcy Rats
VAMPNet
0.1984 = 0.0196
0.3886 + 0.0206
100
30.0%
SDE-Net
0.2700 = 0.0153
0.3731 + 0.0285
18.4%
16.5%
1.5%
PESLA
0.1031 0.0125
0.1796 1 0.0234
250
400
6.0 820
1000
Code Har
Figure 3: Visualization of the results on the 2D Prinz Potential. (a): potential surface and short sam- ple trajectory (left), codebook distinguished by color and shape (center), and comparison between estimated energy and ground truth (right). Bluc/red represents low/high values in the heatmap: (b): measures of marginal and transition JS divergence for unfolded sequences at the lag time of 100; (c): The impact of the preset number of codewords on evolution prediction accuracy (top), energy estimation accuracy (center), and codeword occupancy (bottom).
4.2 2D PRINZ POTENTIAL
We first apply PESLA to the 2D particle movement system asymmetric potential energy 2024). The particle displacement is governed by the (X,)dd dr + odW, where the potential energy function (an
surface (Mardt et al., 2018; Federici et al. stochastic differential equation as dX,
2x 321 consists of four interconnected (left). A total of 10 trajectories with 100K time steps are generated and details on the generation and preprocessing can be found in Appendix[A.4 A.4 The results of energy estimation and evolution prediction are presented in Figure 2a (bottom) and Figure 3b respectively, where PESLA significantly outperforms the baseline methods in both tasks.
V, defined by V(x) = (zf TO low-energy regions, as shown Figure 3a
2 2x + 3x
in
Figure 3a (center) visualizes the adaptive codebook learned by PESLA from historical trajecto- ries, with different codewords distinguished by color and shape, representing their mapped regions in the original state space. The varying density of codewords at different locations directly re- flects PESLA's adaptive scaling partitioning. We emphasize that the adaptive codebook captures the dynamical knowledge the energy landscape, which fundamentally different from the simple equidistant grid-based binning approach. At macro level, the codebook divides the plane into an approximate 2x2 region corresponding to the four potential wells. The high energy barriers between potential wells serve as the boundaries of the four codeword regions. Low-energy wells are assigned more codewords (e.g., the bottom-left well), suggesting that the model recognizes the importance of low-energy regions as long-term dynamic attractors and allocates more attention them, which aligns with the higher accuracy observed in low-energy regions shown in Figure 3a (right). We present the codebooks of multiple independent experiments Appendix B demonstrating that this is not coincidental phenomenon. At finer level, multiple codewords are assigned to the center of each potential well, while the outer areas are divided into mapped regions approximately perpen- dicular the equipotential lines. This indicates that the random walk behavior induced by noise is effectively captured through the differences between codewords.
We examine the sensitivity of PESLA to data size and hyperparameters. Figure 2a (top) illustrates the impact of training data size on performance when all energy estimation algorithms are required to estimate across full-space samples. Due to the coarse-graining of adaptive codebook, PESLA
7

---

Published as a conference paper at ICLR 2025
Robustness
(a)
Fitness Landscape
Codebook
Energy (pr -0,9289)
720
g 0.35
53
59
0.30
20
45
60
60
103
22
Code Nar
0
X
(b)
Tajettay
Ful Space
MJS 4
TJS 4
40
66
60
100
NeuralMJP
0.4267 0.0717
0.4981 + 0.0919
Code Num
T-IB
0.3376 0.0291
0.3720 1 0.0518
100.0%
Occuparcy Rats
Coousttd Count
82
VAMPNet
0.4355 0.0553
0.5182 + 0.0728
8 3
SDE-Net
0.5612 = 0.0582
0.6417 + 0.0702
448%
PESLA
0.3111 = 0.0397
0.3277 = 0.0424
25
85
100
Code Ham
Figure 4: Visualization of the results on the Ecological Evolution. (a): fitness landscape and short sample trajectory (left), codebook distinguished by color and shape (center), and energy landscape fitted by RANSAC regression from estimated energy (right). Blue/red represents low/high values in the heatmap: (b): measures of marginal and transition JS divergence for unfolded sequences the lag time 7 of 10; (c): The impact of the preset number of codewords on evolution prediction accuracy (top), energy estimation accuracy (center), and codeword occupancy (bottom).
maintains optimal performance even with reduced data. In contrast, baseline methods are signifi- cantly limited by insufficient sample coverage in the state space. As shown in Figure 2a (top). as the data volume increases, the performance of MSM improves due the enhanced sample coverage. Figure [c reports the robustness of PESLA concerning the preset number of codewords. Although the number of preset codewords can be continuously increased, the actual number of occupied code- words converges automatically, and the accuracy of energy and evolution predictions reaches its peak.
4.3 ECOLOGICAL EVOLUTION
We examine the strong selection weak mutation system within eco-evolutionary dynamics, which is widely studied in ecology to understand the adaptive evolution of populations in specific environ- ments (Kryazhimskiy al. 2009; Bank et al. 2016). The fixation probability of candidate state (new mutation) with fitness J, is governed by the Kimura formula derived from the Wright-Fisher 2,,
model, given by Pi-j the selection coefficient. Sella & Hirsh (2005) have mathematically demonstrated that the loga- rithmic fitness of such evolutionary systems aligns with the energy of thermodynamic systems. We simulate KK trajectories, each with 100 time steps, under the two-locus setting where each locus has 64 possible mutation types as our dataset. Figure 2b (bottom) and Figure 4b respectively report PESLA's superior predictive performance for fitness and system evolution.
N where N represents the population size and sx()) = 4
In ecology, fitness measures the relative advantage of genotype and is negatively correlated with energy (Sella & Hirsh 2005). PESLA estimates the energy function of genotypes with correla- tion coefficient close to -I. The predicted energy is fitted with RANSAC regression model (see Appendix A.5) and visualized in Figure 44a (right). The distribution pattern of codewords within the codebook indicates that PESLA successfully identifies the set of genotypes with high fitness in eco-evolutionary dynamics. Moreover, since the genotype space is characterized by the Hamming distance, states in the same row column of the codebook are more likely to be mapped to the same codeword (Figure 4a (center)). This indicates that the adaptive codebook incorporates knowledge of system dynamics rather than relying simple equidistant grid binning.
8

---

Published as a conference paper at ICLR 2025
(a)
Homeodomain
BBL
BBA
NTL9
A3D
TRC
TIC
TIC
TIC
1100
1100
1100
(b)
T1C
0
TICO
T1C0
T100
1100
T100
(c)
0.5041
0.5
4
Energy (Pro
Enwoyy Pred)
0.0
Energy 1..
12
3
10 Energy (Ret)
Energy (Ref)
Energy (Ro0)
Energy (Ref)
Enengy (Rc0)
Figure 5: Visualization of the results on the Protein Folding. (a): reference energy landscapes of
each protein; (b): adaptive codebooks of each protein; (c): trajectory correlation coefficients between predicted and reference energy.
We also examine the impact data size and the preset number codewords on PESLA in this system. For the energy estimation within the sample-covered region (measured by Pr). PESLA shows minimal sensitivity to data size, as shown in Figure 2b)(t))) The influence of the preset number of codewords is similar to that observed in last case, which validates PESLA's robustness.
4.4 PROTEIN FOLDING
We apply PESLA to the folding data of five fast-folding proteins simulated by the Anton supercom- puter (Lindorff-Laut Larsen al., 2011). Each protein has two folding trajectories of equal length, used for model training and testing, respectively. Due to the lack of true energy, we estimate the reference energy using Time-lagged Independent Component Analysis (TICA) and the Markov State Model (MSM) based on the complete dataset (three times larger than the training data), consistent with pre- vious studies (Majewski al. 2023 Mardt al.) 2018). For each protein, the lag time used in TICA processing and experiments based the mean transition path time reported by Hondorffi-arreen ct al. (2011). The reference energy distribution on the 2D principal component plane identified by TICA shown in Figure 5a with implementation details provided in Appendix A.4 Each protein features a varying number (1 to 4) of fow-energy regions with different distributions, posing chal- lenges for energy estimation. Figure 5b) shows PESLA's partitioning of the state space for each pro- tein on the TICA principal component plane, demonstrating that PESLA differentiates low-energy, high-energy, and unknown energy regions with varying codeword aggregation rates. This automatic scaling ensures that PESLA's energy predictions remain consistent with reference values (as shown in Figure (cc). even in challenging protein folding problems. Additionally, PESLA achieved the best performance in the evolution prediction task (see Appendix

---

Published as a conference paper at ICLR 2025
5
RELATED WORK
5.1 ENERGY ESTIMATION
Estimating the energy landscape is crucial problem across multiple disciplines. The most funda- mental approach involves collecting data through modern sequencing techniques and manual exper- iments. (Sarkisyan et al. (2016) measured tens of thousands of Aequorea victoria (avGFP) deriva- tive genotypes to construct the local fitness landscape of green fluorescent proteins. Chen et al.] (202)) analyzed the fitness of all single mutations in VIM-2 3-lactamase across 64-fold range of ampicillin concentrations. Additionally, [aang et al. (2024)) conducted high-throughput functional genomics on Salmonella identify gene networks related adaptive effects. There have been many similar efforts (Starr et al. 2018). However, these manual experiments are often associated with high operational costs, making machine learning promising solution to improve this process in data-driven manner (Rupp et al. [012 Han et al. 2023). Zhang et al. (2018) introduce the deep potential molecular dynamics method, using neural networks model interatomic forces and potential energy. To mitigate overfitting issues in deep neural networks, Aggazadeh et al. (2021)
apply sparse recovery algorithms from coding theory for spectral regularization. Zhang et al. (022)
employ high-speed atomic force microscopy to collect data for training U-net model to predict the energy landscape of spatial angles on the DHR10-miee 10-micaN protein. Additionally, Tonner et al.
(202)) and Skwara al. (2023) offer interpretable predictions of mutation effects and population functions through hierarchical Bayesian modeling and polynomial regression, respectively. More recently, Du et al. (2024) developed graph neural network model intermolecular interactions, predicting Gibbs free energy in solute-solvent interactions. Despite these efforts, these models often depend on true energy values or molecular force fields as supervisory signals. In contrast these methods, our PESLA does not require supervisory signals for energy; instead, it learns to estimate energy through self-supervised evolution prediction task. An additional benefit of this approach is that the predicted energy effectively enhances the accuracy of evolution prediction.
5.2 EVOLUTION PREDICTION
Predicting the evolution of stochastic dynamical systems is challenging due the unknown un- derlying energy landscape. Vlachas et al. (2022) employ dimensionality reduction techniques to construct reduced-order models that capture essential macroscopic information, thereby simplifying the analysis of large-scale systems. To handle the challenge of modeling long-term dynamics, ap- proaches such as learning inne-invariant representations have been explored (Federici et al.) 2024] Kostic et al. 2024b Li et al.) 2023). Furthermore, Kostic al. (2022) 2024a) extend Koopman op- crator theory to map system states into a Hilbert space, facilitating the learning and interpretation of nonlinear dynamics. Wu et al. (2018) and Seifner & Sánchez (2023) represent stochastic dynamical processes as discrete state transitions within a Markov process framework. In contrast to existing methods, PESLA utilizes energy landscape knowledge to guide system dynamics modeling.
6 CONCLUSION
In this paper, we propose the PESLA method to estimate the energy landscape from historical evolu- tion trajectories in self-supervised manner. By integrating adaptive codebook learning and graph neural Tokker-Plande equation, PESLA collaboratively models the energy landscape and system dy- namics, even with limited observational data. We introduce physics-inspired regularization help PESLA move beyond the reliance on thermodynamic equilibrium sampling, Experimental results across various systems demonstrate that PESLA outperforms state-of-thed methods in both energy estimation and evolution prediction. PESLA does not require supervisory signals for energy, making a powerful data-driven tool for understanding and predicting stochastic dynamical systems.
Limitations and Future work This work focuses on estimating the energy landscape of a class of energy-drives evolutionary systems. However, when a system is driven by non-conservative forces, an energy landscape does not exist, 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. Future work will need to explore additional model designs to accommodate the dynamics of time-varying landscapes, where E(*) needs to adapt to F.(,, t).
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported in part by the National Natural Science Foundation of China under U23B2030, 62171260, and 922701 14. We sincerely appreciate the inspiration and valuable insights from discussions with Dr. Jiliang Hu.
REFERENCES
Amirali Aghazadeh, Hunter Nisonoff, Orhan Ocal, David H Brookes, Yijie Huang, O Ozan Koy- luoglu, Jennifer Listgarten, and Kannan Ramchandran. Epistatic net allows the sparse spectral regularization of deep neural networks for inferring fitness functions. Nature communications, 12 (1):5225, 2021.
Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Zugner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion mod- els and force fields for coarse-grained molecular dynamics. Journal of Chemical Theory and Computation, 19(18):6151-613 2023.
Claudia Bank, Sebastian Matuszewski, Ryan T Hietpas, and Jeffrey D Jensen. On the (un) pre- dictability of a large intragenic fitness landscape. Proceedings of the National Academy of Sci- ences, 11((4)): 14001-14455--099 2016.
Zachary D Blount, Richard E Lenski, and Jonathan B Losos. Contingency and determinism in evolution: Replaying life's tape. Science, 362(6415) 2018.
Ben Chamberlain, James Rowbottom, Maria I Gorinova, Michael Bronstein, Stefan Webb, and Emanuele Rossi. Grand: Graph neural diffusion. In International conference on machine learn- ing. pp. 1407-1418, PMLR, 2021a.
Benjamin Chamberlain, James Rowbottom, Davide Eynard, Francesco Di Giovanni, Xiaowen Dong, and Michael Bronstein. Beltrami flow and neural diffusion on graphs. Advances in Neural Infor- mation Processing Systems, 14:,,99-1 2021b.
Patrick Charbonneau, Jorge Kurchan, Giorgio Parisi, Pierfrancesco Urbani, and Francesco Zamponi. Fractal free energy landscapes in structural glasses. Nature communications, 5(1):3725, 2014.
Hongyi Chen, Jingtao Ding, Yong Li, Yue Wang, and Xiao-Ping Zhang. Social physics informed diffusion model for crowd simulation. In Proceedings of the AAAI Conference on Artificial Intel- ligence, volume 38, pp. 482, 2024.
JZ Chen, DM Fowler, and N Tokuriki. Environmental selection and epistasis in empirical phenotype-eniioomen--iie
landscape. Nature Ecology Evolution, 6(4):427-438, 2022.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science advances, e6 2017.
Shui-Nee Chow, Wen Huang, Yao Li, and Haomin Zhou. Fokker-plande equations for free energy functional or markov process on graph. Archive for Rational Mechanics and Analysis, 203: 969-1008, 2012.
Alejandro Couce, Anurag Limdi, Melanie Magnan, Siân V Owen, Cristina M Herren, Richard Lenski, Olivier Tenaillon, and Michael Baym. Changing fitness effects of mutations through long-term bacterial evolution. Science, 383(668) ))eadd1417, 2024.
Jingtao Ding, Chang Liu, Yu Zheng, Yunke Zhang, Zihan Yu, Ruikun Li, Hongyi Chen, Jinghua Piao, Huandong Wang, Jiazhen Liu, et al. Artificial intelligence for complex network: Potential, methodology and application. arXiv preprint arXiv:2402.16887 2024.
Wenjie Du, Shuai Zhang, Di Wu, Jun Xia, Ziyuan Zhao, Junfeng Fang, and Yang Wang. Mmgnn: A molecular merged graph neural network for explainable solvation free energy prediction. In Pro- ceedings of the Thirty-Thrrd Third International Joint Conference on Artificial Intelligence, pp. 5808- 5816, 2024.
11

---

Published as a conference paper at ICLR 2025
Marco Federici, Patrick Forré, Ryota Tomioka, and Bastiaan S Veeling. Latent representation and simulation markov processes via time-lagged information bottleneck. In The Tweifth Interna- tional Conference on Learning Representations, 2024.
Sam F Greenbury, Ard A Louis, and Sebastian E Ahnert. The structure of gemotype-phenotype maps makes fitness landscapes navigable. Nature Ecology & Evolution, 6(1 6(11):1742-1752 2022.
Xingyue Guan, Qian- Quan-Yuan Tang, Weitong Ren, Mingchen Chen, Wei Wang. Peter G Wolynes, and Wenfei Li. Predicting protein conformational motions using energetic frustration analysis and alphafold2. Proceedings of the National Academy of Sciences, 121(35):c2400662121, 2024.
Barbara A Han, Kush R Varshney, Shannon LaDeau, Ajit Subramaniam, Kathleen C. Weathers, and Jacob Zwart. synergistic future for ai and ecology. Proceedings of the National Academy of Sciences, 120(38):c222028812) 2023.
Moritz Hoffmaan,, Martin Scherer, Tim Hempel, Andreas Mardt, Brian de Silva, Brooke E Hu- sic, Stefan Klus, Hao Wu, Nathan Kutz, Steven I. Brunton, et Deeptime: a python library for machine learning dynamical models from time series data. Machine Learning: Science and Technology, 3(1):015009, 2021.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool. Russ Bates, Augustin Zidek, Anna Potapenko, al. Highly accurate protein structure prediction with alphafold. nature, 596(7873) 2021.
Hanna Kamyshanska and Roland Memisevic. The potential energy of an autoencoder. IEEE trans- actions on pattern analysis and machine intelligence, 3716)122) 2014.
Lingkai Kong, Jimeng Sun, and Chao Zhang. Sde-net: Equipping deep neural networks with un- certainty estimates. In International Conference on Machine Learning, pp. 5405-5415. PMLR, 2020.
Vladimir Kostic, Pictro Novelli, Andreas Maurer, Carlo Ciliberto, Lorenzo Rosasco, and Massimil- iano Pontil. Learning dynamical systems via koopman operator regression in reproducing kernel hilbert spaces. Advances in Neural Information Processing Systems, 35:4017-4031, 2022.
Vladimir Kostic, Karim Lounici, Pictro Novelli, and Massimiliano Pontil. Sharp spectral rates for koopman operator learning. Advances in Neural Information Processing Systems, 36, 2024a.
Vladimir R Kostic, Pietro Novelli, Riccardo Grazzi, Karim Lounici, and Massimiliano Pontil. Learn- ing invariant representations of hime-nomogeneous stochastic dynamical systems. In The Twelfth International Conference Learning Representations, 2024b.
Sergey Kryazhimskiy, Gašper Tkačik, and Joshua B Plotkin. The dynamics of adaptation cor- related fitness landscapes. Proceedings of the National Academy of Sciences, 106(44):18638- 18643, 2009.
Sergey Kryazhimskiy, Daniel P Rice, Elizabeth R Jerison, and Michael M Desai. Global epista- sis makes adaptation predictable despite sequence-level stochasticity. Science, 344(6191):1519- 1522, 2014.
Ruikun Li, Huandong Wang, and Yong Li. Learning slow and fast system dynamies via automatic separation of time scales. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 4380-4390, 2023.
Ruikun Li, Huandong Wang, Jinghua Piao, Qingmin Liao, and Yong Li. Predicting long-term dy- namics of complex networks via identifying skeleton in hyperbolic space. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 1655-1666, 2024.
Zequn Lin, Zhaofan Lu, Zengru Di, and Ying Tang. Learning noise-induced transitions by multi- scaling reservoir computing. Nature Communications, 15(1 15(1),6584, 2024.
Kresten Sindorff-Larsen Stefano Piana, Ron O Dror, and David E Shaw. How fast-folding proteins fold. Science, 334(6055):517-5220 2011.
12

---

Published as a conference paper at ICLR 2025
Maciej Majewski, Adrià Pérez, Philipp Thölke, Stefan Doemr Nicholas E Charron, Toni Giorgino, Brooke E Husic, Cecilia Clementi, Frank Noć, and Gianni De Fabrituis. Machine learning coarse- grained potentials of protein thermodynamics. Nature communications, 14(1):5739, 2023.
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets for deep learning of molecular kinetics. Nature commumications, 9(1):5, 2018.
Frank Noé, Christof Schüitte, Eric Vanden-Eijnden, Lothar Reich, and Thomas R Weikl. Construct- ing the equilibrium ensemble of folding pathways from short off-equilibrium simulations. Pro- ceedings of the National Academy of Sciences, 106(45):10111-9916, 2009.
Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 147, 2019.
Christoffer Norn, Basile IM Wicky, David Juergens, Sirui Liu, David Kim, Doug Tischer, Brian Koepnick, Ivan Anishchenko, Foldit Players, David Baker, et al. Protein sequence design by conformational landscape optimization. Proceedings of the National Academy of Sciences, 118 (11):2201722811) 2021.
Michael Packer and David R Liu. Methods for the directed evolution of proteins. Nature Reviews Genetics, 16(7):379-394, 2015.
Andrei Papkou, Lucia Garcia-Pastor, José Antonio Escudero, and Andreas Wagner. rugged yet easily navigable fitness landscape. Science, 382(6677))cadh38660 2023.
Frank Poelwijk, Daniel J Kiviet, Daniel M Weinreich, and Sander Tans. Empirical fitness land- scapes reveal accessible evolutionary paths. Nature, 445(7126):383-386, 2007.
Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating diverse high-fidelity images with vq-vae-2. Advances in neural information processing systems, 32, 2019.
H Risken. The lokker-planck equation, 1996.
Matthias Rupp, Alexandre Tkatchenko, and accurate modeling of molecular atomization energies with machine learning. Physical review letters, 10855300, 2012.
Müller, and O Anatole Von Lilienfeld. Fast
Karen S Sarkisyan, Dmitry A Bolotin, Margarita Meer, Dinara R Usmanova, Alexander S Mishin, George V Sharonov, Dmitry N Ivankov, Nina Bozhanova, Mikhail Baranov, Onuralp Soyle- mez, et al. Local fitness landscape of the green fluorescent protein. Nature, 533(7603):397--40,, 2016.
Issei Sato and Hiroshi Nakagawa. Approximation analysis of stochastic gradient langevin dynam- ics by using ookker-planck equation and ito process. In International Conference on Machine Learning. pp. 982 990. PMLR, 2014.
Patrick Seifner and Ramsés J Sánchez. Neural markov jump processes. In International Conference on Machine Learning, pp. 30523-30552. PMLR, 2023.
Guy Sella and Aaron E Hirsh. The application of statistical physics to evolutionary biology. Pro- ceedings of the National Academy of Sciences, 102(27):9541-9444 2005.
Abigail Skwara, Karna Gowda, Mahmoud Yousef, Juan Diaz-Colunga, Arjun S Raman, Alvaro Sanchez, Mikhail Tikhonov, and Seppe Kuehn. Statistically learning the functional landscape of microbial communities. Nature ecology & evolution, 7(11):1823-18 1131882338111-111111111111 2023.
Tyler N Starr, Julia M Flynn, Parul Mishra, Daniel NA Bolon, and Joseph W Thornton. Pervasive contingency and entrenchment in a billion years of hsp90 evolution. Proceedings of the National Academy of Sciences, 115(17):443-4444 4458, 2018.
Vincent Thibeault, Antoine Allard, and Patrick Desrosiers. The low-rank hypothesis of complex systems. Nature Physics, 2022):294302, 2024.
13

---

Published as a conference paper at ICLR 2025
Peter D Tonner, Abe Pressman, and David Ross. Interpretable modeling of landscapes with state-of-the-an predictive power. Proceedings of the National Academy of Sci- ences, (19(6))e221100))) 14021111110211199 2022.
Aaron Van Den Oord, Oriol Vinyals, al. Neural discrete representation learning. Advances in neural information processing systems, 30, 2017.
Pantelis R Vlachas, Georgios Arampatzis, Caroline Uhler, and Petros Koumoutsakos. Multiscale simulations of complex systems by learning their effective dynamics. Nature Machine Intelli- gence, (4)::59-3666 2022.
Benjamin X Wang, Dmitry Leshchiner, Lijuan Luo, Miles Tuncel, Karsten Hokamp, Jay CD Hinton, and Denise M Monack. High-throughput fitness experiments reveal specific vulnerabilities of human-adapted salmonella during stress and infection. Nature Genetics, pp. 1-12, 2024a.
Hanchen Wang. Tianfan Fu, Yuanqi Du, Wenhao Gao, Kexin Huang, Ziming Liu, Payal Chandak, Shengchao Liu, Peter Van Katwyk, Andreea Deac, et al. Scientific discovery in the age of artificial intelligence. Nature, 620(7972):47 2023.
Huandong Wang. Huan Yan, Can Rong, Yuan Yuan, Fenyu Jiang, Zhenyu Han, Hongjie Sui, De- peng Jin, and Yong Li. Multi-scale simulation of complex systems: perspective of integrating knowledge and data. ACM Computing Surveys, 56(12):1-38, 2024b.
Eli Weinstein, Alan Amin, Jonathan Frazer, and Debora Marks. Non-identifiability and the blessings of misspecification in models of molecular fitness. Advances in neural information processing systems, 35:5484-5497, 2022.
Hao Wu, Andreas Mardt, Luca Pasquali, and Frank Noe. Deep generative markov state models. Advances in Neural Information Processing Systems, 31, 2018.
Tao Wu, Xiangyun Gao, Feng An, Xiaotian Sun, Haizhong An, Zhen Su, Shraddha Gupta, Jianxi Gao, and Jürgen Kurths. Predicting multiple observations in complex systems through low- dimensional embeddings. Nature Communications, 15(1):2242, 2024.
Yuan Yuan, Jingtao Ding, Jie Feng, Depeng Jin, and Yong Li. Unist: A prompt-empowered universal model for urban spatio-temporal prediction. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pp. 4106, 2024a.
Yuan Yuan, Chonghua Han, Jingtao Ding, Depeng Jin, and Yong Li. Urbandit: A foundation model for open-world urban spatio-temporal learning. arXiv preprint arXiv:2411.12166 2024b.
Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and Weinan E. Deep potential molecular dynamics: scalable model with the accuracy of quantum mechanics. Physical review letters, 123(14): 2018.
Shuai Zhang. Robbie Sadre, Benjamin A Legg, Harley Pyles, Talita Perciano, E Wes Bethel, David Baker, Oliver Rübel, and James De Yoreo. Rotational dynamics and transition mech- anisms of surface-adsorban proteins. Proceedings of the National Academy of Sciences, 119(16): e2020242119, 2022.
14

---

Published as a conference paper at ICLR 2025
A EXPERIMENTAL DETAILS
We provide details on the experimental systems, baselines, evaluation metrics, and implementation to ensure clarity and reproducibility of the reported results.
A.1 INTRODUCTION AND IMPLEMENTATION OF BASELINES
Our baselines cover both the energy estimation and evolution prediction tasks. For the energy esti- mation task, we have:
Markov State Model (MSM) is commonly used method for estimating the relative energy of system states based on statistical probabilities. It discretizes the system using equidistant grid binning and then calculates the negative log of the frequency distribution for all states in the dataset as the reference energy. This approach is often limited by the inefficiency of Monte Carlo sampling. When the dataset fails cover the entire sample space, some state frequencies become zero, making it impossible to infer unobserved sam- ples from the existing data. In our experiments, we used nearest-neighhor interpolation to compute the full-space energy correlation coefficient PF for quantitative evaluation.
Autoencoder Potential Energy (APE). Kamyshanska Memisevic (2014) demonstrate that autoencoder can estimate the energy of a sample by treating the reconstruction error as a proxy for energy, where a lower reconstruction error indicates that the sample lies in nigh-probability, low-energy region of the learned manifold, while higher error corresponds to higher energy. For an autoencoder with sigmoid activations, with weights W, hidden biases by. and reconstruction biases b,., the energy function is given by:
F(x) Σ log(1 + exp(W x bn.)) || b,.11 + const,
where W x represents the linear combination of inputs, bng is the hidden bias term for the k-th hidden unit, and b, is the reconstruction bias.
For the evolution prediction task, we compare PESLA with:
Neural MJP. /Seifner& & Sánchez (2023) introduce Neural MJP as an alternative variational inference algorithm for Markov jump processes, which relies on neural ordinary differen- tial equations in the form of the master equation. Neural MJP predefines the number of discrete states and encodes observed states into one-hot vector representing the discrete state distribution as the starting point for state evolution. The key difference between these predefined discrete states and PESLA's codewords is that Neural MJP does not character- ize them by energy but instead relies on a black-box neural network to fit the transition probabilities. The number of preset discrete states is treated as hyperparameter.
T-IB (Federici et al., 2024) captures time-invariant representations continuous dynami- cal systems using representation learning objective derived from information bottleneck theory and models state transitions the representation space through a conditional flow model. This efficient representation allows T-IB filter out high-frequency fluctuations as noise and model long-term dynamics over extended time spans.
VAMPnet (Mardt et al., 2018) captures the dynamics of molecular systems by directly learning transformation from molecular configurations to Markov state model using deep neural network that maximizes a variational score. This end-to-end approach allows it to identify slow dynamical processes and long-timescale kinetics effectively.
SDE-Net (Kong et al., 2020) explicitly models the drift and noise diffusion terms of stochastic dynamical systems by parameterizing these two mechanisms within neural dif- ferential equation, enhancing the representational capacity of the neural network. However, NeE-Net does not treat the dynamical system Markov process, making it challenging to capture transition characteristics between metastable states. Despite this limitation, serves as a benchmark for all models.
15

---

Published as a conference paper at ICLR 2025
A.2 ARCHITECTURE PESLA MODEL
We summarize all components of the PESLA model and the parameter shapes of each component Table [ where is the dimension the observed state, K is the preset number codewords, and is the proportion of activated codewords.
Table 1: Module and layer specifications.
Module
Layer name
Parameter shape
i. Adaptive Codebook Learning
Encoder =
Layer-FC
[D, 64, 32]
Codebook C'
[K, 32]
Layer-FC
[32, 64, 64]
Gaussian Decoder ?
Linear (mu)
[64, D]
Linear-Sigmoid (std)
[64, D]
ii. Graph Neural Planer Equation
Energy Function E(())
Layer-FC
[32, 64, 1]
Positional Encoding
[rK, 3]
Probability Encoder +
GCN-FCI
[3+1+1, 64]
GCN-FC2
[64, 64]
Linear (q)
64]
Neighborhood Attention W
Linear (k)
64]
Coefficient Vector 3g
[64.]
GCN-FC1
[64, 64]
Probability Decoder
GCN-FC2
[64, 1]
A.3 GRID SEARCHING FOR HYPERPARAMETERS
To ensure fair comparison across all models, we used the same batch size, optimizer, and learning rate decay strategy during training, conducting grid search only on the learning rate and model- specific hyperparameters to achieve optimal performance. The range and targets of the hyperparam- eter search are detailed in Table
Table 2: Hyperparameter search settings.
Model
Learning Rate
Specific Hyperparameters
Description
NJu
0.01 0.0001
10 1000
The preset number of discrete states Information bottleneck coefficient Output dimensionality of the encoder Intensity coefficient of noise term The preset number of codewords
T-IB
0.01 0.0001
0.01 1.0
VAMPNet
0.01 0.0001
10 1000
SDE-Net
0.01 0.0001
0 1.0
PESLA
0.01 0.0001
10 1000
A.4 DATA GENERATION OR PREPROCESSING
Consistent with previous studies, we simulate the 2D Prinz potential and ecological evolution sys- tems to obtain the datasets. For protein folding, we use the official data provided by the authors, as described in Table 3 For the first two systems, the training and testing sets split in 7:3 ratio. For the protein data, each protein has two trajectories equal length, one used for training and the other for testing.
16

---

Published as a conference paper at ICLR 2025
Figure 6: Visualization of the adaptive codebooks from five independent experiments PESLA on the 2D Prinz potential.
The folding data for the five proteins includes 3D spatial coordinates of 28 to 73 alpha-carbon atoms. We performed TICA implemented by Deeptime (Hoffmann et al. 202T), extracting the linear projections of the top two principal time components for each protein as preprocessing. which were then used for estimating reference energy and testing all models.
Table 3: Protein folding dataset.
Homeodomain
BBL
BBA
NTL9
A3D
C-atom Num
52
28
39
73
Trajectory Length (Ju8)
100
100
200
300
300
Time unit (ns))
10
10
10
10
10
A.5 EVALUATION METRICS CALCULATION
For the evolution prediction task, we test each model using the following steps:
1. Randomly initialize M initial states and predict N future steps using the model to obtain XTM;
2. Discretize each dimension into a K x K state space using a uniform grid, resulting in finite set of discrete states;
3. Compute the marginal and transition probability distributions for each state across all M trajectories;
4. Calculate the Jensen- Sansen-Shannon divergence of the marginal and transition probabilities for each trajectory and take the average.
We set K to 5, 8, and 8 for the 2D Prinz potential, ecological evolution, and protein folding systems, respectively.
For the trajectory energy correlation reported in Figure 4a (right), we used the Random Sample Consensus (RANSAC) regression algorithm fit the maximum likelihood expression Epred ((Errue) of the true energy and PESLA's predicted energy. We then mapped the data from Figure 4a (left) using and visualized in Figure 4a (right).
B SUPPLEMENTARY EXPERIMENTAL RESULTS FOR 2D PRINZ POTENTIAL
We provide the codebooks from five independent experiments on the 2D Prinz potential in Figure showing that PESLA consistently learns similar codeword distribution patterns.
C SUPPLEMENTARY EXPERIMENTAL. RESULTS FOR PROTEIN FOLDING
We provided a supplementary comparison of the predictive performance of all models using the BBA protein example, shown in Table
17

---

Published as a conference paper at ICLR 2025
Table Predictive performance of all models on the BBA protein data with lag time of 0.7 pss and 8x8 grid discretization. All experiments are run 10 times to obtain statistical values.
MJS L
TJS
NeuralMJP
0.0231 = 0.0048
0.3637 = 0.0051
T-IB
0.0388 + 0.0053
0.4327 : 0.0083
VAMPNet
0.0411 + 0.0021
0.4566 + 0.0228
SDE-Net
0.0561 = 0.0181
0.5502 = 0.0611
PESLA
0.0207 = 0.0061
0.2468 = 0.0215
D SUPPLEMENTARY EXPERIMENTAL RESULTS FOR ADDITIONAL EXPERIMENTS
D.I INTERPRETABILITY
Our PESLA synergistically estimates energy and predicts trajectories to simultaneously improve the accuracy of both. Although the quality of evolution prediction directly influences the precision of energy estimation, it remains unclear how the accuracy of energy estimation, in turn, impacts evo- lution prediction. Here, investigate how the correlation between the estimated energy landscape and the true energy landscape influences the evolution prediction. Specifically, we aim to clarify the degree of correlation required between the predicted energy and the true energy to ensure accurate evolution prediction.
We disable the energy prediction module of PESLA, replacing the predicted energy of each code- word with a dummy energy value. When the Pearson correlation coefficient, denoted as p, equals 1.0, the dummy energy is derived from the mean true energy values of all samples within the region of each codeword. We gradually introduce noise to the dummy energy to reduce its correlation with the true energy, illustrated in the first row Figure 7 Subsequently, we train PESLA under various dummy energy conditions and evaluate the prediction error. As shown in Figure 7 the pre- diction error progressively increases as the correlation coefficient between the dummy energy and the true energy decreases. When the correlation coefficient drops below 0.5, PESLA's predictive performance begins lag behind the optimal baseline algorithm (NeuralMJP). This indicates that the quality of evolution prediction is directly influenced by the accuracy of energy estimation.
D.2 CONSISTENCY
Although the degree of discretization of the state space depends the predefined number of code- words, a robust prediction model should yield consistent energy landscapes across different settings. We evaluate the correlation between energy values predicted by PESLA under various hyperparam- eter settings (predefined number of codewords K)) and random seeds. As shown in the correlation matrix in Figure the energy landscapes identified by PESLA remain consistent not only across par- allel experiments with different random seeds but also across different choices of hyperparameters K.
D.3 NOISE ROBUSTNESS
Considering that real-world trajectory data is usually noisy and sparse, the robustness of predic- tive model to noise and limited data determines its practical utility. As reported in Figure 2a of the main text, PESLA outperforms all baselines when available data is reduced. Here, we further evaluate PESLA's robustness to noisy data. Specifically, we add Gaussian noise of varying strength to the dataset, with a noise amplitude equal to the original data magnitude when the strength is set to 1.0. The results in Figure 9 indicate that PESLA remains sufficiently robust to noise until the noise strength exceeds 0.6. PESLA'S robustness to noise can be attributed to its adaptive codebook learning model, which incorporates compact representation of the original state space, PESLA inherently possesses the ability to filter uncertainties such noise-related errors.
aeduced-orde approach. By identifying low-dimensional,
18

---

Published as a conference paper at ICLR 2025 5
Figure 7: Evolution prediction accuracy as a function of the corre lation coefficient P between the dummy energy and the true energy on the 2D Prinz potential. The skyblue and red dashed lines are the optimal baseline and PESLA performances the main text, respectively. All experiments are run 10 times to obtain statistical values.
Lower Correlation
0.0
PESLA
0.
PES
0.1
0.2
0.2
0.3
03
04
0.5
0.6
40
0.7
0.8
2
7
0.
Optim
Oytima
10
0.175
0.50
0.25
0.00
0,30
0.20
0.15
19

---

Published as a conference paper at ICLR 2025
(a)
(b)
Figure 8: Correlation matrix energy values predicted by PESLA for BBA protein at different (a)
preset numbers of codewords K and (b) random seeds.
0.4
MJS
TJ5
Pr
1.0
Energy Correlation
15 Divergence
0.3
0.8
0.6
0.4
0.1
-0.2
0.0
0.2
0.4
0.6
0.8
1.0
Noise Strength
Figure 9: Energy and evolution prediction accuracy on the 2D Prinz potential as functions of noise strength. All experiments are run 10 times to obtain statistical values.
D.4 TRANSFERABILITY
Here, we explore the transferability of PESLA. Although the energy landscapes of different evo- lutionary systems are inherently distinct, PESLA's modules can be partially reused in similar state spaces by learning a generalized spatial mapping mechanism. We proceed by evaluating PESLA's transferability across five different protein-folding datasets.
Due to differences sequence length and arrangement, the folding processes of different proteins occur within their unique energy landscapes, which means that the energy function E(() and transi- tion model need to be specifically trained for each type of protein. However, the mapping functions and between the observed space X' and the latent space S have the potential for transferabil- ity. By learning a universal encoder 3. decoder Ω and codebook C, it is promising to project the structures of various proteins onto a unified latent space.
To test such transferability, we conduct cross-protein experiments for each protein. Specifically, for a given protein i, we train the encoder E, decoder Ω and codebook using data from the other four proteins. We then freeze their parameters and use single folding trajectory of protein i to train the energy function E(*) and the Graph Neural Fokker-Planck equation. Finally, we evaluate the accuracy of energy and evolution predictions unseen folding trajectories of protein i. The results are presented in Table 5 Although the predictive performance all cross- roos--rootein transfer experi- ments is lower than that achieved by training specific proteins, the average transfer performance P: for energy prediction across all proteins reaches over 80% of the performance of specifically trained models. Additionally, the transfer performance of evolution prediction significantly better
20

---

Published as a conference paper at ICLR 2025
than that of the optimal baseline under the same training settings. We believe that the number of available proteins increases, the model's transferability will show promising improvement. The validation experiments in this section demonstrate its feasibility.
Table 5: Comparison of mean Pt. MJS, and TJS metrics for encoder E, decoder Ω and codebook C trained on specific protein data versus cross- cross-protein data for Homeodomain, BBL,, BBA, NTL9, and A3D, All experiments are run 10 times to obtain mean values.
Homeodomain
BBL
BBA
MJS
TJS
Pt
MJS
TJS
Pt
MJS
TJS
PESLA-speciffe
0.9341
0.0203
0.2342
0.9014
0.0200
0.2322
0.9179
0.0207
0.2468
PESLA-croSs
0.8583
0.0875
0.3510
0.7014
0.0775
0.4362
0.6665
0.1055
0.4065
NeuralMJP-cross
0.5837
0.7724
0.5303
0.6703
0.4382
0.5651
NTL9
A3D
Pt
MJS
TJS
Pt
MJS
TJS
PESIA--ppecffic
0.8867
0.0167
0.2625
0.8186
0.0414
0.3055
PESLA-crOss
0.7443
0.1089
0.4597
0.6235
0.2068
0.6034
NeuralM.JP-cross
0.3539
0.4525
0.7340
0.7983
D.5 SCALABILITY
To investigate the relationship between the size of the state space and the codebook size, we evaluate the impact of the preset number of codewords K energy and evolution prediction in protein folding datasets with varying numbers of alpha-C atoms, as shown Figure T0 As K increases from 10 to 1000, the relative performance of the model improves. For the BBA protein, which has only 28 alpha-C atoms, the prediction accuracy for energy reaches over 90% of the performance observed at K 1000 when K 100. For larger proteins, such as A3D, the model's predictive performance converges at K = 500. In fact, protein size increases the complexity of the state space, thereby adding to the modeling challenge. For larger proteins or other systems with complex state spaces, the codebook size needs to be sufficiently large to ensure PESLA's modeling capacity. For the systems studied in this paper, we recommend setting the preset number of codewords K to 1000. For other unfamiliar systems, starting with relatively large K value is generally advisable.
Furthermore, as analyzed in Appendix F the time complexity for training and inference grows sub- linearly with the increase in K. Since only limited number of codewords are activated in the preset codebook, the size of the energy landscape constructed by PESLA is at most K. Consequently, when modeling the state transition distribution over the landscape using the Graph Neural Fokker-Planck equation, the number of codewords to be considered does not necessarily increase lincarly with K. This design allows users to efficiently explore and select appropriate values for K.
E ABLATION STUDIES
PESLA comprises multiple loss function terms and submodules. Here, we introduce additional ab- lation studies elucidate the individual contribution of each component to the overall performance.
As summarized in Section 3.3 the training process involves five loss function terms: Lyees Leg, Ltatents Leode: and Lphy- The Lree and Lug terms jointly guide the adaptive codebook learning module, while Titaeet and Lende direct the synergistic learning process for energy and evolution prediction. Additionally, the Tppyy term incorporates physical knowledge to further inform energy estimation. Beyond the essential loss terms, individually evaluate the performance impact of the auxiliary terms, Ltatent and Lphy- The results are presented Table 6
We firstly remove the Liatent term, training the prediction module solely with the Loode term. The results indicate that, without the predictive constraint from the latent space, the accuracy of evo- lution prediction deteriorates, and the precision of energy estimation is also affected. Next, upon removing the Lphy term, PESLA's performance in both energy and evolution prediction declined significantly. This suggests that, although self-supervised learning on the evolution prediction task
21

---

Published as a conference paper at ICLR 2025
0.3
025
000
107
157
20'
10
101
(a) Evolution Prediction
(b) Energy Prediction
303
200 400004
200
207
100
103
300
Nemerdoman
(c) Minimum codebook for evolution
(d) Minimum codebook energy
Figure 10: Mean (a) MJS and (b) Pt for proteins of different sizes function of codebook size K, normalized by the metric at K 1000. (c) and (d) report the minimum codebook size required achieve 90% prediction performance for different proteins. All experiments are run 10 times to obtain mean values.
can drive energy estimation, physical knowledge remains crucial for guiding this joint optimization task effectively.
In Section 3.2 we enhance the model's capability by encoding one-hot probability vectors as initial conditions for the Graph Neural Fokker-Planck equation. Here, validate this design. By deacti- vating the encoder and decoder V, we require the neural Fokker-Planck equation to directly model the diffusion of the probability vector. As shown in Table when + and V are deactivated, PESLA's performance in both energy and evolution prediction deteriorates, confirming the importance of the high-dimensional encoded space for effective graph neural diffusion modeling.
Table 6: 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 obtain statistical values.
2D Prinz potential
Pt 0.9290 E 0.0342
Pf
MJS
TJS
PESLA
0.7419 0.0318
0.1031
0.0125
0.1796 0.0234
w/o Liphy
0.0641 + 0.0182
0.003 0.0928
0.1435 - 0.0102
0.2559 + 0.0358
w/o Ltatent
0.8089 + 0.0672
0.7192 + 0.0291
0.1270 + 0.0334
0.2010 + 0.0327
w/o &&
0.8994 + 0.0477
0.6925 0.0903
0.1675
0.0089
0.3535 = 0.0122
Ecological
Evolution
P: 0.9067 0.0100
Pf
MJS
TJS 0.3277 + 0.0424
PESLA
-0.7582
0.0241
0.3111 0.0397
w/o Lphy
00.021 + 0.0281
-0.002 = 0.0817
0.4455 = 0.0865
0.4683 + 0.0257
w/o Ltatent
0.8982 0.0071
-0.6912 0.0182
0.3228
0.0441
0.3441 = 0.0232
w/o &
0.8980 + 0.0075
-0.7018 + 0.0202
0.3564 + 0.0236
0.4685 + 0.0227
F COMPUTATIONAL CoST
We denote the sample size and the preset number of codewords as N and K, respectively. The training process of PESLA consists of two modules: adaptive codebook learning and the graph neu- ral Fokker-Planck equation. The former includes encoding and decoding each sample, as well as codeword matching operations. The computational complexity of encoding and decoding is O(N), while codeword matching, which involves similarity calculations with each codeword, has a time complexity of O(NK). In the second module, all computations occur the codeword topology
22

---

Published as a conference paper at ICLR 2025
with size of Orr)), where r is the proportion of activated codewords. Since the encoding, decod- ing, and diffusion processes for the probability vector, as described in Equations 4 and 5, involve operations the entire topology, the computational complexity is O(rNK). Therefore, the over- all time complexity during training is O(N + NK + rNK) = ((NK).
Once training is complete, the model only needs to retain the activated codewords, resulting in an inference time complexity of Or.NK). In the worst case, r = 100%, meaning a 100% codeword reported in Section 4.3, for discrete state space of size 10,000,
activation ratio. However, fewer than 100 codewords are typically sufficient for reliable prediction. Thus, r is usually a low value, making the model's inference cost manageable. Additionally, our encoder E employs MLP architecture, with time complexity that scales linearly with the dimension of the observed state.
We evaluate the training time of all algorithms across three datasets, with batch size and epochs uniformly set to 128 and 10, respectively, to ensure fairness. The experiments are conducted on a hardware platform equipped with an Intel i 14600KF CPU and an NVIDIA RTX 4060Ti GPU. As shown in Figure 11 the total runtime of PESLA's two phases is shorter than that of the optimal baseline, NJU reported in main text. Additionally, PESLA's time bottleneck clearly concentrated in phase 1 (adaptive codebook learning), with phase 2 (graph neural Fokker-Planck equation) accounting for less than half the total training time, aligning with the conclusions of previous analysis.
3511.45
PESLA phase2
Training trow Dlgg scale
20"
NeunPPP
317.20
SDENut
Till VAMPree
1.0-
78:18
76.59
54.83
15.56
29.12
10'
4:13
3.18
3.06
Ecological
Figure 11: Total training time of all algorithms across the three datasets. All experiments are run 10 times obtain mean values.
G DEGRADED OBSERVATION
In Equation 2 of the main text, the system's intrinsic state evolving the energy landscape is mapped to the observation space via the observation function g. which serves as the input to PESLA. In the main experiments, the observation state 2; retains the primary information of s;; however, this may not hold under certain lossy observation functions. In complex systems modeling, the time- delay embedding method reconstructs the manifold of system's evolution by embedding multi-step trajectories of high-ddieesionnll system in a limited-dimensionall space (Wu et al. (024). Here, we supplement a of degraded observation experiments to verify that PESLA can leverage similar idea, modeling the energy landscape using multi-step historical observations.
Table 7: Energy prediction as function of lookback steps.
lookback
4
5
Pt
0.7087
0.7338
0.7916
0.8037
0.8056
Specifically, we applied an observation function g(r, v) [cos(? # sin []]
x
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. We tested the perfor- mance of PESLA's energy prediction function of the historical observation step length, as shown in Table The prediction performance is poor when using only a single observation step, as the one-dimensional degraded observation lacks complete information about the energy. As the input
23

---

Published as a conference paper at ICLR 2025
lookback steps increase, the model gains access the system's past evolution trajectories, improv- ing prediction performance to over 85% of that under lossless observations. The results indicate that degraded observations be mitigated by incorporating multi-step historical trajectories, aligning with the consensus in the field.
24


Reviews:
[{"review_id": "7c77LpI4as", "reviewer": "Reviewer_c7fR", "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.", "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?\n\n* Scalability. To model a bigger protein, how large does the codebook need to be? What about training and inference cost for longer sequences?\n\n* Spell out KNN (k-nearest neighbor?)\n\n* How are encoder Ξ, decoder Ω implemented? GCNN was mentioned but I still have no idea what's in the model.\n\n* There is no ablation tests, so it's hard to see what the various loss terms or model designs offer. \n\n* There is no discussion of the limitations of the method.\n\n* 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.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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?\n\n* Overall, the writing of the paper is too hand-waving. As a referee I don't like guesswork.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Cutgd1EKge", "reviewer": "Reviewer_ciD9", "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.", "questions": "* to address (3) above, please provide better justifications for the main model and the associated training losses\n\n* to address (4) and (5), please compare with published results on recovering the energy and fitting new trajectories of new molecules or systems", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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).\n\n(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.\n\n(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).\n\n(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.\n\n(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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IySuNSelcc", "reviewer": "Reviewer_jDjA", "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.", "questions": "See Weaknesses", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\n2. 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?\n\n3. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "0rkTngqrz6", "reviewer": "Reviewer_iFFC", "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.", "questions": "- What is the computational complexity of the proposed methods and its comparison to existing methods? \n\n- Are the estimate energy landscape consistent under different hyperparameters?\n\n- Does the granularity of the discretization of the energy landscape influence the accuracy of system evolution estimation?\n\n- Considering that real-world trajectory data is usually noisy and sparse, how robust is the method to such data?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is clearly written. \n\n- The method is trained in a self-supervised manner, hence it can conduct estimation without knowing true energy landscape.\n\n- 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.\n\n- 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.\n\n- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)