Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference

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.

Paper

Similar papers

Peer review

Reviewer GBML6/10 · confidence 3/52024-06-14

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.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

I think the authors have adequately addressed the limitations.

Reviewer e5if7/10 · confidence 4/52024-06-15

Summary

This 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. The paper shortly validates their method in a didactic experimental setup.

Strengths

This paper is very well written, the presentation of theoretical results, their intuition, and how these are used is for a large part excellent. The 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.

Weaknesses

### Summary I 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. Although 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. ### Major comments - 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. - 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. - The interpolation in Eq. 7) only presents the mean-representation. There is no such discussion on the covariance. - 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. - 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. ### Technical comments - 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)$? - 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). - 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? - Missing proof of Lemma 3, no reference to appendix. ### Minor comments - Line 166 typo: 'that that' - 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. - Line 215, refer back to Eq.3 for the definition of $c$ - Line 231, proof of lemma 2 is in appendix A.3, not A.2. - 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. - 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). - 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. - Figures 5, 6 are not chronologically ordered/ numbered (6 appears before 5). - 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). - 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.

Questions

.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

1) 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. 2) 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.

Reviewer 4BoG6/10 · confidence 2/52024-07-14

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?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

N/A

Reviewer 4JuK6/10 · confidence 2/52024-07-15

Summary

This 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.

Strengths

This 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.

Weaknesses

The 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?

Questions

The 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? I 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". Addressing the above two concerns can help understand the contributions better.

Rating

6

Confidence

2

Soundness

4

Presentation

2

Contribution

3

Limitations

Yes, the authors have adequately addressed the limitations.

Reviewer e5if2024-08-08

A solid paper given the finetuned presentation

I want to thank the authors for addressing most of my comments and the additional experiment in the global response. --- I'm happy that the authors aim to improve upon so many points, all of my concerns are answered. To respond to a few answers in particular: > Would it be more clear if we separated Lemma 2 into two results? The main problem I had was that I couldn't exactly deduce what the contribution was, so this concern is about presentation. The sentence leading up to Lemma 2 also read: "Our main result is that the posterior distribution over waypoint representations has a closed form solution ...", which added to this confusion. I think a minor rephrasing and 1-2 additional sentences that: a) discuss that the result in lemma 2 is a known consequence of the Gaussians, b) cite prior work, and c) that your contribution is that you tie these puzzle pieces together so that you can use this for intermediate state-planning, would clear this concern. This also goes for lemma 3. > When the matrix $A$ is learned on the datasets used in the paper, we empirically observe that it is often close to a rotation matrix This is an interesting insight of the method, could the authors add an explicit part to the paper (either main or appendix) that discusses this? --- I will raise my score from 4 to 7 and contribution from 2 to 3. This is a good paper with an interesting message that I think should be explored more in reinforcement learning literature. I think these types of model-features could potentially improve the scalability of planning algorithms when used in combination with neural networks (e.g., as in EfficientZero).

Authorsrebuttal2024-08-12

Dear Reviewer, > This is a good paper with an interesting message that I think should be explored more Thanks for the kind words :) > I think a minor rephrasing and 1-2 additional sentences that ... > add an explicit part to the paper that discusses the rotation matrix ... We will make these changes.

Reviewer 4JuK2024-08-12

Thanks for the authors' clarifications and I have no more concerns. This is a solid paper overall, though there is room to enhance the clarity of the writing. Based on these considerations, I have increased my score.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC