Latent Representation and Simulation of Markov Processes via Time-Lagged Information Bottleneck
Authors: Marco Federici
Patrick Forré
Ryota Tomioka
Bastiaan S. Veeling
Venue: ICLR 2024 poster (2024)
Keywords: Markov Processes
Information Theory
Information Bottleneck
Latent Simulation
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We introduce a tractable objective to represent, simplify and simulate Markov processes based on the Information Bottleneck principle. We motivate our model both theoretically and empirically on synthetic data and molecular dynamics.

Abstract: Markov processes are widely used mathematical models for describing dynamic systems in various fields. However, accurately simulating large-scale systems at long time scales is computationally expensive due to the short time steps required for accurate integration. In this paper, we introduce an inference process that maps complex systems into a simplified representational space and models large jumps in time. To achieve this, we propose Time-lagged Information Bottleneck (T-IB), a principled objective rooted in information theory, which aims to capture relevant temporal features while discarding high-frequency information to simplify the simulation task and minimize the inference error. Our experiments demonstrate that T-IB learns information-optimal representations for accurately modeling the statistical properties and dynamics of the original process at a selected time lag, outperforming existing time-lagged dimensionality reduction methods.

LATENT REPRESENTATION AND SIMULATION OF MARKOV PROCESSES VIA TIGGEL INFORMATION BOTTLENECK
Marco Federici
Patrick Forré
Ryota Tomioka Microsoft Research A14Science
AMLab
A14Science Lab, AMLab University of Amsterdam p. d. forre@uva.n.
University of Amsterdam
zyoto@microosot.. com
m. Federici@uva.nn.
Bastiaan S. Veeling"
Microsoft Research A14Science bavveiing@nicrrofftt.co com
ABSTRACT
Markov processes are widely used mathematical models for describing dynamic systems in various fields. However, accurately simulating large-scale systems at long time scales is computationally expensive due to the short time steps required for accurate integration. In this paper, introduce inference process that maps complex systems into simplified representational space and models large jumps in time. To achieve this, we propose Time-lagged Information Bottleneck (T-IB), principled objective rooted in information theory, which aims to capture relevant temporal features while discarding high-frequency information simplify the simulation task and minimize the inference error. Our experiments demonstrate that T-IB learns information I representations for accurately modeling the statistical properties and dynamics of the original process at selected time lag, outperforming existing time-lagged dimensionality reduction methods.
INTRODUCTION
Markov processes have long been studied in the literature (Norris) 1997 Ethier & Kurtz 2009). as they describe relevant processes in nature such as weather, particle physics, and molecular dynamics. Despite being well-understood, simulating large systems over extensive timescales remains a challenging task. In molecular systems, analyzing meta-stable molecular configurations requires unfolding simulations over several milliseconds (T A (0)3s), while accurate simulation necessitates integration steps on the order of femtoseconds (To = 10 15 The time required to simulate 10 steps determined by the time of a single matrix multiplication, which takes on the order of milliseconds on modern hardware, resulting in simulation time of multiple years.
Deep learning-based approximations have shown promising results in the context of time series forecasting (Staudemeyer & Morris) 2019 Lim & Zohren 2021). including applications in weather forecasting |Viillette et al. (2020), sea surface temperature prediction (Ham et al., 2019 Gao ct al. 2022). and molecular dynamics (Sidky al. 2020: Klein al., 2023] SChreiner et al. 2023). Mapping observations into dower-dimensional spaces has proven to be an effective method for reducing computational costs. Successful examples in molecular dynamics include learning system dynamics through coarse-grained molecular representations. (Wang et al. 2019a Köhler et al. 2023
Arts al. 2023), linear (Koopman, 1931 Molgedey & Schuster 1994) and non-linear (Wehmeyer & Noč 2018; Mardt et al. 2018 Sidky et al. 2020) projections of molecular features.
Modern deep representation learning methods have proven effective in creating representations for high-dimensional structured data, including images (Hjelm et al. 2019 Chen et al. 20200), audio
"Conresonding author.
Work partially done during an internship Microsoft Research, A14Science.

---

yt-r
Polle-bbag-
-
max T(me ri 1
Gpl and me)
Xt
Pyl py(x((x))
min (annBKKLWA /xLLL//////////////0//////
0,0
Z
Yt
Figure I: The Time-lagged Information Bottleneck objective aims to maximize the mutual information between sampled representations zr-r; Z₁ at temporal distance while minimizing mismatch between the encoding distribution Pe( pz(z, z, xx) and the learned variational transitional distribution (z( (zz-r). This results in minimal representations capturing dynamics at timescale т or larger, which can be used predict properties of interest yı, such inter-atomic distances, over time.
(van den Oord et al.; 2018 Saeed et al. 2021), ). text (Devlin et al. 2018 Radford et al.
2018), and
graphs (Veličković et al. 2018:] Wang al. 2022). These methods often aim to capture relevant
information while reducing the complexity of the data. In this context, information theory provides
a compelling direction for further analysis (Wennekers & Ay 2003 Gao et al., 2022 Lozano
Durán & Arranz 2022). In particular, the information bottleneck principle (Tishby et al., 2000
Tishby & Zaslavsky 2015) suggests that an optimal representation should retain relevant information
while discarding unnecessary features. Applying this principle to the context of Markov process simulations has the potential to simplify the modeling task, reduce computational complexity, and aid in identifying the salient characteristies that define the relevant dynamics.
In this paper, we make the following contributions: (i) we introduce a probabilistic inference scheme for Markov processes, Latent Simulation (I.S), and characterize the inference by defining Time- lagged InfoMax (T-InfoMax) as general family of principled training objectives. (ii) We propose Time-lagged Information Bottleneck (T-IB, Figure T novel objective that follows the T-InfoMax principle to preserve system dynamics while discarding superfluous information to simplify modeling tasks. (iii) We empirically compare the performance of models trained using the T-InfoMax and T-IB objectives on synthetic trajectories and molecular simulations, showcasing the importance of the T- InfoMax principle and the advantages of the proposed T-IB method for both representation learning and latent simulation inference compared to other models in the literature.
METHOD
We delve into the problem of efficiently representing and simulating Markov processes starting by defining Latent Simulation as an inference procedure and characterizing the corresponding error (section 2. Next, in section 2.2 we analyze the problem of capturing system dynamics from an
perspective, defining and motivating Time-Lagged InfoMax: family of objectives that minimizes the latent simulation error. Finally, we introduce Time-lagged Information Botteneck (section (.3) as an extension of T-InfoMax that aims to simplify the representation space. A schematic representation of our proposed model is visualized in Figure
2.1 LATENT SIMULATION
Consider sequence of T random variables, denoted [xi] which form homogeneous Markov Chain. This chain models a dynamical process of interest, such as molecular dynamics, global climate systems, or particle interactions. Let y, represent specific (noisy) property of x, that we aim to model over time. Formally, we define y: = =f f(,, <.). where f X x 3 + Y is some function and & is temporally uncorrelated noise. Examples of such properties could include the energy or momentum of a particle, the meta-stable state molecular structure, and the amount of rainfall. Each of these properties y, can be derived from more comprehensive high-dimensional state description x,-
2

---

(a)
Yott
(b)
Y++7
(x,) and visualized together with the Simulation steps
X,
11S
Z
7
Figure 2: Graphical models for the joint distribution of a sequence, targets, and representations, 2a shows the data-generating process in which red arrows denote computationally expensive simulation steps. [bb represents the corresponding Variational Latent Simulation, in which the transitions are modeled in the latent space. Gray dashed lines indicate distributions that are not used for inference.
Given an initial observation Xs, the joint distribution the sequence of K future targets [ys+krlk- 1 at some lag time 0 can be expressed as:
p[(y++rTlEe [x,== xx)=-(--.
II
p(x (*+ 1yr) ply
p(y s+r (X dx,+r
dx
Transition
Prediction
Each transition distribution p(X(+7(x)) may necessitate J integration steps at a finer timescale Tu т. Given the sequential nature simulation, generating trajectories over extended time horizons may require substantial computational resources. To mitigate the challenges of simulating large- scale system dynamies, we adopt two modeling strategies: (i) rather than modeling the transition distribution in the original space X, we learn a time-independent encoder pe (2,l xx) that maps into simpler representation space Z: and (ii) we directly model the dynamics for larger jumps > To. We refer to the process of unfolding simulations in the latent representation space as Latent Simulation (LS). The joint distribution for trajectories of targets unfolded using LS starting from X, is defined as:
K
KK ps (lys+orl (xx) : /f... po(zg x.) II p(zo+kr Zx (k- p( (2+tkt dz, Kr+Kr.
(2)
Encoding
transition
Latent prediction
Unfolding LS requires access to the latent transition p(zz+T|zi, (Z;) and predictive p(yr|ze) distribu- tions, which are generally intractable for an arbitrary choice of encoding distribution (o ( To
circumvent this intractability, we introduce variational transition and variational target predictive distributions, denoted as 4s (2r (2r/zz-7) and qu (yr(zz), respectively. The resulting joint inference distri- bution for the future targets [y. [ys+kr]k=:) unfolding from the initial observation x,, referred as the Variational Latent Simulation distribution g'' (lys+ K
graphical model for the genereteeeringgggg process in Figure 2
The Kullback-Leible (KL) divergence, which quantifies the discrepancy between the ground truth and the variational latent simulation distributions, can be upper-bounded as follows:
KT (T((p(y (TT(ppyy+++Tll + 114 1S (ly + J xx,) (T- K
[p LS ((y k=1 (x.)
Variational Simulation
Latent Simulation
K
KL KL (
Z++
(zo+kr 12. 1)r)) + KL (L( /z+ 1lq (Y Z zo+kr
Variational transition gap
Variational latent prediction gap
The upper bound consists the latent simulation error and the sum of the variational gaps for both the latent transition and target predictive distributions. Unfortunately, terms on the right side of equation 3 are intractable. To address this, we propose a two-step optimization procedure: (i) first learn encoding distribution pe (z (x:) that minimizes the latent simulation error, effectively capturing the dynamical properties of the system in the representation; then (ii), assuming fixed (optimal) encoding distribution, we optimize the variational latent transition qo (z (2:) and
3

---

predictive Go (y.(z)) distributions using maximum likelihood to minimize their respective variational gaps. In the following sections, we will focus on step (i), analyzing the problem of learning representations that preserve dynamical properties from information theoretical perspective. Additional details about this two-step procedure are available in Appendix C.1
2.2 TEMPORAL INFORMATION ON MARKOV CHAINS
crucial prerequisite for ensuring that the latent simulation process does not introduce any error is to guarantee that each representation Z, is as informative as the original data x; for the prediction of any future target of interest Yt+r- 2 less predictive than x; for yt+r, the statistics for the corresponding predictive distribution p(y++> (2)) would deviate from those based on the original data p(y++r Xz). This first requirement can be expressed by equating mutual information" thatti that x; and ze share with Yt+ri (Xx; Yt+r) = ( yr+r). We will refer to this requirement as sufficiency of ze for Y1+r, Sufficiency is achieved only when X₂ and 2 yield identical predictive distributions for the future target, i.e., p(y++r/xx) p(y++r/zi) (z.).
Secondly, we introduce the concept of autoinformation. Autoinformation at given lag time T defined as the mutual information between the current observation x; and its corresponding future Xt+r- Formally, (T == [x;; xt+r). This concept extends the statistical notion of autocorrelation, which measures the linear relationship between values of variable at different times
(Brockwell & Davis 2002), to include nonlinear relationships (Chapeau-Blondcau Blondcau 2007 von Wegner
al. 2017).
Since 2, is derived from Xt, the autoinformation for x; sets an upper-bound for the autoinformation for z;: Ax;;)) 2 T[( (proof Appendix B.3 We refer to the difference between the
two values as the autoinformation gap AIG(Z(;T = AI (xx:T) - Al (Z;; r) and we say that Z
preserves autoinformation whenever autoinformation gap zero.
Lemma 1. Autoinformation and Sufficiency (proof in Appendix B.5
A representation Ze preserves autoinformation at lag time T if and only if is sufficient for any target yt+r- Conversely, whenever 2, does not preserve autoinformation for lag time is always possible to find a target YE+r for which 2z is not sufficient:
AIG(ZeT) =
I(xx))+7) I(ze:(++r) Vyt+r = f(x
In simpler terms, a representation that preserves autoinformation encapsulates all dynamic properties of the original data for the temporal scale т. As result, the representation z, can replace x, in predicting any future properties at time t +.
For a temporal sequence [x=]?_a. we define the autoinformation at lag time 7 as the average autoinfor- mation between all pairs of elements in the sequence that are time-steps apart: AT((x))'=>:>> Ex~U(e, -r) [AT [xx;:)) r)], where (s,, T - r) refers to a uniform distribution. If p(xx) is stationary,
the amount of autoinformation for sequence [xo]]_ is equivalent to autoinformation at any point X4. Using this definition, we can show:
Lemma 2. Autoinformation and Markov Property (proof in Appendix B.6
If a sequence of representations (ze)ts of a homogeneous Markov chain [x, tms preserves autoinfor- mation lag time т, then any of its sub-sequences of elements separated by time-steps must also form a homogeneous Markov chain:
AIG(ZZ)] :T) = - + [z is homogeneous Markov Chain,
for every s' E [ r), K [0, L(T' s///T].
Building on this, we further establish that dynamics at a predefined timescale 7 also encode informa- tion relevant larger timescales:
Lemma 3. Slower Information Preservation (proof in Appendix B.8
Any sequence of representations |ze)!=, that preserves autoinformation lag time T must also preserve autoinformation at any larger timescale r':
AIG([z.)] = -
TT.
We refer the reader Appendix A for further details on the notation.
4

---

By synthesizing the insights from Lemma and we infer that any representation preserving autoinformation at lag time 7 captures the dynamical properties of the system across timescales T' that equal or larger than Specifically. we conclude that: (i) Z, can replace x, in predicting any Y++7' (Lemma Lemma 3): (ii) any sequence of representations (us+kr'k== will form homogeneous Markov Chain (Lemma 2 Lemma 3). Furthermore, we establish an upper bound for the expected Latent Simulation error equation using the autoinformation gap:
Ex (L(p([y<+++1 / ((
K ([y++krle kk= [xx) S K
ALG((Z))Z.iT)
Latent Simulation error for K simulations steps with Tag time 7
Antoinformation gap for lag time
with f ~ ss. + T 1) and T := + (K + 1)T 1. In words, the latent simulation error is upper-bounded by the product of the number of simulation steps and the autoinformation gap. A full derivation is reported in Appendix B.9
Given that the autoinformation between elements of the original sequence fixed, we can train repre- sentations that minimize the autoinformation gap at resolution by maximizing the autoinformation between the corresponding representations at the same or higher temporal resolution. We refer to this training objective as Time-lagged InfoMax (T-InfoMax):
CT-inliauux ((x,) T; 0) AIG([2))),
-E,-U,,TT->) (T<;; Z(++)]-
Among the various differentiable methods for maximizing mutual information in the literature (Poole
eet al. 2019 Hjelm et al.] 2019 [Song & Ermon (020), focus on noise contrastive methods
due to their flexibility and computational efficiency (van den Oord et al.) 2018 Chen et al.
2020). Therefore, we introduce an additional critic architecture Fe : Z
R with parameters
to define an upper-bound on the T-InfoMax loss:
B log
Pe(zi,
CT-InfoMax (1x1)== ([x 0; < CT-InfoMax
Fel Pe( +)
==1
B 1==
In this equation, 1; is sampled uniformly in the interval (-, 2t, and -t,-7 are the representations of Fe; and Tre-7 encoded via (2)))))). and B denotes the mini-batch size. We refer the reader to Appendix C.2 for additional discussion regarding the proposed approximations.
2.3 FROM TIME-LAGGED INFOMAX TO TIME-LAGGED INFORMATION BOTTLENECK
In the previous section, we emphasized the importance of maximizing autoinformation for accurate latent simulation. However, it is also critical to design representations that discard as much irrelevant information possible. This principle, known as Information Bottleneck (Tishby et al. 22000), aims to simplify the implied transition p (2 4 and predictive p(y: ( distributions to ease the variational fitting tasks, decreasing their sample complexity. In dynamical systems, the information that z, retains about X, can be decomposed into the autoinformation at the lag time 7 and superfluous information:
Txx;; Ze
AT(Z(-ri7)
T(x;; (t-7)
Total Information
Autointormation at lag time 7 Supertiuous information
As shown in Appendix B.TT. superfluous information consists of time-independent features and dynamic information for temporal scales smaller than
Incorporating sufficiency from equation 4 with the minimality of superfluous information we obtain a family of objectives that denote Time laaged Information Bottleneck (T-IB):
(8) Here, 3 is a hyperparameter that trades off sufficiency (maximal autoinformation, 8 0) and minimality (minimal superfluous information, - +c)). Given that superfluous information can not be computed directly, we provide a tractable upper bound based on the variational latent transition distribution qs (2z/z,,7) Together with equation 6 this defines tractable T-IB InfoNCE objective:
T-B ([xx) : 3:0) =L [T-InfoMax ([x:)== ([x,
77;0) + 3 TE, [[xx;21/==<>]
Fel (
3 log Po (xt.)
L TB InfoNCE F
B: 8:0, ob
Σ log
EE Pelzi
go Zz, -7,
in which the encoder pe pe (z,) xx) is parametrized using Normal distribution with learnable mean and standard deviation as in Alemi et al. (2016); Federici et al. (202)). Details on the upper bound in equation 9 are reported in Appendix
5

---

3 RELATED WORK
cnoomattion-theorette methods have gained traction in fluid mechanics, offering valuable insights into
energy transfer mechanisms (Betchov 1964 Cerbus & Goldburg 2013 Lozano-Durán & Arranz
2022). Measures like Transfer Entropy (Schreiber (000) and Delayed Mutual Information (Materassi
et al. 2014) closely align with the concept of Autoinformation, which is central in this work. However, previous literature predominantly focused on designing localized reduced-order models (Lozano)
Durán & Arranz 2022) by factorizing spatial scales and independent sub-system components, rather
than learning flexible representations that capture dynamics at the desired temporal scale. Moreover, the theory and application of these principles have largely been confined to discrete-state systems 2002) and model selection tasks (Akaike 1974 Burnham & Anderson 2004)).
(Kaiser & Schreiber
widely used approach dynamical system representation involves measuring and maximizing linear autocorrelation (Calhoun et al., 2001] [Pérez-Heráándee et al.) 2013 Wiskott & Sejnowski 2002). In particular, (Sidky et al. (2020) proposes a latent simulation inference that leverages linear correlation maximization, coupled with mixture distribution for latent transitions. As shown in Appendix D autocorrelation maximization can be also interpreted as autoinformation maximization constrained to jointly Normal random variables (Borga 2001I. However, the lincar restriction requires high-dimensional embeddings (Kantz & Schreiber [003 von Wegner et al. (0177). and may introduce training instabilities for non-linear encoders (Mardt et al. 2018 Wu & Noē 2020, Lyu el. 2022). In
this work, we prove that the requirement of linear transitions is not necessary to capture slow-varying signals, demonstrating the benefits of modern non-linear mutual information maximization strategies.
The proposed T-InfoMax family also generalizes existing models based on the reconstruction of future states [Wehmeyer& & Noë 2018 Hernández al. 2018). On one hand, these approaches
proven to maximize mutual information (Barber & Agakov 2003 Poole et al., 2019). on the other their effectiveness and training costs are contingent on the flexibility of the decoder architectures (Chen et al. 2019). For this reason, we chose to maximize autoinformation using contrastive methods, which rely on more flexible critic architecture (van den Oord et al. 2018 Hjelm et al. 2019 Chen
et al. 2020) instead of decode While contrastive methods have already been applied to temporal
series (van den Oord et al., 2018 Opolka et al. 2019 Gao & Shardt 2022, Yang el al. 2023). our work additionally provides formal characterization of InfoMax representations of Markov processes.
Another key contribution of our work lies in the introduction of an explicit bottleneck term to remove superfluous fast features. The proposed T-IB approach builds upon [ang et al. (2019b), which first proposes a reconstruction- cimeemiioonalitt-eddeinn linear encoder instead of flexible deep neural architecture to implicitly reduce information. Wang & Tiwary (202T) later developed a related bottleneck objective, focusing on future target reconstruction instead of autoinformation maximization and using marginal prior for compression. Although less reliant on the decoder architecture, this objective is not guaranteed to produce accurate simulation for arbitrary targets, as demonstrated in Appendix D.3
information bottleneck objective for molecular time series, utilizing
4 EXPERIMENTAL RESULTS
We perform experiments on (i) a controlled dynamical system consisting of non-linear mixing of slow and fast processes, and (ii) molecular simulations of peptides. Our goal is, primarily, to examine the effect of the information maximization strategy (lincar vs. contrastive) and the impact of the bottleneck regularization on the trajectorics unfolded using LS. We further aim to validate our theory by estimating autoinformation and superfluous information for the models considered in this analysis.
Models We analyze representations obtained using correlation maximization methods based either on linear projections (TICA) (Molgedey & Schuster, 1994) or non-linear encoders (VAMPNet) (Mardi eet al. 2018) against and non-linear autoinformation maximization (T-InfoMax) and corresponding bottleneck (T-IB) based InfoNCE. The regularization strength / is selected based on the validation scores We use a conditional Flow++ architecture (Ho et al. 2019 to model the variational transition distribution qu (Ze|Zze))) This is because of the modeling flexibility, the tractability of the likelihood, and the possibility of directly sampling to unfold latent simulations. Multi-layer perceptrons (MLPs)
3
We refer the reader to Appendix D.2 for further details.
AAblaiion studies the effect of and the effect of stochastic encoder can be found in Appendix F11

---

(a)
Free Energy
Slow Modes
Fast Modes
Autoinformation
Slow Modes
Fast Modes
14
wew.ommmaiinnni (nats]
TICA
A5
T-BB T-InfoMax VAMPNer
-14
02 X[
10
60
10
10 Lag (sieps)
x[0]
Xd
(b)
AIG(( 61) 1 I(2:: y1) 1
1(2 yi
MJSL L
TJSØ512
Original (x)
1.31 1 0.00 1.30 + 0.00
4.1 4.7
4.0 1.2
Constant (2,=c)
0.88 0.00
13.6 17.9
105.9
A
TICA (ID)
0.10 0.00
1.16 = 0.00 0.06 + 0.00
14.3 11.2
17.9
VAMPNet
0.05 + 0.02
0.99
= 0.03 0.00 + 0.00
20.5 21.0
10.4
T-InfoMax
0.03 = 0.00
1.29 = 0.02 0.30 4 0.07
13.0
0.5
T-IB (3 0.01)
0.05 0.00
1.30 0.01 0.00 0.00 2.4
2.4 0.3
Figure 3: Visualization of the results the Prinz 2D dataset. 3a free energy and short sample trajectory (left), samples colored by the slow and fast mode index (center), and autoinformation for the full process and its components several lag times (right). 3bj measures of autoinformation gap, mutual information between the representation and the discrete fast and slow modes in nats, and value of marginal and transition JS divergence for unfolded sequences in milli-nats. 3cj trajectories encoded in the latent space 2, through various trained models. Quantitative and qualitative results confirm that T-IB uniquely captures relevant (slow) information while discarding irrelevant (fast) components. This results in more accurate LS as measured by the marginal and transition JS.
are used to model Te (2)) mapping the representations z, into the logits a categorical distribution over the target yz. For all objectives, we use the same encoder, transition, and predictive architectures.
Training We first train the parameters 0 of the encoder po using each objective until convergence. Note that T-IB also optimizes the parameters of the transition model (Z; Zz-r) during this step (as shown in equation 9. Secondly, we fix the parameters 0 and fit the variational transition (2 (2,-)) and predictive [y( distributions. This second phase is identical across all the models, which are trained until convergence within maximum computational budget (50 epochs) with the AdamW optimizer (Loshchilov & Hutter (019) and early stopping based on the validation score. Standard deviations are obtained by running experiments for each tested configuration with different seeds. Additional details on architectures and optimization can be found in Appendix E.2
Quantitative evaluation We estimate the autoinformation of the representations Al([z]] 7) at several lag time using SMILE (Song & Ermon 2020 and measure the amount of information that the representations contain about the targets of interest yz):ye) using difference of discrete entropies: (y:)) H (y+lzz) (Poole et al. 2019 McAllester & Stratos 2020 Given initial system state x, of a test trajectory {x, and the sequence of corresponding targets [un] e we use the trained encoder, transition, and prediction models to unfold trajectories [ūs+kr k=] g ([y++k kr -=== (x.) that cover the same temporal span as the test trajectory (K- = [(T' s)/T]). Similarly to previous work (Arts al. 2023). for evaluation purposes, we consider only discrete targets y, so that we can estimate the marginal and transition probabilities for the ground truth and unfolded target trajectories by counting the frequency of each target state and the corresponding transition matrix (Figure (5a). We evaluate the fidelity of the unfolded simulation by considering the Jensen-Shannon divergence (JS) between the ground truth and unfolded target marginal (MJ and target transition distribution for several T' > (TJS@r'). Further details on the evaluation procedures are reported in Appendix E.3
2D Prinz Potential Inspired by previous work (Mardt et al. 2018 Wu et al. 2018 we design a 2D process consisting of fast x and slow xf components obtained from 2 independent simulations on the 1D Prinz potential (Prinz et al. 2011). This potential energy function consists of four interconnected low-energy regions, which serve as the discrete targets y and yř. The two components mixed through linear projection and tanh non-linearity produce a 2D process consisting of total of 4 (fast) x 4 (slow) modes, visualized in Figure (3a] We generated separate training, validation, and test trajectories of 100K steps each. The encoders po (2: x;) consist of simple MLPs and 2, is fixed
7

---

Alanine Dipeptide
Peptide
TICA
VAMPNet
T-InfoMax
T-IB
(3-0.1)
N (nats)
1.00
TICA
Alanine D.P.
VAMPHA
(8-001)
0.50
THB (9:0.1)
025
T-B (8-1)
10"
Chignolin
Chignolin
Villin
Villin
(a)
Figure Comparison 2D representations for Alanine Dipeptide, Chignolin, and Villin simulations. visualizations are colored by molecular configuration clusters yr obtained by clustering torsion angles (Alanine Dipeptide) and TICA projections (Chignolin, Villin). 4b) corresponding values of autoinformation (y-axis) at multiple lag times (x-axis). An optimal representation should maximize autoinformation at the trained lag time т (indicated by the dashed vertical line) while minimizing information faster processes (to the left of the dashed line). Correlation maximization methods struggle to capture all relevant dynamics in larger systems, while T-IB regularization can effectively regulate the amount fast information in 2z- Visually this results in simpler clustered regions.
4a
to be 2D. As shown in the autoinformation plot in Figure 3a (on the right), the choesen train lag time (T = 64, vertical dashed line), the fast components are temporally independent, and all the relevant information is given by the slow process: (x, 64) s Ax;;;; 64) > AI(x! 64) Rs 0. Therefore, information regarding x can be considered superfluous (equation 7 and should be discarded.
Figure 3e visualizes the representations obtained with several models colored by the slow (left col- umn) and fast (right column) mode index y" and y!. We can visually observe that our proposed T-IB model preserves information regarding the slow process while removing all information regarding the irrelevant faster component. This is quantitatively supported by the measurements of mutual informa- tion reported in Table 3b which also reports the values of marginal and transition JS divergence for the unfolded slow targets trajectories 1o:1 We observe that the latent simulations unfolded from the T-IB representations are statistically more accurate, improving even upon trajectories unfolded by fitting the transition distribution directly in the original space x,. We believe this improvement is due to the substantial simplification caused by the T-IB regularization.
Molecular Simulations We analyze trajectories obtained by simulating Alanine Dipeptide and two fast-folding mini-proteins, namely Chignolin and Villin (Lindorff-Larsee et al. 2011 in water solvent. We define disjoint train, validation, and test splits for each molecule by splitting trajectories into temporally distinct regions. Encoders pe (z) xx)) employ a TorchMD Equivariant Transformer architecture (Thölke & Fabritiis| (2022) for rotation, translation, and reflection invariance. Following previous work (Köhler et al.) 2023). TICA representations are obtained by projecting invariant features such as inter-atomic distances and torsion angles. Following Arts et al. (2023), the targets y, are created by clustering 32-dimensional TICA projections using K-means with centroids. Further details the data splits, features and targets can be found in Appendix(E.12 E.1.2
In Figure we show 2D representations obtained by training the encoders on the molecular trajecto- ries (Figure 4a). and the corresponding measure of autoinformation (Figure 4b)) at several time scales (x-axis), while FigureB3reppe transition and marginal JS for trajectories unfolded on larger latent spaces (16D for Alanine simulations and 32D for Chignolin and Villin). While previous work demon-
8

---

Ground Truth
TICA
VAMPNet
T-InfoMax
T-IB
301
63
4M
53
94
14
62
3M
14
430
30
(8)
4.11
Alanine Dipeptide
Chignolin
Villin
MJS 4 TJS@256 ps
MJS
TJS051.2 ns
MJS 4
TJS@51.2nS J
(b) TICA
2.3 1.8
14.9
8.1 7.8
11.2 12.1
12.1=104
1.71
VAMPNet
220 350
150 - 21
35 - 30
106 100
96 64
130 110
T-InfoMax
3.9 1.6
13.7 + 2.3
11.0 4.0
3.2
35.2
37.8
9.1 0.9
T-IB
0.2
12.9 + 1.2
5.5 + 1.9
3.6
22.5 + 17.9
0.9
Figure 5: Evaluation of the statistical fidelity of unfolded molecular trajectories. 5a visualization of transition matrices for ground-truite and VLS target trajectories for different models on Villin at 51.2 ns. [5] corresponding values of marginal and transition JS on Alanine Dipeptide, Chignolin and Villin. LS based on T-IB representations consistently results in lower simulation error, improving upon linear methods and unregularized T-InfoMax models.
strated that a linear operator can theoretically approximate expected system dynamics on large latent spaces (Koopman 1931 Mezič 2005). we note that models trained to maximize linear correlation (TICA, VAMPNet) face difficulties in extracting dynamic information in low dimensions even with non-linear encoders. Moreover, our empirical observations indicate that higher-dimensional represen- tations obtained with VAMPNet yield transition and prediction distributions that are more difficult to fit (see Table 5J and Appendix F) resulting in less accurate unfolded target trajectories. Methods based on non-linear contrastive T-InfoMax produce more expressive representations in low dimensions. The addition of a bottleneck term aids in regulating the amount of information on faster processes (Figure 4b) left of the dashed line). As shown in Sauurreaaaan Table 5b) T-IB consistently improves the transition and marginal statistical accuracy when compared to the unregularized T-InfoMax coun- terpart. Results for additional targets and train lag times are reported in Appendix F We estimated that training and unfolding Villin latent simulations of the same length of the training trajectory with T-IB take approximately 6 hours on a single GPU. In contrast, running molecular dynamics on the same hardware takes about 2-3 months. Further details the run times can be found in Appendix G
CONCLUSIONS
In this work, we propose an inference scheme designed to accelerate the simulation of Markov processes by mapping observations into representation space where larger time steps can be modeled directly. We explore the problem of creating such a representation from an information- theoretic perspective, defining a novel objective aimed at preserving relevant dynamics while limiting superfluous information content through an Information Bottleneck. We demonstrate the effectiveness of our method from both representation learning and latent inference perspectives by comparing the information content and statistics of unfolded trajectories on synthetic data and molecular dynamics.
Limitations and Future work The primary focus of this work is characterizing and evaluating the dynamic properties of representations. Nevertheless, modeling accurate transition in the latent space remains a crucial aspect, and we believe that more flexible classes of transition models could result in higher statistical fidelity at the cost of slower sampling. Another challenging aspect involves creating representations of systems with large autoinformation content (e.g. chaotic and unstable systems). This is because the variance modern mutual information lower bounds increases exponentially with the amount of information to extract (McAllester & Stratos 2020). To mitigate this issue and validate the applicability of our method to other practical settings, future work will consider exploiting local similarity and studying the generalization capabilities of models trained on multiple systems and different simulation conditions. We further aim to evaluate the accuracy of long-range unfolded trajectories when only collections of shorter simulations available during training time.

---

ACKNOWLEDGMENTS
We thank Frank Noć, Rianne van den Berg, Victor Garcia Satorras, Chin-Wei Huang, Marloes Arts, Wessel Bruinsma, Tian Xie, and Claudio Zeni for the insightful discussions and feedback provided throughout the project. This work was supported by the Microsoft Research PhD Scholarship Programme.
10

---

REFERENCES
H. Akaike. A new look at the statistical model identification. IEEE Transactions on Automatic Control, 19(6):716-723) 723, 1974. doi: 1011097AAC.1.99411007755
Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy, Deep variational information bottleneck. CoRR, abs/1612.00410, 2016. URL
Galen Andrew, Raman Arora, Jeff A. Bilmes, and Karen Liveseu. Deep canonical correlation analysis. In Proceedings of the 30th International Conference on Machine Learning, ICML 2013, Atlanta, GA, USA, 16-21 June 2013, volume 28 of JMLR Workshop and Conference Proceedings, pp. 1247-1255. JMLR.org, 2013. URL htt.httpppppooceeiiiiie html
Marloes Arts, Victor Garcia Satorras, Chin-Wei Huang, Daniel Zügner, Marco Federici, Cecilia Clementi, Frank Noé, Robert Pinsler, and Rianne van den Berg. Two for one: Diffusion models and force fields for boars--rrained molecular dynamics. CoRR, abs/2302. 2023. doi: 10.48550/arxxw22222.00000 URL httph::/ wwt.oo...g/////aa./..///// arxiv. 2302. 00600
David Barber and Felix V. Agakov. The im algorithm: variational approach information maximization. In NIPS 2003, 2003.
R. Betchov. Measure of the Intricacy of Turbulence. The Physics of Fluids, 160 08 1964. ISSN 0031-9171. doi: 10. 10.1063/1.1111356.
1711356
Magnus Borga. Canonical correlation: A tutorial, Jan 2001.
Peter Brockwell and Richard Davis. Introduction to time series and and forecasting. Springer, 2002.
Kenneth Burnham and David Anderson. Model selection and multimodel inference. A Practical Informattin-thhorrtiic Approach, 01 2004. doi: 10. 00.1007/978--0338--2455----
Vince Calhoun, Tülay Adali, Godfrey Pearlson, and J.J. Pekar. Spatial and temporal independent component analysis of functional mri data containing pair of task-related waveforms. Human brain mapping, 13:43-53, 06 2001.
doi: 1002/hbb. 1024.
R. T. Cerbus and W. L. Goldburg. Information content of turbulence.. Phys. Rev. E, 88:053012, Nov 2013. doi: 10. 10.110//hysReeE....5552. URL httl:/ccas//ID...a...oro//// PhysRevF. 88. 053012
François Chapeau- Blondeau. Autocorrelation versus entropy-based autoinformation for measuring dependence random signal. Physica A: Statistical Mechanics and its Applications, 380:1- 18, 2007. ISSN 0378-4371. doi: http:://bo.org100000066./.aaaaap
Chapeau-Blondeau.
www.0000.......7. 02.077....7.00000000 URL https?
// w/www. sciencedi rect. cttp://sceeeee/rrrtce//////////0
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey E. Hinton. A simple framework for contrastive learning of visual representations. Proceedings of the 37th International Conference Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 1597-1607, PMLR, 2020. URL httl:ttppppprceeeinie.
mlr. wttpprss//111//che221111
Wei Chen, Hythem Sidky, and Andrew L. Ferguson. Capabilities and limitations of time-lagged autoencoders for slow mode discovery in dynamical systems. CoRR, abs/1906.00325, 2019. URL
httt:://arivv org/ abs (00os//996...0 0932.
Jacob Devlin, Ming-Wee Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding, arXiv preprint arXiw:1810.04805. 2018.
Stewart N Ethier and Thomas Kurtz, Markov processes: characterization and convergence. John Wiley & Sons, 2009.
11

---

Marco Federici, Anjan Dutta, Patrick Forré, Nate Kushman, and Zeynep Akata. Learning robust representations via multi view information bottleneck. In 8th International Conference on Learning Representations. ICLR 2020, Addis Ababa. Ethiopia, April 26-30, 2020. OpenReview.net, 2020. UJR.hhttss://ppeneeeee
net/ d= Xww.:/m
Marco Federici, David Ruhe, and Patrick Forré. On the Effectiveness of Hybrid Mutual Information Estimation. arXiv e-prints, art. arXiv:2306.0060 006608 June 2023. doi: 10.48550//rX1..2330.000
Ian S. Fischer. The conditional entropy bottleneck. Entropy, 22(9):999, 2020. e22090999. URL
doi: 10.3390/
Xinrui Gao and Yuri A.W. Shardt. Evolve-infomax: A new criterion for slow feature analysis of non-
linear dynamic system from an informatiio-thheorttiaal perspeccivee. PPACCPPpersOnnteee 55(20):
43-48, 2022. ISSN 2405-8963. doi: http://wo..ory/oo..00fi..eeoo..2 069.. URL https: // www. www.sccieneeiicoomm rect com http:///.............../////////
10th
Vienna International Conference on Mathematical Modelling MATHMOD 2022.
Zhihan Gao, Xingjian Shi, Hao Wang. Zhu, Yuyang Bernie Wang, Mu Li, and Dit-Yan Yeung. Earthformer: Exploring space-time transformers for earth system forecasting. Advances in Neural Information Processing Systems, 35:25390-25403, 2022.
Yoo-Geun Ham, Jeong-Hwan Kim, and Jing-Jia Luo. Deep learning for multi-year enso forecasts.
Nature, 573(77755:568-577 57088555727700 2019.
Carlos Hernández, Hannah K Wayment-Stecke Mohammad M Sultan, Brooke Husic, and Vijay Pande. Variational encoding of complex dynamics. Physical Review E, 97(6):062412, 2018.
R. Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchildon, Karan Grewal, Philip Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL [att::ttp:////eeeee// net//Eoumm? id= JOcKX
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow-based generative models with variational dequantization and architecture design. CoRR, abs/1902.000255, 2019, (ttp: org./ abs 00275
Moritz Hoffmann, Martin Scherer, Tim Hempel, Andreas Mardt, Brian de Silva, Brooke Husic, Stefan Klus, Hao Wu, Nathan Kutz, Steven I. Brunton, et al. Deeptime: a python library for machine learning dynamical models from time series data. Machine Learning: Science and Technology, 3(1)0015009, 2021.
A. Kaiser and T. Schreiber. Information transfer in continuous processes. Physica D: Nonlinear Phenomena, 166(1):43-62, 2002. ISSN 0167-2789. doi: https://do..or//10 1016/80066--27999000044433 URL httpt::: hww.:////ww..iie.crreeeeetee wwwp:reetttto/cciieireceee asticLee/11//0117777900000
Holger Kantz and Thomas Schreiber. Nonlinear Time Series Analysis. Cambridge University Press, edition, 2003. doi: 10.1017/CB0978051 1755798.
Leon Klein, Andrew Y. K. Foong, Tor Erlend Fjelde, Bruno Mlodozeniec, Mare Brockschmidt, Sebastian Nowozin, Frank Noé, and Ryota Tomioka. Timewarp: Transferable acceleration of molecular dynamics by learning time-coarsened dynamics. CoRR, abs/2302.01 2023. doi: 0..45500/axxiv.3322.00 170. U /o wttpo///1.....///2222.. 2302. 01170
B. O. Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 17(5):315-3188 318, 1931. doi: 10. 10.1073/pnas. 17.5.315. URLhttps
17.5.3................
//www, , pnw.ppass..... wwwp://.///////////o/ooooo
10w.1 http://aaa............ 5.15
Jonas Köhler, Yaoyi Chen, Andreas Krämer, Cecilia Clementi, and Frank Noé. Flow-matching Efficient coarse- coarse-graining graining of molecular dynamics without forces. Journal of Chemical Theory and Computation, 19(3):442-95 952, 2023, doi: 10. jcc.. URL battps://doii. /rr/ 10. 1e21021/acoo jote. 3c0001
PMID: 36668906.
12

---

Bryan Lim and Stefan Zohren. Time-series forecasting with deep learning: survey. Philosophical Transactions of the Royal Society A, 379(2194):20200099, 2021.
Kresten Lindorff-Larsen, Stefano Piana, Ron O. Dror, and David E. Shaw. How fast-folding proteins fold. Science, 334(6555):517-520) 2011. doi: 10. 10.1126/science, h 1208351. URL https
www.sciience.rrg.gooo
wwwp:///ccccccecceeccceeeeeeeeec 1208351
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. httttttss//::peeeee:
Adrián Lozano-Durde and Gonzalo Arranz. thommattin-thhoretii formulation of dynamical systems: Causality, modeling, and control.
Phys. Rev. Res., 4:023195, Jun 2022. doi:
10. 0010099Phhseeersssecc440411111 URL https: aps .epps.rrg/ooooo /10. 1103/
PhysRevRessarch....00 023195
Qi Lyu, Xiao Fu, Weiran Wang, and Songtao Lu. Understanding latent correlation-basee multiview learning and supersuper An identifiability perspective. The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.nnt, 2022.
URL /LLtps://ppezeee
net forum? id 5F 5FUq05QRc5b
Andreas Mardt, Luca Pasquali, Hao Wu, and Frank Noé. Vampnets: Deep learning of molecular kinetics. Nature Communications, 9, 01 2018. doi: 10. 10.1038/s414677 40088/41447-001--033888-
Massimo Materassi, Giuseppe Consolini, Nathan Smith, and Rossana De Marco. Information theory
analysis of cascading process in a synthetic model of fluid turbulence. Entropy, 1623):1222-1286. 2014. doi: 10.3390/e 100399/e66011100 1272. URL httpp:://dot
David McAllester and Karl Stratos. Formal limitations on the measurement of mutual information. In Silvia Chiappa and Roberto Calandra (eds.), The 23rd International Conference on Artificial Intelligence and Statistics, AISTATS 2020, 26-28 August 2020, Online [Palermo, Sicily, Italy], volume 108 of Proceedings of Machine Learning Research, pp. 875 884. PMLR, 2020. URL html
http://proceecinngs mww.
Igor Mezić. Spectral properties of dynamical systems, model reduction and decompositions. Nonlin- ear Dynamics, :1:0-99355 2005.
L. Molgedey and H. G. Schuster. Separation of mixture of independent signals using time delayed correlations. Phys. Rev. Lett. 72:3634-3637, Jun 1994. doi: 10. hww.://Phyyseeeee........3.. URL 76.3634
hittps://iime aps org. /t1/10.111111111/
Frank Noé and Feliks Nüske. A variational approach to modeling slow processes in stochastic dynamical systems. Multiscale Model, Simul., 11(2)6335-55, 2013. doi: 10. 10858616.
URL
J. R. Norris. Markov Chains. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 1997. doi: 10.1017/CB0978051 1810633..
Felix L. Opolka, Aaron Solomon, Catalina Cangea, Petar Velickovic, Pictro Liò, and R. Devon Hjelm. Spatio-temporal deep graph infomax. CoRR, abs/1904.06316, 2019. URL http:://axii:
orq/ 3xg/abs/1904.665 06316
Ben Poole, Sherjil Ozair, Aäron van den Oord, Alexander A. Alemi, and George Tucker. On variational bounds mutual information. Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, ICML 2019, 9- 15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 5171-5180. PMLR, 2019. URL hhttp://proceedinggs htt.:gg.................. www.ps/sssv1 pool el 9a. html
J. H. Prinz, H. Wu, M. Sarich, B. Keller, M. Senne, M. Held, J. D. Chodera, C. Schütte, and F Noé. Markov models molecular kinetics: generation and validation. The Journal of Chemical Physics, (34411):744 2011 1, doi: http://doo.org/11.10631.3333333
13

---

Guillermo Pérez-Hernánee Hernandez, Fabian Paul, Toni Giorgino, Gianni De Fabritiis, and Frank Noé. Identification of slow molecular order parameters for Markov model construction. The Journal of Chemical Physics, 139 07 2013. ISSN 0021-9606. doi: 10.1063/1.4811444 URL
ht..ps://dof org/ 1063/ 4881//4...1183
Alec Radford, Karthik Narasimhan, Tim Salimans, and llya Sutskever. Improving language under- standing by generative pre-training, OpenAl Blog, 2018.
Aaqib Saced, David Grangier, and Neil Zeghidour. Contrastive learning of general-purpor audio representations. In ICASSP 2021-2021 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3875-3879. IEEE, 2021.
Thomas Schreiber. Measuring information transfer. Phys. Rev. Leut, 85:461 Jul 2000. doi: 10. 10.1100/phy URL http:://iix.a....rgggooooooo PhysRevLett 85. 461
Mathias Schreiner, Ole Winther, and Simon Olsson. Implicit Transfer Operator Learning: Multiple Time-Resolution Surrogates for Molecular Dynamics. arXiv e-prints, 2023.
David E. Shaw, Peter J. Adams, Asaph Azaria, Joseph A. Bank, Brannon Batson, Alistair Bell, Michael Bergdorf, Jhanvi Bhatt, J. Adam Butts, Timothy Correia, Robert M. Dirks, Ron O. Dror, Michael P. Eastwood, Bruce Edwards, Amos Even, Peter Feldmann, Michael Fenn, Christopher H. Fenton, Anthony Forte, Joseph Gagliardo, Gennette Gill, Maria Gorlatova, Brian Greskamp. J. P. Grossman, Justin Gullingsrud, Anissa Harper, William Hasenplaugh. Mark Heily, Benjamin Colin Heshmat, Jeremy Hunt, Douglas J. Terardi, Lev Iserovich, Bryan L.. Jackson, Nick P. Johnson, Mollie M. Kirk, John I.. Klepeis, Jeffrey S. Kuskin, Kenneth M. Mackenzie, Roy J. Mader, Richard McGowen, Adam McI McLaughlin, Mark A. Moraes, Mohamed H. Nasr, Lawrence J. Nociolo, Lief O' O'Donelll Andrew Parker, Jon L. Peticolas, Goran Pocina, Cristian Predescu, Terry Quan, John K. Salmon, Carl Schwink, Keun Sup Shim, Naseer Siddique, Jochen Spengler, Tamas Szalay, Raymond Tabladillo, Reinhard Tartler, Andrew G. Taube, Michael Theobald, Brian Towles, William Vick, Stanley C. Wang, Michael Wazlowski, Madeleine J. Weingarten, John M. Williams, and Kevin A. Yuh. Anton 3: twenty microseconds of molecular dynamics simulation before lunch. In Bronis R. de Supinski, Mary W. Hall, and Todd Gamblin (eds.), International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St. Louis, Missouri, USA, November 14-19, 2021, pp. 1. ACM, 2021. doi: 10.1145/3458887734883977 3487397
URL https: / /tt::://dooooooo http://oooooogg///////
Hythem Sidky, Wei Chen, and Andrew I.. Ferguson. Molecular latent space simulators. Chem. Sci., 11:9459 9467, 2020. doi: 10. D0.003//DDSC03335H. URLpttpp://dx doi, ott.//..g/00....
DOSC03635H
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators. In 8th International Conference Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL hattps://operreviie
nett/Eorun?1d==3xx62211
Ralf C Staudemeyer and Eric Rothstein Morris. Understanding Istm-a tutorial into long short-term memory recurrent neural networks. arXiv preprint 1rXiv:190..09586, 2019.
Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular
potentials. In International Conference on Learning Representations, 2022. URL https
//openreview net pppparrevioo..eetffoooooooooee--
Rw
Naftali Tishby and Noga Zaslavsky. Deep learning and the information bottleneck principle. In 2015 ieee information theory workshop (itw), pp. 1-5. IEFF, 2015.
Naftali Tishby, Fernando C. N. Pereira, and William Bialek. The information bottleneck method. CoRR, physics/0004057, 2000. httppppp/aara......ra/a//////llp
Aäron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. CoRR, abs//807.0n 007403348 2018. URl.http://arxxvv.oo
14

---

Mark Veillette, Siddharth Samsi, and Chris Mattioli. Sevir: A storm event imagery dataset for deep learning applications radar and satellite meteorology. Advances Neural Information Processing Systems, 33:22009-2233 22019,, 2020.
Petar Veličković, William Fedus, William L Hamilton, Pietro Liò, Yoshua Bengio, and R Devon Hjelm. Deep graph infomax. arXiv preprint arXiv: 10341, 2018.
F. von Wegner, E. Tagliazucchi, and H. Laufs. Information-thheoretical analysis of resting state eeg microstate sequences markovianity, non-stationariiy and periodicities. Neurolmage, 158:99- 111, 2017. ISSN 1053-8119. doi: wttpa///wc.rereerr...aap//////./ // wwwww www..cciencd......c rect com
/tt/htpss
Jif Vymětal and Jiří Vondrášek. Metadynamics as a tool for mapping the conformational and free- energy space of peptides the alanine dipeptide case study. The Jourmal of Physical Chemistry B. I 114(16):5632-6444 2010. doi: 10001//pp000900 URLhttps:
0p150950w PMID: 20361773.
Dedi Wang and Pratyush Tiwary. State predictive information bottleneck. The Journal of Chemical Physics, 154(13):13411, 04 2021. ISSN 0021-9606. doi: 10.1063/5.0038198 URL https
//doi.org//0.10633333 10331.00313335513311333
Jiang Wang, Simon Olsson, Christoph Wehmeyer, Adrià Pérez, Nicholas E Charron, Gianni De Fabri- tiis, Frank Noé, and Cecilia Clementi. Machine learning of coarse-grained molecular dynamics force fields. ACS central science, 5(5):755-761, 2019a.
Yihang Wang, João Marcelo Lamim Ribeiro, and Pratyush Tiwary. Past-future information bottle- neck for sampling molecular reaction coordinate simultancously with thermodynamics and kinet- ies. Nature Communications, 10, 2019b. URL htttt::t///att www.nt semanticscholar.ccom org
CozpussD:119991653
Yuyang Wang, Jianren Wang, Zhonglin Cao, and Amir Barati Farimani. Molecular contrastive learning of representations via graph neural networks. Nature Machine Intelligence, 4(3):279-287, 2022.
Christoph Wehmeyer and Frank Noé. Time-lagged autoencoders: Deep learning of slow collective variables for molecular kinetics. The Journal of Chemical Physics, 148(24):241703, 2018
Thomas Wennekers and Nihat Ay. Temporal infomax on markov chains with input leads finite state automata. Neurocomputing, 52-54:431-453 2003. ISSN 0925-2312. doi: https://doi.on org/ 110016//099252231222208888838882 URL httpt:tpt/////t wciisiiiiii enciecdedii. rect come.ccoo//// scsssiisss snwpsccieee//e///////////ee/cccc A.ticllepp11/0922533100002 2003.
Computational Neuroscience: Trends in Research
Laurenz Wiskott and Terrence J. Sejnowski. Slow feature analysis: Unsupervised learning of invariances. Neural Comput., 14(4):715-770, 2002. doi: 10. 10.1662009977600211333333 URL
Hao Wu and Frank Noé. Variational approach for learning markov processes from time series
data. J. Nonlinear Sci,, 30(1):23-66 2020. doi: 10.11077:00033--110005556 URL https
//doi.or//10 /00 1007/s00332-00 9- 9-09567-v
Hao Wu, Andreas Mardt, Luca Pasquali, and Frank Noé. Deep generative markov state models.
CoRR, abss/88050066l 2018. holLhtt::////
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp. 3-19, 2018.
Yuncong Yang. Jiawei Ma. Shiyuan Huang, Long Chen, Xudong Lin, Guangxing Han, and Shih-Fu Chang. Tempelr: Temporal alignment representation with contrastive learning. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL http:tt:: http://ptweweeeeeeeeeeeeeeeeeeee htd.:////dd..-0000mm0002 /ttp://///tttttr///////eeeeeeeee
15

---

A NOTATION
Throughout the paper we use the following notation to indicate mutual information:
[(x; z) KL(p(x, z) ()(()(())))
log p(x,2) 2)
= -p(x.))
(10)
p())
To improve readability we omit the subscript for the expectation. Unless otherwise specified, expectations are computed with respect to the ground-truth distribution p(x,
Similarly, we leave the expectation for conditional KL-divergence implicit:
pxx ((x1z)
(L(p(xxz))l(((()) := Ep(x.e) log
(11)
We will use e to indicate a stochastic source external to the system, i.e. T(E; ) = with as placeholder for any variable (or combination of variables) in the system (excluding itself).
B PROOFS
We start by introducing the assumptions and properties that used throughout the section. Then, we list proofs for the statements the main text.
B. GENERAL ASSUMPTIONS
As preliminary step for proving the statements in Section
2 we clarify our general assumptions
(A.1) z, is representation of x,, With this statement, we signify that z, can be expressed as noisy function of x,: 2; (xx 6z). This implies that 2, is conditionally independent of any other variable of the system when x, is observed: [(z (xx,, 0, in which is placeholder for other variables (or combinations thereof) in the system.
(A.2) is representation of x,,
Analogously the previous assumption, we assume that the target of interest can be expressed as noisy function of x,. Therefore we will assume the same corresponding conditional independence.
(A.3) [x.)r= form homogeneous Markov Chain.
This assumption can be expressed in terms of conditional independence between past events [x) and future x when the current event x is observed:
[[x))= m-1 ;xm-1/X O,
for any E 1,
B.2 PROPERTIES
For completeness, here we list propertics of mutual information that are used to prove the statements in the following sections. Let a, b. c. d be random variables with some joint distribution p(a, b, c, d).
(P.1)
Non-negativity of (conditional) mutual information. Mutual information (and conditional mutual information) is non-negative:
T(a; blc) 2 0
(12)
(P.2) Chain rule of mutual information. Mutual information (and conditional mutual information) can be factorized as follows:
T(ab; c/d) = [(a; c|d) + I(b; c/ad)
7 7(b; cld) T(a; c(bd)
(13)
16

---

(P.3) Data processing inequality (DPI).
Mutual information (and conditional mutual information) between two random variables cannot increase by applying functions to either argument (on the left side of the conditioning). In this paper, we will use a slightly more general version of DPI in which we also consider noisy functions (with independent noise):
[(a; b.c) 2 (()(c)
(14)
B.3 AUTOINFORMATION AND DATA PROCESSING INEQUALITY
Here we demonstrate that the autoinformation in the original space (x) is an upper bound for the autoinformation of the representation (Z;
Statement. The autoinformation for x, upper-bounds the autoinformation for the corresponding representation 2;
Al(x;T) 2 Al(z,T)
(15)
Proof.
AI(x:T) I(x X(+T)
P.2)
(x:Z)) X(+r) - (2r; Xx4rlx1)
(X/26X444)
P21
[z ((+) + [(x X1+ (21)
P2)
(Ze) x Z I(z;; Zetr (X++)) + I(x;; xe+7 (2,))
(2,; x, x ) + [(x;; X (2:)
(Zii (Zi7Zi47) ) (Z1; ; XL++ (1++) ) + (xx; Xx+7 (z))
AT(2; T) + [22;; 1 221r) + T(x;; Xx: (2x).
(16)
Using property (P.1) we infer AI(x;; 7) 2 AI(Z;;
Remark 4. The autoinformation gap upper bounds the amount of information that x, and X share whenever one of the two corresponding representations observed:
AIG(ze:) 2 ((xx:x++11
(17)
and
AIG(Z,;)) 2 X(x,; Xe1 (211))
(18)
Proof. Re-arranging the terms in equation 16 we can also characterize the autoinformation gap as:
AT(x;:) AT(2:T) [(2, X (2(47) + 7 7(x;; xxx7;x,+771
(19)
[(x; 2. (2:)) + / (xx; X (())))
(20)
The expression in the second line be derived by symmetry. Statement follows from (P.1)
B.4 AUTOINFORMATION OF SEQUENCES
Statement. A sequence of representations preserves autoinformation at 7 if and only if all the pairs of its elements at temporal distance T preserve autoinformation
AIG([Z][_.:T)= = + AIG(zm: T) = Vm E [s, T
(21)
Proof. We prove the two directions of the implication separately:
Proof by contradiction. Assume
17

---

(T.1) 3m E | - 7] for which the autoinformation in X is strictly larger than the autoinformation the corresponding representation Zm
Ax T) > AT Al(zm T)
The autoinformation gap between the two sequences can be written as:
AIG(ZZ)E.iT) E, (xGT)
(z T)]
1
(AI(xmii) - Al (zmit))
(22)
We derived that the autoinformation gap must be strictly positive, which results a contra- diction.
If we assume that mutual information is the same for all the pairs, clearly their average is also the same.
B.5 AUTOINFORMATION AND SUFFICIENCY
Statement. Whenever preserves autoinformation, 2, is sufficient for any (noisy) function of Xetri
AIG(ZØT) > I(x6y14)) I(2a)144) Vyr4+ [(xeer,5))
Proof. We address the directions of the implication in Lemmal separately:
We start by assuming that 2, preserves information AT(x;;)) AT(Z;;
I 1(4
P2
(xizi; - [ Not
(x1215 ye+r)
(Zi5 yetr) + (xxx V (z))
P33
(2i, Yetr) + I [xxe Xetr (2:)
Since we showed [xxx)t++7) S I(=r) and using DPI we have I(xx)t+77 2 [(2 Ve4r), we must conclude that [(xx; [2z Y147)
= We prove the second direction of the double implication by contradiction.
(T.1) Let yetr be noisy function of x,+r for which z, is not sufficient:
T(x.)x T(2E)111))
18

---

then
AT(z;;]
T(ze; tt)
(Z) ( M (1+44)47)
111
(zz) Yes7
2.
(4:: (t17) + [z, 2x (8c1T)
1 1 + (Z Zi+T (
P.3
(x + (x1; X (
P2
P2
[xx + (xxi ((+T)
(A22)
(23)
We derived that the AlG(Z(T) 0, which contradicts the premises, concluding the proof.
B.6 MARKOV PROPERTY
Statement. Sequences of representations of homogeneous Markov Chain that preserve information at some lag time T also form a homogeneous Markov Chain at temporal resolution T:
AIG([Z)) ;T) = = [zy+erlkez is homogeneous Markov Chain,
(24)
with s' E T r). K & [(T s
Proof. In order to prove that [2.++ A form a homogeneous Markov Chain, we first show that [Zy+trlk= satisfies the Markov property. This be shown by upper-bounding the amount of
K
information that the past [2=45r]]== carries about the next representation 2.4(641)y whenever the current representation 2 kr observed:
P.3
[(zz:45715_0 : (2/+((+1)r
12
(x, x. (x+((+1)rr (2(4k))
P22
([xe+)r]) k-1 k-1
Z x -(++1))))))) 1(zv+kri (x
([x,+jr)) ( (Z (
P2
(Xo1tkri Xx (Ze' Xx
+ [[xx (/Xks++r])j x ( (x4
A.3)
(Xo44 Xx (Zor4br X (x
P.3)
(xxx kei X ( (krsikri2e Ro (
A Al
AIG((=++rr:) O,
(25)
for any s' E - K. K< (T - s///J. and k E [1,K-1
Using the results from B.4 and the premise that the autoinformation gap is zero, we can conclude that the conditional mutual information in the previous equation must be zero, and the Markov property holds. Furthermore since both p(x,(x,->) and p(z, (x,) are time-independent, we must conclude
that plzz/z, )) must satisfy the same property. Therefore, we conclude that homogeneous Markov Chain.
forms a
19

---

B.7 BOUNDS ON THE AUTOINFORMATION GAP
Statement. For any T' T O, the autoinformation gap for z, lag time I' is upper-bounded by the sum of the autoinformation gap for Z, at lag time T and the autoinformation fap for Zitr-r at lag time T:
AIG(ZeT) & AIG(ZAT) + AIG(Zu+r-riT)
(26)
Proof. Let T' T O. The autoinformation for x, T' can be written
AI Ax;7'')
(x+7))+7))
(x X (rt) T T(z;; (x))
P.1]
(x/Z(iXe
P22)
(zriXe (r)) + (xx;) X4 ir (2))
(P3)
5 ((+rr) + (X1; Xi+r (2))
(Z x + AALG(Zz; T)
P.
(Z X 2 (Zii Zi-r (x+r)) + AIG(Z
(Zii X
AlG(ZiiT)
Zii 2x47) (11; Z(+r') AIG(Z(T)
P3
5
(zc )(tr)) + 1 1(x X X ((tr)) AIG(Z(i)
3
(zei ((t7))+)) + [(x X Betr') + AlG(Zz:T)
Al (Z(i)) + AIG((z; T) + AIG AZ
(27)
Re-arranging the terms, we have:
AIG(Zat) & AIG(Z.:) + AIG(
(28)
Note that whenever z, is sampled from the equilibrium, have Al(z.T') 2A1(2,7)).
B.8 SL.WERR INFORMATION PRESERVATION
Statement. a sequence of preseev autoinformation lag time т, then it preserves autoinformation for any +' 2
AIG( - AIC((Z))(.:7'>
(29)
Proof. Using the result from B.7 we can express the autoinformation gap 7' as:
AIG([Z)) :T')
B |AT + AIG (244
AIG([z,)1 |T= :T) + (I :T).
(30)
Since both 12-'rr and [z.]--. can infer that the right side of equation 30 must be zero. Furthermore, since A :T 2 0,
sub-sequences of [z, and AIG([2)) - ;T)
must conclude (TG(()))".77)) 0.
T"
20

---

B.9 LATENT SIMULATION ERROR AND AUTOINFORMATION
Statement. The average latent simulation error introduced by unfolding K steps using latent simulation starting from X2 with t (s, (ts.s+T-l)) is upper-bounded, by K times the autoinformation gap gap.for the sequence (zi)ee with T (K 1)T
E
(L(p([y xx)llp ([y ([y++xrlee hn-1 [xx))
K (IG((Z)
(31)
Latent Simulation error for K. steps of starting from
Astoinformation for lag time
Proof. We start with the following bound:
KL(p([yes xx)llp ([yettr (x.) (L(p([x(<*r]] (x:)|lp (xx (xx)), (32)
which holds because of assumption (A.2) and the data processing inequality. Secondly, we upper- bound the right-most term as a sum autoinformation:
(L( 1 [x)]ll ((x(((((x)++ 1k xx)) < (L [X++TT: Zitho K
(x))ll (z), [X+trr Zitkr 1a=1 (x.)
log pa (Z, xx)
II
Ti= pl p( Xx+ 1)+ po (Zetkr Xetker
Do 22 xx)
[ p (Ze+ kr Ze+ (xttrr (eekr)
K
K
log p ((t++r (e+ 1)r ( (Z X(tkT)
E
p(Z (z+(k-1)7) p Z(thT)
k=1
K
log p(x++br ()) ((ZL))T)
E
plze+ke Ze4(8-13r) pl (
K
E log p(xc+ter / X2+ 1y
p(Z)
+ E log p(ze+kr Ret
p(x++kT)
A[(x++AT/T) Al(e+tTTiT)
ALG(Z.+AIi))
(33)
Lastly, consider the average error when t ~ U(s, +-
E, KL(p(ly-<rrkt (x))ll((((( ([yctrrk [x:)]
KL(p[[xHurI (x))llbs ((x (xi))
AIG
R-0
AIG(Z
KT
AIG[[Z
(IG((z)et((tI>)r ++(K+1)7-1
;T)
K AIG((Z)[[a:T)
(34)
with == (K 1)r 1. This concludes the proof.
21

---

B.10 SPPER-BOUND FOR THE VARIATIONAL LATENT SIMULATION ERROR
Hereby, we outline the steps to obtain the expression reported in Equation S
(I(p([ys+*r s+ 1x, LS ((y ([y+&T])ke =+ (x,))
= E log 1 1K (x..) p0 (1V 1K (x.) pl.8 (oyo((yrlll) 1k (x) ql ([y++trle K
(x,)
p ([y &ot k= (x,)
w K
KKLPp(VV+trrlee xx)llp' ([yotarl xx)) + E
log
(35)
g ( 1k=1 Xx]
Focusing on the second term:
E log p (1y + x,)
E log p ([y ([o+trr Zz krr
k K
(x,)
9 (1V s+ k=1 xx,)
I/.S [vy - 2 k=1 (x,)
2utkr 2,+ )p(y -+ zg+kr)
E
log
1= kr Zy+ (1+ M4 (v o k
p(z,
Zo 13r)
ply M for 22
E
log
E
log
go (Zothr (o+ 1)r
Io (Y BotkT)
(L(p(Zz=* (zo b2 (6)(()11))))
(e+trr)
(36)
B.11 SCPERFLUOS INFORMATION DECOMPOSITION
The total amount of information that a representation z, contains about the original data can be de-composed using the chain rule of mutual information as follows:
I (x,; z
P.2)
(x 2.)
[[ze Z, (x,)
A.1
(x,2(-ri2
(Z
(xx Zi (Z(-7)
(37)
Superfluous Information
We can further factorize superfluous information by considering the immediate past x;-1 as follows:
I(xe; Zz 2t
P.2)
1: - [2r 1;21(xe)
Superiluous Information
A.11
[x,X,-12 1: (Z(-r)
(x,-1 2, (2) + T(x;; Zu / [,_ 2(-T)
P22
(x,-1i 2, [2(-7 + (x, : 2i (x--1) - (//1ri
(A.1)
(x,-1; Z, | + / (x,Z,-ri - (x()
(x,-1;
(x,; (x-1)) + [(z 2 X2-1)
T(x;- 2z
(2--7 |
I(x:: Zz (x-1)
(38)
Dynamic Information faster than
Time-independent information
in which the last step follows from:
IA..
(Ze: 24
Xx-1)
(X4;
X
X
(39)
22

---

Note that ((x,-1; refers to the information that z, conveys about the immediate past x,-1 when the past representation 2 is observed. This quantity is positive whenever 2 contains infor- mation regarding processes that are faster than i.e. not predictable from the past representa- tion 24 but can be inferred from 24- The second term [xx;; 2,|x, 1) refers to the information that 24 contains about processes that appear amme-independent the highest available time-resolution (T = 1). This component includes both time-independent noise and other time- depeedennd processes that appear uncorrelated at the observed temporal resolution. These two last components are indistin- guishable without having access to higher-resolution sequences.
C COMPUTATION AND
C.I A TWO-STEP MINIMIZATION PROCEDURE
Consider the terms on the right side of expression We use
( (0) (L(p((y++4T))) [xx (Ly (yy K
(x.))
(40)
C (0, Σ 1 (o0t000)00)
(41)
K
(0,)))) = Σ (T.(p(ys:17124 1a
(42)
for notation brevity to underline the dependencies with the parameters 0, y for the encoder, variational transition, and variational predictive distributions respectively. The joint optimization can then be written as:
min C (2) C min C (()(()) min (0, Φ) + min (2,(,, v) 0,0.0
- cLS (0) + min C () + min C C²0 v).
(43)
With 0 = arg min, CLS(
The upper bound in equation 43 is still tight for flexible variational transition and prediction distribu- tion. For a fixed f, the variational transition and predictive gaps depend uniquely on the variational
parameters & and vp which can be optimized by minimizing the negative ligliielihood:
arg min L C²0,,) = arg min
E log go Re -1)>]]
(44)
1=1
arg min L C²(,, arg min
E log qs (Y (x
(45)
C.2 CONTRASTIVE LEARNING ON MARKOV PROCESSES Consider the expression reported in equation c
e F(
LDfONCE: ( -E log
T.)
(46)
Exirpizi) (e
B
F Fel
log
(47)
B 2
Felse
Focusing the denominator in equation 46
note that estimating the partition function would Ti N
require sampling z from p(z:). If the dataset consists of multiple trajectories x
(4)
p(xx)]_>>
then this would require considering the representation of x from multiple trajectories at the given time f. Since we considering time-independent homogeneous processes, even when the dataset
23

---

consists of single trajectory [2:] we can approximate samples from p(x,) by considering any at in the same sequence, with t' T)). This approximation is accurate whenever p(x,) approaches the equilibrium distribution and the trajectory [æ=] is long enough to obtain de-correlated samples. In case multiple trajectories are available at training time, this approach would benefit from creating mati-batches inputs xi" (and corresponding representations z 2 that are sampled from distinct trajectories:
MHlfoMaxx [1) N
Pe(e 20
(4-))
(T;0,5 ~
log
(48)
B
Fel 1 x
1=1
2===
C.3 SUPERFLUOUS INFORMATION UPPER-BOUND
Computing superfluous information would require access to the true transition distribution p(z, (z=-)) Using standard variational inference, can define variational upper-bound based on the variational transition distribution instead:
log p(zz X Z1 Zi-T) p( the 7)
[[ 1(xx; Ze (x1))) Superfloous information
p(zz X1, 2s ks_r) ) Uo (Z) )
log
no (2, Fh 7) p( 2
KL (z, x5, 114 (((-))
K L (L(p(zi(zz-r) (go (2 (z-r))
Variational transition
$ [K<pp(zz
(49)
The Expected value of KL-divergence between encoding and transition distribution can be estimated using sampled representations:
po(2)
K (po((Zz)) (zelz,->) ze log
(50)
go g k=--))
with and 2 as representations sampled from pa la:)) and pa (2t respectively, and Ie: -i-y as samples from the process at temporal distance т. Notably, this procedure is similar to the one used to enforce a bottleneck in /Ficcher (2020).
D COMPARISON WITH THE LITERATURE
D.1 LINEAR CORRELATION MAXIMIZATION AND MUTUAL INFORMATION
conventional and successful approach mutual information maximization is the maximization of linear autocorrelation (Andrew et al. 2013 Noč & Nüske 2013). This can be expressed
arg max T 2.) subject to Cov Novlz,-,,25 z Covlzz, z,]
(51)
Here the maximization of the covariance trace equivalent to the maximization of the sum of its D eigenvalues X, where D denotes the dimensionality of the representation.
variety of surrogates maximize the sum of squared eigenvalues (Mardt et al. 2018 Wu Noč 2020)
the squared Euclidean distance in the representation space (Lyu al. [0222] Wiskott & Sejnowski 2002). This objective can also be equivalently interpreted as maximizing mutual information for jointly Normal random variables (Borgal 2001 with linear encoders.
Assume that the representations 2, and 2, are jointly Normal distributed:
[ [
S S.
[2, r. 2r] ~ NN(,, with S
(52)
S
24

---

Ag + Aa
1122log11 Ap) 1/2log11 -A3]
02
04
0.0
1.0
Figure Visualization of several objectives as a function of the eigenvalues X and x2 of St The vertical lines for d1 and d2 correspond the eigenvalues of Σ determined by the original covariance E.-r,- Note that whenever 2, is a linear projection of xx, Å and X are constrained to be in the shaded region determined by d, and d2. As result, all objectives are optimal for Å₃ dj and xa d2, which corresponds to a projection onto the principal components.
In this instance, autoinformation can be directly computed as follows:
det S det St.
AIv(Z log det S
det St.
log det (S S, n. S.-T.)
1
log det (I A)
2
1
= log det (U A)U?)
2
log det (I - A))
=
2
D
log (1 Xd)
(53)
2
2=1
In which A : 1/2 S.. " and UAUT refers to its eigendecomposition. and
A the corresponding eigenvalues. Under the assumption that S. and St, restricted to be identity matrices. the expression for A simplifies to A = S
As illustrated Figure for any linear encoder in the form z, W WX, maximizing the sum of the eigenvalues of A, the sum of their squared values, or the expression in equation 53 is equivalent. This is true because under the constraint S., St-r.e- upper- bounded by the eigenvalues of E,.-,E,-ra. with E, = Cov(x,,,,
the eigenvalues of S St-r, are
Note that although the correlation matrix does capture linear relation between Z;-r and z, does not consider higher-order interaction between the representations. This is limiting factor especially for low-dimensional representations because of the expressive power of linear transformations. This phenomenon can be clearly observed by comparing the autoinformation plots in Figure 4b (2D representations) and Figure 13 (16/32 dimensional representations). The autoinformation extracted by representations that linear correlation maximization (TICA and VAMPNet) strongly depends on the number of dimensions of the representation 2,, The effect on methods that rely on non-linear contrastive mutual information maximization methods is much more moderate, making them more flexible and suitable for 2D visualizations.
25

---

N+++
Ya+Kr
(b)
Y+tr
YoSKT
3/
3
X
X
Xa+ Kr
x.
X
K
0
0
Z
Figure 7: Viable inference schemes for maximal future state information: (x; T) - T(z x The difference with the Latent Simulation inference scheme lies in the lack of the conditional independence / (2,;x++77 -x++>/Z++> = O. Note that modeling (x, (.--) is generally more difficult than modeling latent transitions go 2, (Zt-)).
D.2 MAXIMIZING INFORMATION WITH RESPECT TO FUTURE STATES
Several models in the literature consider the reconstruction of future states as the training objective (Wehmeyer& & Noë 2018 Wang et al.) (019b).. This objective can be interpreted as the maximization of the mutual information between the current representation Z; and the future state x, Poole et al. (2019):
max [(, max H(x++7) H (x (2,)
H(x+++)
min ) (ail))))) |- log qol (x 0.0
(54)
in which go (x(x++l2z)) refers to the decoder that predicts the future states given the current representa- tion.
note that autoinformation in 2, is always smaller or equal to [(zz; Xt-r), which we will refer to as future predictive information:
AI(x; 7) = [(xx;x++) 2 (z, x++7) 2 I (z,; Ze47) AIz
(55)
Preserving autoinformation is stronger condition than having maximal future predictive information:
ATG(Z;7) = 0 - AT(x;7) - T [(2 x x 0.
(56)
This is because the additional condition / (2,; X1+7 (2447) is required:
AIG(z;:T) = AI(x;; T) [(zz; 2e+r)
(x ((z; x(+T)
+
(57)
Missing future predictive information
This additional condition not directly required prove the results of Lemma Lemma
2 and
Lemma 3 which can be extended to the condition AT(x; r) [(z,; X )= 0. However, the lack of the conditional independence / [[z, X 7 would result in difference inference scheme, in which instead of approximating transitions directly in the latent space, each step would require modeling transitions from z, to x, as shown in Figure 7a Alternatively, one could model latent transitions 9o iz, but the predictive target distribution would depend on both the current and future representation, as shown Figure 7b Both inference schemes and the maximization of X(1;;
computationally expensive than the proposed T-InfoMax training procedure and Latent Simulation inference.
D.3
STATE PREDICTIVE INFORMATION BOTTLENECK AND TARGET SUFFICIENCY
Wang & Tiwary 2021) introduce State Predictive Information Bottleneck (SPIB) objective aiming to create a representation z, that is sufficient for the next target Yi+r while compressing information:
SPIB (0; 3,T) EE, [[ ) (T(xx;2.)).
(58)
26

---

Although this objective seems natural for training effective representations, ciency for a given target Yt+r a necessary but not sufficient condition for autoinformation preser- vation. As result, representation that is optimal according to the SPIB objective may introduce inference error even when the true latent transition p(z,|z, bhe +) and latent future predictive p(y:: rl2)) distributions available inference time, shown in the following example.
show that suffi-
Consider a dynamic system in which each state is described by particle position, velocity, and acceleration governed by simple time-discrete update:
Fr +
X
Vi =
V ra = (x nc),
(59)
in which the acceleration at each time step is sampled from a limeiindeeendent Normal distribution m: NO, 1) and D, refers to the function used unroll the true system dynamics at the time scale Clearly, the system is an instance of homogenous Markov process.
We are interested in predicting the particle position y: rt- Clearly, since the next position depends solely the current position and the current velocity, have that representation that contains only velocity and position information is sufficient for the next target prediction:
SPID x( T(yeiri (PID) with z
r
(60)
V+
On the other hand, representation that maximizes autoinformation (and is optimal according to Equation 8) must also contain information regarding the acceleration since current acceleration is predictive for the future velocity:
r;
T(x;; XELT) = [2, T-IB
21 :2: T-IB) >
[z,,
SPIB, SPIB
2
with 2;
T-IB
V:
(61)
[aa.
Note that representation that is optimal according to SPIB would instead explicitly discard accelera- tion because of the compression regularization:
I [xx z! 18) ((
(62)
Since both representations are sufficient for yı-r, they yield the same predictive distribution for the next target:
ply 2 SPIB)
ply (z:) p(y+++xx))
(63)
However, if we look at the predictive distribution at times larger than т, we observe some discrepancies. In particular, we can show that:
ply+ 2r 2 SPID) # p(ye12r bz 1B) p(y+ (x,),
(64)
In which the first inequality follows from the fact that 2, SPIB does not contain knowledge about the acceleration, while the second inequality follows from Lemma Lemmma Therefore showed that latent simulation performed on representations that are optimal according to the SPIB objective (and not according to T-IB) introduces inference error for time scales larger than т. The intuition is that sufficiency for the next target Yttr does not guarantee transfer of the Markov property from the original space x, the representation z, That requirement is satisfied only whenever the representation 2, preserves autoinformation, shown in Lemma T+ Lemma
27

---

E EXPERIMENTAL DETAILS
We include additional details regarding the training data, architectures, and optimization procedure to ensure the reproducibility of the reported results.
E.
DATA
E.I.1
PRINZ 2D
3
y= 0
y=
v
3=
0.5
Potential v(x)
0.0
2
-0.5
-1.0
-1.0
0
2000
4000
6000
8000
10000
X
Figure 8: Left: visualization of the ID Prinz potential, and the corresponding regions used to define the discrete targets y. Right: Visualization of the ID slow component æ colored by y².
The Prinz 2D trajectories consist of sequences of 100K data points generated by diffusing a point
particle into a potential V(x) :- 4 (x + 0.8e 50.² + 0.2e 50(==0.5)* + 0.5e -00(==00.))) with an Euler-Maruyama integrator following the update:
T(+1 = x (z,) + vh mps
(65)
in which = 10-4 refers to the integrator step and is standard Normal uncorrelated noise. We performing 160 integration steps in-between consecutive timesteps, while [x:]]]] l=
generate x
is generated by considering 5 integration steps. The Deep Time package (Hoffmann et al.) 2021) is used produce the slow and fast trajectories, and the corresponding potential V(x) is visualized in Figure The fast and slow independent components then mixed as follows:
tanh(x;; xx) 22
(66)
ttanh(x] -
produce the trajectories visualized in Figure
E.1.2 MOLECULAR DATA
Trajectories We analyze trajectories obtained by simulating Alanine Dipeptide, Chignolin, and Villin LLindorf-Laaree et al. 2011). For Alanine Dipeptide, the three splits correspond to separate simulations of 150K//00K//000 frames respectively. In contrast, for Chignolin and Villin simulation, a single trajectory is split into 3 temporally disjoint parts: 334.743/100K/100F frames for Chignolin, frames for Villin. Each observation x, consists of the of the Euclidean
and 427. 177.90//100K/1000 coordinates of all the atoms and one-hot corresponding to the atomic number for the Alanine Dipeptide trajectories. The input data for the mini-proteins, on the other hand, consists of a coarse- grained representation indicating the 3D location of the amino acids in the protein chain (10 for Chignolin and 35 for Villin), along with one-hot encoding for the amino acid type.
Targets Targets for the Alanine Dipeptide molecules generated by clustering torsion angles and & into 6 regions, corresponding to the known meta-stable states. For the Chignolin and Villin molecules, generate targets y, by clustering the 32D invariant TICA projections obtained by
28

---

Chignelin CC Clusters
Chignalin RMSD Clusters
12
10
Free
20
20
25
30
0.0
0.
044
0.6
0.
Distance (A)
RORD (4]
Villin CC Clusters
Villin RMSD Clusters
14
11
Free
Dec aecannei.. (4)
Figure 9: Visualization of the 1D free energy induced by the distribution of the distances between first and last C-alpha atoms in the chain (CC) and Root Mean Squared Distance (RMSD) for the molecules of Chignoling and Villin. Vertical dashed lines different discretized regions. The thresholds are set to [7, 15]. [1.3, 5] Angstroms for Chignolin CC and RMDS respectively, while the values of [6.5, 15] and [2.5, 4.2] are used for Villin.
used denote the margin between the
following the same procedure described in Köhler et al. 2023) using KMeans with centroids, as depicted in Figure 4a| We produce additional sets of targets by considering the distance between the first and last C-alpha carbon atoms in the amino acid sequence (CC, clusters), and the Root Mean Squared Distance (RMSD, 3 clusters) from the stable folded configuration. The thresholds used to create the clusters are visualized together with the corresponding free energy in Figure 9
Lag time We decide on training lag time 7 for each molecule that long enough capture relevant meta-stable state transitions, see Figure 4b and Figure T3, We focus a time scale on which most the dynamic information still present while modeling transitions that are orders of magnitudes larger compared to the original simulations. We used train lag time of 16 ps on Alanine Dipeptide simulations, while 3200 ps was used for the Chignolin and Villin simulations. The same value of is used both train the encoder (step 1) and the transition model (step 2).
E.2 ARCHITECTURE: AND OPTIMIZATION
The models used for the experiments reported in this paper are described detail in the following sections. The experiments reported in this paper required a total of 25 days of computation on A100 GPUs. This estimation includes model development, hyper- parameter tuning, and evaluation.
E.2.1 ENCODER TRAINING
We train each encoder for a maximum of 50 epochs with mini-batches of size 512 using the Adam W (Loshchilov & Hutter 2019) optimizer. To prevent overfitting, validation loss. Following previous work (Chen et al. (020). the models are trained with an initial learning rate of 10-6, which is gradually increased up x 10-1 over the course of epochs with
early stopping based on the
29

---

linear schedule. The learning rate is then decreased to the initial value using cosine schedule over the following 45 epochs.
For the Prinz 2D experiments, encoders consist of MLPs with two hidden units of size 64 and 2D output. In the molecular settings, each encoder architecture consists of a TorchMD Equivariant Transformer (Thölke & Fabritiis| (022) followed by global mean pooling and linear layer to produce a rotation, translation, reflection, and permutation invariant representation for each molecule. We use total of layers of 32 hidden units with 8 heads each for the Alanine Dipeptide experiment. The more challenging Chignolin and Villin molecules use encoders consisting of 5 layers with 64 hidden units and 8 projection heads instead. For the evaluation of the quality of unfolded trajectories, we use 16-dimensional representation for Alanine Dipeptide. A total of 32 dimensions are used for Chignolin and Villin.
TICA Temporal Independent Component Analysis consists of linear projection of the input data onto the principal temporal components. As result, po(z,(x,, consists of simple linear projection instead of a neural network that has been optimized using the Deeptime python library (Hoffmann)
et al. 2021). For the Prinz2D experiments, we apply TICA directly to the original sequence [x,]=_. to project each data point x, onto the principal temporal component Z₁. For the Alanine Dipeptide Experiments, the TICA representations correspond directly to the torsion angles determined by the carbon skeleton (2 angles), which are commonly used in literature to describe the configuration of this small molecule (Vymétal & Vondrášek, 2010; Mardt et al.] (018)) The representations for Chignolin and Villin are produced following the same procedure described in detail in Köhler et al. (2023), in which torsion angles and inter-atomic distances are projected onto the principal temporal components.
VAMPNet We train the encoder po(2: x,) using VAMP-2 score (Mardt et al. 2018
Wu & Noé
2020) using the implementation from the Deeptime python library (Hoffmann et al. 2021) The VAMPNet model was originally designed for estimating dominant spectral components of molecular simulations. However, Sidky et al. (2020) has shown the effectiveness of VAMPNet for Latent Simulation inference.
T-InfoMax As a representative of non-linear mutual information maximization methods, we consider the popular InfoNCE method (van den Oord et al.) 2018 Chen al. 2020). Following the literature (van den Oord et al., 20188) Poole et al.) 2019), we model the log-ratio between joint and product distribution with a separable architecture:
(()) = ge. (z..)) ge,
"T.. (z.),
(67)
in which Is, and It. neural networks mapping the latent representations into 128-dimensional normalized vector. The two architectures have distinct weights with one hidden layer of 256 units and group normalization (Wu & He 2018) before the ReLU non-linearity.
T-IB Analogously to the T-InfoMax counterpart, the Time-lagged Information Bottleneck objective makes use of InfoNCE for time-lagged information maximization, with an additional regularization term modulated by the hyper-parameter as shown in equation 9 Following (Federici et al.) 2020) we first train the encoder with an initial value of 3 = =10 for 5 epochs. This is to prevent the representation from collapsing into a constant at the beginning of training. Secondly, the regularization strength is gradually increased up to the final desired value over the course of 30 epochs. We empirically observed that the T-IB models benefit from the use of a stochastic encoder encoders is comparable to the corresponding deterministic counterpart.
po(2,(x.) = (22, [xx)),), ag (x(x)))) The parameter vectors (x,( and (x(x.) are obtained using two linear projection heads on top of the encoder features, as result, the size of the stochastic
We initialize the architectures with a value of ox(()) x 1 to reduce the amount of Gaussian addi- tive noise in the initial part of the training. Empirical results showed that the additional stochasticity produces smooth transitions between different levels of regularization strength. This is in contrast with the sharp regime changes observed with deterministic encoders (Figure 11 Top). We believe that this is due to the fact that the addition of Gaussian noise allows the encoder to destroy superflu- ous information locally when necessary.
30

---

E.2.2 TRANSITION AND PREDICTION TRAINING
The variational transition and prediction models (qo (Z.(Zz)))-)) and qu (z.(y))) respectively) are jointl trained on the embedding produced by the encoder trained in the previous step. The training procedure uses mini-batches of size 512 with A and a fixed learning rate of x 10 over a total of 50 epochs. Contrary to the previous step, we did not observe any overfitting with only marginal improvements in the training and validation scores by the end of the training procedure.
Transition The transition model consists of conditional Flow++ layers (Ho et al. 2019) due to their flexibility, sampling speed, and ability to model correlated distributions. The transitions for Prinz2D and Alanine Dipeptide representations consist of 3 flow layers. Each layer is composed of conditional mixture of logistics CDF coupling transformation consisting of neural network with two hidden layers of 64 hidden units, which maps the representations Zz-r into the parameters of mixture 16 logistics distributions. An architecture of 5 layers is used to learn the more challenging transition distributions for Chignolin and Villin. To prevent numerical overflows while unfolding long simulations, we clip samples to be in the interval 10%, 10%.
Prediction Each feature predictor used in this work consists of simple 1-hidden layer MLP with 128 hidden units mapping the representation 2, into the logits for the variational predictive distribution qo (y,lz.).
E.3 EVALUATION
We focus our evaluation on two main aspects. First, we analyze the amount of autoinformation that several models extract from the molecular data to better understand which temporal characteristics of the molecular process are successfully captured. The second aspect involves the evaluation of the fidelity of trajectories unfolded using the Variational Latent Simulation process.
E.3.1 MUTUAL INFORMATION
Autoinformation We estimate autoinformation for evaluation purposes using SMILE (Song & Ermon 2020) on the trained representations z, with clipping interval of [-5, 5]. The ratio estimation architecture consists of an initial projection head g Z - IR 128 with one hidden layer of 256 units and output h, := g(z) with dimension of 128. Pairs of the 128-dimensional feature vectors h, at different temporal resolutions are then concatenated and fed into second MLPr, r, R x R2288 - IR with 64 hidden units and output, which corresponds to the estimated log-ratio value. Each pair of h. ho+r is fed into distinct architecture T, for each r. This setup allows to estimate autoinformation at several time-lags at once to produce the plots visualized in Figure (b) Figure [3 and Figure 14a Each dot in the figure corresponds to the expected output of one ratio estimation model r, (g(( (z,). (2),, g(2(+)) on the entirety of the training set. The ratio estimation models are for at most 20 epochs using carly stopping based the validation loss. Note that samples from the marginal distribution used to estimate the value of the partition function are sampled by sampling x with uniform probability using the same strategy described in Appendix C.2 Estimation is performed
using the Torch-Mist package Federici et al. 2023
Target Information Following Poole et al. 2019); McAllester & Stratos 2020); Song & Ermon
(2020), we estimate the amount of target information in the representations as a difference of cross- entropies:
y.) = Hy.)) ) Hy, S Hyy)) IE[- log do (y<lz))).
(68)
in which the marginal entropy H for the discrete targets y: is estimated by counting the frequency of each class, while the expected cross entropy E log 9a (y, (y,1z,)] evaluated using the trained predictor qu (y:lz)) on the entirety of the test trajectory and computing the corresponding expected log-
likelihood. Note that with [Zz; Y:) we implicitly refer to the expected amount of target information over entire trajectory rather than the amount of information estimated specifically at the time-step t.
E.3.2 UNFOLDING TRAJECTORIES
Accurately estimating a measure divergence between joint distributions when only samples are accessible is generally challenging task due the number of samples required for reliable
31

---

estimation. For this reason, instead of considering continuous multi multi-dimennsonal targets y.. we focus our attention on discrete targets. The targets our experiments are designed to capture properties of interest of the trajectories
Our evaluation procedure can be described steps:
1. First we encode the initial (unobserved) test state 2, into the latent configuration 2, using by sampling
(o(z|(x,) Starting from 2.. sample total of 256 trajectories
((()
| A=1
from the variational transition model go (z,)z sequentially for total temporal duration which is comparable to the time-span covered by the test trajectories T s s KT. Using
the prediction model qu (y))z) we then sample a target ((() for each sampled z(() obtaining K
256 sequences of targets U+AT k-1
-(1) 2. We count the number of transitions from each discrete state 5) to the following Yetkr for various numbers of steps k, effectively creating a series of transition count matrix C
and C.. respectively for (,(0 V+Kr k=1 and [yil? The 256 count matrices for the unfolded
256 O We normalize each row trajectories are then averaged to produce Čer 1/256 Σ
of ÖLr and Ck. to estimate the transition probability matrices Thr and TAT Analogously, we count the number of times that each state is visited to determine the normalized counts m and m using the ground truth and unfolded trajectories respectively.
3. We compute the Jensen- Jensen-Shannow divergence between each row of Ter and TKr, then average the values obtained for each row into a single number, representing the average -haannnn divergence. With this last step, we obtain one value of transition Jensen- Shannon divergence JSJS) for each chosen number of unfolding steps & (see Figure 15). The values for each row are averaged using the same weighting instead of the relative state probability to accentuate errors when transitioning from rare states. Analogously we compute the value of marginal IS (M by computing the divergence between the probability distribution induced by m and m.
F ADDITIONAL RESULTS
In this section, report additional ablation studies and the performance of the models considered in this analysis for different sets targets.
F.1 T-IB REGULARIZATION STRENGTH AND TRAIN LAG TIME
2.0
1.00
1.8
5.0
0.50
Iagal
0.5
SUW
fact
0.0
0.00
Ragu wriatiionn
Regudarzatoon
Figure 10: Visualization of the effect of the regularization strength on Autoinformation. information regarding slow and fast modes, transition JS, and amount of superfluous information on the Prinz 2D dataset. All representations are trained using - = 64. Representations trained with 3 < 0.01 tend to contain information regarding the fast mode and higher autoinformation at small temporal scales, while strong regularization 3 > 0.1 results in representations that contain too little information. Note that the best performance in terms of transition JS divergence is achieved by the representation that contains the least information regarding y and most about yf, which corresponds to the most compressed sufficient representation.
32

---

Alanine Dipeplide
Chignolin
Vilin
0.3
0.4
Allxkl (nets)
0.3
0.2
stochslli
0.2
Tree
Falco
0.1
0.1
(a)
10
10
10"
10
10
10"
10
10
10"
Alanine Dipeptide
Chignalin
Villin
KL.py (2)0 -11 [nate]
10
80
30
20
10 10"
10
10
10"
(b)
Regularizatio f
Rogularization
Regulartzatinn
Figure 11: Visualization of the effect of the regularization strength on the representations produced with T-IB on molecular simulations with fixed train lag time т. [aa estimated autoinformation (y- axis) for the three molecules as a function of the training regularization strength / (x-axis). Stochastic encoders (in blue) show much smoother interpolation. [Tb] the amount of superfluous information (y-axis, Equation 49) as function of the regularization strength.
== 0.4 ns
T = 3.2 ns
T = 25.6 ns
T 204.8 ns
Autoinformation
Train
0.4 ms
3.2 ms
25
204 8
(a)
10
10"
10
(b)
Lag time t [ps)]
Figure 12: Visualization of the effect of the train lag time on 2D T-IB representations of Villin trained with 3 0.01. 12a representation of the test trajectories for models trained with several lag times, colored by the clustered TICA embedding, reported in Figure 4a autoinformation plot in which the dashed vertical lines correspond to the respective training lag times. Note that, motivated in Section 2.2 representation trained with higher temporal resolution also captures slower processes at the cost of introducing more information into the representation. This can be clearly seen by observing the number of distinct clusters emerging the visualized representations.
12b
corresponding
Figure 10 reports the effect of the regularization strength for T-IB representations of the Prinz 2D data. Consistently with the hypothesis, the best-performing model is the one that produces minimal sufficient representations the training time scale 7 64. This corresponds regularization strength of 3 0.01.
Figure compares the effects of several regularization strengths, demonstrating the differences between deterministic and stochastic encoders. Deterministic encoders (in yellow) tend to sharply transition from a fully informative representation (on the left) to a constant uninformative representa- tion (on the right). A secondary advantage of using a stochastic encoder is the possibility to compute an upper bound of superfluous information thanks to the expression for the density pe (z(2(()) This is generally not possible for deterministic encoder for which K LLopo (9g))ll9 (z, (2,r)) can be evaluated only up to a constant. Regularization strength 3 for T-IB is selected based on validation
performance: 3 = 0.01 for Alanine Dipeptide and Villin; 8 = 0.001 for Chignolin.
33

---

In our experiments molecular data, observed that small values of B can have substantial impact on reducing the amount of superfluous information contained in the representations, with only moderate impact autoinformation. We believe the possible reduction of autoinformation for larger 3 is due to the fact that processes faster than r cannot always be fully disentangled. This includes processes that contain lots of information at smaller time scales, but are only marginally informative for events that are 7 time-steps apart. Reducing information regarding faster processes can drop the amount of superfluous information in the representation but still decrease autoinformation whenever the faster process can not be temporally disentangled. Nevertheless, regularization strength in the order of 10-3 reduces the amount of superfluous information by substantial factor (10x) with little to effect the amount of extracted autoinformation at T.
Figure 12b shows the effect of the train lag time selection on T-IB models trained with , 0.01 and a 2-dimensional representation space for the Villin trajectory. Smaller train lag time corresponds to higher information content and more complex representations, while larger train time scales associated with simpler representations that are not suitable for unfolding simulation at higher temporal resolution. Note that the larger the training lag-time the longer the training trajectories need to be.
F.2 AUTOINFORMATION FOR LARGER REPRESENTATIONS
Plots in Figure 4b Figure 4a and Figure B confirm that with an appropriate regularization strength, T-IB model preserves the maximum amount of autoinformation the training timescale while decreasing autoinformation for smaller lag times (left of the dashed lines).
Note that the autoinformation plot all the models considered in this analysis matches for large time scales. This suggests that all the corresponding representations preserve autoinformation at large lag times while still differing in the amount of superfluous information at faster scales and the representation structure. The perfect overlap is also justified by Lemma 3 which guarantees that representations that preserve autoinformation at some lag time 7 must also preserve information at larger lag times.
Encoders trained with the VAMPNet objective on complex systems tend to preserve autoinformation only for slower processes. We further observe that VAMPNet models tend to become less numeri- cally stable for increasing representation size, while methods based on non-linear autoinformation maximization are less affected by this hyperparameter choice.
F.3
EVALUATING STATISTICS FOR MULTIPLE TARGETS AND TIME-STEPS
Figure 15
Figure 14b and Table report the values of average Jensen-Shannon divergence for transi- tion distribution for different targets y.. We observe that the T-IB model consistently outperforms the other models for transition matrices computed based on different objectives and several lag times.
One of the main challenges for the evaluation of statisties of slow processes (large transition times in Figure 15) lies in the limited amount of test time frames. We observed that, for large time intervals, the estimation the ground-truth transition distribution from rare states may be too noisy to produce accurate measures of Jensen-Shannon divergence. As result, the values reported for large transition times -xiis become dependent on the specific test trajectory used for evaluation. Nevertheless, we believe that the relative comparison between the performance of different models may still represent their ability to match the original statistics. More accurate quantitative analysis in this regime would require access much longer molecular simulations.
34

---

Alanine Dipeptide
Chignolin
Villin
1.0
1.5
0.8
Model
0.8
1.0
4
TICA
VAMPNet
0.5
2
T-InfoMax
0.2
T-IB (8~0.001)
0.0
10
10"
10"
10"
10"
10
Lag r [ps]
Lag [Ps]
Lag time [ps]
Figure 13: Autoinformation plot for high dimensional representations (16 for Alanine Dipeptide, 32 for Chignolin and Villin). Shaded regions indicate the standard deviation measured across 3 seeds and the dashed vertical line indicates the lag time at which the representations are trained. Representations trained with the VAMPNet objective generally less consistent (higher variance) across different seeds. T-IB produces sufficient representations (maximal autoinformation at the training time scale) while minimizing the autoinformation for smaller scales.
Autoinformation
Slow Mode
Model
Alliz.) [nats]
T [nats]
Original
0.01
TICA (1D)
VAMPNet
T-InfoMax
T-IB (8 0.01)
0.001
10
10
10"
10
10
10
(a)
Lag time
(b)
Transition steps
Figure 14: Measurements of autoinformation and transition JS estimated for several time scales. [4a values of autoinformation estimated at several lag times 7 for representations trained with r 64. 14b; values of transition JS estimated at several time scales 7' from unfolded trajectories. T-IB contains the least autoinformation at small time scales while preserving information at the train lag time or larger. At the same time, T-IB results the smaller TJS at all the considered time scales. The measure of standard deviation is obtained by considering 3 seeds for each model.
Chignolin
Villin
CC Cluster
RMSD Cluster
CC Cluster
RMSD Cluster
MJS
TJS051.2
MJS
TJS@51.2
MJS TJS851.2
MJS
TJS85122
TICA
7.5 7.6
5.2 f
7.4
3.7
1.7 1.5
6.1 2.4
7.3
5.3 3.3 3.7
VAMPNet
30 20
103
31.9 21.8
117 102
63
57
71
40 +
T-IooMaxx
1.7
3.8
4.8 1.6
0.7
2.1 0.5
1.9
5.8
6.3323
8.5
T-IB
2.3
1.1
2.9 2.2
4.1
0.8 0.3
4.4
1.7
4.1 1.6
Table Values of marginal (MJS) and transition (T'SS Jensen-Shannon divergence for trajectories unfolded on latent spaces obtained with different models for the prediction of the CC and RMSD cluster targets described in Appendix 1.1 The regularized T-IB model consistently outperforms the corresponding unregularized counterpart ("IIInfoMax) at the considered time scale.
G SIMULATION TIME
G.1 MOLECULAR DYNAMICS SIMULATION
According to the data reported in Shaw et al. 2021 the 64-node Anton supercomputer can simulate up to 250 microseconds per day for system consisting of ~ 10" atoms, which is similar to the total atoms in the Villin and Chignolin simulations. On the other hand, a single A100 GPU can simulate only up to 1.5 microseconds each day. The estimate is based on the data reported in Table III in
35

---

TICA Cluster
Alanine Dipeptide
Chignolin
Villin
Model
10
TICA
VAMPNet
TJS
T-InfoMax
10
T-IB
10
10
10
10
10
10"
10"
10
Transition time r' [ps]
Transition time r [ps]
Transitinn time r' [ps]
CC Cluster
RMSD Cluster
Chignolin
Villin
Chignolin
Villin
0.1
0.1
0.1
0.01
0.01
0.01
10
10
10
10
10"
10
10
10
10
10"
10"
10"
Transtion time [p4]
[so]
[pa]
Transition time
Figure 15: Measure of the average Jensen-Shannol divergence (y-axis) for the unfolded transition matrix for several discrete targets y, as a function of the number of unfolding steps (x-axis).
Shaw et al. (2021) and the simulation condition described in the supplementary material provided by Lindorff Saandorff-Lasent Larsen et al. (201T). Therefore, simulating a time jump of r 3.2 nanoseconds would require approximately 200s on an A100 GPU and about second on Anton 3.
G.2 LATENT SIMULATION
Unfolding one transition step using the Flow++ transition model used in our experiments requires approximately 100 milliseconds on a single A100 GPU. As result, for the reported Chignolin and Villin experiments, our estimated acceleration is about factor x x1000 compared to molecular simulations on the same hardware and x x10 for the highly specialized Anton 3 supercomputer. The total simulation time to produce new molecular simulation of the same length as the training one (I' ~ 100 microseconds) is approximately 3 months on single A100 GPU, hour with Latent Simulation on the same GPU, and 10 hours on 64-nodes Anton 3.
Note that total time require to unfold latent simulation T decreases increase the lag time
=T = T
in which the cost &S is determined by the size of the latent space, the transition model, the prediction model, and the hardware. As shown in Table for our experiments, the prediction time is negligible when compared to the cost of unfolding latent transitions. This is because the prediction model consists of simple MLP. Whenever the target of interest y. is also high-dimensional, the prediction cost may increase significantly. However, it is reasonable to assume both t/. and th to require in the order of 100 milliseconds for most tasks of interest.
The Latent Simulation process is also highly parallelizable. As shown in Table 2 is possible to simultaneously unfold more than 10a trajectories
single A100 GPU with little overhead.
is important to note that learning encoder, transition, and prediction models for Latent Simulation still require several ground truth trajectories of length T >= r, and the unfolded Latent Simulations are approximations of the molecular dynamics. This is because we do not directly represent the water molecules around the proteins nor the single atoms composing the amino acids.
36

---

Trajectories
100 Trajectories
1000 Trajectories
10000 Trajectories
Transition Prediction
124 + 3 0.690 0.001
128 + 1
130 +
166 = 0.739 + 0.003
0.700 0.001
0.732 = 0.002
Table 2: Estimations for the time (in milliseconds) required to unfold one step of parallel Latent Simulation of Chignoling and Villin single A100 GPU. The estimates refer to the time required to produce samples from the conditional distribution qu Zz >/Zz) and qu. (Yx z,) for given ž,. Note that
the cost of conditioning and sampling the transition model dominates the one of making predictions.
Details on the architectures for the transition prediction models are described in Appendix
TICA [8]
VAMPnet [10ss]
T-InfoMax s10's]
T-IB [10's
Alanine Dipeptide
1.04 + 0.01
2.39 + 0.02
2.63 + 0.08
2.78 + 0.02
Chignolin
42.2 0.2
2.8 + 0.3
3.0 + 0.3
3.5 + 0.3
Villin
60 +
12.5 = 0.3
12.8 = 0.2
13.3 + 0.4
Table 3: Training time required to train the encoder architectures on the Alanine Dipeptide, Chignolin, and Villin data. The measurements are reported seconds for the TICA experiments, and 10 seconds for the other models relying on TorchMD encoders.
Data
Training Time [10's]
Alanine Dipeptide Chignolin Villin
1.7 = 0.1
4.5 + 0.3
4.5 + 0.5
Table 4: Estimated training time required fit the transition and predictive model for a fixed representation. The estimates also include the time required to unroll and evaluate latent simulation for validation purposes.
G.3 TRAINING TIME
We report the training time corresponding to all the models reported in our experimental section by differentiating the time required to train the encoder (step 1) from the training of transition and prediction model (step 2) described in Section 2.
Table 3 reports the total time required to train the encoder architectures with the TICA, VAMPnet, T-InfoMax and T-IB objectives. The training time for TICA substantially shorter since relies on linear mapping instead of flexible TorchMD architecture. The variance of the time estimates is computed over three runs per experiment.
The training time for the second step equivalent for all models since the same transition and prediction architecture are fit to each representation using maximum likelihood. Train time is not influenced by the encoder (linear or Deep NN) since we encode and store the entire dataset to disk at the end of step 1. As a result, the total cost depends solely on the dataset size and size of the latent representation, as reported Table 4
The total training time (step 1 step 2) for the T-IB model on Villin amounts to approximately 5 hours. Unfolding a latent simulation the same length of the training trajectory requires another hour, bringing the total 6 hours. Even by including the training time, Latent Simulation is 100x to x faster than running molecular dynamics on comparable hardware.
37


Reviews:
[{"review_id": "8NS5Ul0IDu", "reviewer": "Reviewer_H2G8", "summary": "Markov processes can require computing many intermediate steps before reaching the final state which may be expensive. This paper proposes modeling state in latent space and computing large steps in time directly. The transition and emission distributions are approximated using variational inference. The optimization scheme is two step: first the encoding distribution $q(z | x)$ is learned and then the rest is learned. The approach ensures that the latent space contains information about the original state through the mutual information and \"autoinformation\". They introduce a loss T-InfoMax that maximizes mutual information on latent state across some resolution gap. Finally, they introduce an information bottleneck to discard unnecessary information such as time dependent information, resulting in a loss T-IB. All the mutual information terms are approximated, some are learned with neural contrastive estimation.", "questions": "- How to choose $\\tau$ and once chosen, does this parameter remain fixed and the model has to be retrained for different choice?\n\n- Can you provide run times?\n\n- How does the model scale with the amount of data available for training?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is nicely written and the method is explained well.\n\nThe information theoretic approach is a good motivator for what follows in the paper. The method itself is going through a couple of departures from the pure mutual information and get increasingly complex and further approximated.\n\nThe theoretical results in 2.2 are nice and follow the story, although I did not check the proofs in detail.\n\nThe results look good although some additional discussion is missing.", "weaknesses": "The method has many moving parts and some choices are not well motivated. For example, different choices of architectures for different distributions, training with contrastive estimation. A study that sees which part is contributing to the performance would be nice.\n\nThe training requires two step procedure.\n\nAs far I can see the authors do not provide actual run times. What should be compared is the total time to obtain the training data, train the model and perform the simulation.\n\nThere are two experiments, one with synthetic data which is a toy example and one with molecular simulation which again seems small in scale.", "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": "3glZoTWP6j", "reviewer": "Reviewer_mH7G", "summary": "The paper addresses the challenge of accurately simulating large-scale Markov processes over extended time scales, a computationally expensive task. To overcome this, the authors introduce an inference scheme based on the Time-lagged Information Bottleneck (T-IB) framework. T-IB seeks to simplify the simulation process by mapping complex systems into a reduced representational space that captures relevant temporal features while discarding high-frequency information. Through experiments, the paper demonstrates the efficacy of T-IB in creating information-optimal representations for modeling the statistical properties and dynamics of the original process at a selected time lag.", "questions": "See Weaknesses", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The introduction of the Time-lagged Information Bottleneck (T-IB) framework, rooted in information theory, offers a novel perspective on data-driven model reduction of complex stochastic systems. This framework is developed based on Tiwary et al.'s past-future Information Bottleneck (IB) theory and leverages recent advancements in machine learning to achieve improved estimation of time-lagged mutual information. The experiments showcases the robustness and effectiveness of T-IB in learning information-optimal representations for Markov process simulation.", "weaknesses": "1. Comparison with Past-Future IB Method: The paper builds upon the past-future IB method proposed by Tiwary et al., but it introduces an optimization objective for I(z_t, z_t+tau) instead of I(z_t, x_t+tau). While theoretically, optimizing I(z_t, z_t+tau) is expected to yield a lower bound on I(z_t, x_t+tau), the paper lacks a comprehensive discussion of the differences between the two approaches. A deeper exploration of the distinctions and the implications of this choice is needed to provide a clearer understanding.\n\n2. Choice of Estimation Method: In the machine learning domain, there are various methods for estimating lower bounds on mutual information. The paper selects a specific method, but it would benefit from a more thorough justification for this choice. Providing a comparative analysis of the selected method against other available options would enhance the paper's robustness.\n\n3. Complexity of Additional Terms in Formula (8): The introduction of two additional terms in Formula (8) may appear somewhat perplexing. From an information bottleneck theory perspective, the goal is to discard as much irrelevant information as possible. However, dimensionality reduction inherently leads to information loss. The rationale and impact of introducing these terms need further clarification. Moreover, the experimental results suggest that the inclusion of these terms does not significantly affect the dimensionality reduction outcome, requiring a more detailed explanation.\n\n4. Quantitative Evaluation of Dimensionality Reduction: The paper lacks a robust quantitative evaluation of the dimensionality reduction results. While the field may lack a general method for quantitatively assessing dimensionality reduction outcomes, the authors could consider exploring quantitative evaluation approaches, such as measuring mutual information under different tau conditions, to provide a more concrete assessment of the method's performance.\n\n5. Comparison with VAMPnet: VAMPnet, as mentioned in the weaknesses, primarily serves the purpose of extracting dominant or slow components and estimating kinematic properties. It may not be a direct dimensionality reduction method. Comparing VAMPnet with the proposed approach for dimensionality reduction may not be entirely fair or informative.", "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": "LLkdeOgHNE", "reviewer": "Reviewer_Yi96", "summary": "The paper aims to solve the problem of accurately and efficiently sampling large-scale systems at long time scales. The solution is based on information theory where the aim is to learn latent representations that capture the dynamical properties of the original process at at a selected time-lag by capturing relevant temporal features and discarding irrelevant high frequency information. The final objective is a trade-off between sufficiency (how informative is the learned representation of the dynamic properties of interest) and minimality of irrelevant information.", "questions": "1: The related works section currently reads like a list of related works. It will be far better to re-write this section as a discussion of how current literature relates to this proposed work and how this proposed work builds on top of this existing literature, and how it is better than its closest related works.\n\n2: Can the authors comment on \"Since $z_t$ is derived from $x_t$, the autoinformation for $x_t$ sets an upper-bound for the autoinformation for $z_t$. Given that the encoder is time-independent, how can we expect the autoinformation between $x_t$ to be related in any way to the autoinformation in $z_t$?\n\n3: The paper sets up the reader to expect a solution that is faster than current available methods for simulation. I was expecting experiments measuring how long it takes for the simulations to be run compared to existing methods. Can the authors provide that?\n\nIf above questions are addressed I'll be happy to re-visit my score.\n\n[UPDATE: authors addressed above points sufficiently. Score has been updated.]", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "To the best of my knowledge the paper is original and presents novel ideas. The quality and clarity of the paper overall is also commendable, with well explained notations and good structuring of the paper. The idea is significant since efficiently and accurately simulating large scale systems on long time scales is a critical problem in various fields. The experiments seem to largely support the claims made in the paper.", "weaknesses": "This is a thorough work on its own but a few things can be improved, These are detailed in the \"Questions\" below. Generally this pertains to the way related works section is written and also a few claims that don't seem to be backed up.", "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)