Title: Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference
Venue: NeurIPS 2024
Accepted: yes
Keywords: contrastive learning

prediction

planning

inference

time-series
Abstract: Given time series data, how can we answer questions like ``what will happen in the future?'' and ``how did we get here?'' These sorts of probabilistic inference questions are challenging when observations are high-dimensional. In this paper, we show how these questions can have compact, closed form solutions in terms of learned representations. The key idea is to apply a variant of contrastive learning to time series data. Prior work already shows that the representations learned by contrastive learning encode a probability ratio. By extending prior work to show that the marginal distribution over representations is Gaussian, we can then prove that joint distribution of representations is also Gaussian. Taken together, these results show that representations learned via temporal contrastive learning follow a Gauss-Markov chain, a graphical model where inference (e.g., prediction, planning) over representations corresponds to inverting a low-dimensional matrix. In one special case, inferring intermediate representations will be equivalent to interpolating between the learned representations. We validate our theory using numerical simulations on tasks up to 46-dimensions.

Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference

Abstract

Given time series data, how can we answer questions like 'what will happen in the future?' and 'how did we get here?' These sorts of probabilistic inference questions are challenging when observations are high-dimensional. In this paper, we show how these questions can have compact, closed form solutions in terms of learned representations. The key idea is to apply a variant of contrastive learning to time series data. Prior work already shows that the representations learned by contrastive learning encode a probability ratio. By extending prior work to show that the marginal distribution over representations is Gaussian, we can then prove that joint distribution of representations is also Gaussian. Taken together, these results show that representations learned via temporal contrastive learning follow a Gauss-Markov chain, a graphical model where inference (e.g., prediction, planning) over representations corresponds to inverting a low-dimensional matrix. In one special case, inferring intermediate representations will be equivalent to interpolating between the learned representations. We validate our theory using numerical simulations on tasks up to 46-dimensions. 1

1 Introduction

Probabilistic modeling of time-series data has applications ranging from robotic control [1] to material science [2], from cell biology [3] to astrophysics [4]. These applications are often concerned with two questions: predicting future states (e.g., what will this cell look like in an hour), and inferring trajectories between two given states. However, answering these questions often requires reasoning over high-dimensional data, which can be challenging as most tools in the standard probabilistic toolkit require generation. Might it be possible to use discriminative methods (e.g., contrastive learning) to perform such inferences?

Many prior works aim to learn representations that are easy to predict while retaining salient bits of information. For time-series data, we want the representation to remain a sufficient statistic for distributions related to time - for example, they should retain bits required to predict future states (or representations thereof). While generative methods [5-8] have this property, they tend to be computationally expensive (see, e.g., [9]) and can be challenging to scale to high-dimensional observations.

ψ

A

small dim

easy planning

large dim

hard planning

In this paper, we will study how contrastive methods (which are discriminative, rather than generative) can perform inference over times series. Ideally, we want representations of observations x to be a sufficient statistic for temporal relationships (e.g., does x ′ occur after x ?) but need not retain other information about x (e.g. the location of static objects). This intuition motivates us to study how contrastive representation learning methods [10-14] might be used to solve prediction and planning problems on time series data. While prior works in computer vision [10, 15] and natural language processing (NLP) [16] often study the geometry of learned representations, our results show how geometric operations such as interpolation are related to inference. Our analysis will focus on a regularized version of the symmetrized infoNCE objective [17], generating positive examples by sampling pairs of observations from the same time series data. We will study how representations learned in this way can facilitate two inference questions: prediction and planning. 2 As a stepping stone, we will build upon prior work [20] to show that regularized contrastive learning should produce representations whose marginal distribution is an isotropic Gaussian distribution.

The main contribution of this paper is to demonstrate how intermediate and future time steps in a time series can be inferred easily using contrastive representations. This inference problem captures a number of practical tasks: interpolation, in-filling, and even planning and control, where the intermediate steps represent states between a stand and goal. While ordinarily these problems require an iterative inference or optimization procedure, with contrastive representations this can be done simply by inverting a low-dimensional matrix. In one special case, inference will correspond to linear interpolation. Our first step is to prove that, under certain assumptions, the distribution over future representations has a Gaussian distribution, with a mean that is a linear function of the initial state representation (Lemma 1). This paves the main to our main result (Theorem 2): given an initial and final state, we show that the posterior distribution over an intermediate state representations also follows a Gaussian distribution. Said in other words, the representations follow a Gauss-Markov chain, 3 wherein any joint or conditional distribution can be computed by inverting a low-dimensional matrix [23, 24] (See Fig. 1). In one special case, inference will correspond to linearly interpolating between the representations of an initial state and final state. Section 5 provides numerical experiments.

2 Related Work

Representations for time-series data. In applications ranging from robotics to vision to NLP, users often want to learn representations of observations from time series data such that the spatial arrangement of representations reflects the temporal arrangement of the observations [10, 16, 25, 26].

Ideally, these representations should retain information required to predict future observations and infer likely paths between pairs of observations. Many approaches use an autoencoder, learning representations that retain the bits necessary to reconstruct the input observation, while also regularizing the representations to compressed or predictable [6, 27-31]. A prototypical method is the sequential VAE [5], which is computationally expensive to train because of the reconstruction loss, but is easy to use for inference. Our work shares the aims of prior prior methods that attempt to linearize the dynamics of nonlinear systems [32-37], including videos [38, 39]. Our work aims to retain uncertainty estimates over predictions (like the sequential V AE) without requiring reconstruction. Avoiding reconstruction is appealing practically because it decreases the computational requirements and number of hyperparameters; and theoretically because it means that representations only need to retain bits about temporal relationships and not about the bits required to reconstruct the original observation.

Contrastive Learning. Contrastive learning methods circumvent reconstruction by learning representations that merely classify if two events were sampled from the same joint distribution [17, 40, 41]. When applied to representing states along trajectories, contrastive representations learn to classify whether two points lie on the same trajectory or not [10, 25, 42-44]. Empirically, prior work in computer vision and NLP has observed that contrastive learning acquires representations where interpolation between representations corresponds to changing the images in semantically meaningful ways [16, 45-49].

Our analysis will be structurally similar to prior theoretical analysis on explaining why word embeddings can solve analogies [50-52]. Our work will make a Gaussianity assumption similar to Arora et al. [51] and our Markov assumption is similar to the random walks analyzed in Arora et al. [51], Hashimoto et al. [53]. Our paper builds upon and extends these results to answer questions such as: 'what is the distribution over future observations representations?' and 'what is the distribution over state (representations) that would occur on the path between one observation and another?' While prior work is primarily aimed at explaining the good performance of contrastive word embeddings (see, e.g., [51]), we are primarily interested in showing how similar contrastive methods are an effective tool for inference over high-dimensional time series data. Our analysis will show how representations learned via temporal contrastive learning (i.e., without reconstruction) are sufficient statistics for inferring future outcomes and can be used for performing inference on a graphical model (a problem typically associated with generative methods).

Goal-oriented decision making. Much work on time series representations is done in service of learning goal-reaching behavior, an old problem [54, 55] that has received renewed attention in recent years [56-63]. Some of the excitement in goal-conditioned RL is a reflection of the recent success of self-supervised methods in computer vision [64] and NLP [65]. Our analysis will study a variant of contrastive representation learning proposed in prior work for goal-conditioned RL [42, 43]. These methods are widespread, appearing as learning objectives for learning value functions [26, 43, 66-73], as auxiliary objectives [71, 74-79], in objectives for model-based RL [32, 80-82], and in exploration methods [83, 84]. Our analysis will highlight connections between these prior methods, the classic successor representation [85, 86], and probabilistic inference.

Planning. Planning lies at the core of many RL and control methods, allowing methods to infer the sequence of states and actions that would occur if the agent navigated from one state to a goal state. While common methods such as PRM [87] and RRT [88] focus on building random graphs, there is a strong community focusing on planning methods based on probabilistic inference [19, 89, 90]. The key challenge is scaling to high-dimensional settings. While semi-parametric methods make progress on this problem this limitation through semi-parametric planning [91-93], it remains unclear how to scale any of these methods to high-dimensional settings when states do not lie on a low-dimensional manifold. Our analysis will show how contrastive representations may lift this limitation, with experiments validating this theory on 39-dimensional and 46-dimensional tasks.

3 Preliminaries

Our aim is to learn representations of time series data such that the spatial arrangement of representations corresponds to the temporal arrangement of the underlying data: if one example occurs shortly after another, then they should be mapped to similar representations. This problem setting arises in many areas, including video understanding and reinforcement learning. To de-

fine this problem formally, we will define a Markov process with states x t indexed by time t : 4 p ( x 1: T | x 0 ) = ∏ T t =0 p ( x t +1 | x t ) . The dynamics p ( x t +1 | x t ) tell us the immediate next state, and we can define the distribution over states t steps in the future by marginalizing over the intermediate states, p t ( x t | x 0 ) = ∫ p ( x 1: t | x 0 ) d x 1: t -1 . A key quantity of interest will be the γ -discounted state occupancy measure, which corresponds to a time-averaged distribution over future states:

Contrastive learning. Our analysis will focus on applying contrastive learning to a particular data distribution. Contrastive learning [10, 40, 94] acquires representations using 'positive' pairs ( x, x + ) and 'negative' pairs ( x, x -) . While contrastive learning typically learns just one representation, we will use two different representation for the two elements of the pair; that is, our analysis will use terms like ϕ ( x ) , ψ ( x + ) and ψ ( x -) . We assume all representations lie in R k .

̸

The aim of contrastive learning is to learn representations such that positive pairs have similar representations ( ϕ ( x ) ≈ ψ ( x + ) ) while negative pairs have dissimilar representations ( ϕ ( x ) = ψ ( x -) ). Let p ( x, x + ) be the joint distribution over positive pairs (i.e., ( x, x + ) ∼ p ( x, x + ) ). We will use the product of the marginal distributions to sample negative pairs ( ( x, x -) ∼ p ( x ) p ( x ) ). Let B be the batch size, and note that the positive samples x + j at index j in the batch serve as negatives for x i for any i = j . Our analysis is based on the infoNCE objective without resubstitution [10, 11]:

̸

̸

̸

We will use the symmetrized version of this objective [17], where the denominator is the sum across rows of a logits matrix and once where it is a sum across the.

While contrastive learning is typically applied to an example x and an augmentation x + ∼ p ( x | x ) of that same example (e.g., a random crop), we will follow prior work [10, 42] in using the time series dynamics to generate the positive pairs, so x + will be an observation that occurs temporally after x . While our experiments will sample positive examples from the discounted state occupancy measure ( x + ∼ p t + ( x t + | x ) ) in line with prior work [43], our analysis will also apply to different distributions (e.g., always sampling a state k steps ahead).

While prior work typically constrains the representations to have a constant norm (i.e., to lie on the unit hypersphere) [10], we will instead constrain the expected norm of the representations is bounded, a difference that will be important for our analysis:

Because the norm scales with the dimension of the representation, we have scaled down the left side by the representation dimension, k . In practice, we will impose this constraint by adding a regularization term λ E p ( x ) [ ∥ ψ ( x ) ∥ 2 2 ] to the infoNCE objective (Eq. 2) and dynamically tuning the weight λ via dual gradient descent.

3.1 Key assumptions

This section outlines the two key assumptions behind our analysis, both of which have some theoretical justification. Our main assumption examines the distribution over representations:

Assumption 1. Regularized, temporal contrastive learning acquires representations whose marginal distribution representations p ( ψ ) ≜ ∫ p ( x ) 1 ( ψ ( x ) = ψ ) d x is an isotropic Gaussian distribution:

In Appendix A.1 we extend prior work [20] provide some theoretical intuition for why this assumption should hold: namely, that the isotropic Gaussian is the distribution that maximizes entropy subject to an expected L2 norm constraint (Eq. 3) [95-97]. Our analysis also assumes that the learned representations converge to the theoretical minimizer of the infoNCE objective:

Assumption 2. Applying contrastive learning to the symmetrized infoNCE objective results in representations that encode a probability ratio:

This assumption holds under ideal conditions [98, 99] (see Appendix A.5), 5 but we nonetheless call this an 'assumption' because it may not hold in practice due to sampling and function approximation error. This assumption means the learned representations are sufficient statistics for predicting the probability (ratio) of future states: these representations must retain all the information pertinent to reasoning about temporal relationships, but need not retain information about the precise contents of the observations. As such, they may be much more compressed than representations learned via reconstruction.

Combined, these assumptions will allow us to express the distribution over sequences of representations as a Gauss-Markov chain. The denominator in Assumption 2, p ( x ) , may have a complex distribution, but Assumption 1 tells us that the distribution over representations has a simpler form. This will allow us to rearrange Assumption 2 to express the conditional distribution over representations as the product of two Gaussian likelihoods. Note that the left hand side of Assumption 2 already looks like a Gaussian likelihood.

4 Contrastive Representations Make Inference Easy

In this section, our main result will be to show how representations learned by (regularized) contrastive learning are distributed according to a Gauss-Markov chain, making it straightforward to perform inference (e.g., planning, prediction) over these representations. Our proof technique will combine (known) results about Gaussian distributions with (known) results about contrastive learning. We start by discussing an important choice of parametrization (Section 4.1) that facilitates prediction (Section 4.2) before presenting the main result in Section 4.3.

4.1 A Parametrization for Shared Encoders

This section describes the two encoders ( ψ ( · ) , ϕ ( · ) ) to compute representations of x and x + . While prior work in computer vision and NLP literature use the same encoder for both x and x + , this decision does not make sense for many time-series data as it would imply that our prediction for p ( x t | x 0 ) is the same as our prediction for p ( x 0 | x t ) . However, the difficulty of transiting from x 0 to x t (e.g., climbing to the peak of a mountain) might be more difficult than the reverse (e.g., sledding down a mountain). Our proposed parametrization will handle this asymmetry.

We will treat the encoder ψ ( · ) as encoding the contents of the state. We will additionally learn a matrix A so that the function ψ ↦→ Aψ corresponds to a (multi-step) prediction of the future representation.

ψ

(

x

0

)

φ

(

x

0

) =

Aψ

(

x

0

)

A

ψ

(

x

t

)

To map this onto contrastive learning, we will use ϕ ( x ) ≜ Aψ ( x ) as the encoder for the initial state. One way of interpreting this encoder is as an additional linear projection applied on top of ψ ( · ) , a design similar to those used in other areas of contrastive learning [41]. Once learned, we can use these encoders to answer questions about prediction (Section 4.2) and planning (Section 4.3).

4.2 Representations Encode a Predictive Model

Given an initial state x 0 , what states are likely to occur in the future? Answering this question directly in terms of high-dimensional states is challenging, but our learned representations provide a straightforward answer. Let ψ 0 = ψ ( x 0 ) and ψ t + = ψ ( x t + ) be random variables representing the representations of the initial state and a future state. Our aim is to estimate the distribution over these future representations, p ( ψ t + | ψ 0 ) . We will show that the learned representations encode this distribution.

Lemma 1. Under the assumptions from Section 3, the distribution over representations of future states follows a Gaussian distribution with mean parameter given by the initial state representation:

The main takeaway here is that the distribution over future representations has a convenient, closed form solution. The representation norm constraint, c , determines the shrinkage factor c c +1 ∈ [0 , 1) ; highly regularized settings (small c ) move the mean closer towards the origin and decrease the variance, as visualized in Fig. 3. Regardless of the constraint c , the predicted mean is a linear function ψ ↦→ c c +1 Aψ . The proof is in Appendix A.2. The proof technique is similar to that of the law of the unconscious statistician.

4.3 Planning over One Intermediate State

We now show how these representations can be used for a specific type of planning: given an initial state x 0 and a future state x t + , infer the representation of an intermediate 'waypoint' state x w . The next section will extend this analysis to inferring the entire sequence

ψ

(

x

0

)

Aψ

(

x

0

)

c

c

+1

Aψ

(

x

0

)

of intermediate states. We assume x 0 → x w → x t + form a Markov chain where x w ∼ p ( x t + | x 0 = x 0 ) and x t + ∼ p ( x t + | x 0 = x w ) are both drawn from the discounted state occupancy measure (Eq. 1). Let random variable ψ w = ψ ( x w ) be the representation of this intermediate state. Our main result is that the posterior distribution over waypoint representations has a closed form solution in terms of the initial state representation and future state representation:

Theorem 2. Under Assumptions 1 and 2, the posterior distribution over waypoint representations is a Gaussian whose mean and covariance are linear functions of the initial and final state representations:

The proof (Appendix A.3) uses the Markov property together with Lemma 1. The main takeaway from this lemma is that the posterior distribution takes the form of a simple probability distribution (a Gaussian) with parameters that are linear functions of the initial and final representations.

We give three examples to build intuition:

Example 1: A = I and the c is very large (little regularization). Then, the covariance is Σ -1 ≈ 2 I and the mean is the simple average of the initial and final representations µ ≈ 1 2 ( ψ 0 + ψ t + ) . In other words, the waypoint representation is the midpoint of the line ψ 0 → ψ t + .

Example 2: A is a rotation matrix and c is very large. Rotation matrices satisfy A T = A -1 so the covariance is again Σ -1 ≈ 2 I . As noted in Section 4.2, we can interpret Aψ 0 as a prediction of which representations will occur after ψ 0 . Similarly, A -1 ψ t + = A T ψ t + is a prediction of which representations will occur before ψ t + . Theorem 2 tells us that the mean of the waypoint distribution is the simple average of these two predictions, µ ≈ 1 2 ( A T ψ t + + Aψ 0 ) .

Example 3: A is a rotation matrix and c = 0 . 01 (very strong regularization). In this case Σ -1 = 0 . 01 0 . 01+1 A T A + 0 . 01+1 0 . 01 I ≈ 100 I , so µ ≈ 1 100 ( ψ 0 + ψ t + ) ≈ 0 . Thus, in the case of strong regularization, the posterior concentrates around the origin.

4.4 Planning over Many Intermediate States

This section extends the analysis to multiple intermediate states. Again, we will infer the posterior distribution of the representations of these intermediate states, ψ w 1 , ψ w 2 , · · · . We assume that these states form a Markov chain.

Theorem 3. Given observations from a Markov chain x 0 → x 1 · · · x t + , the joint distribution over representations is a Gaussian distribution. Using ψ 1: n = ( ψ w 1 , · · · , ψ w n ) to denote the concatenated representations of each observation, we can write this distribution as

where Σ -1 is a tridiagonal matrix

This distribution can be written in the canonical parametrization as Σ = Λ -1 and µ = Σ η . Recall that Gaussian distributions are closed under marginalization. Thus, once in this canonical parametrization, the marginal distributions can be obtained by reading off individual entries of these parameters:

The key takeaway here is that this posterior distribution over waypoints is Gaussian, and it has a closed form expression in terms of the initial and final representations (as well as regularization parameter c and the learned matrix A ).

In the general case of n intermediate states, the posterior distribution is

where ψ w 0 = ψ 0 and ψ w n +1 = ψ t + . This corresponds to a chain graphical model with edge potentials f ( ψ, ψ ′ ) = e -1+ 1 c 2 ∥ c c +1 Aψ -ψ ′ ∥ 2 2 .

Special case. To build intuition, consider the special case where A is a rotation matrix and c is very large, so c c +1 A T A + c +1 c ≈ 2 I . In this case, Σ -1 is a (block) second difference matrix [100]:

The inverse of this matrix has a closed form solution [101, Pg. 471], allowing us to obtain the mean of each waypoint in closed form:

where λ ( i ) = i n +1 . Thus, each posterior mean is a convex combination of the (forward prediction from the) initial representation and the (backwards prediction from the) final representation. When A is the identity matrix, the posterior mean is simple linear interpolation between the initial and final representations!

5 Numerical Simulation

We include several didactic experiments to illustrate our results. All results and figures can be reproduced by running make in the source code: https://github.com/vivekmyers/contrastive_ planning . The expected compute time is a few hours on a A6000 GPU. Figures in this section show error across different training and dataset split seeds.

5.1 Synthetic Dataset

To validate our analysis, we design a time series task with 2D points where inference over intermediate points (i.e., in-filling) requires nonlinear interpolation. Fig. 4 (Top Left) shows the dataset of time series data, starting at the origin and spiraling outwards, with each trajectory using a randomly-chosen initial angle. We applied contrastive learning with the parametrization in Section 4.1 to these data and used the learned representations to solve prediction and planning problems (see Fig. 4 for details). Note that these predictions correctly handle the nonlinear structure of these data - states nearby the initial state in Euclidean space that are not temporally adjacent are assigned low likelihood.

5.2 Solving Mazes with Inferred Representations

Our next experiment studies whether the inferred representations are useful for solving a control task. We took a 2d maze environment and dataset from prior work (Fig. 5, Left ) [102] and learned encoders

DATASET

start (xo)

end (xT)

BACKWARD PREDICTION

PLANNING

start (xo)

end (xT)

2

4

6

8

10

0

0

.

2

0

.

4

0

.

6

0

.

8

1

Initial

/lscript

2

Distance to Goal

Success Rate

Maze Planning Success Rates

contrastive (ours)

VIP

PCA

autoencoder

no planning

start

★

from this dataset. To solve the maze, we take the observation of the starting state and goal state, compute the representations of these states, and use the analysis in Section 4.3 to infer the sequence of intermediate representations. We visualize the results using a nearest neighbor retrieval (Fig. 5, Left ). Figure 7 contains additional examples.

Finally, we studied whether these representations are useful for control. We implemented a simple proportional controller for this maze. As expected, this proportional controller can successfully navigate to close goals, but fails to reach distant goals (Fig. 5, Right ). However, if we use the proportional controller to track a series of waypoints planned using our representations (i.e., the orange dots shown in Fig. 5 (Left) ), the success rate increases by up to 4 . 5 × . To test the importance of nonlinear representations, we compare with a 'PCA' baseline that predicts waypoints by interpolating between the principal components of the initial state and goal state. The better performance of our method indicates the importance of doing the interpolation using representations that are nonlinear functions of the input observations. While prior methods learn representations to encode temporal distances, it

MSE of waypoint prediction for door opening

0

2

4

6

8

Waypoint MSE

contrastive planning

(

ours

)

VIP planning

PCA planning

AE planning

no planning

x

0

xT

contrastive

planning

(ours)

x

0

xT

PCA planning

x

0

xT

VIP planning

is unclear whether these methods support inference via interpolation. To test this hypothesis, we use one of these methods ('VIP' [69]) as a baseline. While the VIP representations likely encode similar bits as our representations, the better performance of the contrastive representations indicates that the VIP representations do not expose those bits in a way that makes planning easy.

5.3 Higher dimensional tasks

In this section we provide preliminary experiments showing the planning approach in Section 4 scales to higher dimensional tasks. We used two datasets from prior work [102]: door-human-v0 (39-dimensional observations) and hammer-human-v0 (46-dimensional observations). After learning encoders on these tasks, we evaluated the inference capabilities of the learned representations. Given the first and last observation from a trajectory in a validation set, we use linear interpolation (see Eq. 7) to infer the representation of five intermediate waypoint representations.

We evaluate performance in two ways. Quantitatively , we measure the mean squared error between each of the true waypoint observations and those inferred by our method. Since our method infers representations, rather than observations, we use a nearest-neighbor retrieval on a validation set so that we can measure errors in the space of observations. Qualitatively , we visualize the high-dimensional observations from the validation trajectory using a 2-dimensional TSNE [103] embedding, overlying the infer waypoints from our method; as before, we convert the representations inferred by our method to observations using nearest neighbors.

We compare with three alternative methods in Fig. 6. To test the importance of representation learning, we first naïvely interpolate between the initial and final observations ('no planning'). The poor performance of this baseline indicates that the input time series are highly nonlinear. Similarly, interpolating the principle components of the initial and final observations ('PCA') performs poorly, again highlighting that the input time series is highly nonlinear and that our representations are doing more than denoising (i.e., discarding directions of small variation). The third baseline, 'VIP' [69],

learns representations to encode temporal distances using approximate dynamic programming. Like our method, VIP avoids reconstruction and learns nonlinear representations of the observations. However, the results in Fig. 6 highlight that VIP's representations do not allow users to plan by interpolation. The error bars shown in Fig. 6 ( Top Right ) show the standard deviation over 500 trajectories sampled from the validation set. For reproducibility, we repeated this entire experiment on another task, the 46-dimensional hammer-human-v0 from D4RL. The results, shown in Appendix Fig. 8, support the conclusions above. Taken together, these results show that our procedure for interpolating contrastive representations continues to be effective on tasks where observations have dozens of dimensions.

6 Discussion

Representation learning is at the core of many high-dimensional time-series modeling questions, yet how those representations are learned is often disconnected with the inferential task. The main contribution of this paper is to show how discriminative techniques can be used to acquire compact representations that make it easy to answer inferential questions about time. The precise objective and parametrization we studied is not much different from that used in practice, suggesting that either our theoretical results might be adapted to the existing methods, or that practitioners might adopt these details so they can use the closed-form solutions to inference questions. Our work may also have implications for studying the structure of learned representations. While prior work often studies the geometry of representations as a post-hoc check, our analysis provides tools for studying when interpolation properties are guaranteed to emerge, as well as how to learn representations with certain desired geometric properties.

Limitations. Our analysis hinges on the two assumptions mentioned in Section 3.1, and it remains open how errors in those approximations translate into errors in our analysis. One important open question is whether it is always possible to satisfy these assumptions using sufficiently-expressive representations.

Acknowledgments

We thank Seohong Park, Gautam Reddy, Chongyi Zheng, and anonymous reviewers for feedback and discussions that shaped this project. This work was supported by Princeton Research Computing resources at Princeton University. This work was partially supported by AFOSR FA9550-22-1-0273.

References

Evangelos Theodorou, Jonas Buchli, and Stefan Schaal. A Generalized Path Integral Control Approach to Reinforcement Learning. Journal of Machine Learning Research , 11:3137-3181, 2010.

Hannes Jónsson, Greg Mills, and Karsten W Jacobsen. Nudged Elastic Band Method for Finding Minimum Energy Paths of Transitions. In Classical and Quantum Dynamics in Condensed Phase Simulations , pp. 385-404. World Scientific, 1998.

Wouter Saelens, Robrecht Cannoodt, Helena Todorov, and Yvan Saeys. A Comparison of Single-Cell Trajectory Inference Methods. Nature Biotechnology , 37(5):547-554, 2019.

Steven R Majewski, Ricardo P Schiavon, Peter M Frinchaboy, Carlos Allende Prieto, Robert Barkhouser, Dmitry Bizyaev, Basil Blank, Sophia Brunner, Adam Burton, Ricardo Carrera, et al. The Apache Point Observatory Galactic Evolution Experiment (APOGEE). Astronomical Journal , 154(3):94, 2017.

Shengjia Zhao, Jiaming Song, and Stefano Ermon. Towards Deeper Understanding of Variational Autoencoding Models. arXiv:1702.08658, 2017.

Yizhe Zhu, Martin Renqiang Min, Asim Kadav, and Hans Peter Graf. S3vae: Self-Supervised Sequential VAE for Representation Disentanglement and Data Generation. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 6538-6547. 2020.

Alireza Makhzani, Jonathon Shlens, Navdeep Jaitly, Ian Goodfellow, and Brendan Frey. Adversarial Autoencoders. arXiv:1511.05644, 2015.

Vincent Dumoulin, Ishmael Belghazi, Ben Poole, Alex Lamb, Martin Arjovsky, Olivier Mastropietro, and Aaron Courville. Adversarially Learned Inference. In International Conference on Learning Representations . 2016.

Ali Razavi, Aaron Van den Oord, and Oriol Vinyals. Generating Diverse High-Fidelity Images With Vq-VAE-2. Neural Information Processing Systems , 32, 2019.

Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation Learning With Contrastive Predictive Coding. arXiv:1807.03748, 2018.

Kihyuk Sohn. Improved Deep Metric Learning With Multi-Class N-Pair Loss Objective. Neural Information Processing Systems , 29, 2016.

Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved Baselines With Momentum Contrastive Learning. arXiv:2003.04297, 2020.

Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive Multiview Coding. In Computer Vision-ECCV 2020: 16th European Conference , pp. 776-794. 2020.

Zhirong Wu, Yuanjun Xiong, Stella X Yu, and Dahua Lin. Unsupervised Feature Learning via Non-Parametric Instance Discrimination. In IEEE Conference on Computer Vision and Pattern Recognition , pp. 3733-3742. 2018.

Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A Simple Framework for Contrastive Learning of Visual Representations. In International Conference on Machine Learning , pp. 1597-1607. 2020.

Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient Estimation of Word Representations in Vector Space. arXiv:1301.3781, 2013.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning Transferable Visual Models From Natural Language Supervision. In International Conference on Machine Learning , pp. 8748-8763. 2021.

Matthew Botvinick and Marc Toussaint. Planning as Inference. Trends in Cognitive Sciences , 16(10):485-488, 2012.

Hagai Attias. Planning by Probabilistic Inference. In International Workshop on Artificial Intelligence and Statistics , pp. 9-16. 2003.

Tongzhou Wang and Phillip Isola. Understanding Contrastive Representation Learning Through Alignment and Uniformity on the Hypersphere. In International Conference on Machine Learning , pp. 9929-9939. 2020.

George E Uhlenbeck and Leonard S Ornstein. On the Theory of the Brownian Motion. Physical Review , 36(5):823, 1930.

George EP Box, Gwilym M Jenkins, Gregory C Reinsel, and Greta M Ljung. Time Series Analysis: Forecasting and Control . John Wiley & Sons, 2015.

Dmitry M Malioutov, Jason K Johnson, and Alan S Willsky. Walk-Sums and Belief Propagation in Gaussian Graphical Models. Journal of Machine Learning Research , 7:2031-2064, 2006.

Yair Weiss and William Freeman. Correctness of Belief Propagation in Gaussian Graphical Models of Arbitrary Topology. Neural Information Processing Systems , 12, 1999.

Rui Qian, Tianjian Meng, Boqing Gong, Ming-Hsuan Yang, Huisheng Wang, Serge Belongie, and Yin Cui. Spatiotemporal Contrastive Video Representation Learning. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 6964-6974. 2021.

Benjamin Eysenbach, Ruslan Salakhutdinov, and Sergey Levine. C-Learning: Learning to Achieve Goals via Recursive Classification. In International Conference on Learning Representations . 2020.

Siddharth Karamcheti, Suraj Nair, Annie S. Chen, Thomas Kollar, Chelsea Finn, Dorsa Sadigh, and Percy Liang. Language-Driven Representation Learning for Robotics. arXiv:2302.12766, 2023.

Seongmin Park and Jihwa Lee. Finetuning Pretrained Transformers Into Variational Autoencoders. arXiv:2108.02446, 2021.

Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova. BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding. In NAACL-Hlt , pp. 4171-4186.

2019.

Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, et al. UniMASK: Unified Inference in Sequential Decision Problems. arXiv:2211.10869, 2022.

Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio. A Recurrent Latent Variable Model for Sequential Data. In Neural Information Processing Systems , volume 28. 2015.

Rui Shu, Tung Nguyen, Yinlam Chow, Tuan Pham, Khoat Than, Mohammad Ghavamzadeh, Stefano Ermon, and Hung Bui. Predictive Coding for Locally-Linear Control. In International Conference on Machine Learning , pp. 8862-8871. 2020.

Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to Control: A Locally Linear Latent Dynamics Model for Control From Raw Images. Neural Information Processing Systems , 28, 2015.

Ershad Banijamali, Rui Shu, Hung Bui, Ali Ghodsi, et al. Robust Locally-Linear Controllable Embedding. In International Conference on Artificial Intelligence and Statistics , pp. 17511759. 2018.

Brandon Cui, Yinlam Chow, and Mohammad Ghavamzadeh. Control-Aware Representations for Model-Based Reinforcement Learning. In International Conference on Learning Representations . 2020.

Tung D Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Temporal Predictive Coding for Model-Based Planning in Latent Space. In International Conference on Machine Learning , pp. 8130-8139. 2021.

Tung Nguyen, Rui Shu, Tuan Pham, Hung Bui, and Stefano Ermon. Non-Markovian Predictive Coding for Planning in Latent Space. 2020.

Ross Goroshin, Michael F Mathieu, and Yann LeCun. Learning to Linearize Under Uncertainty. Neural Information Processing Systems , 28, 2015.

Dinesh Jayaraman and Kristen Grauman. Slow and Steady Feature Analysis: Higher Order Temporal Coherence in Video. In IEEE Conference on Computer Vision and Pattern Recognition , pp. 3852-3861. 2016.

Michael Gutmann and Aapo Hyvärinen. Noise-Contrastive Estimation: A New Estimation Principle for Unnormalized Statistical Models. In International Conference on Artificial Intelligence and Statistics , pp. 297-304. 2010.

Xinlei Chen and Kaiming He. Exploring Simple Siamese Representation Learning. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 15750-15758. 2021.

Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, Sergey Levine, and Google Brain. Time-Contrastive Networks: Self-Supervised Learning From Video. In IEEE International Conference on Robotics and Automation , pp. 1134-1141. 2018.

Benjamin Eysenbach, Tianjun Zhang, Sergey Levine, and Russ R Salakhutdinov. Contrastive Learning as Goal-Conditioned Reinforcement Learning. Neural Information Processing Systems , 35:35603-35620, 2022.

Mengda Xu, Zhenjia Xu, Cheng Chi, Manuela Veloso, and Shuran Song. Xskill: Cross Embodiment Skill Discovery. In Conference on Robot Learning , pp. 3536-3555. 2023.

Laurenz Wiskott and Terrence J. Sejnowski. Slow Feature Analysis: Unsupervised Learning of Invariances. Neural Computation , 14(4):715-770, 2002.

Jia-Wei Yan, Ci-Siang Lin, Fu-En Yang, Yu-Jhe Li, and Yu-Chiang Frank Wang. SemanticsGuided Representation Learning With Applications to Visual Synthesis. In International Conference on Pattern Recognition , pp. 7181-7187. 2021.

Alon Oring, Zohar Yakhini, and Yacov Hel-Or. Autoencoder Image Interpolation by Shaping the Latent Space. In International Conference on Machine Learning , pp. 8281-8290. 2021.

Ying-Cong Chen, Xiaogang Xu, Zhuotao Tian, and Jiaya Jia. Homomorphic Latent Space Interpolation for Unpaired Image-to-Image Translation. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 2403-2411. 2019.

Xiaofeng Liu, Yang Zou, Lingsheng Kong, Zhihui Diao, Junliang Yan, Jun Wang, Site Li, Ping Jia, and Jane You. Data Augmentation via Latent Space Interpolation for Image Classification. In International Conference on Pattern Recognition , pp. 728-733. 2018.

Omer Levy and Yoav Goldberg. Linguistic Regularities in Sparse and Explicit Word Representations. In Computational Natural Language Learning , pp. 171-180. 2014.

Sanjeev Arora, Yuanzhi Li, Yingyu Liang, Tengyu Ma, and Andrej Risteski. A Latent Variable Model Approach to Pmi-Based Word Embeddings. Transactions of the Association for Computational Linguistics , 4:385-399, 2016.

Carl Allen and Timothy Hospedales. Analogies Explained: Towards Understanding Word Embeddings. In International Conference on Machine Learning , pp. 223-231. 2019.

Tatsunori B Hashimoto, David Alvarez-Melis, and Tommi S Jaakkola. Word Embeddings as Metric Recovery in Semantic Spaces. Transactions of the Association for Computational Linguistics , 4:273-286, 2016.

Allen Newell, John C Shaw, and Herbert A Simon. Report on a General Problem Solving Program. In IFIP Congress , volume 256, p. 64. 1959.

John E Laird, Allen Newell, and Paul S Rosenbloom. Soar: An Architecture for General Intelligence. Artificial Intelligence , 33(1):1-64, 1987.

Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Reinforcement Learning via Sequence Modeling. Neural Information Processing Systems , 34:15084-15097, 2021.

Elliot Chane-Sane, Cordelia Schmid, and Ivan Laptev. Goal-Conditioned Reinforcement Learning With Imagined Subgoals. In International Conference on Machine Learning , pp. 1430-1440. 2021.

Cédric Colas, Tristan Karch, Olivier Sigaud, and Pierre-Yves Oudeyer. Intrinsically Motivated Goal-Conditioned Reinforcement Learning: A Short Survey. Preprint , 2021.

Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking Goal-Conditioned Supervised Learning and Its Connection to Offline RL. In International Conference on Learning Representations . 2021.

Yecheng Jason Ma, Jason Yan, Dinesh Jayaraman, and Osbert Bastani. How Far I'll Go: Offline Goal-Conditioned Reinforcement Learning via f -Advantage Regression. arXiv:2206.03023, 2022.

Yannick Schroecker and Charles Isbell. Universal Value Density Estimation for Imitation Learning and Goal-Conditioned Reinforcement Learning. arXiv:2002.06473, 2020.

Michael Janner, Qiyang Li, and Sergey Levine. Offline Reinforcement Learning as One Big Sequence Modeling Problem. Neural Information Processing Systems , 34:1273-1286, 2021.

Joey Hejna, Jensen Gao, and Dorsa Sadigh. Distance Weighted Supervised Learning for Offline Interaction Data. arXiv:2304.13774, 2023.

Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-Resolution Image Synthesis With Latent Diffusion Models. In IEEE/CVF Conference on Computer Vision and Pattern Recognition , pp. 10684-10695. 2022.

OpenAI. GPT-4 Technical Report. arXiv:2303.08774, 2023.

Chongyi Zheng, Benjamin Eysenbach, Homer Rich Walke, Patrick Yin, Kuan Fang, Ruslan Salakhutdinov, and Sergey Levine. Stabilizing Contrastive RL: Techniques for Robotic Goal Reaching From Offline Data. In International Conference on Learning Representations . 2024.

Stephen Tian, Suraj Nair, Frederik Ebert, Sudeep Dasari, Benjamin Eysenbach, Chelsea Finn, and Sergey Levine. Model-Based Visual Planning With Self-Supervised Functional Distances. In International Conference on Learning Representations . 2020.

Alekh Agarwal, Nan Jiang, Sham M Kakade, and Wen Sun. Reinforcement Learning: Theory and Algorithms. CS Dept , pp. 10-4, 2019.

Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang. VIP: Towards Universal Visual Reward and Representation via Value-Implicit Pre-Training. In International Conference on Learning Representations . 2022.

Yecheng Jason Ma, Vikash Kumar, Amy Zhang, Osbert Bastani, and Dinesh Jayaraman. LIV: Language-Image Representations and Rewards for Robotic Control. International Conference on Machine Learning , 2023.

Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3M: A Universal Visual Representation for Robot Manipulation. In Conference on Robot Learning , pp. 892-909. 2023.

Bo Liu, Yihao Feng, Qiang Liu, and Peter Stone. Metric Residual Network for Sample Efficient Goal-Conditioned Reinforcement Learning. In AAAI Conference on Artificial Intelligence , volume 37, pp. 8799-8806. 2023.

Tongzhou Wang, Antonio Torralba, Phillip Isola, and Amy Zhang. Optimal Goal-Reaching Reinforcement Learning via Quasimetric Learning. In International Conference on Machine Learning . 2023.

Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-Efficient Reinforcement Learning With Self-Predictive Representations. In International Conference on Learning Representations . 2020.

Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Avila Pires, Yash Chandak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, et al. Understanding Self-Predictive Learning for Reinforcement Learning. In International Conference on Machine Learning , pp. 33632-33656. 2023.

Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling Representation Learning From Reinforcement Learning. In International Conference on Machine Learning , pp. 9870-9879. 2021.

Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Information Prioritization Through Empowerment in Visual Model-Based RL. In International Conference on Learning Representations . 2021.

Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised State Representation Learning in Atari. Neural Information Processing Systems , 32, 2019.

Pablo Samuel Castro, Tyler Kastner, P. Panangaden, and Mark Rowland. MICo: Improved Representations via Sampling-Based State Similarity for Markov Decision Processes. In Neural Information Processing Systems . 2021.

Raj Ghugare, Homanga Bharadhwaj, Benjamin Eysenbach, Sergey Levine, and Russ Salakhutdinov. Simplifying Model-Based RL: Learning Representations, Latent-Space Models, and Policies With One Objective. In International Conference on Learning Representations . 2022.

Cameron S. Allen. Learning Markov State Abstractions for Deep Reinforcement Learning. In Neural Information Processing Systems . 2021.

Bogdan Mazoure, Benjamin Eysenbach, Ofir Nachum, and Jonathan Tompson. Contrastive Value Learning: Implicit Models for Simple Offline RL. In Conference on Robot Learning , pp. 1257-1267. 2023.

Zhaohan Guo, Shantanu Thakoor, Miruna Pîslar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al. Byol-Explore: Exploration by Bootstrapped Prediction. Neural Information Processing Systems , 35:3185531870, 2022.

Yilun Du, Chuang Gan, and Phillip Isola. Curious Representation Learning for Embodied Intelligence. IEEE/CVF International Conference on Computer Vision , pp. 10388-10397, 2021.

Peter Dayan. Improving Generalization for Temporal Difference Learning: The Successor Representation. Neural Computation , 5:613-624, 1993.

André Barreto, Will Dabney, Rémi Munos, Jonathan J Hunt, Tom Schaul, Hado P van Hasselt, and David Silver. Successor Features for Transfer in Reinforcement Learning. Neural Information Processing Systems , 30, 2017.

Lydia E Kavraki, Petr Svestka, J-C Latombe, and Mark H Overmars. Probabilistic Roadmaps for Path Planning in High-Dimensional Configuration Spaces. IEEE Transactions on Robotics and Automation , 12(4):566-580, 1996.

Steven M LaValle, James J Kuffner, BR Donald, et al. Rapidly-Exploring Random Trees: Progress and Prospects. Algorithmic and Computational Robotics: New Directions , 5:293-308, 2001.

Sep Thijssen and H. J. Kappen. Path Integral Control and State-Dependent Feedback. Physical Review E , 91(3):032104, 2015.

Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model Predictive Path Integral Control Using Covariance Variable Importance Sampling. arXiv:1509.01149, 2015.

Kuan Fang, Patrick Yin, Ashvin Nair, Homer Rich Walke, Gengchen Yan, and Sergey Levine. Generalization With Lossy Affordances: Leveraging Broad Offline Data for Learning Visuomotor Tasks. In Conference on Robot Learning , pp. 106-117. 2023.

Ben Eysenbach, Russ R Salakhutdinov, and Sergey Levine. Search on the Replay Buffer: Bridging Planning and Reinforcement Learning. In Neural Information Processing Systems , volume 32. 2019.

Tianjun Zhang, Benjamin Eysenbach, Ruslan Salakhutdinov, Sergey Levine, and Joseph E Gonzalez. C-Planning: An Automatic Curriculum for Learning Goal-Reaching Tasks. In International Conference on Learning Representations . 2021.

Nikunj Saunshi, Orestis Plevrakis, Sanjeev Arora, Mikhail Khodak, and Hrishikesh Khandeparkar. A Theoretical Analysis of Contrastive Unsupervised Representation Learning. In International Conference on Machine Learning , pp. 5628-5637. 2019.

Claude Elwood Shannon. A Mathematical Theory of Communication. Bell System Technical Journal , 27(3):379-423, 1948.

E. T. Jaynes. Information Theory and Statistical Mechanics. Physical Review , 106(4):620-630, 1957.

Keith Conrad. Probability Distributions and Maximum Likelihood. 2010.

Zhuang Ma and Michael Collins. Noise Contrastive Estimation and Negative Sampling for Conditional Models: Consistency and Statistical Efficiency. In Empirical Methods in Natural Language Processing . 2018.

Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On Variational Bounds of Mutual Information. In International Conference on Machine Learning , pp. 51715180. 2019.

Nick Higham. What Is the Second Difference Matrix? https://nhigham.com/2022/01/ 31/what-is-the-second-difference-matrix/ , 2022.

Morris Newman and John Todd. The Evaluation of Matrix Inversion Programs. Journal of the Society for Industrial and Applied Mathematics , 6(4):466-476, 1958.

Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for Deep Data-Driven Reinforcement Learning. arXiv:2004.07219, 2020.

Laurens Van der Maaten and Geoffrey Hinton. Visualizing Data Using T-SNE. Journal of Machine Learning Research , 9(11), 2008.

A Proofs

This section contains the proofs omitted from the main text.

A.1 Marginal Distribution over Representations is Gaussian

Recall Assumption 1, which states that the marginal distribution over representations is Gaussian.

Assumption 1. Regularized, temporal contrastive learning acquires representations whose marginal distribution representations p ( ψ ) ≜ ∫ p ( x ) 1 ( ψ ( x ) = ψ ) d x is an isotropic Gaussian distribution:

We will motivate this statement by connecting the optimal contrastive infoNCE objective to the maximum entropy marginal distribution over the representations.

The infoNCE objective (Eq. 2) can be decomposed into an alignment term and a uniformity term [20], where the uniformity term can be simplified as follows:

̸

̸

The derivation above extends that in Wang and Isola [20] by considering a Gaussian distribution rather than a von Mises Fisher distribution. We are implicitly making the assumption that the marginal distributions satisfy p ( x ) = p ( x -) . This difference corresponds to our choice of using a negative squared L2 distance in the infoNCE loss rather than an inner product, a difference that will be important later in our analysis. A second difference is that we do not use the resubstitution estimator (i.e., we exclude data point x i from our estimate of ˆ p GMM when evaluating the likelihood of x i ), which we found hurt performance empirically. The takeaway from this identity is that maximizing the uniformity term corresponds to maximizing (an estimate of) the entropy of the representations.

We next prove that the maximum entropy distribution with an expected L2 norm constraint is a Gaussian distribution. Variants of this result are well known [95-97], but we include a full proof here for transparency.

Lemma 4. The maximum entropy distribution satisfying the expected L2 norm constraint in Eq. (3) is a multivariate Gaussian distribution with mean µ = 0 and covariance Σ = c · I

Proof. We start by defining the corresponding Lagrangian, with the second constraint saying that p ( x ) must be a valid probability distribution.

We next take the derivative w.r.t. p ( x ) :

Setting this derivative equal to 0 and solving for p ( x ) , we get

We next solve for λ 1 and λ 2 to satisfy the constraints in the Lagrangian. Note that x ∼ N ( µ = 0 , Σ = c · I ) has an expected norm E [ ∥ x ∥ 2 2 ] = c · k , so we must have λ 1 = -1 2 c . We determine λ 1 as the normalizing constant for a Gaussian, finally giving us:

corresponding to an isotropic Gaussian distribution with mean µ = 0 and covariance Σ = c · I .

A.2 Proof of Lemma 1

Below we present the proof of Lemma 1.

Lemma 1. Under the assumptions from Section 3, the distribution over representations of future states follows a Gaussian distribution with mean parameter given by the initial state representation:

Proof. Our proof technique will be similar to that of the law of the unconscious statistician:

In (a) we applied Bayes' Rule and removed the denominator, which is a constant w.r.t. ψ t + . In (b) we factored the joint distribution, noting that ψ t + and ψ 0 are deterministic functions of x t + and x 0 respectively, so they are conditionally independent from the other random variables. In (c) we used Assumption 2 after solving for p ( x t + | x 0 ) = p ( x t + ) e -1 2 ∥ Aψ ( x 0 ) -ψ ( x ) ∥ 2 2 . In (d) we noted that when the integrand is nonzero, it takes on a constant value of e -1 2 ∥ Aψ 0 -ψ t + ∥ 2 2 , so we can move that constant outside the integral. In (e) we used the definition of the marginal representation distribution (Eq. 6). In (f) we used Assumption 1 to write the marginal distributions p ( ψ t + ) and p ( ψ 0 ) as Gaussian distributions. We removed the normalizing constants, which are independent of ψ t + . In (g) we completed the square and then recognized the expression as the density of a multivariate Gaussian distribution.

A.3 Proof of Theorem 2: Waypoint Distribution

Theorem 2. Under Assumptions 1 and 2, the posterior distribution over waypoint representations is a Gaussian whose mean and covariance are linear functions of the initial and final state representations:

Proof.

In line (a) we used the definition of the conditional distribution and then simplified the numerator using the Markov property. Line (b) uses the Lemma 1. Line (c) completes the square, the details of which are below:

where Σ -1 = c c +1 A T A + c +1 c I and µ = Σ( A T ψ t + + Aψ 0 ) . Above, we have used const. = to denote equality up to an additive constant that is independent of ψ w .

A.4 Proof of Theorem 3: Planning over Many Intermediate States

Theorem 3. Given observations from a Markov chain x 0 → x 1 · · · x t + , the joint distribution over representations is a Gaussian distribution. Using ψ 1: n = ( ψ w 1 , · · · , ψ w n ) to denote the concatenated representations of each observation, we can write this distribution as

where Σ -1 is a tridiagonal matrix

Proof. We start by recalling that the waypoints form a Markov chain, so we can express their joint density as a product of conditional densities:

The aim of this lemma is to express the joint distribution over multiple waypoints, given an initial and final state representation:

where

In (a) we applied Bayes' rule and removed the denominator because it is a constant with respect to ψ 1: n . In (b) we applied the Markov assumption. In (c) we used Lemma 1 to express the conditional probabilities as Gaussians, ignoring the proportionality constants (which are independent of ψ . In (d) we simplified the exponents, removing terms that do not depend on ψ 1: n .

A.5 Formalizing Assumption 2

Assumption 2 relates the learned contrastive critic to a log-likelihood ratio between the positive and negative data distribution.

Assumption 2. Applying contrastive learning to the symmetrized infoNCE objective results in representations that encode a probability ratio:

We can justify this assumption by analyzing the general solution to the symmetrized version of the Oord et al. [10] infoNCE objective, which we do in Lemma 5. Applying this lemma to our representation learning objective (2) for sufficiently large batch size B then yields Eq. (5), with the function approximator ∥ ϕ ( x ) -ψ ( x + ) ∥ 2 ≈ f ( x, x + ) .

Lemma 5. The solution to the optimization problem

satisfies for some C .

Proof of Lemma 5. We first break down the LHS and RHS of Eq. (2):

̸

̸

̸

̸

for arbitrary c 1 ( x ) , c 2 ( x + ) .

For any C , when c 1 ( x ) = Cp ( x ) and c 2 ( x + ) = Cp ( x + ) ,

It follows that Eq. (12) maximizes both J 1 and J 2 , and is precisely the optimal solution Eq. (9) for Eq. (8).

What does C represent? From Eq. (9), we can connect C to the mutual information I ( x, x + ) :

Proof of Lemma 6. We can first consider J 1 without loss of generality. Denoting

̸

̸

we take the functional derivative:

̸

̸

We now use the following result from Ma and Collins [98]:

Lemma 6. The optimal solutions f 1 and f 2 for J 1 and J 2 satisfy

̸

̸

̸

̸

̸

̸

̸

start

★

goal

start

★

start

★

goal

start

★

goal

start

★

(a) Contrastive representations

(b) PCA representations

Xo

Xo

̸

This is zero when i.e.,

as in Eq. (10), and Eq. (11) follows similarly, exchanging x and x + .

B Additional Experiments

Fig. 7 visualizes the inferred waypoints from the task in Fig. 5. Fig. 8 visualizes the representations learned on a 46-dimensional robotic hammering task (see Section 5.3).

B.1 Stock Prediction

We show results on a stock opening price task in Fig. 9.

0

500

1000

40

60

price (USD)

A

0

500

1000

20

40

60

AAL

0

500

1000

80

100

120

140

160

180

200

AAP

0

500

1000

60

80

100

120

140

160

180

AAPL

0

500

1000

40

60

80

100

120

price (USD)

ABBV

0

500

1000

60

80

100

120

ABC

0

500

1000

40

60

ABT

0

500

1000

80

100

120

140

160

ACN

0

500

1000

days

40

60

80

100

120

140

160

180

200

price (USD)

ADBE

0

500

1000

days

40

60

80

100

ADI

0

500

1000

days

40

ADM

0

500

1000

days

60

80

100

120

ADP

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: The claims in the abstract and introduction are supposed by the Lemmas in the main text.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: See 'Limitations' in Sec. 6, and discussion following the Assumptions in the main text.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: Proofs for each Lemma are provided in either the main text or the appendices. Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: All results and figures in paper can be reproduced by running make in the linked code.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Code has been released. The datasets used have been released by prior work. Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Details are in the appendix and/or the supplemental code.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: See Sec. 5.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: See Sec. 5

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We are unaware of any violations.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: This paper provides theoretical analysis of representations learned by contrastive learning. This analysis suggests new ways of using these representations, for both beneficial and harmful purposes.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: N/A

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: See citations to the D4RL dataset.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [No]

Justification: N/A

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: N/A

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: N/A

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: SUMMARY:
This paper presents a variant of contrastive learning to time series data. Contrastive learning has been used in many areas but the insight from this paper is joint distribution of representations is also Gaussian. In addition, the experimental results show that their theory can be applied to tasks up to 46 dimensions.

STRENGTHS:
1. The problem studied in this paper is fundamental.
2. This paper provides a strong theoretical analysis.
3. The writing of this paper is clear. It is easy for me to follow the main idea of this paper.

WEAKNESSES:
1. Using contrastive learning to address various problems is not novel. However, this paper provides some new insights such as representations learned via temporal contrastive learning following a Gauss-Markov chain.
2. The evaluation section is too simple, and it is possible to add more evaluation. Since it is an ML conference, more experimental results can make the paper more solidated.

QUESTIONS:
1. As for section 5 Numerical Simulation, is it possible for the authors to use real datasets? For example, the authors can collect some datasets from the Traffic and Stock areas and evaluate the solution mentioned in this paper over them.
2. In addition, I think efficiency is also important for high-dimensional tasks. Could the authors show the running time of each experiment?

Minor:
1. As for citation, the template of NeurIPS is [10, 12] instead of (10;12). I am not sure whether you have changed the template, but it would be great if you could fix it.

LIMITATIONS:
I think the authors have adequately addressed the limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-06-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper shows that learning contrastive representations with the infoNCE objective combined with an L2 constraint on these representations permit Gaussian distributions on the representations. The authors then show for timeseries data, that, when we make a specific Markov chain assumption (a kind of linear state-space model) then the distribution over learned representations permits a nice Gaussian form for all future or past representations. As a consequence of this, planning can be performed through simple linear interpolation within the learned space.\n\nThe paper shortly validates their method in a didactic experimental setup.\n\nSTRENGTHS:\nThis paper is very well written, the presentation of theoretical results, their intuition, and how these are used is for a large part excellent.\n\nThe experiments nicely contribute to the main point of the method, that, within a contrastive learned space we can do planning in an effectively closed form rather than using an expensive dynamic programming method. The visualizations nicely show how the authors' method differ from the baselines.\n\nWEAKNESSES:\n### Summary\n\nI find the discussion and narrative of this paper too weak to confirm the novelty and contributions of the method. The authors make assumptions on their model such that it satisfies a nice practical property: planning through interpolation. However, the theoretical results that are presented should be well known consequences of Gaussian assumptions for linear state space models. The presentation makes this insight almost seem profound, but the weak discussion of related work makes me skeptical of the contribution. \n\nAlthough this paper presents a nice observation, taking all listed points of improvement into account, I don't think this paper is sufficient yet for an accept. But I'm open to improve the score if the authors address my concerns.\n\n### Major comments\n - Lemma 2 is stated as the main result. Although this way of defining waypoint conditionals is an interesting perspective on planning, the result should be a well-known consequence of using Gaussians (The same goes for Lemma 3). This makes it unclear what the novelty/ contribution of this result is. The authors need to improve this by discussing prior work/ results after introducing their model.\n - The title of the paper: inference via interpolation, is too hidden within the main text (at the end of 4.4 in Eq. 7). I would strongly suggest turning this result into a corollary (indeed with proof referring to Newman and Todd 1958 Pg. 471), then moving this result all the way to the top of section 4, and restructuring the text to move towards this result. \n - The interpolation in Eq. 7) only presents the mean-representation. There is no such discussion on the covariance.\n - The experiment of figure 5 is not well-designed. The introduced method makes use of strong assumptions that allow it to make use of linear-interpolation for planning towards future or past states. Still, the authors compare this to baselines where (also without discussion) we suddenly assume that these have the same functionality, why is this a reasonable assumption? As a result, a baseline without planning outperforms the baselines **with** planning.\n - Discussion Line 347 (last sentence) strongly oversells the presentation of the paper. The interpolation property is a nice consequence of the method's assumptions (Eq 3. and As.1, 2) and the contrastive learning. However, there is no discussion on how and why other methods cannot achieve similar things. My suggestion is to drop the 'when' comment, and only keep the 'how', because that's what the paper ultimately shows.\n\n### Technical comments\n - Eq 2. is confusing since it is first written that $(x, x^+) \\sim p(x, x)$ but in the Eq. it appears as $p(x)$. The term $p(x, x)$ is wrong, write this as $p(x, x^+)$ and keep its use consistent. Perhaps use $(x^+, x) \\sim p(x^+ | x)p(x)$ and $x^- \\sim p(x)$?\n  - Assumption 2, phrasing this result as an assumption weakens the story of the authors, especially since this is the equation that the authors' rely on for later analysis. It's usual to make assumptions or simplifications of the real world, so the fact that approximation error or model misspecification might not satisfy the result should not take away from it. I suggest making this a proposition and writing an assumption around this that discusses approximation error, please directly refer to the proof afterwards (citation + where to find it).\n  - Lemma 3, $\\eta$ isn't clearly shown to have zero entries apart from the first and last index. Why not use the same definition as in Appendix A.4?\n  - Missing proof of Lemma 3, no reference to appendix.\n\n### Minor comments\n - Line 166 typo: 'that that'\n- The two representations are $\\phi$ and $\\psi$ composed with $A$. Calling this two separate representations is confusing, there is only 1 learned representation and $A$ is used to learn the transport from embedding $\\psi(a)$ to $\\psi(b)$. Maybe state this upfront, i.e., in the first sentence line 187 say $\\psi(x)$ and $\\phi(x) = A\\psi(x)$, that helps when reading chronologically.\n - Line 215, refer back to Eq.3 for the definition of $c$\n - Line 231, proof of lemma 2 is in appendix A.3, not A.2.\n - The three examples in 4.3 for lemma 2 are nice to have for an unfamiliar audience, but I don't see how this contributes to the result since $A$ is learned, and will likely never satisfy trivial solutions.\n - Line 268, the interpolation result is nice, but the linear interpolation insight is not particularly profound (the text makes it seem as if it were so).\n - The start and end symbols in Figure 4 are too small to see and overlap with the black dots too much, even at 500% zoom.\n - Figures 5, 6 are not chronologically ordered/ numbered (6 appears before 5).\n - The choice of baselines and experiments are somewhat limited. I would have liked a comparison to a decoding method which the authors criticize in the introduction paragraph 2 (e.g., using the DreamerV2/ PlaNet RSSM model).\n - There is still space left on page 9, the authors could have spent more space on discussing the limitations of their method, which right now is rather shallow.\n\nQUESTIONS:\n.\n\nLIMITATIONS:\n1) the paper concludes that it provides a kind of 'recipe' for how to get models that have a linear interpolation property for planning, but aside from the model that the authors designed in a very specific way this is not discussed.\n\n2) the experiments are somewhat trivial since the authors expected baselines, that were not designed for a particular feature (planning through linear interp.) to suddenly have this feature.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-15"}

comments: SUMMARY:
This paper proposes to use contrastive learning technique in time series feature learning, which are helpful for improving inference performance later on.

STRENGTHS:
The method proposed here is simple but effective as shown in the numerical experiments. Employment of contrastive learning in time series control task has the potential as a complementary technique to the generative modeling methods that are currently popular.

WEAKNESSES:
Seeing the analysis of the learned feature unrolling as a Gauss-Markov process reminds me of the very classic Kalman filter. In fact, although Kalman filter is not designed to be learned in the discriminative way that this paper proposes, the equations in this paper share fair bit of similarity to those in Kalman filter. I would like to see some discussion on the difference that separates from this work from the prior time-series modeling with similar formulations.

QUESTIONS:
Does such a contrastive learning method for time series require sampling from the environments repeatedly from any given starting state? If so, how is the applicability of this method on the environments that are only sampled once?

LIMITATIONS:
N/A
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-14

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the challenges of probabilistic inference with high-dimensional time series data. It provides compact, closed form solutions in terms of learned representations by extending a prior work to show that the learned representations can be modeled as a Gauss-Markov chain. It demonstrates the effectiveness of this theoretical result via numerical experiments on 39-dimensional and 46-dimensional tasks.\n\nSTRENGTHS:\nThis paper clearly states the backgrounds and motivations. It builds upon prior work to derive a theoretical result that is useful for managing high-dimensional data in practical applications. It enhances reader understanding through examples that build intuition and illustrations that visualize methods. It also has comprehensive numerical experiments to not only validate theoretical results but also demonstrate the applicability to real-world applications.\n\nWEAKNESSES:\nThe writing of this paper can be improved, particularly in the preliminary section, which is hard to follow due to a lack of self-contained explanations for some notations and definitions. For example, why use product of the marginal distributions to sample negative pairs and what does ≈ mean for similarity? What is the infoNCE objective without resubstitution, particularly what is f and B in (2)? Why the subscript of expectation in (2) seems to not match with the sampling rule (x, x+) ∼ p(x, x) mentioned earlier? What does TSNE in the experiment section stand for?\n\nQUESTIONS:\nThe usefulness of the geometric properties remains somewhat abstract. While the special case provides some insights, it is not fully clear what it actually means for practical applications. Could you elaborate on how the geometric properties are used in practical scenarios and whether it has been used in your numerical experiments? Why linear interpolation can be used for your high-dimensional tasks while nonlinear interpolation is required for your synthetic task?\n\nI suggest separating Assumption 2 into the assumption inherited from the prior work and a statement of its results. Additionally, the term \"prior work\" is used throughout the paper, creating some ambiguity. It would enhance clarity by specifying which particular prior work is the one that shows \"the representations learned by contrastive learning encode a probability ratio\".\n\nAddressing the above two concerns can help understand the contributions better.\n\nLIMITATIONS:\nYes, the authors have adequately addressed the limitations.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-15"}