<figure><img src="image_1.png" /></figure>

HAL
open science

# VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang

<figure><img src="image_2.png" /></figure>

## To cite this version:

Yongchao Huang. VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models. 2026. hal-05468475

HAL Id: hal-05468475

https://hal.science/hal-05468475v1

Preprint submitted on 20 Jan 2026

HAL is a multi-disciplinary open access archive for the deposit and dissemination of scientific research documents, whether they are published or not. The documents may come from teaching and research institutions in France or abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est destinée au dépôt et à la diffusion de documents scientifiques de niveau recherche, publiés ou non, émanant des établissements d'enseignement et de recherche français ou étrangers, des laboratoires publics ou privés.

<figure><img src="image_3.png" /></figure>

Distributed under a Creative Commons CC BY 4.0 - Attribution - International License
# VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang *

01/01/2026

## Abstract

Joint Embedding Predictive Architectures (JEPA) offer a scalable paradigm for self-supervised learning by predicting latent representations rather than reconstructing high-entropy observations. However, existing formulations rely on *deterministic* regression objectives, which mask probabilistic semantics and limit its applicability in stochastic control. In this work, we introduce *Variational JEPA (VJEPA)*, a *probabilistic* generalization that learns a predictive distribution over future latent states via a variational objective. We show that VJEPA unifies representation learning with Predictive State Representations (PSRs) and Bayesian filtering, establishing that sequential modeling does not require autoregressive observation likelihoods. Theoretically, we prove that VJEPA representations can serve as sufficient information states for optimal control without pixel reconstruction, while providing formal guarantees for collapse avoidance. We further propose *Bayesian JEPA (BJEPA)*, an extension that factorizes the predictive belief into a learned dynamics expert and a modular prior expert, enabling zero-shot task transfer and constraint (e.g. goal, physics) satisfaction via a Product of Experts. Empirically, through a noisy environment experiment, we demonstrate that VJEPA and BJEPA successfully filter out high-variance nuisance distractors that cause representation collapse in generative baselines. By enabling principled uncertainty estimation (e.g. constructing credible intervals via sampling) while remaining likelihood-free regarding observations, VJEPA provides a foundational framework for scalable, robust, uncertainty-aware planning in high-dimensional, noisy environments.

## Contents

<table><tr><td>1 Introduction</td><td>5</td></tr><tr><td>2 Related Work</td><td>6</td></tr><tr><td>2.1 Predictive State Representations</td><td>6</td></tr><tr><td>2.2 World Models and Latent Dynamics</td><td>6</td></tr><tr><td>2.3 Active Inference</td><td>7</td></tr><tr><td>2.4 Joint Embedding Predictive Architectures (JEPAs)</td><td>7</td></tr><tr><td>2.5 JEPA World Models and Latent-Space Planning</td><td>8</td></tr><tr><td>3 JEPA: Background and Deterministic Formulation</td><td>9</td></tr></table>

<footnote>*Author email: yongchao.huang@abdn.ac.uk</footnote>
<table>
  <tbody>
    <tr>
      <td>4</td>
      <td>VJEPA: Variational Formulation of JEPA</td>
      <td>10</td>
    </tr>
    <tr>
      <td>4.1</td>
      <td>Probabilistic Predictive Model</td>
      <td>11</td>
    </tr>
    <tr>
      <td>4.2</td>
      <td>Inference Model from the Target Encoder</td>
      <td>11</td>
    </tr>
    <tr>
      <td>4.3</td>
      <td>Variational Objective</td>
      <td>12</td>
    </tr>
    <tr>
      <td>4.4</td>
      <td>Schematic Overview</td>
      <td>12</td>
    </tr>
    <tr>
      <td>4.5</td>
      <td>Training Procedure</td>
      <td>13</td>
    </tr>
    <tr>
      <td>4.6</td>
      <td>Prediction and Uncertainty Propagation</td>
      <td>14</td>
    </tr>
    <tr>
      <td>4.7</td>
      <td>Collapse Avoidance</td>
      <td>15</td>
    </tr>
    <tr>
      <td>5</td>
      <td>Time-indexed VJEPA as a Latent Space Dynamical System</td>
      <td>16</td>
    </tr>
    <tr>
      <td>5.1</td>
      <td>Background: State-Space Models</td>
      <td>17</td>
    </tr>
    <tr>
      <td>5.2</td>
      <td>From JEPA Masking to Temporal Dynamics</td>
      <td>17</td>
    </tr>
    <tr>
      <td>5.3</td>
      <td>Latent Transition Model Induced by JEPA/VJEPA</td>
      <td>18</td>
    </tr>
    <tr>
      <td>5.4</td>
      <td>Sequential but Not Autoregressive</td>
      <td>18</td>
    </tr>
    <tr>
      <td>5.5</td>
      <td>Clip-Level Prediction (V-JEPA) vs. Sequential Dynamics</td>
      <td>19</td>
    </tr>
    <tr>
      <td>6</td>
      <td>JEPA/VJEPA for Control</td>
      <td>20</td>
    </tr>
    <tr>
      <td>6.1</td>
      <td>Background: POMDPs and Belief-State Control</td>
      <td>20</td>
    </tr>
    <tr>
      <td>6.2</td>
      <td>JEPA/VJEPA Latent State as a Predictive Information State</td>
      <td>21</td>
    </tr>
    <tr>
      <td>6.3</td>
      <td>Predictive Sufficiency for Control (Formalization)</td>
      <td>22</td>
    </tr>
    <tr>
      <td>6.4</td>
      <td>Latent-Space Planning Objective</td>
      <td>23</td>
    </tr>
    <tr>
      <td>6.5</td>
      <td>Sampling vs. Point Prediction Planning</td>
      <td>24</td>
    </tr>
    <tr>
      <td>6.6</td>
      <td>Algorithm: VJEPA Sampling-Based Latent MPC</td>
      <td>24</td>
    </tr>
    <tr>
      <td>6.7</td>
      <td>When is MAP Control Optimal?</td>
      <td>25</td>
    </tr>
    <tr>
      <td>7</td>
      <td>Information-Theoretic Analysis</td>
      <td>27</td>
    </tr>
    <tr>
      <td>7.1</td>
      <td>Variational Maximization of Predictive Information</td>
      <td>27</td>
    </tr>
    <tr>
      <td>7.2</td>
      <td>Predictive Information Bottleneck</td>
      <td>28</td>
    </tr>
    <tr>
      <td>7.3</td>
      <td>Contrast with Generative Objectives</td>
      <td>29</td>
    </tr>
    <tr>
      <td>8</td>
      <td>Bayesian JEPA (BJEPA)</td>
      <td>30</td>
    </tr>
    <tr>
      <td>8.1</td>
      <td>Bayesian Factorization via Product of Experts</td>
      <td>31</td>
    </tr>
    <tr>
      <td>8.2</td>
      <td>The BJEPA Architecture</td>
      <td>31</td>
    </tr>
    <tr>
      <td>8.3</td>
      <td>Training and Inference</td>
      <td>34</td>
    </tr>
    <tr>
      <td>8.4</td>
      <td>Sequential BJEPA as Latent Bayesian Filtering</td>
      <td>36</td>
    </tr>
    <tr>
      <td>8.5</td>
      <td>MAP Optimality under BJEPA</td>
      <td>37</td>
    </tr>
    <tr>
      <td>8.6</td>
      <td>Energy-Based Priors and Latent Optimization</td>
      <td>38</td>
    </tr>
    <tr>
      <td>8.7</td>
      <td>Algorithm: Gradient-Based Planning with BJEPA</td>
      <td>38</td>
    </tr>
    <tr>
      <td>8.8</td>
      <td>Modularity, Transfer, and Continual Learning</td>
      <td>39</td>
    </tr>
    <tr>
      <td>9</td>
      <td>Toy Experiment: a "Noisy TV" Linear System</td>
      <td>40</td>
    </tr>
    <tr>
      <td>9.1</td>
      <td>Environment Setup</td>
      <td>40</td>
    </tr>
    <tr>
      <td>9.2</td>
      <td>Model Architectures and Evaluation Alignment</td>
      <td>41</td>
    </tr>
    <tr>
      <td>9.3</td>
      <td>Results and Analysis</td>
      <td>41</td>
    </tr>
  </tbody>
</table>
<table>
  <tbody>
    <tr>
      <td><b>10 Discussion</b></td>
      <td><b>43</b></td>
    </tr>
    <tr>
      <td>10.1 Unifying JEPA with Probabilistic State-Space Models</td>
      <td>43</td>
    </tr>
    <tr>
      <td>10.2 Nuisance Invariance: The “PCA vs. CCA” Distinction</td>
      <td>43</td>
    </tr>
    <tr>
      <td>10.3 Applications and Broader Impact</td>
      <td>44</td>
    </tr>
    <tr>
      <td>10.4 Limitations and Future Directions</td>
      <td>45</td>
    </tr>
    <tr>
      <td><b>11 Conclusion</b></td>
      <td><b>45</b></td>
    </tr>
    <tr>
      <td><b>A Equivalence Between Squared Loss and Gaussian Likelihood</b></td>
      <td><b>52</b></td>
    </tr>
    <tr>
      <td><b>B Relationships Between Information-Theoretic Quantities</b></td>
      <td><b>53</b></td>
    </tr>
    <tr>
      <td>B.1 Cross-Entropy, Entropy, and KL Divergence</td>
      <td>53</td>
    </tr>
    <tr>
      <td>B.2 Mutual Information and Entropy</td>
      <td>53</td>
    </tr>
    <tr>
      <td>B.3 Chain Rule for Mutual Information</td>
      <td>54</td>
    </tr>
    <tr>
      <td>B.4 Symmetry Properties</td>
      <td>54</td>
    </tr>
    <tr>
      <td>B.5 Evidence Lower Bound (ELBO)</td>
      <td>55</td>
    </tr>
    <tr>
      <td>B.6 Non-Negativity Properties</td>
      <td>55</td>
    </tr>
    <tr>
      <td>B.7 Summary of Relationships</td>
      <td>56</td>
    </tr>
    <tr>
      <td><b>C Compare VJEPA with Hidden Markov Models (HMMs)</b></td>
      <td><b>56</b></td>
    </tr>
    <tr>
      <td>C.1 Classical HMMs</td>
      <td>56</td>
    </tr>
    <tr>
      <td>C.2 VJEPA as “Latent-Only” Predictive Modeling</td>
      <td>56</td>
    </tr>
    <tr>
      <td>C.3 Similarities: Latent-State Semantics, Dynamics, and Beliefs</td>
      <td>57</td>
    </tr>
    <tr>
      <td>C.4 Differences</td>
      <td>57</td>
    </tr>
    <tr>
      <td>C.5 Sequential VJEPA as an “HMM Without Emissions”</td>
      <td>57</td>
    </tr>
    <tr>
      <td><b>D Kalman-Filter-Like Interpretation of VJEPA</b></td>
      <td><b>58</b></td>
    </tr>
    <tr>
      <td>D.1 The Linear-Gaussian Benchmark</td>
      <td>58</td>
    </tr>
    <tr>
      <td>D.2 VJEPA as Generalized Latent Prediction</td>
      <td>58</td>
    </tr>
    <tr>
      <td>D.3 Amortized Correction</td>
      <td>59</td>
    </tr>
    <tr>
      <td><b>E A Particle-Filter Interpretation of VJEPA</b></td>
      <td><b>59</b></td>
    </tr>
    <tr>
      <td><b>F Detailed Proof of Theorem 1</b></td>
      <td><b>61</b></td>
    </tr>
    <tr>
      <td><b>G Detailed Proof of Theorem 3</b></td>
      <td><b>63</b></td>
    </tr>
    <tr>
      <td><b>H Proof of Proposition 1 (Invariance to Nuisance)</b></td>
      <td><b>66</b></td>
    </tr>
    <tr>
      <td>H.1 Case 1: Generative / Reconstruction Objective</td>
      <td>66</td>
    </tr>
    <tr>
      <td>H.2 Case 2: VJEPA Prediction Objective</td>
      <td>67</td>
    </tr>
    <tr>
      <td><b>I Derivation of the BJEPA Predictive Factorization</b></td>
      <td><b>67</b></td>
    </tr>
    <tr>
      <td><b>J Toy Experiment Details</b></td>
      <td><b>68</b></td>
    </tr>
    <tr>
      <td>J.1 Environment and Data Generation</td>
      <td>68</td>
    </tr>
    <tr>
      <td>J.2 Model Architectures and Training Strategy</td>
      <td>69</td>
    </tr>
    <tr>
      <td>J.2.1 Generative Baselines (Pixel Reconstruction)</td>
      <td>69</td>
    </tr>
    <tr>
      <td>J.2.2 Joint-Embedding Architectures (Latent Autoregression)</td>
      <td>70</td>
    </tr>
    <tr>
      <td>J.2.3 Temporal Alignment in Training and Evaluation</td>
      <td>71</td>
    </tr>
    <tr>
      <td>J.2.4 Linear Probe Evaluation Protocol</td>
      <td>72</td>
    </tr>
  </tbody>
</table>
<table><tr><td>J.3</td><td>Optimization and Hyperparameters</td><td>73</td></tr><tr><td>J.4</td><td>Results</td><td>74</td></tr><tr><td>J.5</td><td>Computational Resources</td><td>77</td></tr><tr><td>J.6</td><td>Code Availability</td><td>77</td></tr></table>
# 1 Introduction

Self-supervised representation learning has traditionally followed two dominant paradigms: (i) generative modeling, which optimizes likelihoods over high-dimensional observations such as pixels or tokens [39, 63, 79, 61], and (ii) contrastive learning, which discriminates between positive and negative pairs [28, 80, 19, 31]. Both approaches suffer from fundamental limitations: generative models must model nuisance variability and high-entropy details irrelevant for downstream tasks [3, 33, 47], while contrastive methods require carefully engineered negative sampling schemes and can introduce representational bias [78, 4, 19].

Joint Embedding Predictive Architectures (JEPA) offer a third paradigm [5, 8]. Rather than reconstructing observations or contrasting samples, JEPA learns by predicting *representations* of missing or future data [5]. This removes the need for observation-level likelihoods and explicitly biases learning toward predictable, task-relevant structure [8, 43, 70]. Recent instantiations such as I-JEPA [5] and V-JEPA [8, 9] have demonstrated strong empirical performance on vision and video tasks. Early on, LeCun [44] proposed *Hierarchical JEPA* (H-JEPA), a conceptual framework for multi-level, multi-timescale predictive models that emphasizes learning latent dynamics without observation reconstruction as a basis for model-predictive control and long-horizon reasoning.

More recently, several works have begun to instantiate JEPA-style predictors as *world models* for planning and control. Bardes et al. [6] show that large-scale JEPA pretraining on internet video, combined with lightweight action-conditioned prediction, enables zero-shot robotic planning via latent-space model-predictive control. Terver et al. [76] systematically study action-conditioned *JEPA World Models* (JEPA-WMs), demonstrating that planning directly in representation space can solve navigation and manipulation tasks without observation reconstruction and without training on reward-labeled data. Destrade et al. [21] further study *value-guided JEPA planning*, which shapes the representation space so that distances approximate goal-conditioned value, and improves optimization-based planning.

Despite these advances, existing JEPA-based world models remain largely *deterministic* and are trained using regression-based objectives [5, 8]. As a result, the underlying probabilistic semantics of JEPA predictions, the representation of uncertainty over future latent states, and the conditions under which the learned representation constitutes a sufficient information state for planning and control remain unformalized. In particular, current approaches implicitly assume predictive sufficiency, but do not characterize when or why this assumption should hold.

In this work, we ask and answer several foundational questions about JEPA:

* *Implicit Probabilistic Model:* what probabilistic objective does deterministic JEPA implicitly optimize, and how can it be generalized to handle uncertainty?

* *Dynamics and Control:* can JEPA be formalized as a latent dynamical system where the learned representation serves as a *sufficient information state* for optimal control, without reconstructing observations?

* *Bayesian Connection:* how does JEPA relate to Bayesian filtering and Predictive State Representations (PSRs), and can this connection enable the injection of structural priors or constraints into latent planning?

* *Nuisance Invariance:* does introducing temporal structure force JEPA to adopt autoregressive observation likelihoods, or can it remain likelihood-free to avoid modeling high-variance nuisance distractors?
We address these questions by formulating JEPA as a *probabilistic predictive state-space model*, and showing that sequential predictive modeling does not require autoregressive likelihood factorization and that JEPA representations can serve as sufficient information states for control and planning.

**Contributions.** We propose *Variational JEPA (VJEPA)*, a probabilistic formulation that explicates the implicit generative model of JEPA. We theoretically establish VJEPA as a predictive state-space model, proving that its representations serve as *sufficient information states* for optimal control and providing formal guarantees for collapse avoidance. We further extend this to *Bayesian JEPA (BJEPA)*, a modular framework that factorizes predictive belief into learned dynamics and structural priors, enabling the insertion of constraints and physical laws, as well as zero-shot task transfer via Product of Experts. Finally, we empirically demonstrate in a noisy environment that VJEPA and BJEPA successfully filter out high-variance distractors, validating that principled uncertainty-aware world modeling is achievable without autoregressive observation reconstruction.

# 2 Related Work

## 2.1 Predictive State Representations

Predictive State Representations (PSRs) were introduced by Sutton et al. [46] and independently formalized by Singh et al. [70]. They replace latent states with sufficient statistics of future observations, representing the state as

$$s_t = p(o_{t+1:t+H} | o_{\le t}),$$

where $H$ denotes the prediction horizon. Sutton et al. [46] emphasized the theoretical connection between PSRs and history-based models, showing how state can be represented purely in terms of observable predictions. Singh et al. [70] developed the system-dynamics matrix formalism and demonstrated that PSRs generalize $n$-th order Markov models and hidden-state models such as HMMs or POMDPs. Later, we will see that VJEPA can be viewed as a neural, amortized PSR that compresses predictive distributions into a latent embedding $Z_t$.

## 2.2 World Models and Latent Dynamics

Model-based latent dynamics have been widely explored in reinforcement learning. PlaNet [30] introduced a probabilistic latent dynamics model that predicts future latent states $z_{t+1}$ conditioned on actions $a_t$ and reconstructs observations $x_t$, optimizing an evidence lower bound (ELBO) over pixels. Dreamer [29] extended PlaNet by integrating actor-critic reinforcement learning in the learned latent space, enabling policy optimization without decoding full trajectories to pixels. Formally, these methods learn

$$p(z_{t+1} | z_t, a_t), \quad p(x_t | z_t),$$

where the observation likelihood $p(x_t | z_t)$ is explicitly modeled.

In contrast, JEPA [5], V-JEPA (video-JEPA) [8] and VJEPA (this work) discard observation reconstruction, learning predictive latent representations that compress future information into $Z_t$, emphasizing predictive sufficiency rather than pixel fidelity.

MuZero [68] also learns latent dynamics without reconstructing observations, but it still relies on value and policy heads for planning. VJEPA differs by learning latent dynamics purely through representation prediction, without access to rewards, policies, or observation likelihoods, making it a self-supervised model suitable for both control and general representation learning.
## 2.3 Active Inference

Active Inference is a theoretical framework proposed by Friston and colleagues in which perception, action, and learning are cast as variational inference in a latent variable model of the world [23, 24, 25]. The central idea is that biological agents act to minimize variational free energy, which provides an upper bound on sensory surprise (negative log-evidence). Formally, variational free energy is defined as

$$ \mathcal{F}(q) = \mathbb{E}_{q(z)} [\log q(z) - \log p(x, z)], $$

where $z$ denotes latent (hidden) states of the environment, $x$ denotes sensory observations, and $q(z)$ is an approximate posterior belief. Minimizing $\mathcal{F}(q)$ drives $q(z)$ towards the true Bayesian posterior $p(z|x)$, thereby implementing perception as approximate Bayesian inference. This formulation was introduced and developed in early work by Friston [23] and extended to perception-action loops in subsequent studies [24].

Within Active Inference, belief updating corresponds to approximate Bayesian filtering over latent states, often implemented as gradient descent on free energy. More recent work has emphasized a distinction between perception and action: while perception minimizes current free energy, action selection is driven by the minimization of *expected free energy* over future trajectories, enabling agents to plan by imagining future outcomes and selecting policies that reduce expected surprise and uncertainty [25].

VJEPA aligns with Active Inference at the level of *latent belief propagation*: both frameworks aim to learn compact latent representations that summarize past observations and support prediction of future states. However, VJEPA deliberately omits modeling sensory likelihoods $p(x|z)$, preferences, or policy evaluation. Instead, it focuses exclusively on learning predictive latent dynamics through representation prediction, yielding a self-supervised model of belief dynamics rather than a full Active Inference agent.

## 2.4 Joint Embedding Predictive Architectures (JEPAs)

JEPA was introduced to self-supervised learning by Assran et al. [5] in the Image-based JEPA (I-JEPA) model, where representations of held-out image regions are predicted from encoded context patches. I-JEPA demonstrated that prediction in a latent representation space can produce highly semantic features without pixel reconstruction or handcrafted augmentations, and scales to large vision models with competitive performance [5].

Subsequent work extended the JEPA paradigm to video. Bardes et al. proposed Video JEPA (V-JEPA) for self-supervised learning from video by predicting masked spatio-temporal regions in a learned latent space [8]. While V-JEPA improves temporal representation learning relative to image-only JEPA, it retains the deterministic regression objective.

Beyond these instantiations, the JEPA framework has inspired a variety of architectural and domain adaptations. For example, MC-JEPA [11] incorporates motion and content prediction objectives for richer video representation [11], S-JEPA applies the JEPA principle to skeletal action recognition [1], and Point-JEPA adapts JEPA to point cloud self-supervised learning [67]. Other variants explore JEPA in graph domains [71] and in joint vision-language models [18]. Additional work (e.g. C-JEPA [56], DSeq-JEPA [32]) proposes contrastive or sequential extensions that address specific limitations such as representation collapse or prediction ordering.

Despite this growing body of JEPA-based methods and applications, to our knowledge there has been no formal variational or Bayesian JEPA that explicitly models predictive uncertainty and the underlying probabilistic generative process. Some work incorporates auxiliary losses to stabilize training (e.g. VICReg-style regularization [10] in C-JEPA [56]) or implicitly addresses uncertainty
through architectural choices, but these do not derive from a principled likelihood-based objective. In contrast, our VJEPA formulates JEPA as a variational predictive model, learning an explicit predictive distribution over future latent states with likelihood and KL regularization, enabling a formal probabilistic interpretation and control-theoretic extensions not addressed by prior JEPA work.

## 2.5 JEPA World Models and Latent-Space Planning

The use of JEPA as a foundation for world modeling and planning was first articulated at a conceptual level by LeCun [44]. In his proposal of *Hierarchical JEPA* (H-JEPA), LeCun argues that autonomous agents should learn multi-level, multi-timescale predictive representations of the world that support planning and model-predictive control without relying on observation reconstruction or explicit reward modeling. H-JEPA provides a normative blueprint for predictive world modeling and hierarchical control, but does not specify a concrete training objective, probabilistic semantics, or learning algorithm.

Building on this vision, more recent work has developed concrete, action-conditioned JEPA-based world models that enable planning directly in representation space. Terver et al. [76] introduce *JEPA World Models* (JEPA-WMs), which combine pretrained visual encoders with action-conditioned JEPA predictors trained via a deterministic representation prediction loss. Planning is performed by sampling or optimizing action sequences whose predicted latent trajectories minimize a goal-conditioned cost in embedding space. Their results demonstrate that latent-space planning without observation reconstruction can match or outperform reconstruction-based world models, particularly when reward signals are unavailable or sparse.

Several subsequent works explore extensions of JEPA-based planning. Destrade et al. [21] propose *value-guided JEPA planning*, in which the latent representation space is shaped so that distances approximate goal-conditioned value functions, improving the optimization landscape for planning. More recently, Bardes et al. [6] demonstrate that large-scale JEPA pretraining on internet video, followed by lightweight action-conditioned prediction, enables zero-shot robotic manipulation via latent-space model-predictive control.

While these works provide strong empirical evidence that JEPA representations can support planning and control, they are uniformly based on *deterministic* predictors trained via regression objectives. As a result, uncertainty over future latent states, belief propagation, and the conditions under which the learned representation constitutes a sufficient information state for optimal control remain implicit and unformalized. In particular, these approaches assume predictive sufficiency of the latent state without providing a probabilistic model or formal guarantees.

Our work complements JEPA-based world models by providing a *variational* formulation of JEPA that makes the underlying predictive distribution explicit. By modeling uncertainty over future latent states, VJEPA enables belief propagation, distributional planning, and formal connections to predictive state representations and Bayesian filtering, which are capabilities that not addressed by existing deterministic JEPA world models.

Taken together, the above lines of work leave open a gap that is not fully addressed by existing approaches. Predictive State Representations provide a principled account of state as predictive statistics, but are not designed as scalable, amortized representation learners for high-dimensional sensory data. Latent world models such as PlaNet, Dreamer, and MuZero learn predictive dynamics suitable for control, but rely on observation reconstruction, reward supervision, or policy optimization objectives that entangle representation learning with task-specific signals. Active Inference offers a unifying theoretical framework for perception and action, but requires explicit modeling of
sensory likelihoods, preferences, and policies, and is not typically instantiated as a self-supervised representation learning algorithm. Finally, existing JEPA formulations learn predictive representations using deterministic regression objectives, which obscure the underlying probabilistic structure and preclude explicit modeling of uncertainty over future latent states.

VJEPA occupies a complementary point in this design space by extending JEPA with an explicit probabilistic formulation. It learns predictive latent states as amortized representations while modeling a full predictive distribution over future embeddings, rather than a single point estimate. This enables principled uncertainty estimation, belief propagation, and compatibility with Bayesian filtering and control, while retaining JEPA's core advantages: avoiding observation reconstruction, reward supervision, and autoregressive likelihood factorization. By focusing exclusively on predictive sufficiency through variational representation prediction, VJEPA bridges predictive state representations, latent world models, and belief propagation frameworks, while remaining a purely self-supervised model that is directly applicable to downstream control and planning.

## 3 JEPA: Background and Deterministic Formulation

Joint Embedding Predictive Architectures (JEPA) were introduced as a self-supervised learning framework that learns representations by predicting embeddings of missing or future data, rather than reconstructing observations or discriminating between positive and negative samples [5, 8]. The central design principle of JEPA is to encourage representations that capture *predictable, task-relevant structure* while discarding nuisance variability that is difficult or unnecessary to predict.

Given an input $x$ (e.g. an image or a video clip), JEPA first partitions the input into two disjoint subsets: a *context* $x_C$ and a *target* $x_T$. The context contains the information available to the model, while the target represents missing or held-out content that must be predicted at the representation level. This partitioning can correspond to spatial masking (as in I-JEPA [5]) or temporal masking (as in video JEPA [8]), but the formulation itself is agnostic to the specific modality. The context is mapped to a latent representation using a trainable encoder

$$Z_C = f_{\theta}(x_C), \qquad (1)$$

which produces a compact embedding summarizing the available information. The target is encoded using a separate encoder

$$Z_T = f_{\theta'}(x_T), \qquad (2)$$

whose parameters $\theta'$ are updated as an exponential moving average (EMA) of the context encoder parameters[55, 54, 45, 26]:

$$\theta' \leftarrow \tau\theta' + (1 - \tau)\theta. \qquad (3)$$

This asymmetric design, inspired by prior work [26], stabilizes training and prevents representational collapse by ensuring that the target embeddings evolve more slowly than the predictor.

To predict the target representation, JEPA employs a predictor network

$$\hat{Z}_T = g_{\phi}(Z_C, \xi_T), \qquad (4)$$

where $\xi_T$ encodes structural (side) information about the target, such as spatial location or temporal index. The predictor is tasked with mapping the context representation to the target embedding without direct access to the target observations themselves.
With these settings, training is then performed by minimizing a regression loss between the predicted and target embeddings:

$$ \begin{aligned} \mathcal{L}_{\text{JEPA}} &= \|\hat{Z}_T - Z_T\|^2 \\ &= \|g_{\phi}(Z_C, \xi_T) - Z_T\|^2 \\ &= \|g_{\phi}(f_{\theta}(x_C), \xi_T) - f_{\theta'}(x_T)\|^2 \\ &= \|g_{\phi}(f_{\theta}(x_C), \xi_T) - f_{\tau\theta'+(1-\tau)\theta}(x_T)\|^2, \end{aligned} \quad (5) $$

where the target encoder parameters are updated as per Eq. (3). Unlike generative models, this objective does not require reconstructing observations or modeling pixel-level likelihoods; instead, it enforces consistency between representations, encouraging the latent space to organize itself around predictable relationships in the data.

Although the JEPA loss is typically presented as a mean-squared error, it admits a natural probabilistic interpretation. Specifically, minimizing $\mathcal{L}_{\text{JEPA}}$ is equivalent to maximizing the log-likelihood of $Z_T$ under an isotropic Gaussian distribution centered at $\hat{Z}_T$ with fixed variance (see Appendix A a brief derivation):

$$ \mathcal{L}_{\text{JEPA}} \propto -\log p(Z_T | Z_C), \quad p(Z_T | Z_C) = \mathcal{N}(Z_T | g_{\phi}(Z_C, \xi_T), \sigma^2 I), \quad (6) $$

This implicit assumption of a simple predictive distribution motivates the variational extension of JEPA introduced in the next section, in which predictive uncertainty is modeled explicitly.

# 4 VJEPA: Variational Formulation of JEPA

Eq. (5) shows that the deterministic JEPA objective is equivalent to maximum likelihood under an isotropic Gaussian predictive model with fixed variance (Eq. (6)). This suggests a natural generalization: instead of predicting a single target embedding $\hat{Z}_T$ (i.e. a point estimate), we can model an explicit predictive distribution over possible target embeddings, enabling uncertainty-aware prediction and multi-modal futures. Importantly, as in deterministic JEPA, prediction is conditioned not only on the context representation but also on structural (side) information describing the target.

**From deterministic JEPA to VJEPA.** VJEPA can be understood as a minimal probabilistic extension of deterministic JEPA. Rather than altering the architectural principles of JEPA, VJEPA replaces point prediction in representation space with distributional prediction while preserving the same context-target structure, conditioning mechanism, and training asymmetry. As a preview, we summarise the key differences in Table.1; these features will be progressively introduced in later sections.
<table><caption>Table 1: Comparison between deterministic JEPA and its variational extension VJEPA. VJEPA preserves the structural design of JEPA while endowing it with an explicit probabilistic semantics.</caption><thead><tr><th>Feature</th><th>JEPA</th><th>VJEPA</th></tr></thead><tbody><tr><td>Target prediction</td><td>Z&#770;<sub>T</sub> = g<sub>&phi;</sub>(Z<sub>C</sub>, &xi;<sub>T</sub>)</td><td>p<sub>&phi;</sub>(Z<sub>T</sub> | Z<sub>C</sub>, &xi;<sub>T</sub>)</td></tr><tr><td>Prediction type</td><td>Point estimate</td><td>Predictive distribution</td></tr><tr><td>Training objective</td><td>MSE (equally Gaussian likelihood)</td><td>Negative log-likelihood + KL</td></tr><tr><td>Uncertainty modeling</td><td>Implicit (fixed variance)</td><td>Explicit (learned variance / distribution)</td></tr><tr><td>Multi-modal futures</td><td>Not supported</td><td>Supported</td></tr><tr><td>Collapse avoidance</td><td>Architectural (EMA, asymmetry)</td><td>Objective-level (predictive mismatch + KL reg)</td></tr><tr><td>Probabilistic interpretation</td><td>Implicit Gaussian</td><td>Explicit variational model</td></tr><tr><td>Compatibility with filtering / control</td><td>Limited</td><td>Natural (belief propagation)</td></tr></tbody></table>

## 4.1 Probabilistic Predictive Model

Let $x$ be an input sample and $(x_C, x_T)$ be a context/target partition. As in JEPA, the context is encoded as deterministic $Z_C = f_\theta(x_C)$, see Eq. (1). In addition, let $\xi_T$ denote side information specifying the structure of the target (e.g. spatial location, temporal index, or masking pattern, etc), as used in Eq. (4). VJEPA introduces a probabilistic predictive model over the target representation:

$$p_{\phi}(Z_T | Z_C, \xi_T), \qquad (7)$$

which generalizes the deterministic predictor $g_\phi(Z_C, \xi_T)$ by allowing the predicted target representation to be stochastic, e.g. $p_\phi(Z_T | Z_C, \xi_T) = N(\mu_\phi(Z_C, \xi_T), \Sigma_\phi(Z_C))$; if $p_\phi(Z_T | Z_C, \xi_T) = g_\phi(Z_C, \xi_T)$ is chosen, VJEPA reduces to the deterministic JEPA. The probabilistic, predictive model Eq. (7) can be further decomposed, e.g. formulated as a Bayesian posterior (see later Section.8).

Optionally, one may also posit an observation model

$$p_{\psi}(x_T | Z_T), \qquad (8)$$

yielding the joint factorization

$$p_{\phi, \psi}(Z_T, x_T | x_C, \xi_T) = p_{\phi}(Z_T | Z_C, \xi_T) p_{\psi}(x_T | Z_T). \qquad (9)$$

In this text, in the spirit of JEPA, we *do not* optimize an observation-level reconstruction objective over $p_\psi(x_T | Z_T)$; the learning signal is provided solely through the predictive distribution in representation space, $p_\phi(Z_T | Z_C, \xi_T)$.

## 4.2 Inference Model from the Target Encoder

To train the predictive model without reconstructing $x_T$, we require a representation-space “posterior” for $Z_T$ given the target observation. We define an amortized inference distribution using the (EMA) target encoder:

$$q_{\theta'}(Z_T | x_T), \qquad (10)$$

which may be instantiated as a diagonal Gaussian whose mean and variance are parameterized by $f_{\theta'}(x_T)$. The EMA update in Eq. (3) ensures that the inference target evolves slowly, which stabilizes training and prevents degenerate fixed points.
## 4.3 Variational Objective

Given $(x_C, x_T, \xi_T)$, VJEPA trains the predictive model $p_\phi(Z_T | Z_C, \xi_T)$ to match the target-encoder distribution $q_{\theta'}(Z_T | x_T)$. A simple and effective objective is the *regularized negative log-likelihood*:

$$ \mathcal{L}_{\text{VJEPA}} = \mathbb{E}_x \mathbb{E}_{Z_T \sim q_{\theta'}(\cdot | x_T)} \left[ -\log p_\phi(Z_T | Z_C, \xi_T) \right] + \beta \mathbb{E}_x \text{KL}(q_{\theta'}(Z_T | x_T) \| p(Z_T)), \quad (11) $$

where $p(Z_T)$ is a fixed prior (e.g. $\mathcal{N}(0, I)$) and $\beta > 0$ controls the strength of regularization.

The first term trains a probabilistic predictor in *representation space*, conditioned on both the available context and the structural specification of the target. The second term prevents trivial or degenerate solutions by encouraging the inferred target representation distribution to remain non-collapsed and well-calibrated relative to the prior. Together, this objective (i) introduces explicit predictive uncertainty, (ii) permits multi-modal futures through expressive $p_\phi(\cdot | Z_C, \xi_T)$, and (iii) yields collapse-avoidance guarantees under mild conditions (see later Section 4.7).

**Deterministic JEPA as a special case.** If we set

$$ q_{\theta'}(Z_T | x_T) = \delta(Z_T - f_{\theta'}(x_T)) $$

and

$$ p_{\phi}(Z_T | Z_C, \xi_T) = \mathcal{N}(Z_T | g_{\phi}(Z_C, \xi_T), \sigma^2 I) $$

with fixed $\sigma^2$, then minimizing Eq. (11) reduces (up to additive constants) to the squared-loss JEPA objective in Eq. (5).

## 4.4 Schematic Overview

The complete VJEPA architecture, illustrating the probabilistic predictor, target sampling, and the variational loss components, is schematically shown in Fig.1.
<figure><img src="image_4.png" /><figcaption>Figure 1: **Variational JEPA (VJEPA) architecture.** The context $x_C$ is encoded into $Z_C$, conditioning a probabilistic predictor $p_\phi(Z_T | Z_C, \xi_T)$. The target $x_T$ is encoded into a distribution $q_{\theta'}$ (typically via EMA), from which a sample $Z_T$ is drawn. Training minimizes the sum of a negative log-likelihood term (red box) and a KL regularization term (blue box). $f_{\theta'}$ can output sufficient statistics (e.g. mean and covariance for Gaussians) for parameterizing $q_{\theta'}$.</figcaption></figure>

## 4.5 Training Procedure

Algorithm 1 summarizes VJEPA training. In practice, the reparameterization trick may be used when $q_{\theta'}$ is continuous, and $p_\phi$ may be implemented as a Gaussian with learned mean and covariance (diagonal or structured), a mixture model, or a normalizing flow, depending on the desired expressivity. A concrete example is given in Section.9.

<caption>**Algorithm 1 VJEPA Training**</caption>

**Require:** Dataset $\{x^{(i)}\}_{i=1}^N$, context encoder $f_\theta$, target encoder $f_{\theta'}$, predictive model $p_\phi$, prior $p(Z)$, EMA rate $\tau$, regularization weight $\beta$
1: **for** each minibatch $\{x\}$ **do**
2: &nbsp;&nbsp;&nbsp;&nbsp;Sample a context-target partition $(x_C, x_T)$ and corresponding target structure $\xi_T$
3: &nbsp;&nbsp;&nbsp;&nbsp;Encode context: $Z_C \leftarrow f_\theta(x_C)$
4: &nbsp;&nbsp;&nbsp;&nbsp;Compute inference distribution from target: $q_{\theta'}(Z_T | x_T)$
5: &nbsp;&nbsp;&nbsp;&nbsp;Sample target representation (e.g. using reparameterization): $Z_T \sim q_{\theta'}(Z_T | x_T)$
6: &nbsp;&nbsp;&nbsp;&nbsp;Evaluate predictive distribution: $p_\phi(Z_T | Z_C, \xi_T)$
7: &nbsp;&nbsp;&nbsp;&nbsp;Compute variational loss: $\mathcal{L} = -\log p_\phi(Z_T | Z_C, \xi_T) + \beta \text{KL}(q_{\theta'}(Z_T | x_T) || p(Z_T))$
8: &nbsp;&nbsp;&nbsp;&nbsp;Update $\theta$ and $\phi$ via gradient descent on $\mathcal{L}$
9: &nbsp;&nbsp;&nbsp;&nbsp;Update target encoder via EMA: $\theta' \leftarrow \tau\theta' + (1-\tau)\theta$
10: **end for**

**Remark.** Algorithm 1 reduces to standard JEPA training when the inference distribution $q_{\theta'}(Z_T | x_T)$ collapses to a Dirac delta, $q_{\theta'}(Z_T | x_T) = \delta(Z_T - f_{\theta'}(x_T))$, and the predictive distribution is
chosen as an isotropic Gaussian with fixed variance, $p_\phi(Z_T | Z_C, \xi_T) = \mathcal{N}(Z_T | g_\phi(Z_C, \xi_T), \sigma^2 I)$. In this case, minimizing the VJEPA objective in Eq. (11) is equivalent (up to constants) to minimizing the squared-loss JEPA objective in Eq. (5).

## 4.6 Prediction and Uncertainty Propagation

After training, VJEPA can be used as a predictive model in representation space. Given a new input $x_C$ and target specification $\xi_T$, we first compute the context representation (i.e. the context encoder Eq. (1))

$$Z_C = f_{\theta}(x_C),$$

and then form the predictive distribution (i.e. the predictor Eq. (7))

$$p_{\phi}(Z_T | Z_C, \xi_T).$$

**Point prediction.** A point estimate of the target representation may be obtained from the predictive distribution $p_\phi(Z_T | Z_C, \xi_T)$ in several standard ways. A common choice is the predictive mean,

$$\hat{Z}_T^{\text{mean}} = \mathbb{E}_{p_\phi(Z_T|Z_C, \xi_T)}[Z_T],$$

which reduces to the deterministic JEPA predictor when $p_\phi$ is Gaussian with fixed variance.

Alternatively, a maximum a posteriori (MAP) estimate may be used:

$$\hat{Z}_T^{\text{MAP}} = \arg\max_{Z_T} p_\phi(Z_T | Z_C, \xi_T),$$

which coincides with the predictive mean for unimodal Gaussian models but differs for multi-modal or heavy-tailed distributions.

More generally, one may consider a *mode prediction*,

$$\hat{Z}_T^{\text{mode}} = \text{mode}[p_\phi(Z_T | Z_C, \xi_T)],$$

which selects a single most probable latent realization. For expressive predictive models (e.g. mixture models or normalizing flows), MAP or mode predictions allow VJEPA to select a representative future consistent with the learned multi-modal predictive structure.

**Distributional prediction and uncertainty representation.** Rather than using a single predicted embedding, VJEPA represents future uncertainty explicitly through the predictive distribution $p_\phi(Z_T | Z_C, \xi_T)$. Downstream tasks may operate directly on this distribution, either by propagating summary statistics (e.g. moments) when $p_\phi$ admits a closed form, or by approximating expectations using samples. In this sense, VJEPA produces a *belief* over future representations rather than a point forecast. A general and flexible way to work with the predictive belief is via Monte Carlo sampling:

$$Z_T^{(m)} \sim p_\phi(Z_T | Z_C, \xi_T), \quad m = 1, \dots, M.$$

where $M$ is the total number of samples. These samples represent alternatively plausible future representations consistent with the observed context and target specification. Expectations of downstream quantities can then be approximated as

$$\mathbb{E}[f(Z_T)] \approx \frac{1}{M} \sum_{m=1}^{M} f(Z_T^{(k)}),$$

allowing uncertainty-aware prediction even when $p_\phi$ is complex (e.g. mixtures or normalizing flows).
**Multi-step prediction and uncertainty propagation.** For sequential prediction, VJEPA propagates uncertainty by integrating the predictive model in latent space. Given a belief $p(Z_t)$, the predictive belief at a future step $t + \Delta$ is

$$p(Z_{t+\Delta}) = \int p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}) p(Z_t) dZ_t, \qquad (12)$$

which defines a belief-state evolution analogous to *Bayesian filtering*. In practice, this integral may be approximated by moment propagation (when available) or by Monte Carlo rollouts:

$$Z_{t+1}^{(m)} \sim p_{\phi}(Z_{t+1} | Z_t^{(m)}, \xi_{t+1}), \quad m = 1, \dots, M,$$

which yields a distribution over latent trajectories rather than a single deterministic path. This form of distributional rollout underlies the planning and control algorithms introduced in later sections, e.g. model predictive control (MPC).

**Connection to ensemble world models.** Sampling-based prediction in VJEPA is closely related to ensemble world models commonly used in model-based reinforcement learning. In ensemble methods, multiple deterministic dynamics models {$f_{\phi(m)}$}$_{m=1}^M$ are trained independently, and uncertainty is approximated by disagreement across model predictions. In contrast, VJEPA represents uncertainty within a single probabilistic model $p_{\phi}(Z_T | Z_C, \xi_T)$, where stochasticity arises from an explicit predictive distribution rather than parameter diversity. Monte Carlo sampling from VJEPA therefore plays a role analogous to sampling from an ensemble, while avoiding the computational overhead and calibration issues associated with maintaining multiple models. This perspective highlights VJEPA as a distributional generalization of ensemble world models operating directly in representation space.

## 4.7 Collapse Avoidance

Here we formalize the claim that VJEPA discourages representational collapse.

**Definition 1 (Collapse).** We say the context representation collapses if there exists a constant $c \in \mathbb{R}^d$ such that

$$f_{\theta}(x_C) = c \quad \forall x_C.$$

**Theorem 1 (No Collapsed Global Optimum under Target Diversity).** Consider the VJEPA objective in Eq. (11). Assume:

(i) **(Target diversity)** There exist target inputs $x_T$ and $x'_T$ such that $q_{\theta'}(\cdot | x_T) \neq q_{\theta'}(\cdot | x'_T)$ (as distributions).

(ii) **(Nontrivial conditioning)** The predictive family {$p_{\phi}(\cdot | Z_C, \xi_T)$} is such that, for some $Z_C \neq Z'_C$ and the same $\xi_T$, the induced distributions can differ: $p_{\phi}(\cdot | Z_C, \xi_T) \neq p_{\phi}(\cdot | Z'_C, \xi_T)$ for some $\phi$.

Then any globally optimal solution of the VJEPA training objective Eq. (11) is non-collapsed, i.e. $f_{\theta}(x_C)$ cannot be constant over all contexts at a global minimum.
**Proof sketch.** If $f_{\theta}(x_C) \equiv c$, then the predictor cannot use the context: for each $\xi_T$ it reduces to an unconditional distribution $p_{\phi}(Z_T | c, \xi_T)$. The best possible choice of such an unconditional predictor (for fixed $\xi_T$) minimizes the cross-entropy with the aggregated target distribution $q_{\theta'}(Z_T | \xi_T) := \mathbb{E}_{x_T|\xi_T}[q_{\theta'}(Z_T | x_T)]$, hence achieves

$$ \inf_{p} \mathbb{E}_{x_T|\xi_T} \mathbb{E}_{Z_T \sim q_{\theta'}(\cdot|x_T)}[-\log p(Z_T)] = H_q(Z_T | \xi_T), $$

which equals $H_q(Z_T | X_T, \xi_T) + I_q(Z_T; X_T | \xi_T)$, where $I(\cdot; \cdot)$ denotes mutual information (see Appendix.B). By Assumption (i), we have $I_q(Z_T; X_T | \xi_T) > 0$ (for at least one $\xi_T$ with nonzero probability), so the collapsed predictor incurs a strictly positive irreducible gap relative to a predictor that can condition on informative $Z_C$. Assumption (ii) ensures the model class can exploit non-collapsed $Z_C$ to reduce the first term in Eq. (11). The KL regularizer further discourages degenerate target encodings by penalizing pathological $q_{\theta'}(\cdot | x_T)$, but is not the main source of the strict gap. A more detailed proof can be found in Appendix.F. □

*Remark 1.* Theorem 1 shows that collapse is not merely discouraged but is fundamentally incompatible with global optimality of the VJEPA objective (Eq. (11)) under target diversity. A collapsed context representation forces the predictor to ignore informative conditioning, yielding an irreducible excess loss relative to non-collapsed solutions. Therefore, collapse avoidance in VJEPA arises from the structure of the predictive objective itself, rather than relying solely on architectural asymmetry, stop-gradient operations, or EMA updates.

# 5 Time-indexed VJEPA as a Latent Space Dynamical System

VJEPA always learns a conditional predictive relation in representation space of the form (Eq. (7)):

$$ Z_T \sim p_{\phi}(Z_T | Z_C, \xi_T), $$

where $Z_C$ is a context representation, $Z_T$ is a target representation, and $\xi_T$ specifies structural information about the target (e.g. spatial location, temporal index, or masking pattern). However, whether this predictive relation constitutes a latent dynamical system depends on how the context-target pair is constructed and how $\xi_T$ is interpreted. JEPA is therefore not inherently a sequential or observation-likelihood model; in particular, JEPA does not inherently impose sequential or Markovian structure. This distinction is important for understanding both the capabilities and the limitations of existing JEPA variants, including Video JEPA [8, 9].

Here we formalize a *conditional dynamical-system* perspective on VJEPA, introducing a *latent dynamical system* in representation space when the context and target are indexed by time. Under this time-indexed formulation, VJEPA defines a latent state, a predictive transition model, and a belief-propagation mechanism, enabling filtering-style updates and control directly in latent space. Table.2 summarizes the different regimes in which JEPA-style objectives<fnref n="1" /> are commonly used.

<footnote n="1">¹The temporal formulation in this section also applied to JEPA.</footnote>
<table><thead><tr><th></th><th>What is Predicted</th><th>Structure Learned</th><th>Dynamical System?</th></tr></thead><tbody><tr><td>Image JEPA (I-JEPA [5])</td><td>Masked image regions</td><td>Spatial / semantic relations</td><td>×</td></tr><tr><td>Clip-level V-JEPA [8]</td><td>Masked video blocks</td><td>Intra-clip spatiotemporal correlations</td><td>×</td></tr><tr><td><strong>Time-indexed JEPA</strong></td><td>Z<sub>t+Δ</sub> from Z<sub>t</sub></td><td>Latent state transitions across time</td><td>✓</td></tr></tbody></table>

Table 2: JEPA always learns conditional predictive relations, but it induces a latent dynamical system only when contexts and targets are explicitly indexed by time and predictions are reused across steps. Clip-level V-JEPA captures short-range correlations but does not, by itself, define a compositional state transition model.

In this section, we focus on the *time-indexed* formulation of JEPA/VJEPA, under which the learned predictive model can be interpreted as a latent state-space model. This perspective will later justify belief propagation, filtering-style updates, and planning directly in representation space, while remaining fully consistent with the original JEPA training objective.

## 5.1 Background: State-Space Models

A (controlled) state-space model posits latent states {$s_t$}$_{t=1}^T$ and observations {$x_t$}$_{t=1}^T$ with a transition model and an observation model:

$$p(s_{1:T}, x_{1:T} | u_{1:T-1}) = p(s_1) \prod_{t=1}^{T-1} p(s_{t+1} | s_t, u_t) \prod_{t=1}^{T} p(x_t | s_t), \quad (13)$$

where $u_t$ is a control/action input. Classical HMMs (see Appendix.C) are the special case where $s_t$ is discrete and transitions/emissions belong to parametric families; modern nonlinear state-space models allow continuous $s_t$ and often use approximate inference.

The defining feature of (13) is that prediction and control can be performed by propagating a *belief* over latent states $p(s_t | x_{\le t}, u_{<t})$ forward through the transition model (also sometimes termed the *transition kernel*).

## 5.2 From JEPA Masking to Temporal Dynamics

JEPA is defined by choosing, for each sample $x$, a context-target partition ($x_C, x_T$) and a target specification $\xi_T$ (e.g. the spatial location, temporal index, or masking pattern, etc). In sequential settings, $x$ can represent a window of a video or a trajectory segment, and the target can correspond to a future time index. Concretely, for a temporal horizon $\Delta$, one may take

$$x_C \equiv x_{\le t} \quad \text{and} \quad x_T \equiv x_{t+\Delta}, \quad \xi_T \equiv (t + \Delta, \Delta),$$

so that predicting the target embedding corresponds to predicting a (arbitrary) future latent state $^2$. Let the JEPA encoder define the latent state

$$Z_t := f_{\theta}(x_{\le t}), \quad (14)$$

where $x_{\le t}$ denotes the available history (or the context crop/patch/window). Similarly, the target encoder produces a representation of the target observation

$$Z_{t+\Delta} := f_{\theta'}(x_{t+\Delta}),$$

used as the training target (with EMA parameters $\theta'$ as in Eq. (3)).

<footnote>$^2$We term the time-indexed JEPA a *sequential* (or *temporal*) JEPA. A characteristic feature of JEPA is that it can be trained to predict representations at a *chosen* future horizon $\Delta$ directly (i.e. skip-time prediction), rather than only learning one-step transitions that must be iterated. This capability is complementary to (not incompatible with) Markovian and diffusion models, and it does not require an autoregressive likelihood factorization over observations.</footnote>
## 5.3 Latent Transition Model Induced by JEPA/VJEPA

In deterministic JEPA, the predictor $g_{\phi}$ induces a latent transition (prediction)

$$ \hat{Z}_{t+\Delta} = g_{\phi}(Z_t, \xi_{t+\Delta}), \qquad (15) $$

which can be read as a deterministic latent dynamics model in representation space. In VJEPA, this becomes an explicit stochastic transition:

$$ Z_{t+\Delta} \sim p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}), \qquad (16) $$

which is exactly the state-transition component of a (latent) state-space model, except that it is learned without requiring an observation likelihood. Note that, JEPA/VJEPA allows $\Delta > 1$, i.e. predicting/transitioning to an arbitrary future (or masked) latent state.

To accommodate control, we let $u_{t:t+\Delta-1}$ denote control inputs between $t$ and $t+\Delta$. Then the most general controlled version of (16) is

$$ Z_{t+\Delta} \sim p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1}), \qquad (17) $$

where $\xi_{t+\Delta}$ can include the time index and horizon $\Delta$.

**Optional observation model.** For completeness, one may optionally posit a decoder (as in Eq. (8))

$$ x_{t+\Delta} \sim p_{\psi}(x_{t+\Delta} | Z_{t+\Delta}), \qquad (18) $$

which would turn the latent process in Eq. (16) into a fully generative state-space model. In this case, the implied generative structure is

$$ Z_t \rightarrow Z_{t+\Delta} \rightarrow x_{t+\Delta}, $$

i.e. future observations are generated from future latent states predicted from the current latent state. This contrasts with autoregressive observation models, which directly factorize

$$ Z_t \rightarrow x_t \rightarrow x_{t+1}, $$

and therefore require modeling the full conditional distribution $p(x_{t+1} | x_{\le t})$.

Importantly, the central JEPA/VJEPA design principle is that learning proceeds without optimizing the observation-level likelihood term $\log p_{\psi}(x_{t+\Delta} | Z_{t+\Delta})$. The decoder is indexed at $t+\Delta$ to mirror the skip-time prediction structure of JEPA (Eq. (16)); when $\Delta = 1$, this formulation reduces to the standard one-step state-space model.

## 5.4 Sequential but Not Autoregressive

We clarify here a common confusion of equating “sequential modeling” with “autoregressive modeling”. Autoregression refers specifically to a likelihood factorization over observations:

$$ p(x_{1:T}) = \prod_{t=1}^{T} p(x_t | x_{<t}), \qquad (19) $$

which underlies a wide class of generative sequence models. This includes classical probabilistic models such as Hidden Markov Models (HMMs, see Appendix.C) [60] and linear Gaussian state-space models (Kalman filters) [37], as well as modern deep generative models such as diffusion
models [34]. Although HMMs and Kalman filters introduce latent states, marginalizing these states still yields an autoregressive likelihood over observations $^3$. Optimizing Eq. (19) forces the model to account for the full entropy of observations (e.g. textures, noise, syntax, etc), including aspects irrelevant to prediction for downstream tasks such as control.

Instead, JEPA/VJEPA defines a predictive model in latent space (Eq. (17)):

$$p(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1}), \qquad (20)$$

and learns representations by matching predicted latent states to target-encoder latents. Importantly, this does not require specifying or factorizing an observation likelihood $p(x_{t+1} | x_{\le t})$. Thus, JEPA can be sequential (it supports multi-step prediction and belief propagation in $Z$) without being autoregressive in $x$.

**Belief propagation in latent space.** Given a belief $p(Z_t)$, sequential prediction (and decision) proceeds by pushing the belief through the latent transition:

$$p(Z_{t+\Delta}) = \int p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1}) p(Z_t) dZ_t, \qquad (21)$$

matching Eq. (12) and making explicit the analogy to filtering and planning.

## 5.5 Clip-Level Prediction (V-JEPA) vs. Sequential Dynamics

It is important to distinguish between latent prediction and latent dynamics. In its current implementation, V-JEPA [8, 9] is trained to predict masked spatio-temporal regions within a fixed video clip, using a large masked context that may already contain information from both past and future frames [8]. As a result, the learned predictor is not explicitly constrained to represent a compositional, step-by-step temporal transition model.

This design choice has been noted in recent discussions [9] of V-JEPA: while V-JEPA captures rich temporal cues and yields strong representations for downstream video understanding tasks, it does not enforce sequential consistency across multiple prediction steps, nor does it learn a Markovian latent transition operator that can be iterated forward in time. In particular, predictions at different horizons are made independently, and the training objective does not require that

$$Z_{t+2} \approx p_{\phi}(\cdot | Z_{t+1}, \xi_{t+2}) \quad \text{with} \quad Z_{t+1} \sim p_{\phi}(\cdot | Z_t, \xi_{t+1}).$$

The dynamical-system perspective adopted here makes this distinction explicit. By parameterizing JEPA/VJEPA with a time-indexed context $Z_t = f_{\theta}(x_{\le t})$ and a horizon-dependent target specification $\xi_{t+\Delta}$, the predictor (Eq. (16)) $p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta})$ can be interpreted as a latent transition model. When this transition is applied recursively, it induces a belief-state evolution analogous to classical state-space models, enabling filtering, multi-step prediction, and planning directly in representation space.

<footnote>$^3$Many classical sequence models are autoregressive at the observation level, even when latent states are introduced. For example, a HMM defines $p(s_{1:T}, x_{1:T}) = p(s_1) \prod_{t=2}^{T} p(s_t | s_{t-1}) \prod_{t=1}^{T} p(x_t | s_t)$, which implies, after marginalizing latent states $\{s_t\}$, $p(x_{1:T}) = \prod_{t=1}^{T} p(x_t | x_{<t})$. Linear Gaussian state-space models (Kalman filters) similarly assume latent dynamics $z_{t+1} = Az_t + \epsilon_t, x_t = Cz_t + \delta_t$, and also induce an autoregressive observation likelihood $p(x_{1:T}) = \prod_t p(x_t | x_{<t})$. Modern generative models such as diffusion models [34] and autoregressive transformers [81] make this factorization explicit. Thus, autoregression is a property of observation likelihood modeling, not of temporal dynamics. JEPA introduces temporal structure by predicting latent representations directly, without defining or factorizing an observation-level likelihood.</footnote>
From this perspective, clip-level V-JEPA can be seen as learning a powerful *one-shot latent predictor*, whereas the sequential formulation studied here extends JEPA/VJEPA into a latent dynamical system. This extension does not contradict the original V-JEPA design, but rather clarifies the additional structural assumptions required to support temporal composition, uncertainty propagation, and control.

**Summary.** Time-indexed JEPA/VJEPA can be viewed as learning a latent state-space model in which: (i) the encoder defines the latent state $Z_t$, (ii) the predictor defines a latent transition model, and (iii) optional observation modeling is decoupled from representation learning. This is the mathematical sense in which JEPA is sequential but not autoregressive.

# 6 JEPA/VJEPA for Control

Having interpreted JEPA/VJEPA as defining a latent dynamical system in representation space, we now show how such latent predictive states can be used for *planning and control*. The key message is classical in partially observable control: *optimal control does not require reconstructing observations, but only a sufficient belief (information state) for predicting future consequences of actions*.

## 6.1 Background: POMDPs and Belief-State Control

A partially<fnref n="4" /> observable Markov decision process (POMDP) is given by

$$ \mathcal{M} = (S, \mathcal{U}, \mathcal{X}, P, R, \gamma), $$

with latent state $s_t \in S$, action $u_t \in \mathcal{U}$, observation $x_t \in \mathcal{X}$, transition kernel $P(s_{t+1} | s_t, u_t)$, reward (utility, or negative cost) $R(s_t, u_t)$, and discount $\gamma \in (0, 1)$. The agent does not observe $s_t$ directly and must act based on history

$$ h_t := (x_{1:t}, u_{1:t-1}). $$

A standard result is that an *information state* (a Bayesian belief state) is sufficient for optimal control [36, 72, 73, 88]:

$$ b_t(s) := p(s_t = s | h_t), $$

then there exists an optimal policy that depends on history only through $b_t$:

$$ \pi^*(u_t | h_t) = \pi^*(u_t | b_t). \qquad (22) $$

Thus, control reduces to propagating and planning in belief space rather than observation space.

The sufficiency of the *belief state* (or *information state*) for optimal control is a foundational result established across the control theory and operations research literature. Åström (1965) [88] originally demonstrated that the conditional probability distribution of the latent state serves as a sufficient statistic for the complete history of actions and observations, thereby transforming the

<footnote n="4">⁴POMDP [48, 36] generalizes an MDP to settings where the agent cannot directly observe the true system state. The environment evolves according to Markovian dynamics in a latent state space, but the agent only receives observations generated probabilistically from the underlying state. As a result, the agent must maintain a *belief state*, which is a probability distribution over latent states, based on the history of observations and actions. Policies in a POMDP therefore map from belief states (or equivalently observation histories) to actions, rather than from fully observed states as in an MDP.</footnote>
partially observable problem into a completely observable Markovian problem over the space of distributions. This theoretical basis was operationally expanded by Smallwood and Sondik (1973) [72] and Sondik (1978) [73], who proved the existence of optimal policies that depend exclusively on this "information vector" for both finite and infinite horizon settings, deriving the associated dynamic programming equations. Synthesizing these classical results for artificial intelligence planning, Kaelbling et al. (1998) [36] re-affirm that because the belief state encapsulates all necessary historical information, the control problem formally reduces to solving a "belief MDP" which is a fully observable Markov decision process defined over the continuous space of probability distributions. These collective work confirm that the optimal policy need only map the current belief to an action $\pi^*(u_t | b_t)$, and justify the reduction of general history-based planning to belief-space planning.

## 6.2 JEPA/VJEPA Latent State as a Predictive Information State

JEPA/VJEPA does not explicitly model $p(x_t | s_t)$ and therefore does not construct the Bayesian belief $b_t$ in $S$. Instead, it learns a latent representation $Z_t$ from available context/history and equips it with a predictive transition model (Section 5):

$$Z_t = f_{\theta}(x_{\le t}), \quad Z_{t+\Delta} \sim p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1}).$$

The control-relevant requirement is not that $Z_t$ reconstructs $x_t$, but that it is *predictively sufficient* for evaluating future outcomes under candidate actions.

**Predictive sufficiency (latent information state).** We call $Z_t$ a *predictive information state* for control (over horizon $H$) if for any action sequence $u_{t:t+H-1}$, the conditional distribution of future task variables depends on history only via $Z_t$. A convenient instantiation is to require that $Z_t$ predicts future latent representations:

$$p(Z_{t+1:t+H} | h_t, u_{t:t+H-1}) = p(Z_{t+1:t+H} | Z_t, u_{t:t+H-1}), \quad (23)$$

at least approximately under the learned model. This factorization is precisely the notion of predictive sufficiency formalized in the next subsection, where we show that such a representation is sufficient for optimal control.

**Why reconstruction is unnecessary.** In model-based control, planning means using a predictive model to compare candidate actions by their anticipated consequences<fnref n="5" />. If the task objective can be expressed in terms of predictive features (for example via a learned cost head on top of $Z$), then planning can be performed entirely in latent space. An observation decoder $p_{\psi}(x_t | Z_t)$ is therefore optional: it may aid visualization or auxiliary tasks, but it is not required for the control computation itself. While JEPA/VJEPA does not *guarantee* that $Z_t$ is a sufficient information state without additional assumptions, it is designed to *encourage* predictive sufficiency by learning an amortized predictive state together with a transition model $p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1})$ across horizons. Under standard realizability conditions, namely that the existence of a true predictive information state $\tilde{Z}_t$ (e.g. a belief state or predictive state representation) that renders future task-relevant variables conditionally independent of history, and that $q_{\theta'}$ and $p_{\phi}$ are sufficiently expressive and correctly specified, a near-optimal VJEPA solution can make $Z_t = f_{\theta}(x_{\le t})$

<footnote n="5">⁵Concretely, a planner simulates candidate futures under different action sequences and selects the action (or sequence) that optimizes an objective such as expected reward or cost.</footnote>
an information-preserving transformation of $\tilde{Z}_t$ (up to approximation error). Consequently, $Z_t$ inherits the control-relevant predictive sufficiency of $\tilde{Z}_t$ [46, 14]. These assumptions are standard in predictive-state and representation-learning analyses, so JEPA/VJEPA can be viewed as replacing explicit belief-state estimation with representation learning while preserving the information needed for optimal control.

## 6.3 Predictive Sufficiency for Control (Formalization)

We now formalize the notion of predictive sufficiency introduced above and state the standard control-theoretic consequence for optimal policies. Consider a POMDP with latent state $s_t$, observation $x_t$, and control $u_t$ (Section 6). Let the agent's history be $h_t = (x_{1:t}, u_{1:t-1})$, and let the (post-action) stage cost be $c_{t+1} = c(s_{t+1}, u_t)$ (or any task cost measurable with respect to the environment's trajectory).

**From latent predictive sufficiency to control-relevant sufficiency.** In control, we care about predicting cost-relevant consequences of actions. If the stage cost can be written as a (possibly learned) function of latent states, e.g. $c_{t+k+1} = \ell(Z_{t+k+1}, u_{t+k})$ (or $\ell(h(Z_{t+k+1}), u_{t+k})$ for a cost head $h$), then the latent factorization in Eq. (23) implies that the conditional law of any cumulative cost over horizon $H$ depends on history only through $Z_t$. This motivates the following definition:

**Definition 2 (Control-Relevant Predictive Sufficiency).** A representation $Z_t = f_{\theta}(x_{\le t})$ is (control-)predictively sufficient over horizon $H$ if for any action sequence $u_{t:t+H-1}$, the conditional distribution of future cumulative cost depends on history only through $Z_t$:

$$p\left(\sum_{k=0}^{H-1} c(s_{t+k+1}, u_{t+k}) \mid h_t, u_{t:t+H-1}\right) = p\left(\sum_{k=0}^{H-1} c(s_{t+k+1}, u_{t+k}) \mid Z_t, u_{t:t+H-1}\right).$$

**Lemma 1 (Latent sufficiency implies cost sufficiency).** Assume the stage cost is measurable with respect to the latent trajectory. If Eq. (23) holds for horizon $H$, then $Z_t$ satisfies Definition 2.

This lemma shows that latent predictive sufficiency implies control-relevant predictive sufficiency whenever the task cost is measurable with respect to the learned latent state.

**Proof sketch.** Eq. (23) implies the conditional distribution of the latent trajectory $Z_{t+1:t+H}$ given $(h_t, u_{t:t+H-1})$ depends on history only via $Z_t$. Any measurable functional of that trajectory, including the cumulative cost, therefore has the same property. $\square$

**Theorem 2 (Sufficiency for Optimal Control from a Predictive Information State).** If $Z_t$ is control-predictively sufficient in the sense of Definition 2 for all horizons $H$ (or for the planning horizon used by the controller), then there exists an optimal policy that depends on history only through $Z_t$:

$$\pi^*(u_t | h_t) = \pi^*(u_t | Z_t).$$

**Proof sketch.** The statement follows the standard POMDP information-state argument: if $Z_t$ renders the conditional law of all future (cost-relevant) consequences of action sequences independent of the full history, then dynamic programming (or finite-horizon rollout evaluation) can be carried out using $Z_t$ as the state variable. Hence an optimal policy can be chosen Markov in $Z_t$, i.e. $\pi^*(u_t | h_t) = \pi^*(u_t | Z_t)$. $\square$
**Relation to VJEPA training.** VJEPA does not, by itself, guarantee the predictive sufficiency property in Definition 2 without additional realizability and optimization assumptions. Rather, it is designed to encourage such sufficiency by learning a latent predictive state together with a predictive model $p_{\phi}(Z_{t+\Delta} | Z_t, \xi_{t+\Delta}, u_{t:t+\Delta-1})$ that supports latent-space planning over multiple horizons. Under standard realizability assumptions, namely: (i) that there exists a true latent predictive information state $\tilde{Z}_t$ (e.g. a Bayesian belief state or a predictive state representation in the sense of PSRs [46, 14]) rendering future task-relevant variables conditionally independent of history; (ii) that the target encoder family $q_{\theta'}(Z_T | x_T)$ is expressive enough to represent the true conditional distribution of future predictive states given observations; and (iii) that the predictive model family $p_{\phi}$ is well specified and contains the true latent transition kernel induced by $\tilde{Z}_t$, then if VJEPA training reaches a (near-)global optimum over a sufficiently rich set of prediction horizons $\Delta$, the learned representation $Z_t = f_{\theta}(x_{\le t})$ becomes an information-preserving (invertible up to approximation error) transformation of $\tilde{Z}_t$. Consequently, $Z_t$ inherits the predictive sufficiency of $\tilde{Z}_t$ for evaluating future costs under candidate action sequences, and Theorem 2 applies. In this sense, VJEPA can be viewed as learning an amortized predictive information state whose sufficiency for control follows from correct specification and optimization, rather than from observation reconstruction or explicit belief-state estimation. As in PSR and belief-state analyses, these assumptions characterize an idealized regime; in practice VJEPA is expected to approximate predictive sufficiency to the extent allowed by data, capacity, and optimization.

## 6.4 Latent-Space Planning Objective

**Planning in POMDPs.** In a POMDP, optimal control cannot be expressed directly as a function of the unobserved latent state $s_t$. Instead, planning must be carried out in terms of the belief state

$$b_t(s) := p(s_t = s | x_{1:t}, u_{1:t-1}),$$

which summarizes all information available to the agent. Given a transition model $P(s_{t+1} | s_t, u_t)$, an observation model $p(x_t | s_t)$, and a reward (or cost) function, the POMDP induces a belief-space MDP, in which actions are selected to maximize expected cumulative reward under the evolution of beliefs. Planning in a POMDP therefore amounts to simulating how candidate action sequences propagate the belief distribution forward in time and selecting the action that optimizes expected return [72, 36, 66]:

$$\pi^*(u_t | h_t) = \pi^*(u_t | b_t).$$

This belief-space formulation is exact but typically intractable in high-dimensional or continuous settings [57, 50, 41, 74].

**Latent belief approximation via VJEPA.** VJEPA replaces the intractable Bayesian belief $b_t$ with a learned latent representation $Z_t = f_{\theta}(x_{\le t})$ that serves as an amortized, predictive belief state<fnref n="6" />. The predictive model $p_{\phi}(Z_{t+1} | Z_t, \xi_{t+1}, u_t)$ then plays the role of a belief transition operator, enabling planning directly in latent space without requiring an explicit observation likelihood.

**Latent-space planning objective.** Let $c(Z_{t+1}, u_t)$ be a stage cost defined on latent states (or via a cost head applied to $Z_{t+1}$). Given the current latent state $Z_t$, the H-step planning problem

<footnote n="6">⁶Here ‘belief’ refers to a predictive distribution over latent representations rather than a posterior over true environment states.</footnote>
is

$$u_{t:t+H-1}^* \in \arg\min_{u_{t:t+H-1}} \mathbb{E} \left[ \sum_{k=0}^{H-1} c(Z_{t+k+1}, u_{t+k}) \middle| Z_t, u_{t:t+H-1} \right], \quad (24)$$

where the expectation is taken under the latent dynamics induced by the VJEPA predictive model<fnref n="7" />:

$$Z_{t+k+1} \sim p_{\phi}(Z_{t+k+1} | Z_{t+k}, \xi_{t+k+1}, u_{t+k}).$$

This objective is formally equivalent to stochastic model predictive control (MPC [17, 62, 53]), but instantiated entirely in representation space.

**Distributional planning.** When $p_\phi$ is stochastic, as in VJEPA, the objective (24) involves an expectation over latent trajectories. This enables distributional planning [53, 69]: risk-neutral control arises naturally, and risk-sensitive criteria [15] (e.g. variance penalties [42, 53] or CVaR [64, 20]) can be incorporated by modifying the functional applied to the trajectory distribution.

## 6.5 Sampling vs. Point Prediction Planning

The choice of how to use the predictive distribution is central in practice.

**MAP/mean (trajectory) planning.** A computationally cheap approximation is to plan using a single representative trajectory, e.g. the predictive mean (or MAP point when well-defined):

$$\hat{Z}_{t+1} = \mu_{\phi}(Z_t, u_t) \quad \text{or} \quad \hat{Z}_{t+1} = \arg\max_{Z_{t+1}} p_{\phi}(Z_{t+1} | Z_t, u_t).$$

This reduces planning to deterministic rollout and is often effective when the predictive distribution is unimodal and uncertainty is weak or approximately action-independent.

**Sampling-based (distributional) planning.** Alternatively, one can approximate (24) by Monte Carlo rollouts:

$$Z_{t+k+1}^{(m)} \sim p_{\phi}(\cdot | Z_{t+k}^{(m)}, u_{t+k}), \quad m = 1, \dots, M,$$

and estimate the expected cost by sample averaging. $M$ is the number of samples drawn at each time step. This is useful in multi-modal dynamics or partial observability, where committing to a single MAP/mean trajectory can ignore low-probability but potentially high-cost outcomes. Sampling also naturally supports robust or risk-sensitive planning by changing how samples are aggregated.

## 6.6 Algorithm: VJEPA Sampling-Based Latent MPC

We now formalize the sampling-based planning procedure discussed in Section 6.5. This procedure, commonly referred to as random shooting or Model Predictive Path Integral (MPPI) control in the literature [86, 87], approximates the minimization of the expected cost (Eq. (24)) by simulating $N$ parallel latent trajectories using the VJEPA predictive model.

<footnote n="7">⁷This is learning-based latent-space MPC, see e.g. [40, 84, 29].</footnote>
**Algorithm 2 VJEPA-Based Model Predictive Control (VJEPA-MPC)**

**Require:** History $x_{\le t}$, Horizon $H$, Number of samples $M$, Cost function $c$

1: **Encode:** Estimate current predictive state:

$$Z_t = f_{\theta}(x_{\le t})$$

2: **for** sample $i = 1$ to $M$ **do**

3: Sample candidate action sequence $u_{t:t+H-1}^{(i)} \sim p(u)$

4: $Z_t^{(i)} \leftarrow Z_t$

5: **for** step $k = 0$ to $H-1$ **do**

6: Sample next latent state (dynamics rollout):

$$Z_{t+k+1}^{(i)} \sim p_{\phi}(Z_{t+k+1} | Z_{t+k}, u_{t+k}^{(i)})$$

7: **end for**

8: Compute cumulative cost for trajectory $i$:

$$J^{(i)} = \sum_{k=0}^{H-1} c(Z_{t+k+1}^{(i)}, u_{t+k}^{(i)})$$

9: **end for**

10: **Select:** Optimal action sequence index $i^* = \arg\min_i J^{(i)}$

11: **Execute:** Apply first action $u_t^{(i^*)}$ to environment

## 6.7 When is MAP Control Optimal?

We formalize a simple sufficient condition under which planning on a point prediction (mean or MAP) is equivalent to planning under the full predictive distribution. Following standard practice in stochastic MPC and belief-space control, we index the stage cost at the *post-action* state: applying action $u_t$ induces a transition $Z_t \to Z_{t+1}$, and the cost evaluates the consequence of this action on the resulting state [53, 62, 12].

**Theorem 3 (Optimality of MAP (Mean) Control under Quadratic Costs).** Consider a one-step control problem with latent predictive model $p_{\phi}(Z_{t+1} | Z_t, u_t)$ and stage cost $c(Z_{t+1}, u_t)$. Assume that for each $(Z_t, u_t)$,

$$p_{\phi}(Z_{t+1} | Z_t, u_t) = \mathcal{N}(Z_{t+1} | \mu_{\phi}(Z_t, u_t), \Sigma),$$

where the covariance $\Sigma$ does not depend on $u_t$. Let the stage cost be quadratic in the next latent state:

$$c(Z_{t+1}, u_t) = (Z_{t+1} - z^*)^\top Q_c(Z_{t+1} - z^*) + r(u_t), \quad Q_c \succeq 0,$$

with arbitrary $r(\cdot)$. Then the action minimizing expected cost,

$$u_t^* \in \arg\min_{u_t} \mathbb{E}[c(Z_{t+1}, u_t) | Z_t, u_t],$$

is equivalently obtained by minimizing the cost at the predictive mean:

$$u_t^* \in \arg\min_{u_t} c(\mu_{\phi}(Z_t, u_t), u_t).$$
Moreover, since a Gaussian distribution has its unique MAP point at the mean, this is equivalently MAP control.

*Proof.* Using the quadratic form and standard Gaussian moment identities,

$$ \mathbb{E}[(Z_{t+1} - z^*)^\top Q_c(Z_{t+1} - z^*)] = (\mu_\phi - z^*)^\top Q_c(\mu_\phi - z^*) + \text{tr}(Q_c \Sigma), $$

where the trace term tr($Q_c \Sigma$) does not depend on $u_t$ by assumption. Hence

$$ \arg\min_{u_t} \mathbb{E}[c(Z_{t+1}, u_t)] = \arg\min_{u_t} ((\mu_\phi - z^*)^\top Q_c(\mu_\phi - z^*) + r(u_t)) = \arg\min_{u_t} c(\mu_\phi(Z_t, u_t), u_t). $$

Finally, for a Gaussian predictive distribution, the MAP point coincides with the mean $\mu_\phi(Z_t, u_t)$. A more detailed proof is provided in Appendix G. □

**Discussion.** Theorem 3 captures a common regime in which point-prediction planning is justified: predictive uncertainty contributes only an additive constant to the objective and therefore does not affect the optimizer. When the predictive covariance depends on the action (risk-sensitive settings), when costs are non-quadratic, or when $p_\phi$ is multi-modal (non-Gaussian), sampling-based rollouts become important, as optimizing a single MAP/mean trajectory may ignore uncertainty and rare but costly outcomes.

**Corollary (Multi-step mean/MAP planning under action-independent covariance).** Consider an H-step planning problem in latent space with stochastic predictive dynamics

$$ Z_{t+k+1} | Z_{t+k}, u_{t+k} \sim \mathcal{N}(\mu_\phi(Z_{t+k}, u_{t+k}), \Sigma_k), \quad k = 0, \dots, H-1, $$

where each covariance $\Sigma_k$ is positive semidefinite and does not depend on the action sequence $u_{t:t+H-1}$. Assume the stage cost is quadratic in the next latent state:

$$ c(Z_{t+k+1}, u_{t+k}) = (Z_{t+k+1} - z^*)^\top Q_k (Z_{t+k+1} - z^*) + r_k(u_{t+k}), \quad Q_k \succeq 0. $$

Define the risk-neutral MPC objective<fnref n="8" />

$$ J(u_{t:t+H-1}) := \mathbb{E} \left[ \sum_{k=0}^{H-1} c(Z_{t+k+1}, u_{t+k}) \middle| Z_t, u_{t:t+H-1} \right]. $$

Assume additionally that the induced state covariances $Cov(Z_{t+k})$ under the rollout are independent of the chosen action sequence (e.g. in linear-Gaussian dynamics with additive, action-independent noise). Then the optimal action sequence

$$ u_{t:t+H-1}^* \in \arg\min_{u_{t:t+H-1}} J(u_{t:t+H-1}) $$

is equivalently obtained by deterministic planning on the mean (or MAP) trajectory:

$$ u_{t:t+H-1}^* \in \arg\min_{u_{t:t+H-1}} \sum_{k=0}^{H-1} c(\bar{Z}_{t+k+1}, u_{t+k}), $$

<footnote n="8">⁸Risk-neutral control refers to optimizing the expected cumulative cost $\mathbb{E}[\sum_t c(Z_t, u_t)]$ without explicit preference for or against uncertainty. Two action sequences with equal expected cost are treated as equally desirable, even if they differ in variance or tail risk. By contrast, risk-sensitive control incorporates uncertainty directly into the objective, for example via variance penalties, exponential utility functions, or tail-risk criteria such as CVaR [64, 20].</footnote>
where the mean trajectory {$\bar{Z}_{t+k}$}$_{k=0}^H$ is defined recursively by

$$
\bar{Z}_t := Z_t, \quad \bar{Z}_{t+k+1} := \mu_\phi(\bar{Z}_{t+k}, u_{t+k}).
$$

Since each Gaussian predictive distribution has its unique MAP point at the mean, this procedure
is equivalently MAP planning at every step.

Proof. Write $Z_{t+k+1} = \bar{Z}_{t+k+1} + \varepsilon_{t+k+1}$, where $\bar{Z}_{t+k+1} := \mathbb{E}[Z_{t+k+1}]$ under the chosen action sequence and $\varepsilon_{t+k+1}$ is a zero-mean deviation with covariance $\text{Cov}(Z_{t+k+1})$. For each $k$, expand the quadratic term:

$$
(Z_{t+k+1} - z^*)^\top Q_k (Z_{t+k+1} - z^*) = (\bar{Z}_{t+k+1} - z^*)^\top Q_k (\bar{Z}_{t+k+1} - z^*) + 2 \varepsilon_{t+k+1}^\top Q_k (\bar{Z}_{t+k+1} - z^*) + \varepsilon_{t+k+1}^\top Q_k \varepsilon_{t+k+1}.
$$

Taking expectations, the cross term vanishes because $\mathbb{E}[\varepsilon_{t+k+1}] = 0$, and

$$
\mathbb{E}[\varepsilon_{t+k+1}^{\top} Q_k \varepsilon_{t+k+1}] = \text{tr}(Q_k \text{Cov}(Z_{t+k+1})).
$$

Therefore,

$$
\mathbb{E} \left[ (Z_{t+k+1} - z^*)^\top Q_k (Z_{t+k+1} - z^*) \right] = (\bar{Z}_{t+k+1} - z^*)^\top Q_k (\bar{Z}_{t+k+1} - z^*) + \text{tr}(Q_k \text{Cov}(Z_{t+k+1})).
$$

Summing over $k$, the expected cost decomposes as

$$
J(u_{t:t+H-1}) = \sum_{k=0}^{H-1} \left( (\bar{Z}_{t+k+1} - z^*)^\top Q_k (\bar{Z}_{t+k+1} - z^*) + r_k(u_{t+k}) \right) + \sum_{k=0}^{H-1} \text{tr}(Q_k \text{Cov}(Z_{t+k+1})).
$$

Under the stated assumptions, the covariance terms are independent of the action sequence and thus constitute a constant offset with respect to optimization over $u_{t:t+H-1}$. Dropping this constant yields deterministic planning on the mean trajectory. Finally, since each conditional Gaussian $N(\mu_\phi(\cdot), \Sigma_k)$ has its MAP point at the mean, this is equivalently MAP planning, i.e.

$$
\arg \min_{u_{t:t+H-1}} J(u_{t:t+H-1}) = \arg \min_{u_{t:t+H-1}} \sum_{k=0}^{H-1} \left( (\bar{Z}_{t+k+1} - z^*)^\top Q_k (\bar{Z}_{t+k+1} - z^*) + r_k(u_{t+k}) \right).
$$

□

7 Information-Theoretic Analysis

In this section, we analyze VJEPA through the lens of information theory, specifically relating
the variational objective to the maximization of predictive mutual information and the Predictive
Information Bottleneck (PIB) principle [77, 13, 3, 2, 83, 82].

7.1 Variational Maximization of Predictive Information

We establish that minimizing the VJEPA loss is equivalent to maximizing a lower bound on the mutual information<fnref n="9" /> between the current latent state $Z_t$ and the future latent state $Z_{t+\Delta}$.

<footnote n="9">⁹In probability theory and information theory, mutual information (MI) measures the mutual dependence between two random variables, quantifying the information obtained about one by observing the other. Let (X, Y) be a pair of random variables over space X × Y with joint distribution P(X, Y) and marginals P_X, P_Y. The mutual information is defined as the KL divergence between the joint and product distributions: I(X; Y) = D_{KL}(P(X,Y) || P_X ⊗ P_Y). It can also be expressed via entropy H(·) and conditional entropy H(·|·) as: I(X; Y) = H(X) - H(X|Y) = H(Y) - H(Y|X). By the properties of KL divergence, I(X; Y) ≥ 0, with equality if and only if X and Y are independent. See Appendix.B for more details.</footnote>
**Theorem 4** (Variational Mutual Information Lower Bound). Let $(Z_t, Z_{t+\Delta})$ be the joint distribution of context and target representations induced by the data and the encoder policies. The mutual information $I(Z_t; Z_{t+\Delta})$ is lower-bounded by the negative cross-entropy (or expected log-likelihood) of the predictive distribution:

$$I(Z_t; Z_{t+\Delta}) \geq \mathbb{E}_{p(Z_t, Z_{t+\Delta})} [\log p_{\phi}(Z_{t+\Delta} | Z_t)] + H(Z_{t+\Delta}),$$

where $H(Z_{t+\Delta})$ is the marginal entropy of the target representations.

*Proof.* By definition, $I(Z_t; Z_{t+\Delta}) = H(Z_{t+\Delta}) - H(Z_{t+\Delta} | Z_t)$. The conditional entropy is defined as $H(Z_{t+\Delta} | Z_t) = -\mathbb{E}[\log p(Z_{t+\Delta} | Z_t)]$. Using the non-negativity of KL divergence, $D_{KL}(p(\cdot|Z_t) \| p_{\phi}(\cdot|Z_t)) \geq 0$, we have:

$$\mathbb{E}[\log p(Z_{t+\Delta} | Z_t)] \geq \mathbb{E}[\log p_{\phi}(Z_{t+\Delta} | Z_t)].$$

Substituting this inequality yields the result. This is known as the Barber-Agakov bound [7, 58]. □

**Implication.** The VJEPA objective (Eq. 11) minimizes $-\log p_{\phi}(Z_{t+\Delta} | Z_t)$. Since $H(Z_{t+\Delta})$ depends only on the target encoder (which evolves slowly via EMA), minimizing the VJEPA loss effectively maximizes the mutual information between the past and the future representations, $I(Z_t; Z_{t+\Delta})$.

## 7.2 Predictive Information Bottleneck

The Information Bottleneck (IB) method [77] provides a general information-theoretic framework for finding a compressed representation $Z$ of an input source $X$ that retains the maximum possible information about a relevant target variable $Y$. Formally, it seeks to minimize the functional

$$\mathcal{L}_{IB} = I(X; Z) - \beta I(Z; Y),$$

where $I(\cdot; \cdot)$ denotes mutual information<fnref n="10" /> and $\beta$ is a Lagrange multiplier controlling the trade-off between complexity (compression of $X$) and accuracy (prediction of $Y$). The Predictive Information Bottleneck (PIB) [13] specializes this principle to temporal data: it aims to extract a summary of the past ($X = x_{\le t}$) that is maximally predictive of the future ($Y = x_{t+\Delta}$), while discarding irrelevant, noisy, or redundant details.

A core distinction of VJEPA is its adherence to this predictive principle. Unlike autoencoders, which effectively maximize $I(Z_t; x_{\le t})$ by reconstructing the full input history, VJEPA acts as a PIB. It seeks to capture information about the future while remaining invariant to non-predictive details of the past.

Let $x_{\le t}$ be the observation history. We posit a generative process where observations consist of a predictable signal $S$ and nuisance noise $N$: $x = S + N$. The future $x_{t+\Delta}$ depends on history only through $S$.

**Proposition 1** (Invariance to Nuisance Variability). Let the input $x$ decompose into a predictive signal $S$ and a nuisance variable $N$ (where $N$ is independent of the future target $Z_{t+\Delta}$). Generative models maximizing $\log p(x | Z)$ must encode both $S$ and $N$ to minimize reconstruction error (since

<footnote n="10">¹⁰The mutual information $I(\cdot; \cdot)$ is a measure of mutual dependence between two random variables defined as the Kullback-Leibler divergence between their joint distribution and the product of their marginals: $I(X; Y) = D_{KL}(P_{(X,Y)} \| P_X \otimes P_Y) = H(X) - H(X|Y)$. See Appendix.B.</footnote>
the entropy $H(x)$ includes $H(N)$). In contrast, the VJEPA objective is invariant to representations that discard $N$ provided they preserve the mutual information with the future, $I(Z_t; Z_{t+\Delta})$. Consequently, VJEPA admits minimal sufficient statistics that filter out nuisance variability while maintaining optimal prediction loss. (See Proof in Appendix.H).

This proposition implies that VJEPA allows the representation to be minimal (compressing away noise), whereas reconstruction-based objectives force the representation to be maximal (retaining noise). This efficiency is critical for control, where acting on noise can lead to instability.

## 7.3 Contrast with Generative Objectives

We can now formally distinguish VJEPA from autoregressive (AR) world models [27, 29] by analyzing their respective training objectives through an information-theoretic lens.

**Derivation of the Generative Penalty.** Standard generative world models are trained to maximize the log-likelihood of future observations $x_{t+\Delta}$ given the current latent state $Z_t$. This corresponds to minimizing the Negative Log-Likelihood (NLL):

$$ \mathcal{L}_{AR} = \mathbb{E}_{data} [-\log p_{\psi}(x_{t+\Delta} | Z_t)]. $$

In the limit of infinite data and a sufficiently expressive model, minimizing the NLL is equivalent to minimizing the conditional entropy of the targets given the representation:

$$ \mathcal{L}_{AR} \approx H(x_{t+\Delta} | Z_t). $$

Using the fundamental identity relating entropy and mutual information, $H(X | Y) = H(X) - I(X; Y)$, we can rewrite the objective as:

$$ \mathcal{L}_{AR} \approx \underbrace{H(x_{t+\Delta})}_{\text{Constant}} - I(x_{t+\Delta}; Z_t). $$

Since the marginal entropy of the data $H(x_{t+\Delta})$ is a constant property of the environment (determined by the dataset), minimizing the generative loss is equivalent to maximizing the mutual information between the latent state and the raw pixels:

$$ \min \mathcal{L}_{AR} \iff \max I(x_{t+\Delta}; Z_t). $$

**Decomposition into Signal and Noise.** Consider the observation $x$ as a composition of task-relevant signal $S$ and nuisance noise $N$ (e.g., “Noisy TV” static, camera grain), such that $x_{t+\Delta} = (S_{t+\Delta}, N_{t+\Delta})$. By applying the chain rule for mutual information<fnref n="11" />, we can decompose the information the representation $Z_t$ holds about the full observation into two components:

$$ I(x_{t+\Delta}; Z_t) = I(S_{t+\Delta}, N_{t+\Delta}; Z_t) = I(S_{t+\Delta}; Z_t) + I(N_{t+\Delta}; Z_t | S_{t+\Delta}). $$

The first term, $I(S_{t+\Delta}; Z_t)$, measures how much the representation captures the predictable, task-relevant signal. The second term, $I(N_{t+\Delta}; Z_t | S_{t+\Delta})$, measures the information the representation retains about the noise, given that the signal is already known. If we assume the noise $N_{t+\Delta}$ is

<footnote n="11">¹¹The chain rule for mutual information states that, for any random variables $X, Y, Z$, the information provided by the pair $(X, Y)$ about $Z$ is: $I(X, Y; Z) = I(X; Z) + I(Y; Z | X)$.</footnote>
independent of the signal $S_{t+\Delta}$ (a common property of sensor noise), this second term identifies the specific capacity the model must allocate to capturing unpredictable nuisance variability.

Substituting this back into the loss function, we obtain the effective objective for AR models:

$$ \mathcal{L}_{\text{AR}} \approx -I(Z_t; S_{t+\Delta}) \underbrace{-I(Z_t; N_{t+\Delta})}_{\text{Penalty term}}. \qquad (25) $$

Importantly, the second term acts as a penalty: the loss increases if the representation $Z_t$ fails to capture information about the noise $N_{t+\Delta}$. To minimize $\mathcal{L}_{\text{AR}}$, the model is mathematically forced to allocate capacity to predict high-entropy, task-irrelevant details.

**The VJEPA Advantage.** In contrast to generative models that operate on pixels, VJEPA defines the prediction target in an abstract representation space $Z_{t+\Delta}$, generated by an encoder $f_{\theta'}$ designed to capture the predictable signal $S$ while discarding nuisance noise $N$. We can formally derive the VJEPA objective by considering the minimization of the latent negative log-likelihood (NLL):

$$ \mathcal{L}_{\text{VJEPA}} = \mathbb{E}[-\log p_{\phi}(Z_{t+\Delta} | Z_t)] \approx H(Z_{t+\Delta} | Z_t), \qquad (26) $$

where the expected NLL corresponds to the conditional entropy of the future latent state. Using the fundamental identity $H(X | Y) = H(X) - I(X; Y)$, we rewrite the objective as:

$$ \mathcal{L}_{\text{VJEPA}} \approx H(Z_{t+\Delta}) - I(Z_t; Z_{t+\Delta}). \qquad (27) $$

Under the assumption of a well-regularized encoder, maintained via the KL-regularization term shown in Fig.1, the marginal entropy $H(Z_{t+\Delta})$ is encouraged to remain high to prevent representation collapse. Consequently, the optimization reduces to maximizing the mutual information in the latent space: $\min_{\mathcal{L}_{\text{VJEPA}}} \Leftrightarrow \max I(Z_t; Z_{t+\Delta})$.

According to the Predictive Information Bottleneck (PIB) principle, an ideal target encoder should filter out unpredictable nuisances such that $Z_{t+\Delta} \approx f_{\theta'}(S_{t+\Delta})$. Substituting this signal-centric target yields the decomposed VJEPA advantage:

$$ \mathcal{L}_{\text{VJEPA}} \approx -I(Z_t; Z_{t+\Delta}) \approx -I(Z_t; S_{t+\Delta}). \qquad (28) $$

Because the target $Z_{t+\Delta}$ has already discarded the pixel-level noise $N_{t+\Delta}$, the VJEPA objective contains no penalty term for ignoring $N$. This allows the world model to filter out high-frequency observation noise before the prediction loss is computed, focusing its limited capacity directly on the causal dynamics relevant for planning and control.

# 8 Bayesian JEPA (BJEPA)

VJEPA models uncertainty directly in *representation space* by learning a conditional distribution over future latent states. For example, the predictive model Eq. (7) can be formulated as

$$ p_{\phi}(Z_T | Z_C) = \mathcal{N}(\mu_{\phi}(Z_C), \Sigma_{\phi}(Z_C)), \qquad (29) $$

which scales naturally to high-dimensional embeddings and aligns with the interpretation of $Z_t$ as a predictive information state.

In this section, we introduce Bayesian JEPA (BJEPA), which extends VJEPA by explicitly factoringizing the predictive model into independent factors: a learned dynamics expert and a constraint-based expert. This enables principled incorporation of goals, constraints, and task structure while preserving JEPA’s core design principles.
## 8.1 Bayesian Factorization via Product of Experts

BJEPA approximates the target posterior distribution using a Product of Experts (PoE) mechanism<fnref n="12" />. Strictly speaking, PoE is the computational operation used to fuse multiple density functions; however, we structure the experts such that their product corresponds to a Bayesian update.

We treat the historical context and the auxiliary constraints as conditionally independent sources of information regarding the future state, combining them via the product of their respective probability densities:

$$p(Z_T | Z_C, \eta) \propto p_{\text{like}}(Z_T | Z_C)p_{\text{prior}}(Z_T | \eta). \quad (30)$$

This formulation allows us to interpret the factors through a Bayesian lens:

* $p_{\text{like}}(Z_T | Z_C)$ acts as the predictive likelihood (or transition model). It is learned from data and encodes the system's typical latent dynamics conditioned on past context.

* $p_{\text{prior}}(Z_T | \eta)$ acts as a latent-space prior. It injects auxiliary information $\eta$, such as goal regions, safety constraints, or physical manifolds, into the inference process.

Mathematically, the product formulation implements a logical AND<fnref n="13" />: the predicted state must satisfy the dynamics and the constraints simultaneously. Just as a Kalman Filter update intersects the prediction with a measurement, BJEPA intersects physical possibility with task necessity. This effectively carves the latent space to the intersection of physically probable and task-compliant states, i.e. we get a predictive distribution that only has high probability in regions that are both physically reachable (dynamics) and closer to the goal (prior). As shown in Appendix I, this formulation<fnref n="14" /> is exact under the assumption that the context $Z_C$ and auxiliary info $\eta$ are conditionally independent given the target $Z_T$, up to a normalization constant and marginal scaling. Thus, BJEPA utilizes the PoE mechanism to realize a Bayesian inference process entirely within the representation space. Compared to a standard autoregressive (AR) model, we would have to search the entire predicted pixel space to find a path to a goal; in BJEPA, the PoE constraints this search within the intersected space<fnref n="15" />.

## 8.2 The BJEPA Architecture

BJEPA builds on the standard JEPA/VJEPA architecture by retaining the context/target encoder structure used for representation learning, and augmenting it with an explicit latent-space prior encoder. Concretely, BJEPA consists of the following components:

* a context (online) encoder $f_{\theta}$ that maps past observations to a latent context representation,

$$Z_C = f_{\theta}(x_C),$$

this is identical to the context encoder used in JEPA/VJEPA (Eq. (1));

<footnote n="12">¹²PoE is a machine learning framework that models a probability distribution by combining the output from several simpler probability models (experts), this is achieved by multiplying their probability distributions and renormalizing.</footnote>

<footnote n="13">¹³Unlike a mixture model, which effectively performs a logical OR.</footnote>

<footnote n="14">¹⁴The full Eq.(30) is: $p(Z_T | Z_C, \eta) = \frac{1}{\mathcal{Z}} (p_{\text{like}}(Z_T | Z_C) \times p_{\text{prior}}(Z_T | \eta))$. In the context of PoE, multiplying probabilities (element-wise multiplication of two PDFs) corresponds to a logical AND operation on constraints; the resulting distribution represents states that are consistent with both the dynamics (history) AND the auxiliary constraints (goals). After multiplication, the resulting function usually does not integrate to unity; a normalization step divides the product by the partition function $\mathcal{Z}$ to ensure the result is a valid probability distribution.</footnote>

<footnote n="15">¹⁵If both $p_{\text{like}}$ and $p_{\text{prior}}$ are Gaussian, the posterior can be derived analytically (also Gaussian), by simply adding the precisions (inverse variances) of the two experts in the latent space.</footnote>
* a target encoder $f_{\theta'}$ that maps future observations to target latent representations,

$$Z_T = f_{\theta'}(x_T),$$

as in standard JEPA training (Eq. (2)); the parameters $\theta'$ are typically updated via an exponential moving average (EMA) of $\theta$ (Eq. (3)) and are not directly optimized by gradient descent;

* a predictive likelihood network parameterizing

$$p_{\text{like}}(Z_T | Z_C), \qquad (31)$$

which generalizes the VJEPA probabilistic predictor (Eq. (7)) and models uncertainty over future latent representations conditioned on context;

* a prior encoder that maps auxiliary information $\eta$, e.g. goals, constraints, demonstrations, physical structure, to a latent-space prior distribution:

$$p_{\text{prior}}(Z_T | \eta), \qquad (32)$$

which is new relative to JEPA/VJEPA and enables Bayesian conditioning in representation space.

The full Bayesian JEPA (BJEPA) architecture can be summarized by the following latent-space mappings:

$Z_C = f_{\theta}(x_C)$, Context / online encoder: what has happened (history)

$Z_T = f_{\theta'}(x_T)$, Target encoder (training only): what usually happens (learned dynamics)

$p_{\text{like}}(Z_T | Z_C) = \text{Likelihood}_{\phi}(Z_C)$, Predictive likelihood

$p_{\text{prior}}(Z_T | \eta) = \text{Prior}_{\phi'}(\eta)$, Latent prior: what should or must or expect to happen (goals/constraints)

$p(Z_T | Z_C, \eta) \propto p_{\text{like}}(Z_T | Z_C)p_{\text{prior}}(Z_T | \eta)$, 'Bayesian posterior'

These components are schematically sketched in Fig.2, in which the combination operator ("Combine") fuses the likelihood and prior distributions. Mathematically, this implements PoE via pointwise multiplication of densities followed by normalization. In the common case where both experts are parameterized as Gaussians, i.e. $p_{\text{like}} = N(\mu_{dyn}, \Sigma_{dyn})$ and $p_{\text{prior}} = N(\mu_{aux}, \Sigma_{aux})$, this operation admits a closed-form solution equivalent to a Kalman update in latent space. The posterior mean $\mu_{post}$ is computed as the precision-weighted average of the expert means: $\mu_{post} = (\Sigma_{dyn}^{-1} + \Sigma_{aux}^{-1})^{-1}(\Sigma_{dyn}^{-1}\mu_{dyn} + \Sigma_{aux}^{-1}\mu_{aux})$. This mechanism ensures that the predicted representation $\hat{Z}_T$ satisfies both the physical dynamics and the task constraints, effectively intersecting the two manifolds in latent space. Note that, all components operate entirely in representation space. Observation reconstruction, reward supervision, and autoregressive likelihood factorization are not required.
<figure><img src="image_5.png" /><figcaption>Figure 2: Bayesian JEPA (BJEPA) architecture. A predictive likelihood models dynamics while a prior encoder maps auxiliary input $\eta$ to a latent constraint. These distributions are fused via a PoE operator to yield a posterior predictive distribution. During training, the prediction is matched against the target encoder output using a JEPA loss.</figcaption></figure>

**VJEPA as a special case of BJEPA** VJEPA is recovered as a special case of BJEPA by choosing a uniform (uninformative / vague) prior:

$$p_{\text{prior}}(Z_T | \eta) \propto 1.$$

In this case, the Bayesian posterior predictor becomes

$$p(Z_T | Z_C, \eta) = p_{\text{like}}(Z_T | Z_C),$$

which reduces exactly to the VJEPA predictive distribution. Thus, BJEPA generalizes VJEPA without altering its training objective.
## 8.3 Training and Inference

BJEPA allows for disentangled learning: the dynamics (what is possible) can be learned from large-scale unlabeled data, while the prior (what is desired) can be specified or learned separately.

**Representation Learning: Training with Structural Priors.** The primary training phase focuses on the context encoder $f_{\theta}$ and the predictive likelihood expert $p_{\text{like}}(Z_T | Z_C)$. While we can assume an uninformative prior<fnref n="16" /> to learn pure data-driven dynamics, BJEPA also allows us to inject structural priors $p_{\text{struct}}(Z_T)$ during training. These priors can enforce desirable properties, such as stationarity, sparsity, or slowness, acting as an information bottleneck that filters out high-entropy nuisance variables (e.g. the "Noisy TV" distractors in later experiment).

The objective minimizes the variational negative log-likelihood of the target representations provided by the target encoder $f_{\theta'}$ (which serves as the "teacher" via EMA). The target encoder $f_{\theta'}$ provides stable regression targets (the "ground truth" latent states) derived from the data, exactly as in the standard VJEPA formulation. We optimize the parameters of the context encoder and the likelihood expert to maximize the predictive likelihood of the target representations while optionally regularizing the dynamics against the structural prior. Adapting Eq. (11), the general BJEPA training objective is:

$$ \mathcal{L}_{\text{BJEPA}} = \mathbb{E}_x \left[ \underbrace{\mathbb{E}_{Z_T \sim q_{\theta'}(\cdot|x_T)} \left[ -\log p_{\text{like}}(Z_T | Z_C) \right]}_{\text{Maximize Predictive Likelihood}} + \beta \underbrace{\text{KL}(q_{\theta'}(Z_T | x_T) \| p_{\text{ref}}(Z))}_{\text{Regularize Information State}} + \gamma \underbrace{\text{KL}(p_{\text{like}}(Z_T | Z_C) \| p_{\text{struct}}(Z_T))}_{\text{Enforce Structural Constraints}} \right]. \quad (33) $$

* **Likelihood ($\mathcal{L}_{\text{VJEPA}}$):** The first two terms (controlled by $\beta$) constitute the standard VJEPA objective (Eq.11). They ensure the learned representation captures the predictable mutual information between the past and future.

* **Structural Regularization ($\gamma$):** The third term penalizes the model if the "physically feasible" distribution ($p_{\text{like}}$) deviates too far from the "task compliant" distribution ($p_{\text{prior}}$), therefore forces the predicted dynamics to adhere to the structural prior. In our "Noisy TV" experiments (Section 9), we use a *static prior* ($p_{\text{struct}} \approx \text{const}$), which penalizes the model for tracking drifting, non-stationary noise, effectively filtering out distractors that violate the prior's structure.

The target encoder parameters are updated asymmetrically: the mean parameters $\mu_{\theta'}$ are updated via EMA as per Eq. (3), while the variance parameters $\sigma_{\theta'}^2$ are learned directly via gradient descent on the KL regularization term (the second term in Eq. (33)). This allows the model to learn appropriate uncertainty estimates for the target representations while maintaining stability through the EMA-updated mean. Alternatively, both the mean and variance parameters of the target encoder can be learned, or set to be constant. $p_{\text{like}}$ can capture the multi-modal distribution of valid future states (the "physics" of the latent space) without collapsing, thanks to the reference prior $p_{\text{ref}}$ (typically a unit Gaussian) and the asymmetry of the target encoder.

Note that we explicitly maximize the log-likelihood of the dynamics expert $p_{\text{like}}(Z_T | Z_C)$ rather than the full posterior $p(Z_T | Z_C, \eta)$ (Eq. (30)). This factorization is deliberate: the likelihood

<footnote n="16">¹⁶We can assume an uninformative (uniform) prior, $p_{\text{prior}}(Z_T | \eta) \propto 1$, allowing the model to learn pure latent dynamics driven solely by observational data. Note that, in a modular architecture such as BJEPA, the training of the prior (the "constraint expert") is often distinct from the training of the dynamics (the "likelihood expert"). If they were to be trained jointly end-to-end, the dynamics model might overfit to specific tasks rather than learning general world physics.</footnote>
expert is responsible for learning how the world actually evolves (data-driven physics), while the prior expert encodes how we expect or want the world to behave (constraints). Optimizing the full, fused posterior directly against the ground truth $Z_T$ can be risky, as it would force the model to pretend that the data satisfies our arbitrary constraints, for example, ignoring "Noisy TV" drift because the prior enforces stationarity (see later the experiment section). Instead, we require the model to fit the observational data (via the negative log-likelihood term) while being penalized for deviating from the structural expectations (via the KL regularization term). This ensures accurate physical modeling while filtering out nuisance variables that violate the structural prior.

**Inference with Task-Specific Priors.** Once the dynamics $p_{\text{like}}$ are trained, we can introduce new, task-specific priors $p_{\text{task}}(Z_T | \eta)$ at inference time (or fine-tuning) to guide the agent, without retraining the underlying physics model. Unlike the likelihood, the prior expert $p_{\text{prior}}(Z_T | \eta)$ encodes task-specific constraints or goals. It is specified or trained in one of two ways, depending on the nature of the auxiliary information $\eta$:

* **Learned Goal Priors (e.g. image goals):** if $\eta$ represents a goal observation (e.g. an image of the desired outcome), the prior encoder can be trained to map $\eta$ to a distribution in latent space. A common and efficient strategy is to reuse the frozen target encoder $f_{\theta'}$ or train a separate lightweight mapper to predict the representation of the goal, e.g.

$$p_{\text{prior}}(Z_T | \eta_{\text{img}}) \approx \mathcal{N}(Z_T | f_{\theta'}(\eta_{\text{img}}), \sigma_{\text{goal}}^2 I).$$

This allows the agent to pursue visual goals without retraining the dynamics model.

* **Analytic / Energy-Based Priors:** for logical constraints or safety regions (e.g. "avoid region $Z_{\text{unsafe}}$"), the prior can be specified analytically as an energy function:

$$p_{\text{prior}}(Z_T | \eta) \propto e^{-E(Z_T; \eta)}.$$

For example, a "stay close to reference trajectory" prior might be defined as a quadratic energy well centered on a reference path.

**Inference: Planning via Bayesian Fusion.** During planning, the target encoder is inactive<fnref n="17" /> (as the future is unknown). Instead, we fuse the learned dynamics with the task-specified prior using the Product of Experts mechanism to obtain the BJEPA posterior:

1. **Likelihood Step:** the learned expert $p_{\text{like}}(Z_T | Z_C)$ predicts where the system will go based on history $Z_C$.

2. **Prior Step:** the prior expert $p_{\text{prior}}(Z_T | \eta)$ specifies where the system should go to satisfy auxiliary information $\eta$ (e.g. goals or safety constraints).

3. **Fusion:** the planner samples or optimizes against the product distribution:

$$p(Z_T | Z_C, \eta) \propto p_{\text{like}}(Z_T | Z_C) \times p_{\text{prior}}(Z_T | \eta).$$

This effectively restricts the search space to the intersection of dynamically feasible and task-compliant trajectories, enabling efficient latent-space planning without requiring the model to have seen the specific task $\eta$ during the representation learning phase.

<footnote n="17">¹⁷The target encoder is not required at inference time unless new observations $x_T$ are incorporated.</footnote>
The planner then samples or optimizes latent states $Z_T$ that maximize the joint posterior:

$$ \log p(Z_T | Z_C, \eta) \cong \log p_{\text{like}}(Z_T | Z_C) + \log p_{\text{prior}}(Z_T | \eta) + \text{const.} $$

This effectively intersects the manifold of physically feasible futures (from the Likelihood) with the manifold of task-compliant futures (from the Prior), enabling zero-shot generalization to new tasks defined by $\eta$.

The distinction between the training and inference phases in BJEPA represents a transition from latent alignment to explicit manifold intersection. During training, the system employs soft fusion, where the task prior $\eta$ acts as a regularizer via a KL divergence term; this biases the dynamics predictor to learn a world model that naturally favors task-compliant regions without strictly forcing the prediction to originate from the fused distribution $p(Z_T | Z_C, \eta)$. In contrast, the inference phase utilizes hard fusion through the PoE mechanism, explicitly intersecting the manifold of physically feasible trajectories with the task-specific prior. This ensures that the final latent plan is strictly confined to the intersection of what the world model deems possible and what the task requires, enabling efficient, zero-shot generalization to new constraints.

**MAP vs. Sampling in the Posterior.** Given the fused posterior in Eq. (30), prediction and planning may formally proceed via:

* **MAP prediction:**

$$ \hat{Z}_T^{\text{MAP}} = \arg \max_{Z_T} \left[ \log p_{\text{like}}(Z_T | Z_C) + \log p_{\text{prior}}(Z_T | \eta) \right]; $$

* **Sampling-based prediction:**

$$ Z_T^{(m)} \sim p(Z_T | Z_C, \eta), \quad m = 1, \dots, M. $$

with $M$ being the total number of samples.

MAP prediction is computationally efficient and sufficient for unimodal posteriors, while sampling becomes essential in multi-modal or risk-sensitive settings where preserving the full uncertainty profile is necessary for robust control.

## 8.4 Sequential BJEPA as Latent Bayesian Filtering

When applied sequentially over time, BJEPA naturally induces a recursive filtering mechanism in representation space. By treating the context $Z_t$ as the current state summary ($Z_t = f_{\theta}(x_{\le t})$ summarizes history) and the auxiliary information $\eta_{t+1}$ as a “virtual observation” or constraint for the next step, the Product of Experts formulation implements a Bayesian filter update.

**The BJEPA Filter Recursion.** Standard Bayesian filtering consists of a prediction step (time update) via a transition model and a correction step (measurement update) via an observation likelihood. BJEPA implements this structure entirely in latent space:

1. **Prediction (Time Update):** The predictive likelihood $p_{\text{like}}(Z_{t+1} | Z_t)$ propagates the state forward based on system dynamics.
2. Correction (Constraint Update): The latent prior $p_{\text{prior}}(Z_{t+1} | \eta_{t+1})$ acts as a measurement likelihood for the auxiliary task information $\eta_{t+1}$, effectively “correcting” the predicted state to satisfy the goal or constraint.

Combining these via the PoE mechanism yields the recursive update for the belief over the future trajectory:

$$p(Z_{t+1} | Z_t, \eta_{t+1}) \propto \underbrace{p_{\text{like}}(Z_{t+1} | Z_t)}_{\text{Dynamics (Prediction)}} \times \underbrace{p_{\text{prior}}(Z_{t+1} | \eta_{t+1})}_{\text{Constraint (Correction)}}. \quad (34)$$

This formulation mirrors Bayesian filtering, but avoids explicit observation likelihoods. It also unifies planning and filtering: planning becomes the task of filtering the state distribution conditioned on the evidence of “optimality” or goal achievement encoded by $\eta$.

**Comparison to Standard Filtering.** In a standard Kalman Filter or HMM, the correction comes from sensory data $x_{t+1}$ via $p(x_{t+1} | Z_{t+1})$. In BJEPA, the correction comes from the intention or constraint $\eta_{t+1}$ via $p_{\text{prior}}$. This allows the agent to reason about the “posterior” future that is consistent with both physics and goals, without needing to generate or score high-dimensional observations.

*Remark 2 (Belief-Space MPC).* This perspective allows BJEPA planning to be interpreted as *Stochastic Model Predictive Control (MPC)* performed via belief propagation. The likelihood expert propagates the physically feasible set (the learned belief), while the prior expert injects goals or constraints and carves out the task-relevant subset. The resulting posterior $p(Z_{t+1} | Z_t, \eta_{t+1})$ represents the optimal control distribution, generalizing the notion of “Control as Inference” to latent representation spaces and yielding a belief-space MPC formulation without explicit state estimation.

## 8.5 MAP Optimality under BJEPA

We now show that the sufficiency of MAP planning (Theorem 3 of VJEPA) extends to the BJEPA architecture, provided the “experts” are well-behaved.

**Theorem 5 (Optimality of MAP Planning under BJEPA).** Consider a one-step planning problem where the belief is formed by a BJEPA posterior:

$$p(Z_{t+1} | Z_t, \eta, u_t) \propto p_{\text{like}}(Z_{t+1} | Z_t, u_t) \times p_{\text{prior}}(Z_{t+1} | \eta).$$

Assume:

1. The likelihood expert is Gaussian with action-independent covariance: $p_{\text{like}} = \mathcal{N}(\mu_{\text{dyn}}(Z_t, u_t), \Sigma_{\text{dyn}})$.

2. The prior expert is Gaussian and independent of the current action: $p_{\text{prior}} = \mathcal{N}(\mu_{\text{aux}}(\eta), \Sigma_{\text{aux}})$.

3. The stage cost $c(Z_{t+1}, u_t)$ is quadratic in $Z_{t+1}$.

Then the action minimizing the expected cost under the full Bayesian posterior is equivalently obtained by minimizing the cost at the posterior MAP state:

$$u_t^* = \arg \min_{u_t} c(\hat{Z}_{t+1}^{\text{MAP}}, u_t).$$
*Proof.* Since the posterior is formed by the product of two Gaussians, it is itself a Gaussian:

$$p(Z_{t+1} | \cdot) = \mathcal{N}(Z_{t+1} | \mu_{\text{post}}, \Sigma_{\text{post}}).$$

The posterior precision (inverse covariance) is the sum of the expert precisions:

$$\Sigma_{\text{post}}^{-1} = \Sigma_{\text{dyn}}^{-1} + \Sigma_{\text{aux}}^{-1}.$$

Since both $\Sigma_{\text{dyn}}$ and $\Sigma_{\text{aux}}$ are independent of $u_t$, the resulting posterior covariance $\Sigma_{\text{post}}$ is also constant with respect to $u_t$. The expected quadratic cost decomposes<fnref n="18" /> into a mean-dependent term and a trace term $\text{tr}(Q\Sigma_{\text{post}})$. Since the trace term is constant with respect to the action, the optimization depends only on the posterior mean. For a Gaussian, the mean coincides with the MAP estimate. □

## 8.6 Energy-Based Priors and Latent Optimization

A key advantage of the Product of Experts formulation is flexibility: the prior $p_{\text{prior}}$ need not be a normalized probability density. It can be specified as an unnormalized energy-based model (EBM):

$$p_{\text{prior}}(Z_T | \eta) \propto \exp(-E_{\phi'}(Z_T; \eta)), \quad (35)$$

where $E_{\phi'}$ is a scalar energy function representing the “cost” of a state $Z_T$ given task $\eta$. This allows for complex, non-convex constraints (e.g. obstacle avoidance potentials or logical constraints) that are difficult to normalize.

Combining a Gaussian predictive likelihood $p_{\text{like}} = \mathcal{N}(\mu_{\phi}, \Sigma)$ with an energy-based prior yields a posterior whose MAP estimate corresponds to a latent-space energy minimization problem:

$$Z_T^{\text{MAP}} = \arg \min_{Z_T} \left[ \underbrace{\frac{1}{2} \|Z_T - \mu_{\phi}(Z_T)\|_{\Sigma^{-1}}^2}_{\text{Dynamics Consistency}} + \underbrace{E_{\phi'}(Z_T; \eta)}_{\text{Task Constraint}} \right].$$

This formulation reveals a deep connection to classical trajectory optimization and control:

* The first term acts as a regularizer anchoring the solution to physically plausible futures (the learned dynamics).

* The second term acts as a task cost or potential field pulling the solution towards the goal.

Since all components in VJEPA/BJEPA are differentiable, this optimization can be solved efficiently via *gradient descent* directly in the representation space, effectively performing *planning as inference* without creating a separate discrete graph or search tree.

## 8.7 Algorithm: Gradient-Based Planning with BJEPA

We synthesize the previous 3 theoretical components, i.e. Bayesian filtering, MAP optimality, and energy-based priors, into a practical control algorithm. Because the BJEPA posterior is fully differentiable, we can implement Model Predictive Control (MPC) via *gradient descent in latent space*, this procedure is detailed in Algorithm 3. At every time step, we optimize a sequence of actions such that the resulting latent trajectory maximizes the joint posterior probability defined by the learned dynamics and the task prior.

<footnote n="18">¹⁸Similar to the proof to Theorem 3 of VJEPA, see Appendix G.</footnote>
**Algorithm 3 BJEPA Gradient-Based MPC**

**Require:** Current history $x_{\le t}$, Task $\eta$, Horizon $H$

**Require:** Differentiable Dynamics $p_{\text{like}}$, Energy Prior $E_{\phi'}$

1: **Initialize:** Estimate current state $Z_t = f_{\theta}(x_{\le t})$

2: **Initialize:** Action sequence $u_{0:H-1}$ (e.g. random or warm-start)

3: **while** optimization budget not exceeded **do**

4: 1. *Forward Rollout (Dynamics Expert)*

5: $\hat{Z}_0 \leftarrow Z_t$

6: **for** $k = 0$ **to** $H-1$ **do**

7: $\hat{Z}_{k+1} \leftarrow \mu_{\text{dyn}}(\hat{Z}_k, u_k)$ {Predict future mean}

8: **end for**

9: 2. *Evaluate BJEPA Posterior (Energy)*

10: $J \leftarrow 0$

11: **for** $k = 1$ **to** $H$ **do**

12: $J_{\text{prior}} \leftarrow E_{\phi'}(\hat{Z}_k; \eta)$ {Task Constraint Energy}

13: $J \leftarrow J + J_{\text{prior}} + \lambda \|u_{k-1}\|^2$ {Add control cost}

14: **end for**

15: 3. *Backward Pass (Optimization)*

16: $\nabla_u J \leftarrow \text{Backpropagate}(J)$

17: $u_{0:H-1} \leftarrow u_{0:H-1} - \alpha \nabla_u J$ {Gradient update}

18: **end while**

19: **Execute:** Apply $u_0$ to environment

20: **Recurse:** Receive $x_{t+1}$, update $Z_{t+1}$, repeat.

This algorithm contrasts with standard trajectory optimization in that the “cost function” is not handcrafted; it is the learned energy prior $E_{\phi'}$. Furthermore, the “dynamics constraints” are provided by the learned likelihood expert $p_{\text{like}}$.

## 8.8 Modularity, Transfer, and Continual Learning

A fundamental advantage of the BJEPA factorization (Eq. 30) is the semantic disentanglement of environmental dynamics from task objectives. This modularity unlocks capabilities that are difficult to achieve in monolithic world models:

**Zero-Shot Transfer via Prior Swapping.** In standard reinforcement learning or world modeling (e.g. Dreamer, MuZero), the reward function is often entangled with the representation or dynamics. To change tasks, one must typically retrain or fine-tune the model weights. In BJEPA, the predictive likelihood $p_{\text{like}}$ learns “objective-agnostic physics.” Once trained, it can be transferred zero-shot to arbitrarily many new tasks simply by swapping the prior expert $p_{\text{prior}}(Z | \eta)$. The agent understands how the world evolves (Likelihood) independently of what it is currently trying to achieve (Prior). This ‘reuse dynamics, adapt priors’ paradigm facilitates zero-shot task transfer.

**Continual Learning without Catastrophic Forgetting.** Because task-specific information is injected at inference time via the prior (or the energy function), the dynamics model does not need to be updated to accommodate new goals. This mitigates catastrophic forgetting: the agent does not overwrite its knowledge of physics to learn a new task. The shared world model accumulates general knowledge, while task specificity is handled by lightweight, modular priors.
**Summary.** BJEPA generalizes JEPA-based world models by introducing an explicit ‘Bayesian posterior’ in representation space. This enables uncertainty-aware planning, structured priors, belief propagation, and principled transfer and continual learning, while retaining JEPA’s core advantages: the avoidance of observation reconstruction, reward supervision, and autoregressive likelihood factorization.

**Comparison with Existing WM Frameworks.** Table.3 contextualizes BJEPA within the landscape of latent world models. While methods such as *Dreamer* handle uncertainty, they rely on computationally expensive observation reconstruction. Conversely, recent JEPA-based world models avoid reconstruction but typically lack explicit uncertainty quantification and rely on monolithic, learned reward functions. BJEPA uniquely occupies the intersection of these capabilities: it is likelihood-free (like JEPA) yet probabilistic (like *Dreamer*), and it replaces monolithic rewards with modular, energy-based priors. This specific combination allows for flexible, zero-shot task specification without retraining the underlying dynamics.

<table><caption>Table 3: Comparison of latent world modeling frameworks. BJEPA is unique in combining likelihood-free learning with explicit uncertainty and modular task specification.</caption><thead><tr><th>Method</th><th>Obs. Reconstruction</th><th>Latent Uncertainty</th><th>Planning Space</th><th>Task Specification</th><th>Reward/Cost</th></tr></thead><tbody><tr><td>Dreamer</td><td>✓</td><td>✓</td><td>Latent</td><td>Monolithic</td><td>Learned</td></tr><tr><td>JEPA-WM</td><td><b>X</b></td><td><b>X</b></td><td>Latent</td><td>Monolithic</td><td>Learned</td></tr><tr><td>VJEPA (Ours)</td><td><b>X</b></td><td>✓</td><td>Latent</td><td>Implicit</td><td>None</td></tr><tr><td><b>BJEPA (Ours)</b></td><td><b>X</b></td><td>✓</td><td><b>Latent</b></td><td><b>Modular Prior</b></td><td><b>Energy / Prior</b></td></tr></tbody></table>

**Note on “likelihood-free”:** In generative world models (e.g. *Dreamer*, VAEs), the model maximizes the likelihood of observed data $max \log p_{\psi}(x | z)$, requiring a decoder to reconstruct every detail (e.g. texture, noise) from the latent state. In contrast, JEPA architectures are “likelihood-free” regarding observations because they never reconstruct the input. The loss operates entirely in representation space ($Z$), optimizing prediction of abstract features rather than pixel-level reconstruction.

# 9 Toy Experiment: a “Noisy TV” Linear System

We present an analytically tractable experiment designed to stress-test the *nuisance invariance* property (Proposition 1) of the JEPA family (JEPA/VJEPA/BJEPA). We construct an environment where observations contain a low-variance signal embedded in high-variance, yet predictable, nuisance noise (“distractors”). This setup mimics the “Noisy-TV” problem<fnref n="19" />, where a model is presented with a dominant source of high-entropy, task-irrelevant variability that threatens to overwhelm the signal representation. Implementation details can be found in Appendix.J.

## 9.1 Environment Setup

We define a Linear-Gaussian System where the observation dimension $D_x = 20$ is significantly larger than the true state dimension $D_s = 4$. The observation $x_t$ is composed of a signal $s_t$ and a scaled *distractor* $d_t$ projected through fixed matrices $C$ and $D$:

$$x_t = Cs_t + D(\sigma d_t) + \epsilon_t.$$

<footnote n="19">¹⁹The “Noisy-TV” problem originated as a thought experiment in Burda et al. [16] regarding curiosity-driven reinforcement learning. It describes a failure mode where an agent, rewarded for seeking novelty or prediction error, becomes fixated on a source of uncontrollable, unpredictable noise (like a TV displaying static). Because the noise is inherently unpredictable, the agent continuously extracts high intrinsic rewards without making meaningful progress, effectively becoming a “couch potato”.</footnote>
The latent processes evolve as follows:

* *The Signal (s<sub>t</sub>):* evolves according to a stable rotation, s<sub>t+1</sub> = A<sub>rot</sub>s<sub>t</sub> + w<sub>t</sub>, where w<sub>t</sub> ~ N(0, 0.1<sup>2</sup>I). The orthogonal transition ensures the signal remains bounded and maintains constant variance, fixing the theoretical recoverability throughout the sequence.

* *The Distractor (d<sub>t</sub>):* evolves as a “sticky” random walk, d<sub>t+1</sub> = 0.9d<sub>t</sub> + v<sub>t</sub>. We scale the distractor by σ ∈ [0, 8], allowing its variance to grow up to ≈ 64× the signal variance. At σ = 8.0, the SNR drops to -2.2 dB, effectively burying the signal in noise.

## 9.2 Model Architectures and Evaluation Alignment

We compare five models using a latent dimension D<sub>z</sub> = D<sub>s</sub> = 4. All architectures use linear transformations to match the environment. To ensure a fair comparison, the evaluation respects the temporal structure of each model:

* *VAE (Static):* acts as a static model where the latent z<sub>t</sub> represents the current state. It is evaluated against the current signal s<sub>t</sub>.

* *AR (Pixel-Predictive):* the bottleneck z<sub>t</sub> is optimized to predict future pixels x<sub>t+1</sub>. Consequently, z<sub>t</sub> is evaluated on its ability to recover the *future* signal s<sub>t+1</sub>.

* *JEPA/VJEPA (Latent-Predictive):* we extract the *predictor output* (the explicit prediction of Z<sub>t+1</sub>) for the linear probe. These are evaluated against s<sub>t+1</sub>.

* *BJEPA (Bayesian Fusion):* at inference, BJEPA uses *hard fusion* (Product of Experts) to combine dynamics with a task prior. The resulting fused mean μ<sub>post</sub> is evaluated against s<sub>t+1</sub>.

## 9.3 Results and Analysis

Performance is measured via the coefficient of determination (R<sup>2</sup>) using a linear probe fit on training data and evaluated on generalization test data.

**Failure of Generative Models.** As shown in Table 4 and Fig.3, generative models (VAE, AR) suffer catastrophic degradation as noise increases. At scale σ = 8.0, VAE test signal recovery drops to R<sup>2</sup> ≈ 0.50 while its noise recovery remains high (R<sup>2</sup> = 0.62), confirming it has prioritized high-variance distractors. Similarly, the AR model drops to R<sup>2</sup> = 0.578.

**Robustness of Joint-Embedding Architectures.** All joint-embedding architectures (JEPA, VJEPA, BJEPA) demonstrate robustness to the “Noisy TV” distractor, maintaining R<sup>2</sup> > 0.84 even at scale 8.0. Deterministic JEPA achieved the highest accuracy (R<sup>2</sup> = 0.93), though it exhibited momentary instability at scale 3.0 (R<sup>2</sup> = 0.841). Probabilistic VJEPA and BJEPA showed superior training stability across seeds.

**Qualitative Filtering.** Fig.4 illustrates the filtering effect. At Scale 8.0, VAE and AR reconstructions (dashed lines) track high-frequency noise. In contrast, BJEPA and VJEPA (solid lines) successfully act as latent filters, tracking the underlying true signal s<sub>t+1</sub> with high fidelity.
<table><caption>Table 4: Performance metrics ($R^2$) across models at representative noise scales. Tr denotes training set, Te denotes test set. Full table is listed in Appendix.J.</caption><thead><tr><th>Scale (SNR)</th><th>Model</th><th>↑Signal R<sup>2</sup> (Tr/Te)</th><th>↓Noise R<sup>2</sup> (Tr/Te)</th><th>↓Time (Tr/Te)</th></tr></thead><tbody><tr><td rowspan="5">0.0 (inf dB)</td><td>VAE</td><td><u>1.000</u> / <u>1.000</u></td><td>NA / NA</td><td>12.6s / 0.01s</td></tr><tr><td>AR</td><td>0.999 / 0.999</td><td>NA / NA</td><td><b>6.4s</b> / 0.01s</td></tr><tr><td>JEPA</td><td>0.947 / 0.930</td><td>NA / NA</td><td>16.9s / 0.01s</td></tr><tr><td>VJEPA</td><td>0.999 / 0.999</td><td>NA / NA</td><td>13.9s / 0.01s</td></tr><tr><td>BJEPA</td><td>0.987 / 0.981</td><td>NA / NA</td><td>23.5s / 0.01s</td></tr><tr><td rowspan="5">4.0 (3.8 dB)</td><td>VAE</td><td>0.822 / 0.730</td><td>0.512 / 0.458</td><td>12.5s / 0.01s</td></tr><tr><td>AR</td><td>0.839 / 0.756</td><td>0.394 / 0.338</td><td><b>6.5s</b> / 0.01s</td></tr><tr><td>JEPA</td><td><b>0.999</b> / <b>0.999</b></td><td><b>0.004</b> / <b>-0.010</b></td><td>16.6s / 0.02s</td></tr><tr><td>VJEPA</td><td>0.994 / 0.993</td><td>0.025 / -0.007</td><td>13.8s / 0.01s</td></tr><tr><td>BJEPA</td><td>0.920 / 0.899</td><td>0.213 / 0.156</td><td>23.2s / 0.01s</td></tr><tr><td rowspan="5">8.0 (-2.2 dB)</td><td>VAE</td><td>0.613 / 0.499</td><td>0.656 / 0.620</td><td>12.3s / 0.01s</td></tr><tr><td>AR</td><td>0.680 / 0.578</td><td>0.491 / 0.449</td><td><b>7.1s</b> / 0.01s</td></tr><tr><td>JEPA</td><td><b>0.947</b> / <b>0.930</b></td><td><b>0.226</b> / <b>0.183</b></td><td>16.1s / 0.01s</td></tr><tr><td>VJEPA</td><td>0.905 / 0.870</td><td>0.299 / 0.251</td><td>13.4s / 0.01s</td></tr><tr><td>BJEPA</td><td>0.896 / 0.841</td><td>0.292 / 0.238</td><td>23.0s / 0.01s</td></tr></tbody></table>

<figure><img src="image_6.png" /><figcaption>Figure 3: Performance metrics across noise scales. Top Row: Training set $R^2$. Bottom Row: Test set $R^2$ (Generalization). The generative models (VAE, AR) degrade linearly as noise increases, tracking the distractor (Bottom Right). The JEPA-based models (Blue/Cyan/Purple) maintain high signal recovery (Bottom Left) even at high noise scales, demonstrating invariance to nuisance variability.</figcaption></figure>
<figure><img src="image_7.png" /><figcaption>Figure 4: **Latent Reconstructions at varying noise scales.** At $\sigma = 8.0$ (Right), the VAE and AR reconstructions (dashed lines) track the high-frequency noise. In contrast, BJEPA and VJEPA (solid lines) successfully filter the noise and track the underlying true signal (black line).</figcaption></figure>

# 10 Discussion

In this work, we have formalized Joint Embedding Predictive Architectures (JEPA) not merely as a self-supervised training trick, but as a principled probabilistic framework for learning latent dynamical systems. By introducing *Variational JEPA (VJEPA)* and its Bayesian extension *BJEPA*, we bridged the gap between deterministic representation learning and probabilistic world models. Our analysis and experiments support several key conclusions regarding the nature of predictive learning.

## 10.1 Unifying JEPA with Probabilistic State-Space Models

JEPA is probably best understood as a predictive state-space model trained by representation prediction rather than observation likelihood. A central contribution of this work is the decoupling of sequential modeling from autoregressive observation likelihoods. Standard world models (e.g. Dreamer [29]) and generative sequence models (e.g. Transformers [81]) typically conflate the two, forcing the model to predict the next observation $x_{t+1}$ given history. This forces the latent state to account for high-entropy local noise.

VJEPA demonstrates that sequential structure does not harm representation learning; only autoregressive factorization over pixels does. By adopting a *time-indexed* specialization (as in Section.5 and Section.6), where contexts and targets are explicitly associated with temporal indices, VJEPA learns a mathematically rigorous *latent* transition model $p_{\phi}(Z_{t+1}|Z_t)$. Notably, this transition kernel does not need to be autoregressive (as seen in Eq.12) and is capable of belief propagation and filtering without ever estimating a density over pixels. This positions VJEPA as a scalable foundation for model-based control in high-dimensional, noisy environments.

## 10.2 Nuisance Invariance: The “PCA vs. CCA” Distinction

Our "Noisy TV" experiment (Section 9) provides empirical validation for the Information-Theoretic analysis in Section 7. The failure of VAE and AR baselines on high-variance distractors highlights a fundamental distinction in representation learning objectives:

* *Generative Reconstruction (e.g. VAE):* These models implicitly perform a non-linear gener-
alization<fnref n="20" /> of *Principal Component Analysis (PCA)* [65]. They prioritize latent dimensions that explain the maximum variance in the input *x*. If nuisance noise has higher variance than the signal (as in our experiment), a reconstruction objective is mathematically compelled to model the noise.

*   *Predictive Association (VJEPA):* JEPA-based models implicitly perform a non-linear generalization of *Canonical Correlation Analysis*<fnref n="21" /> (CCA [85]) or the *Information Bottleneck*. As established in Theorem 4 (Section 7), minimizing the VJEPA loss maximizes a lower bound on the mutual information between the past and the future ($I(Z_t; Z_{t+\Delta})$). Therefore, VJEPA (implicitly CCA) prioritizes dimensions with the highest predictive mutual information. It will ignore the noise regardless of its magnitude. This "nuisance invariance" property (Proposition 1) makes VJEPA as a *semantic noise filter* for downstream planning.

## 10.3 Applications and Broader Impact

While this work focuses on theoretical foundations, the VJEPA framework enables new capabilities in applied domains where state estimation and planning are required but observation reconstruction is costly or unnecessary.

**Robotics and Embodied AI.** In visual robotic control, $x_t$ represents high-dimensional sensory streams (e.g. camera feeds, lidar). Standard world models (e.g. Dreamer [29]) often spend significant capacity modeling local textures and camera noise. VJEPA allows the agent to learn a *task-specific physics engine* in latent space. For a manipulation task, the context $Z_t$ captures the geometry of objects, while the predictive model $p_\phi(Z_{t+1} | Z_t, u_t)$ predicts the consequences of motor commands. Planning can be performed via the *VJEPA-MPC* algorithm (Algo.2) to minimize a latent distance to a target goal embedding, enabling efficient visual control without pixel-level prediction. Furthermore, the BJEPA extension allows for *zero-shot transfer*: a robot can learn the environmental dynamics (Likelihood Expert) once, and then swap in different goal constraints (Prior Expert) at inference time to solve varied tasks without retraining the underlying model.

**Game Playing and Strategic Reasoning.** In high-dimensional game environments (e.g. *Atari*, *Minecraft*), VJEPA functions as a self-supervised latent simulator. Unlike *MuZero* [68], which couples representation learning with a specific reward signal, VJEPA learns a reward-agnostic model of the game mechanics. The predictive model $p_\phi(Z_{t+1} | Z_t, u_t)$ captures the transition rules and physics of the game world. This enables planning via Monte Carlo Tree Search (MCTS [89]) or trajectory optimization in the latent space, allowing the agent to reason about long-term strategies and effectively transfer the learned dynamics to new in-game tasks or modified rule sets without retraining the visual encoder.

**Language-as-Action Planning.** Current LLMs typically operate by minimizing the surface-level entropy of the next token. VJEPA offers an alternative view where language generation can be framed as *planning in semantic space*. Here, the 'world' is informational or symbolic, and the latent state $Z_t$ represents the current discourse status or reasoning state.

<footnote n="20">²⁰VAEs do not perform standard linear PCA by design, but they often implicitly discover the same principal directions (subspaces) as PCA while learning, particularly due to their regularization and training dynamics [65].</footnote>

<footnote n="21">²¹CCA is a multivariate statistical method that finds the strongest linear relationships between two sets of variables, which creates new composite variables (canonical variates) from each set that are maximally correlated with each other [85].</footnote>
* **Control Inputs:** The tokens (or chunks of text) act as discrete actions driving the system ($u_t \in V$). $u_t$ can represent both generated content and user inputs such as prompts, which function as exogenous control signals that steer the initial semantic trajectory.

* **Latent Dynamics:** The model learns how choosing a word $u_t$ updates the semantic state:

$$Z_{t+1} \sim p_{\phi}(Z_{t+1} | Z_t, u_t).$$

In this framework, a response is not evaluated by its likelihood $\prod p(u_t | u_{<t})$, but by its effect on the latent world. Generation becomes a trajectory optimization problem: identifying a sequence of tokens $u_{t:T}$ that drives the initial thought $Z_t$ to a target state $Z_{\text{goal}}$ (e.g. a state satisfying a logical constraint or answer condition). This allows text generation to leverage planning algorithms such as Beam Search [52], Monte Carlo Tree Search (MCTS [89]), or latent-space Cross-Entropy Method (CEM [51]), which decouples reasoning from surface-form statistics and potentially helps address hallucinations and long-horizon coherence.

## 10.4 Limitations and Future Directions

While VJEPA provides a rigorous probabilistic footing, our experiments revealed specific trade-offs:

**Expressivity of the Predictive Distribution.** In our experiments, the probabilistic VJEPA sometimes slightly underperformed the deterministic JEPA in peak accuracy, or required more careful hyperparameter tuning. We attribute this to the unimodality of the Gaussian assumption used in our implementation of $p_{\phi}$. In environments with complex, multi-modal bifurcations (e.g. a robot encountering an obstacle and choosing left or right), a unimodal Gaussian may average distinct modes, resulting in a blurry or incorrect belief. Future work should explore more expressive predictive heads, such as Gaussian Mixture Models (GMMs) [35] or Latent Diffusion heads, to capture complex aleatoric uncertainty while retaining the VJEPA objective.

**Optimization Dynamics.** We observed that the balance between the predictive loss and the KL regularization terms (controlled by $\beta$ in the VJEPA objective Eq. 11, and $\gamma$ in the BJEPA objective Eq. 33) is important. If regularization is too strong, the representation collapses; if too weak, the variance estimates become uncalibrated. Developing adaptive mechanisms for balancing these terms, similar to KL-balancing in latent overshooting, is a promising direction.

## 11 Conclusion

This work introduces Variational JEPA (VJEPA) and Bayesian JEPA (BJEPA), establishing the first formal probabilistic framework for Joint Embedding Predictive Architectures. By recasting JEPA as a variational inference problem in representation space, we resolve the ambiguity<fnref n="22" /> regarding the probabilistic semantics of predictive embeddings and their suitability for control.

<footnote n="22">²²Standard JEPA trains a deterministic predictor using Mean Squared Error. In stochastic environments with multimodal futures (e.g. a fork in the road), such a predictor converges to the conditional mean, which may be a physically invalid state (e.g. the average of ‘go left’ and ‘go right’ is ‘hit the wall’). This creates ambiguity regarding the semantics of the embedding: does it represent a specific future or an average? VJEPA resolves this by explicitly modeling the predictive distribution. Further, the variational derivation provides the theoretical guarantee that the learned state maximizes predictive mutual information, a necessary condition for it to serve as a sufficient information state for control.</footnote>
**VJEPA: A Probabilistic Foundation.** VJEPA generalizes standard deterministic JEPA by learning an explicit predictive distribution over future latent states. We showed that minimizing the VJEPA objective maximizes a variational lower bound on the predictive mutual information between the past and the future. Unlike generative world models, VJEPA achieves this without an autoregressive observation likelihood, effectively filtering out high-entropy nuisance variables (the "Noisy TV" problem) while retaining the information necessary for dynamics modeling. A key practical advantage of this formulation is the ability to perform principled uncertainty estimation (e.g. constructing credible intervals) by drawing multiple samples ($M \gg 1$) from the predictive distribution during inference$^{23}$, a capability absent in deterministic JEPA.

**BJEPA: Modular and Structural Control.** We extended this VJEPA framework to BJEPA, which factorizes prediction into a learned *Likelihood Expert* (dynamics) and a *Prior Expert* (constraints). This modularity allows for the injection of structural priors, such as stationarity or goal-directedness, directly into the latent inference process via Product of Experts which bridges the gap between "what is physically possible" and "what is task-relevant". This separation enables zero-shot task transfer and robust planning in the presence of non-stationary noise.

**Theoretical and Practical Unification.** Our analysis unifies JEPA with classical ideas from Predictive State Representations (PSRs) and Bayesian filtering. We established that observation reconstruction is theoretically unnecessary for optimal control, provided the latent state captures sufficient predictive mutual information. Furthermore, we demonstrated that:

* *Sequential modeling does not imply autoregression:* one can learn rigorous temporal dynamics $p(Z_{t+1}|Z_t)$ without the computational burden of factorizing the high-dimensional observation density $p(x_{t+1}|x_{\le t})$.

* *Predictive representations suffice for control:* we proved that under standard conditions, an optimal policy can be derived solely from the predictive latent state, rendering pixel reconstruction redundant for planning.

* *Likelihood-free world models can be principled:* VJEPA breaks the historical conflation of probabilistic rigor with pixel-level likelihoods. It replaces heuristic regression losses with a formal variational framework operating entirely in latent space, ensuring principled uncertainty estimation and information maximization without modeling nuisance variability.

By combining the representational efficiency of JEPA with the rigor of probabilistic state-space models, VJEPA positions itself as a foundational framework for scalable, robust, uncertainty-aware world modeling and planning in high-dimensional, possibly noisy, real-world environments.

## References

[1] Mohamed Abdelfattah and Alexandre Alahi. S-jepa: A joint embedding predictive architecture for skeletal action recognition. In *European Conference on Computer Vision (ECCV)*. Springer, 2024.

<footnote>$^{23}$We distinguish between two sampling regimes using $K$ and $M$ to denote total number of samples. *First. Training Sampling* ($K = 1$): drawn from the target encoder distribution $q_{\theta'}(\cdot | x_T)$ during training to generate a concrete regression target. A single sample is sufficient because the stochasticity of mini-batch updates averages out the noise over time. *Second. Inference Sampling* ($M \gg 1$): drawn from the predictor $p_{\phi}(\cdot | Z_C)$ during testing or planning (inference) to estimate uncertainty. Multiple samples are required here to accurately map the distribution's shape and construct credible intervals. This is also a key advantage of VJEPA/BJEPA over deterministic JEPA.</footnote>
[2] Alexander A. Alemi. Variational predictive information bottleneck, 2019.

[3] Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep variational information bottleneck, 2019.

[4] Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning, 2019.

[5] Mahmoud Assran, Quentin Duval, Ishan Misra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a joint-embedding predictive architecture, 2023.

[6] Mido Assran, Adrien Bardes, David Fan, Quentin Garrido, Russell Howes, Mojtaba, Komeili, Matthew Muckley, Ammar Rizvi, Claire Roberts, Koustuv Sinha, Artem Zholus, Sergio Arnaud, Abha Gejji, Ada Martin, Francois Robert Hogan, Daniel Dugas, Piotr Bojanowski, Vasil Khalidov, Patrick Labatut, Francisco Massa, Marc Szafraniec, Kapil Krishnakumar, Yong Li, Xiaodong Ma, Sarath Chandar, Franziska Meier, Yann LeCun, Michael Rabbat, and Nicolas Ballas. V-jepa 2: Self-supervised video models enable understanding, prediction and planning, 2025.

[7] David Barber and Felix Agakov. The im algorithm: a variational approach to information maximization. In *Proceedings of the 17th International Conference on Neural Information Processing Systems*, NIPS'03, page 201-208, Cambridge, MA, USA, 2003. MIT Press.

[8] Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mahmoud Assran, and Nicolas Ballas. Revisiting feature prediction for learning visual representations from video, 2024.

[9] Adrien Bardes, Quentin Garrido, Jean Ponce, Xinlei Chen, Michael Rabbat, Yann LeCun, Mido Assran, and Nicolas Ballas. V-JEPA: Latent video prediction for visual representation learning, 2024.

[10] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Variance-invariance-covariance regularization for self-supervised learning, 2022.

[11] Adrien Bardes, Jean Ponce, and Yann LeCun. Mc-jepa: A joint-embedding predictive architecture for self-supervised learning of motion and content features, 2023.

[12] Dimitri P. Bertsekas. *Dynamic Programming and Optimal Control*, volume 1. Athena Scientific, 4th edition, 2012.

[13] William Bialek, Ilya Nemenman, and Naftali Tishby. Predictability, complexity, and learning. *Neural Comput.*, 13(11):2409-2463, November 2001.

[14] Byron Boots, Sajid M Siddiqi, and Geoffrey J Gordon. Closing the learning-planning loop with predictive state representations. *Int. J. Rob. Res.*, 30(7):954-966, June 2011.

[15] V. S. Borkar and S. P. Meyn. Risk-sensitive optimal control for markov decision processes with monotone cost. *Mathematics of Operations Research*, 27(1):192-209, 2002.

[16] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation, 2018.
[17] Eduardo F. Camacho and Carlos Bordons. *Model Predictive Control*. Springer, 2nd edition, 2007.

[18] Delong Chen, Mustafa Shukor, Theo Moutakanni, Willy Chung, Jade Yu, Tejaswi Kasarla, Allen Bolourchi, Yann LeCun, and Pascale Fung. Vl-jepa: Joint embedding predictive architecture for vision-language, 2025.

[19] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, 2020.

[20] Yinlam Chow, Aviv Tamar, Shie Mannor, and Marco Pavone. Risk-sensitive and robust decision-making: a cvar optimization approach. In *Proceedings of the 29th International Conference on Neural Information Processing Systems - Volume 1*, NIPS'15, page 1522–1530, Cambridge, MA, USA, 2015. MIT Press.

[21] Matthieu Destrade, Oumayma Bounou, Quentin Le Lidec, Jean Ponce, and Yann LeCun. Value-guided action planning with jepa world models, 2025.

[22] Katlyn Edwards, Corrie Scalisi, Julianne DeMars-Smith, and Key Lee. Google colab for teaching cs and ml. In *Proceedings of the 55th ACM Technical Symposium on Computer Science Education V. 2*, SIGCSE 2024, page 1925, New York, NY, USA, 2024. Association for Computing Machinery.

[23] Karl Friston. The free-energy principle: a unified brain theory? *Nature Reviews Neuroscience*, 11(2):127–138, 2010.

[24] Karl Friston, Rick Adams, Laurent Perrinet, and Michael Breakspear. Perceptions as hypotheses: saccades as experiments. *Frontiers in Psychology*, 3:151, 2012.

[25] Karl Friston, Thomas FitzGerald, Francesco Rigoli, Philipp Schwartenbeck, and Giovanni Pezzulo. Active inference: a process theory. *Neural Computation*, 29(1):1–49, 2017.

[26] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H. Richmond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukcuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning, 2020.

[27] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In *Proceedings of the 32nd International Conference on Neural Information Processing Systems*, NIPS'18, page 2455–2467, Red Hook, NY, USA, 2018. Curran Associates Inc.

[28] R. Hadsell, S. Chopra, and Y. LeCun. Dimensionality reduction by learning an invariant mapping. In *2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06)*, volume 2, pages 1735–1742, 2006.

[29] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination, 2020.

[30] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels, 2019.

[31] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning, 2020.
[32] Xiangteng He, Shunsuke Sakai, Kun Yuan, Nicolas Padoy, Tatsuhito Hasegawa, and Leonid Sigal. Dseq-jepa: Discriminative sequential joint-embedding predictive architecture, 2025.

[33] Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a constrained variational framework. In *International Conference on Learning Representations*, 2017.

[34] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models, 2020.

[35] Yongchao Huang. Sampling via gaussian mixture approximations, 2025.

[36] Leslie Pack Kaelbling, Michael L. Littman, and Anthony R. Cassandra. Planning and acting in partially observable stochastic domains. *Artificial Intelligence*, 101(1):99-134, 1998.

[37] R. E. Kalman. A new approach to linear filtering and prediction problems. *Journal of Basic Engineering*, 82(1):35-45, 03 1960.

[38] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.

[39] Diederik P Kingma and Max Welling. Auto-encoding variational bayes, 2022.

[40] Torsten Koller, Felix Berkenkamp, Matteo Turchetta, and Andreas Krause. Learning-based model predictive control for safe exploration. In *2018 IEEE Conference on Decision and Control (CDC)*, pages 6059-6066, 2018.

[41] Hanna Kurniawati, David Hsu, and Wee Sun Lee. Sarsop: Efficient point-based pomdp planning by approximating optimally reachable belief spaces. In *Robotics: Science and Systems IV*, pages 65-72. The MIT Press, 06 2009.

[42] Tze Leung Lai, Haipeng Xing, and Zehao Chen. Mean-variance portfolio optimization when means and covariances are unknown. *The Annals of Applied Statistics*, 5(2A), June 2011.

[43] Anders Boesen Lindbo Larsen, Søren Kaae Sønderby, Hugo Larochelle, and Ole Winther. Autoencoding beyond pixels using a learned similarity metric, 2016.

[44] Yann LeCun. A path towards autonomous machine intelligence version 0.9.2, 2022-06-27. *Open Review*, 62(1):1-62, 2022.

[45] Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, 2019.

[46] Michael Littman and Richard S Sutton. Predictive representations of state. In T. Dietterich, S. Becker, and Z. Ghahramani, editors, *Advances in Neural Information Processing Systems*, volume 14. MIT Press, 2001.

[47] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations, 2019.

[48] William S. Lovejoy. Computationally feasible bounds for partially observed markov decision processes. *Operations Research*, 39(1):162-175, 1991.
[49] David J. C. MacKay. *Information Theory, Inference and Learning Algorithms*. Cambridge University Press, New York, NY, USA, June 2002.

[50] Omid Madani, Steve Hanks, and Anne Condon. On the undecidability of probabilistic planning and related stochastic optimization problems. *Artificial Intelligence*, 147(1):5–34, 2003. Planning with Uncertainty and Incomplete Information.

[51] Shie Mannor, Reuven Rubinstein, and Yohai Gat. The cross entropy method for fast policy search. In *Proceedings of the Twentieth International Conference on International Conference on Machine Learning*, ICML’03, page 512–519. AAAI Press, 2003.

[52] Clara Meister, Ryan Cotterell, and Tim Vieira. If beam search is the answer, what was the question? In Bonnie Webber, Trevor Cohn, Yulan He, and Yang Liu, editors, *Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP)*, pages 2173–2185, Online, November 2020. Association for Computational Linguistics.

[53] Ali Mesbah. Stochastic model predictive control: An overview and perspectives for future research. *IEEE Control Systems Magazine*, 36(6):30–44, 2016.

[54] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning, 2016.

[55] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. *Nature*, 518:529–533, 2015.

[56] Shentong Mo and Shengbang Tong. Connecting joint-embedding predictive architecture with contrastive self-supervised learning, 2024.

[57] Christos H. Papadimitriou and John N. Tsitsiklis. The complexity of markov decision processes. *Mathematics of Operations Research*, 12(3):441–450, 1987.

[58] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, *Proceedings of the 36th International Conference on Machine Learning*, volume 97 of *Proceedings of Machine Learning Research*, pages 5171–5180. PMLR, 09–15 Jun 2019.

[59] Martin L. Puterman. *Markov Decision Processes: Discrete Stochastic Dynamic Programming*. Wiley, 2nd edition, 2014.

[60] L.R. Rabiner. A tutorial on hidden markov models and selected applications in speech recognition. *Proceedings of the IEEE*, 77(2):257–286, 1989.

[61] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer, 2023.

[62] James B. Rawlings, David Q. Mayne, and Moritz Diehl. *Model Predictive Control: Theory, Computation, and Design*. Nob Hill Publishing, 2nd edition, 2017.
[63] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models, 2014.

[64] R. Tyrrell Rockafellar and Stanislav Uryasev. Optimization of conditional value-at-risk. *Journal of Risk*, 2:21–41, 2000.

[65] Michal Rolínek, Dominik Zietlow, and Georg Martius. Variational autoencoders pursue pca directions (by accident). In *2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)*, pages 12398–12407, 2019.

[66] Stéphane Ross, Joelle Pineau, Sébastien Paquet, and Brahim Chaib-draa. Online planning algorithms for pomdps. *J. Artif. Int. Res.*, 32(1):663–704, July 2008.

[67] Ayumu Saito, Prachi Kudeshia, and Jiju Poovvancheri. Point-jepa: A joint embedding predictive architecture for self-supervised learning on point cloud, 2025.

[68] Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, and et al. Mastering atari, go, chess and shogi by planning with a learned model. In *Nature*, 2020.

[69] Alexander T. Schwarm and Michael Nikolaou. Chance-constrained model predictive control. *AIChE Journal*, 45(8):1743–1752, 1999.

[70] Satinder Singh, Michael R. James, and Matthew R. Rudary. Predictive state representations: a new theory for modeling dynamical systems. In *Proceedings of the 20th Conference on Uncertainty in Artificial Intelligence*, UAI '04, page 512–519, Arlington, Virginia, USA, 2004. AUAI Press.

[71] Geri Skenderi, Hang Li, Jiliang Tang, and Marco Cristani. Graph-level representation learning with joint-embedding predictive architectures, 2025.

[72] Richard D. Smallwood and Edward J. Sondik. The optimal control of partially observable markov processes over a finite horizon. *Operations Research*, 21(5):1071–1088, 1973.

[73] Edward J. Sondik. The optimal control of partially observable markov processes over the infinite horizon: Discounted costs. *Operations Research*, 26(2):282–304, 1978.

[74] Zachary Sunberg and Mykel Kochenderfer. Online algorithms for pomdps with continuous state, action, and observation spaces, 2018.

[75] Richard S. Sutton and Andrew G. Barto. *Reinforcement Learning: An Introduction*. MIT Press, 2nd edition, 2018.

[76] Basile Terver, Tsung-Yen Yang, Jean Ponce, Adrien Bardes, and Yann LeCun. What drives success in physical planning with joint-embedding predictive world models?, 2026.

[77] Naftali Tishby, Fernando C. Pereira, and William Bialek. The information bottleneck method, 2000.

[78] Michael Tschannen, Josip Djolonga, Paul K. Rubenstein, Sylvain Gelly, and Mario Lucic. On mutual information maximization for representation learning, 2020.

[79] Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural networks, 2016.
[80] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding, 2019.

[81] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023.

[82] Dedi Wang and Pratyush Tiwary. State predictive information bottleneck. *The Journal of Chemical Physics*, 154(13), April 2021.

[83] Y. Wang, J. M. L. Ribeiro, and P. Tiwary. Past–future information bottleneck for sampling molecular reaction coordinate simultaneously with thermodynamics and kinetics. *Nature Communications*, 10:3573, 2019.

[84] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: a locally linear latent dynamics model for control from raw images. In *Proceedings of the 29th International Conference on Neural Information Processing Systems - Volume 2*, NIPS'15, page 2746–2754, Cambridge, MA, USA, 2015. MIT Press.

[85] David Weenink. Canonical correlation analysis. In *Proceedings of the Institute of Phonetic Sciences*, volume 25, pages 81–99, University of Amsterdam, 2003.

[86] Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling, 2015.

[87] Grady Williams, Paul Drews, Brian Goldfain, James M. Rehg, and Evangelos A. Theodorou. Aggressive driving with model predictive path integral control. In *2016 IEEE International Conference on Robotics and Automation (ICRA)*, pages 1433–1440, 2016.

[88] K.J Åström. Optimal control of markov processes with incomplete state information. *Journal of Mathematical Analysis and Applications*, 10(1):174–205, 1965.

[89] Maciej Świechowski, Konrad Godlewski, Bartosz Sawicki, and Jacek Mańdziuk. Monte carlo tree search: a review of recent modifications and applications. *Artificial Intelligence Review*, 56(3):2497–2562, July 2022.

## A Equivalence Between Squared Loss and Gaussian Likelihood

In this appendix, we show that minimizing the JEPA regression loss in Eq. (5) is equivalent to maximizing the log-likelihood of the target embedding under an isotropic Gaussian predictive model (Eq. (6)).

Consider the probabilistic model assumed in Eq. (6):

$$p(Z_T | Z_C) = \mathcal{N} (Z_T | \hat{Z}_T, \sigma^2 I), \quad \hat{Z}_T = g_{\phi}(Z_C, \xi_T).$$

The corresponding negative log-likelihood is given by

$$\begin{aligned} -\log p(Z_T | Z_C) &= -\log \left[ \frac{1}{(2\pi\sigma^2)^{d/2}} \exp \left( -\frac{1}{2\sigma^2} \|Z_T - \hat{Z}_T\|^2 \right) \right] \\ &= \frac{1}{2\sigma^2} \|Z_T - \hat{Z}_T\|^2 + \frac{d}{2} \log(2\pi\sigma^2), \end{aligned}$$

where d is the dimensionality of the embedding space.
Since the second term is constant with respect to the model parameters $(\theta, \phi)$, minimizing the negative log-likelihood is equivalent to minimizing

$$ \|Z_T - \hat{Z}_T\|^2. $$

Thus, minimizing the deterministic JEPA objective $\mathcal{L}_{\text{JEPA}} = \|\hat{Z}_T - Z_T\|^2$ can be translated as maximum likelihood estimation under an isotropic Gaussian predictive distribution with fixed variance $\sigma^2$.

This observation motivates the variational extension introduced in the main text, where the predictive distribution $p(Z_T | Z_C)$ is modeled explicitly and the variance (or more general uncertainty structure) is learned rather than fixed.

# B Relationships Between Information-Theoretic Quantities

In this section, we summarize the fundamental relationships between entropy, cross-entropy, Kullback-Leibler (KL) divergence, and mutual information. These identities underpin the theoretical analysis of VJEPA and BJEPA. For further info, readers are encouraged to refer to standard texts such as [49]. Here, let's denote $p(x)$ as the true (or target) data distribution and $q(x)$ the approximating (or predicted) distribution.

## B.1 Cross-Entropy, Entropy, and KL Divergence

Cross-entropy is the sum of the true entropy and the KL divergence. Intuitively, it represents the average number of bits needed to encode data from $p$ using a code optimized for $q$, which is the optimal bits (entropy of $p$) plus the penalty for using the wrong code (KL divergence):

$$ H(p, q) = H(p) + D_{\text{KL}}(p \| q), \qquad (36) $$

where the terms are defined as:

*   **Cross-Entropy:** $H(p, q) = -\sum_x p(x) \log q(x)$

*   **Entropy:** $H(p) = -\sum_x p(x) \log p(x)$

*   **KL Divergence:** $D_{\text{KL}}(p \| q) = \sum_x p(x) \log \frac{p(x)}{q(x)}$

## B.2 Mutual Information and Entropy

Mutual information (MI) measures the reduction in uncertainty (entropy) of one variable given knowledge of another. It satisfies the following additive identities:

$$ I(X;Y) = H(X) - H(X | Y) \qquad (37) $$

$$ I(X;Y) = H(Y) - H(Y | X) \qquad (38) $$

$$ I(X;Y) = H(X) + H(Y) - H(X,Y) \qquad (39) $$

where $H(X | Y)$ is the conditional entropy. We used the first formula later in our proof of the non-collapsed global optimum associated with the VJEPA objective (Theorem 1, and its detailed proof in Appendix.F).
## B.3 Chain Rule for Mutual Information

The chain rule for mutual information describes how the information provided by a pair of variables $(X, Y)$ about a third variable $Z$ can be decomposed into two distinct stages. The step-by-step expansion is as follows:

1. **Definition of Mutual Information:** The mutual information between $(X, Y)$ and $Z$ is the difference between the marginal entropy of $Z$ and the conditional entropy of $Z$ given $(X, Y)$:

$$I(X, Y; Z) = H(Z) - H(Z | X, Y) \qquad (40)$$

2. **Decomposition of Conditional Entropy:** Using the identity for conditional entropy, we can express $H(Z | X, Y)$ by first conditioning on $X$:

$$H(Z | X, Y) = H(Z | X) - I(Y; Z | X) \qquad (41)$$

Rearranging this reveals that $H(Z | X, Y)$ represents the remaining uncertainty in $Z$ after observing both $X$ and $Y$.

3. **Substitution and Re-grouping:** Substituting the decomposition into the original definition:

$$I(X, Y; Z) = H(Z) - [H(Z | X) - I(Y; Z | X)] \qquad (42)$$

$$I(X, Y; Z) = [H(Z) - H(Z | X)] + I(Y; Z | X) \qquad (43)$$

4. **Final Chain Rule Form:** Recognizing that $H(Z) - H(Z | X)$ is the definition of $I(X; Z)$, we arrive at the chain rule:

$$I(X, Y; Z) = I(X; Z) + I(Y; Z | X) \qquad (44)$$

## B.4 Symmetry Properties

A crucial distinction between these metrics is their behavior under argument swapping:

* **Mutual Information is Symmetric:**

$$I(X; Y) = I(Y; X).$$

Intuitively, the amount of information $X$ provides about $Y$ is identical to the information $Y$ provides about $X$. This symmetry arises from the identity $I(X; Y) = H(X) + H(Y) - H(X, Y)$, where the joint entropy $H(X, Y)$ is invariant to ordering.

* **KL Divergence is Asymmetric:**

$$D_{\text{KL}}(p \| q) \neq D_{\text{KL}}(q \| p) \quad (\text{generally}).$$

Minimizing $D_{\text{KL}}(p \| q)$ (forward KL) encourages the approximation $q$ to cover the support of $p$ (mode-covering), whereas minimizing $D_{\text{KL}}(q \| p)$ (reverse KL) encourages $q$ to seek a single mode of $p$ (mode-seeking).

* **Cross-Entropy is Asymmetric:**

$$H(p, q) \neq H(q, p).$$

Since $H(p, q) = H(p) + D_{\text{KL}}(p \| q)$, the asymmetry of the KL divergence imparts asymmetry to the cross-entropy.
**Common Confusion: Cross-Entropy vs. Joint Entropy.** It is easy to confuse Cross-Entropy $H(p, q)$ with Joint Entropy $H(X, Y)$.

* **Cross-Entropy** $H(p, q)$ is defined between two probability distributions $p$ and $q$:

$$H(p, q) = -\sum_{x} p(x) \log q(x).$$

It is **NOT** symmetric ($H(p, q) \neq H(q, p)$) because the roles of the weighting distribution $p(x)$ and the log-term $q(x)$ are distinct.

* **Joint Entropy** $H(X, Y)$ is defined over two random variables $X$ and $Y$:

$$H(X, Y) = -\sum_{x} \sum_{y} p(x, y) \log p(x, y).$$

It **IS** symmetric ($H(X, Y) = H(Y, X)$) because the joint probability $p(x, y)$ is invariant to ordering.

## B.5 Evidence Lower Bound (ELBO)

In variational inference, we often aim to approximate an intractable posterior $p(z | x)$ using a simpler distribution $q(z | x)$. The Evidence Lower Bound (ELBO) provides a tractable lower bound on the log-likelihood of the data (the evidence) $\log p(x)$. It relates the log-evidence to the KL divergence between the approximate and true posteriors:

$$\log p(x) = \mathcal{L}_{\text{ELBO}}(q) + D_{\text{KL}}(q(z | x) \| p(z | x)), \quad (45)$$

where the ELBO is defined as:

$$\mathcal{L}_{\text{ELBO}}(q) = \mathbb{E}_{q(z|x)}[\log p(x, z) - \log q(z | x)]. \quad (46)$$

Since the KL divergence is non-negative ($D_{\text{KL}} \ge 0$), the ELBO serves as a valid lower bound on the evidence:

$$\log p(x) \ge \mathcal{L}_{\text{ELBO}}(q). \quad (47)$$

Maximizing the ELBO is therefore equivalent to minimizing the KL divergence between the approximate posterior $q(z | x)$ and the true posterior $p(z | x)$, effectively pushing the approximation towards the true distribution.

## B.6 Non-Negativity Properties

An important property of these quantities is their non-negativity (for discrete variables), which ensures they act as valid objective functions or distance/divergence metrics:

1. **Entropy:** $H(X) \ge 0$. It is zero if and only if $X$ is deterministic.

2. **KL Divergence:** $D_{\text{KL}}(p \| q) \ge 0$. This is known as Gibbs' inequality. It is zero if and only if $p = q$ almost everywhere.

3. **Mutual Information:** $I(X; Y) \ge 0$. It is zero if and only if $X$ and $Y$ are independent.

4. **Cross-Entropy:** $H(p, q) \ge H(p)$. Since $D_{\text{KL}} \ge 0$, the cross-entropy is always lower-bounded by the entropy of the true distribution.
## B.7 Summary of Relationships

Table.5 provides a quick reference for these relationships.

<table><caption>Table 5: Summary of relationships and properties of information-theoretic quantities.</caption><thead><tr><th>Quantity</th><th>Relation to others</th><th>Non-negative?</th><th>Symmetric?</th></tr></thead><tbody><tr><td>Cross-Entropy (H(p, q))</td><td>H(p) + D<sub>KL</sub>(p||q)</td><td>Yes (≥ H(p))</td><td>No</td></tr><tr><td>KL Divergence (D<sub>KL</sub>)</td><td>H(p, q) - H(p)</td><td>Yes</td><td>No</td></tr><tr><td>Mutual Information (I(X; Y))</td><td>H(X) - H(X | Y)</td><td>Yes</td><td>Yes</td></tr><tr><td>Mutual Information (I(X; Y))</td><td>D<sub>KL</sub>(p(x, y)||p(x)p(y))</td><td>Yes</td><td>Yes</td></tr></tbody></table>

# C Compare VJEPA with Hidden Markov Models (HMMs)

VJEPA shares a structural resemblance to a Hidden Markov Model (HMM) in that both posit a latent state sequence to explain temporal data. However, the resemblance is only partial: classical HMMs are *generative probabilistic models of observations*, whereas VJEPA is a *representation-predictive* model whose training objective operates entirely in the latent space. Here we clarify the similarities and differences.

## C.1 Classical HMMs

An HMM defines a discrete-time latent Markov process {$s_t$}$_{t \ge 1}$ and a sequence of observations {$x_t$}$_{t \ge 1}$ via the joint factorization:

$$p(s_{1:T}, x_{1:T}) = p(s_1) \prod_{t=2}^{T} p(s_t | s_{t-1}) \prod_{t=1}^{T} p(x_t | s_t), \quad (48)$$

where $p(s_t | s_{t-1})$ is the transition model and $p(x_t | s_t)$ is the emission (observation) model. Learning<fnref n="24" /> typically proceeds by maximum likelihood (e.g. Expectation-Maximization), and inference computes the filtering posterior $p(s_t | x_{\le t})$ or smoothing posterior $p(s_t | x_{1:T})$.

In modern state-space model variants (SSMs), $s_t$ may be continuous and the dynamics nonlinear. However, the defining property remains: the model is trained to explain *observations* through an explicit hidden-to-observation likelihood $p(x_t | s_t)$.

## C.2 VJEPA as “Latent-Only” Predictive Modeling

VJEPA is not an observation-likelihood model. As detailed in Section 4, VJEPA introduces a latent representation for the target region and learns a predictive distribution (Eq. 7):

$$p_{\phi}(Z_T | Z_C, \xi_T),$$

<footnote n="24">²⁴A (discrete) HMM assumes an unobserved state sequence $s_{1:T}$ where each $s_t \in \{1, \dots, N\}$ (the state space) and an observed sequence $x_{1:T}$ where each $x_t \in \{v_1, \dots, v_M\}$ (the observation symbols). It is parameterized by $\lambda = (A, B, \pi)$, where: $A = [a_{ij}]$ with $a_{ij} = \text{Pr}(s_t = j | s_{t-1} = i)$ (transitions); $B = [b_j(k)]$ with $b_j(k) = \text{Pr}(x_t = v_k | s_t = j)$ (emissions); and $\pi = [\pi_i]$ with $\pi_i = \text{Pr}(s_1 = i)$ (initial distribution). The three canonical tasks are: (i) evaluation: compute $\text{Pr}(x_{1:T} | \lambda)$ (e.g. via the Forward algorithm); (ii) decoding: find the most likely hidden state sequence arg $\text{max}_{s_{1:T}} \text{Pr}(s_{1:T} | x_{1:T}, \lambda)$ (e.g. via the Viterbi algorithm); and (iii) learning: estimate $\lambda$ from data (e.g. via the Baum-Welch/EM algorithm). Details can be found in standard course materials, e.g. *Computational Intelligence*, which the author taught at University of Aberdeen.</footnote>
matched against an amortized target-encoder “inference” distribution (Eq. 10):

$$q_{\theta'}(Z_T | x_T). \qquad (49)$$

Training minimizes a representation-space negative log-likelihood plus regularization (Eq. 11), strictly avoiding the computation of a sensory emission likelihood $p(x_T | Z_T)$. While one could optionally introduce a decoder $p_\psi(x_T | Z_T)$, VJEPA does not optimize it; the learning signal flows solely through latent prediction.

## C.3 Similarities: Latent-State Semantics, Dynamics, and Beliefs

Both HMMs and VJEPA posit latent variables that summarize information needed for prediction.

* **Latent State:** In an HMM, $s_t$ is the latent state; in VJEPA, $Z_t$ plays an analogous role as a predictive latent embedding.

* **Transitions:** The VJEPA predictive distribution $p_\phi(Z_T | Z_C, \xi_T)$ acts as a transition model in latent space, mapping a “current” latent summary (context) to a distribution over a “future” latent representation.

* **Filtering:** If $q_{\theta'}(Z_T | x_T)$ has nontrivial variance, VJEPA supports belief-state style reasoning (means, covariances, sampling), similar to probabilistic filtering in HMMs.

## C.4 Differences

(1) **Observations are not modeled.** HMMs are trained to maximize log $p(x_{1:T})$ via the emission likelihood $p(x_t | s_t)$. In contrast, VJEPA focuses exclusively on predictive structure in representation space. It learns a predictive latent process without committing to a generative model of sensory observations, thereby avoiding the modeling of high-entropy noise.

(2) **The “inference model” is defined by the encoder, not the generator.** In an HMM, the posterior $p(s_t | x_{\le t})$ is derived by inverting the generative model Eq. (48). In VJEPA, $q_{\theta'}(Z_T | x_T)$ is an amortized distribution defined by the target encoder (the “teacher”). It serves as the ground truth for the predictor, rather than being an approximation of the predictor’s posterior.

(3) **The Markov property is optional.** The HMM assumes a strict first-order Markov property $p(s_t | s_{\le t}) = p(s_t | s_{t-1})$. VJEPA does not require a particular temporal factorization; it can predict multi-step futures directly ($Z_t \to Z_{t+\Delta}$) without iterating a 1-step transition, and without autoregressive observation factorization (see Section 5).

(4) **Representation geometry is learned, not prescribed.** Classical HMMs typically impose strong inductive biases on the latent space (e.g. discrete states or Gaussian dynamics). VJEPA learns the embedding geometry via neural encoders, making it closer in spirit to continuous Predictive State Representations (PSRs) than to classical HMM fitting.

## C.5 Sequential VJEPA as an “HMM Without Emissions”

Since an HMM is defined by the tuple (Transition, Emission), one can interpret sequential VJEPA as learning a latent predictive process (the Transition component) while deliberately omitting the Emission optimization. This design choice retains probabilistic semantics and uncertainty propagation while discarding the burden of reconstructing irrelevant observation details.
# D Kalman-Filter-Like Interpretation of VJEPA

Although VJEPA does not explicitly model observation likelihoods, its operation in latent space closely mirrors the prediction step of a Bayesian filter. This analogy becomes exact in the limit of linear-Gaussian dynamics, providing a rigorous grounding for VJEPA’s uncertainty estimates and justifying its use in belief-space planning.

## D.1 The Linear-Gaussian Benchmark

Consider a classical linear dynamical system with Gaussian noise. The latent state $s_t$ evolves according to:

$$s_{t+1} = As_t + Bu_t + w_t, \quad w_t \sim \mathcal{N}(0, Q), \qquad (50)$$

where $A$ is the state transition matrix, $B$ is the control matrix, and $Q$ is the process noise covariance.

A Kalman Filter maintains the belief state as a Gaussian distribution $b_t(s_t) = \mathcal{N}(\mu_t, P_t)$. The prediction step (or time update) propagates this belief forward:

$$\mu_{t+1|t} = A\mu_t + Bu_t, \qquad (51)$$

$$P_{t+1|t} = AP_tA^\top + Q. \qquad (52)$$

Here, the mean is shifted deterministically by the dynamics, while the uncertainty (covariance) grows due to the additive process noise $Q$.

## D.2 VJEPA as Generalized Latent Prediction

VJEPA generalizes this structure to non-linear, amortized inference in representation space. If we interpret the context embedding $Z_t$ as a sufficient statistic for the current belief history (i.e. $Z_t \approx \mu_t$), then the probabilistic predictor $p_\phi$ implements a generalized transition update.

Recall the VJEPA predictive distribution<fnref n="25" />:

$$p_{\phi}(Z_{t+1} | Z_t, u_t) = \mathcal{N}(\mu_{\phi}(Z_t, u_t), \Sigma_{\phi}(Z_t, u_t)).$$

Comparing this to the Kalman equations (51) and (52):

1. **Learned Transition ($\mu_\phi \approx A, B$):** The network $\mu_\phi(Z_t, u_t)$ approximates the non-linear flow of the system. Unlike the fixed matrix $A$, it can model state-dependent transitions and complex control interactions.

2. **State-Dependent Uncertainty ($\Sigma_\phi \approx Q$):** The predicted covariance $\Sigma_\phi(Z_t, u_t)$ generalizes the process noise $Q$. Importantly, VJEPA allows this noise to be *heteroscedastic* (state-dependent). For example, the model can predict high uncertainty (large $\Sigma_\phi$) in chaotic regions of the state space and low uncertainty (small $\Sigma_\phi$) in stable regions.

<footnote n="25">²⁵Unlike Theorem 6.7 (MAP/mean predictive control is optimal under quadratic costs), in which we assumed the Gaussian covariance $\Sigma$ of the predictive distribution does not depend on $u_t$, i.e. $p_\phi(Z_{t+1} | Z_t, u_t) = \mathcal{N}(Z_{t+1} | \mu_\phi(Z_t, u_t), \Sigma)$, here we make it a more general case.</footnote>
**Uncertainty Propagation.** In the standard Kalman Filter, the predictive uncertainty $P_{t+1|t}$ depends on the previous uncertainty $P_t$ via $AP_tA^\top$. In standard JEPA, $Z_t$ is typically a point estimate ($P_t \to 0$). However, if we treat the input $Z_t$ as a sample from a belief distribution, the VJEPA predictor naturally propagates this uncertainty via the sampling mechanism discussed in Section 4 (Eq. 12):

$$P_{t+1} \approx \underbrace{\mathbb{E}[\Sigma_{\phi}(Z_t, u_t)]}_{\text{Process Noise Q}} + \underbrace{\text{Var}(\mu_{\phi}(Z_t, u_t))}_{\text{Propagated Uncertainty } AP_tA^\top}.$$

Thus, VJEPA implicitly captures the full Kalman prediction logic: aleatoric uncertainty is output by $\Sigma_{\phi}$, while epistemic/belief uncertainty is propagated by passing a distribution of $Z_t$ through the non-linear mean $\mu_{\phi}$.

### D.3 Amortized Correction

A full Bayesian filter also requires a **correction step** (measurement update) to incorporate new observations $x_{t+1}$:

$$\mu_{t+1} = \mu_{t+1|t} + K_t(x_{t+1} - H\mu_{t+1|t}).$$

VJEPA handles this differently. Instead of an explicit update equation involving observation likelihoods and Kalman gains $K_t$, VJEPA **amortizes** the correction into the context encoder $f_{\theta}$.

$$Z_{t+1} = f_{\theta}(x_{\le t+1}).$$

The encoder learns to map the updated history $x_{\le t+1}$ directly to the posterior latent state $Z_{t+1}$. The VJEPA training objective (minimizing the divergence between the *predicted* belief $p_{\phi}(\cdot|Z_t)$ and the *encoded* belief $q_{\theta'}(\cdot|x_{t+1})$) effectively forces the predictor to be consistent with the encoder's implicit Bayesian updates.

## E A Particle-Filter Interpretation of VJEPA

VJEPA admits a natural *particle filter* interpretation when used sequentially. This allows us to track multimodal beliefs about the latent state without restricting ourselves to Gaussian assumptions, using a method known as Sequential Importance Resampling (SIR). The complete algorithmic flow, illustrating how latent predictions are fused with surrogate measurements, is visualized in Fig. 5.

Let the latent belief at time $t$ be represented by a set of $K$ weighted particles $\{Z_t^{(k)}, w_t^{(k)}\}_{k=1}^K$ approximating the posterior distribution $b_t(Z_t) \approx p(Z_t | x_{1:t})$:

$$b_t(Z_t) \approx \sum_{k=1}^{K} w_t^{(k)} \delta(Z_t - Z_t^{(k)}), \quad \sum_{k=1}^{K} w_t^{(k)} = 1.$$

Given an action $u_t$ (or target specification $\xi_{t+1}$), the filter proceeds in three main stages:

1. **Prediction (Proposal).** We propagate each particle forward through the learned VJEPA dynamics. We employ a *bootstrap filter* design, using the dynamics model itself as the importance sampling proposal distribution:

$$Z_{t+1}^{(k)} \sim p_{\phi}(Z_{t+1} | Z_t^{(k)}, \xi_{t+1}), \quad k = 1, \dots, K.$$

This step draws samples from the predictive belief $p(Z_{t+1} | x_{1:t})$.
**2. Update (Weighting).** To incorporate the new information from observation $x_{t+1}$, we update the importance weights. Since the proposal distribution is the transition prior, the unnormalized weight $\tilde{w}_{t+1}^{(k)}$ is simply proportional to the likelihood of the observation given the particle's state. VJEPA provides a representation-space inference distribution $q_{\theta'}(Z_{t+1} | x_{t+1})$ (Eq. 10), which serves as a surrogate measurement model. We consider two cases for defining this likelihood:

**Option A: Explicit Observation Model.** If one instantiates the optional decoder $p_{\psi}(x_{t+1} | Z_{t+1})$ (Eq. 8), the weight update is the standard observation likelihood:

$$ \tilde{w}_{t+1}^{(k)} \propto w_{t}^{(k)} p_{\psi}(x_{t+1} | Z_{t+1}^{(k)}). $$

**Option B: Latent-Space Update (Likelihood-Free).** Without a decoder, we can derive a pseudo-likelihood in latent space. By inverting the inference encoder using Bayes' rule, we have $p(x | Z) \propto p(Z | x)/p(Z)$. Treating the target encoder $q_{\theta'}$ as an approximation to the true posterior $p(Z | x)$, the weight update becomes:

$$ \tilde{w}_{t+1}^{(k)} \propto w_{t}^{(k)} \frac{q_{\theta'}(Z_{t+1}^{(k)} | x_{t+1})}{p_{\text{ref}}(Z_{t+1}^{(k)})}, $$

where $p_{\text{ref}}(Z)$ is the fixed reference prior (e.g. $N(0, I)$) used in the VJEPA KL regularization term (Eq. 11). Intuitively, this ratio measures how much more probable the particle $Z_{t+1}^{(k)}$ is according to the data-informed encoder compared to the uninformed prior.

**3. Normalize and Resample.** The weights are normalized so they sum to one: $w_{t+1}^{(k)} = \tilde{w}_{t+1}^{(k)} / \sum_{j=1}^{K} \tilde{w}_{t+1}^{(j)}$. To prevent particle degeneracy (where one particle accumulates all the weight), a resampling step is performed, drawing K new particles with replacement from the current set with probabilities proportional to their weights.

This interpretation establishes sequential VJEPA as a full belief-propagation mechanism over latent predictive states, where $p_{\phi}$ provides the dynamics and $q_{\theta'}$ injects measurement information directly in the latent space.
<figure><img src="image_8.png" /><figcaption>Figure 5: Schematic overview of a Particle Filter (Sequential Monte Carlo) algorithm applied within the VJEPA framework. The process begins with a set of weighted particles representing the current belief. The **Prediction** step uses the learned probabilistic dynamics model $p_{\phi}$ as a proposal distribution to propagate particles forward. The **Update** step re-weights these particles based on new observation information, which is incorporated via a surrogate likelihood derived from the target encoder $q_{\theta'}$ (or an optional decoder $p_{\psi}$). Finally, **Resampling** is performed to avoid particle degeneracy, resulting in an updated particle set representing the posterior belief for the next time step.</figcaption></figure>

# F Detailed Proof of Theorem 1

We provide a detailed proof that the VJEPA objective admits no collapsed global optimum under the stated assumptions in Theorem 1.
**Theorem 1: No Collapsed Global Optimum under Target Diversity**

Consider the VJEPA objective

$$ \mathcal{L}_{\text{VJEPA}} = \mathbb{E}_x \mathbb{E}_{Z_T \sim q_{\theta'}(\cdot | x_T)} [-\log p_{\phi}(Z_T | Z_C, \xi_T)] + \beta \mathbb{E}_x \text{KL}(q_{\theta'}(Z_T | x_T) \| p(Z_T)). $$

Assume:

(i) **(Target diversity)** There exist $x_T, x'_T$ such that $q_{\theta'}(\cdot | x_T) \neq q_{\theta'}(\cdot | x'_T)$.

(ii) **(Nontrivial conditioning)** The predictive family $\{p_{\phi}(\cdot | Z_C, \xi_T)\}$ can represent different distributions for different $Z_C$ (for fixed $\xi_T$).

Then no global minimizer of $\mathcal{L}_{\text{VJEPA}}$ satisfies $f_{\theta}(x_C) \equiv c$ for all $x_C$.

*Proof.* The proof follows three steps: we first analyze the structural consequences of a collapsed context encoder, then derive the minimal achievable prediction loss in this regime, and finally demonstrate that non-collapsed representations achieve strictly lower VJEPA objective values.

**Step 1: Consequence of a collapsed context encoder.** Assume, for contradiction, that the context encoder collapses:

$$ f_{\theta}(x_C) = c \quad \forall x_C. $$

Then the predictive model cannot depend on the context content and reduces to

$$ p_{\phi}(Z_T | Z_C, \xi_T) = p_{\phi}(Z_T | c, \xi_T), $$

i.e. an *unconditional* distribution for each $\xi_T$. Under this restriction, the first term of the VJEPA objective becomes

$$ \mathcal{L}_{\text{pred}}^{\text{coll}} = \mathbb{E}_{\xi_T} \mathbb{E}_{x_T | \xi_T} \mathbb{E}_{Z_T \sim q_{\theta'}(\cdot | x_T)} [-\log p_{\phi}(Z_T | c, \xi_T)]. \quad (53) $$

**Step 2: Optimal unconditional predictor.** For a fixed $\xi_T$, define the *aggregated target distribution*

$$ \bar{q}(Z_T | \xi_T) := \mathbb{E}_{x_T | \xi_T} [q_{\theta'}(Z_T | x_T)]. $$

It is a standard result that the distribution minimizing cross-entropy with respect to $\bar{q}$ is $\bar{q}$ itself. Hence,

$$ \inf_{p(\cdot | c, \xi_T)} \mathbb{E}_{x_T | \xi_T} \mathbb{E}_{Z_T \sim q_{\theta'}(\cdot | x_T)} [-\log p(Z_T)] = H(\bar{q}(\cdot | \xi_T)). $$

Therefore, the best achievable collapsed prediction loss equals

$$ \mathcal{L}_{\text{pred}}^{\text{coll}} = \mathbb{E}_{\xi_T} H(Z_T | \xi_T). \quad (54) $$

**Step 3: Decomposition via conditional mutual information.** By standard entropy identities (see Appendix.B),

$$ H(Z_T | \xi_T) = H(Z_T | X_T, \xi_T) + I(Z_T; X_T | \xi_T). $$

Assumption (i) (*target diversity*) implies that

$$ I(Z_T; X_T | \xi_T) > 0 $$

for at least one $\xi_T$ with nonzero probability mass. Hence the collapsed predictor incurs a strictly positive irreducible excess loss.
**Step 4: Advantage of non-collapsed representations.** Consider any non-collapsed encoder producing distinct $Z_C$ values correlated with $x_T$. By Assumption (ii), the predictive family can represent different $p_\phi(Z_T | Z_C, \xi_T)$ and therefore approximate $q_{\theta'}(Z_T | x_T)$ conditionally. In the realizable limit,

$$ \mathcal{L}_{\text{pred}}^{\text{non-coll}} = \mathbb{E}_{\xi_T} \mathbb{E}_{x_T|\xi_T} H(Z_T | X_T = x_T, \xi_T) < \mathcal{L}_{\text{pred}}^{\text{coll}} $$

**Step 5: Role of the KL regularizer.** The KL term

$$ \text{KL}(q_{\theta'}(Z_T | x_T) \| p(Z_T)) $$

is independent of $Z_C$ and therefore does not eliminate the gap identified above. Its role is to prevent pathological collapse of the target encoder, not to create the strict separation between collapsed and non-collapsed optima.

**Conclusion** A collapsed context encoder forces the predictive model to fit a single unconditional distribution to multiple distinct target distributions, incurring an irreducible mutual-information gap. Since a non-collapsed solution can achieve strictly lower loss, no collapsed representation can be globally optimal.

**Interpretation.** Collapse is avoided because predictive uncertainty depends on the target input. VJEPA therefore prevents collapse through information mismatch, not architectural heuristics. Theorem 1 formalizes the intuitive point that if target embeddings vary with the target input, then a collapsed context representation forces the predictor to fit a single unconditional distribution to multiple distinct targets, incurring an irreducible prediction loss. Consequently, any solution satisfying the collapsed mode hypothesis cannot be globally optimal for the VJEPA objective. Under the stated assumptions, preventing collapse is therefore an intrinsic property of the objective itself, rather than a consequence of architectural heuristics or training asymmetries. □

# G Detailed Proof of Theorem 3

We restate Theorem 3 in Section 6.7 for convenience.
**Theorem 3: Optimality of MAP (Mean) Control under Quadratic Costs**

Consider a one-step control problem with latent predictive model $p_{\phi}(Z_{t+1} | Z_t, u_t)$ and stage cost $c(Z_{t+1}, u_t)$. Assume that for each $(Z_t, u_t)$,

$$p_{\phi}(Z_{t+1} | Z_t, u_t) = \mathcal{N}(Z_{t+1} | \mu_{\phi}(Z_t, u_t), \Sigma),$$

where the covariance $\Sigma$ does not depend on $u_t$. Let the stage cost be quadratic in $Z_{t+1}$:

$$c(Z_{t+1}, u_t) = (Z_{t+1} - z^*)^\top Q_c(Z_{t+1} - z^*) + r(u_t),$$

with $Q_c \succeq 0$ and arbitrary $r(\cdot)$. Then the action minimizing expected cost,

$$u_t^* \in \arg\min_{u_t} \mathbb{E}[c(Z_{t+1}, u_t) | Z_t, u_t],$$

is equivalently obtained by minimizing the cost at the predictive mean:

$$u_t^* \in \arg\min_{u_t} c(\mu_{\phi}(Z_t, u_t), u_t).$$

Moreover, since a Gaussian has $\mu_{\phi}(Z_t, u_t)$ as its unique MAP point, this is equivalently MAP control.

*Proof.* Fix a time index $t$ and a given latent state $Z_t$. For each candidate action $u_t$, define the random variable

$$Z \equiv Z_{t+1} \sim \mathcal{N}(\mu, \Sigma), \quad \mu \equiv \mu_{\phi}(Z_t, u_t),$$

where $\Sigma$ is constant with respect to $u_t$ by assumption. We want to analyze this objective

$$J(u_t) := \mathbb{E}[c(Z, u_t) | Z_t, u_t] = \mathbb{E}[(Z - z^*)^\top Q_c(Z - z^*)] + r(u_t). \quad (55)$$

in which the only nontrivial term is the expectation of the quadratic form.

**Step 1: Expand the quadratic form.** Let $d := Z - z^*$. Then

$$d^\top Q_c d = (Z - z^*)^\top Q_c (Z - z^*).$$

Insert and subtract the mean $\mu$:

$$Z - z^* = (Z - \mu) + (\mu - z^*).$$

Therefore,

$$(Z - z^*)^\top Q_c (Z - z^*) = ((Z - \mu) + (\mu - z^*))^\top Q_c ((Z - \mu) + (\mu - z^*)).$$

Expand the product into four terms:

$$((Z - \mu) + (\mu - z^*))^\top Q_c ((Z - \mu) + (\mu - z^*)) = (Z - \mu)^\top Q_c (Z - \mu) + (Z - \mu)^\top Q_c (\mu - z^*) + (\mu - z^*)^\top Q_c (Z - \mu) + (\mu - z^*)^\top Q_c (\mu - z^*).$$
**Step 2: Take expectations term-by-term.** Take $\mathbb{E}[\cdot]$ under $Z \sim \mathcal{N}(\mu, \Sigma)$.

(i) Cross terms vanish. Because $\mathbb{E}[Z - \mu] = 0$, we have

$$ \mathbb{E}[(Z - \mu)^{\top} Q_c(\mu - z^*)] = \mathbb{E}[Z - \mu]^{\top} Q_c(\mu - z^*) = 0. $$

Similarly,

$$ \mathbb{E}[(\mu - z^*)^{\top} Q_c(Z - \mu)] = (\mu - z^*)^{\top} Q_c \mathbb{E}[Z - \mu] = 0. $$

(ii) The constant term remains. Since $(\mu - z^*)^{\top} Q_c(\mu - z^*)$ is deterministic given $u_t$,

$$ \mathbb{E}[(\mu - z^*)^{\top} Q_c(\mu - z^*)] = (\mu - z^*)^{\top} Q_c(\mu - z^*). $$

(iii) The centered quadratic term becomes a trace. Let $\varepsilon := Z - \mu$. Then $\mathbb{E}[\varepsilon] = 0$ and $\mathbb{E}[\varepsilon\varepsilon^{\top}] = \Sigma$.

We claim

$$ \mathbb{E}[\varepsilon^{\top} Q_c \varepsilon] = \text{tr}(Q_c \Sigma). $$

To see this, use the identity for any vector $\varepsilon$ and matrix $Q_c$:

$$ \varepsilon^{\top} Q_c \varepsilon = \text{tr} (\varepsilon^{\top} Q_c \varepsilon) = \text{tr} (Q_c \varepsilon \varepsilon^{\top}), $$

where we used $\text{tr}(a) = a$ for a scalar $a$, and cyclicity of trace $^{26}$. Taking expectations,

$$ \mathbb{E}[\varepsilon^{\top} Q_c \varepsilon] = \mathbb{E}\left[\text{tr}\left(Q_c \varepsilon \varepsilon^{\top}\right)\right] = \text{tr}\left(Q_c \mathbb{E}[\varepsilon \varepsilon^{\top}]\right) = \text{tr}(Q_c \Sigma), $$

where linearity of expectation and linearity of trace justify exchanging $\mathbb{E}$ and $\text{tr}$.

**Step 3: Assemble the expected quadratic cost.** Combining the above,

$$ \mathbb{E}[(Z - z^*)^{\top} Q_c (Z - z^*)] = (\mu - z^*)^{\top} Q_c (\mu - z^*) + \text{tr}(Q_c \Sigma). $$

Hence the full objective in Eq. (55) becomes

$$ J(u_t) = (\mu_{\phi}(Z_t, u_t) - z^*)^{\top} Q_c (\mu_{\phi}(Z_t, u_t) - z^*) + \text{tr}(Q_c \Sigma) + r(u_t). $$

**Step 4: Use the covariance-independence assumption to reduce the argmin.** By assumption, $\Sigma$ does not depend on $u_t$, and therefore $\text{tr}(Q_c \Sigma)$ is a constant with respect to $u_t$. Adding or subtracting a constant does not change the minimizer, so

$$ \begin{aligned} \arg \min_{u_t} J(u_t) &= \arg \min_{u_t} \left[ (\mu_{\phi}(Z_t, u_t) - z^*)^{\top} Q_c (\mu_{\phi}(Z_t, u_t) - z^*) + r(u_t) \right] \\ &= \arg \min_{u_t} c(\mu_{\phi}(Z_t, u_t), u_t), \end{aligned} $$

which proves the mean-planning claim.

<footnote>$^{26}$The cyclic property of the trace says that $\text{tr}(XY) = \text{tr(YX)}$; in our case $X = \varepsilon^T$ and $Y = Q_c \varepsilon$.</footnote>
**Step 5: Mean equals MAP for a Gaussian.** For $Z \sim \mathcal{N}(\mu, \Sigma)$ with $\Sigma > 0$, the log-density is

$$ \log p(Z) = -\frac{1}{2}(Z - \mu)^{\top}\Sigma^{-1}(Z - \mu) + \text{const}, $$

which is strictly concave in $Z$ and uniquely maximized at $Z = \mu$. Thus the MAP point equals the mean:

$$ \arg\max_{Z} p(Z) = \mu. $$

Therefore minimizing cost at the predictive mean is equivalently minimizing cost at the predictive MAP point, i.e. MAP control. $\square$

*Remark 3 (Assumptions matter).* If $\Sigma = \Sigma(u_t)$ depends on the action, then the term $\text{tr}(Q_c\Sigma(u_t))$ is no longer constant and the optimal action generally depends on predictive uncertainty. Likewise, for non-quadratic costs or non-Gaussian/multimodal $p_\phi$, the expected cost typically cannot be reduced to evaluating a single point estimate.

*Remark 4 (Post-action cost indexing).* We index the stage cost as $c(Z_{t+1}, u_t)$ to emphasize that cost is incurred *after* applying the action $u_t$ and observing its effect on the system. This convention is standard in stochastic MPC and belief-space control, where actions are optimized based on predicted future states rather than current ones [53, 62].

Some reinforcement learning and dynamic programming formulations instead write the stage cost as $c(Z_t, u_t)$ [75, 59]. The two conventions are equivalent under expectation for Markov dynamics, since

$$ \mathbb{E}[c(Z_{t+1}, u_t) | Z_t, u_t] $$

can be absorbed into a redefined cost function of $(Z_t, u_t)$. We adopt the post-action form to align with prediction-based control and MPC.

# H Proof of Proposition 1 (Invariance to Nuisance)

We provide a formal derivation showing that generative objectives force the representation to encode nuisance information, whereas VJEPA objectives do not.

**Setup.** Let the observation at time $t$ be a tuple $x_t = (s_t, n_t)$, where:

* $s_t$ is the **signal**: it contains information relevant for predicting the future target $Z_{t+\Delta}$.

* $n_t$ is the **nuisance**: it is independent of the future target, i.e. $I(n_t; Z_{t+\Delta} | s_t) = 0$.

Let $Z_t$ be the representation of $x_t$.

## H.1 Case 1: Generative / Reconstruction Objective

Consider a generative world model (e.g. VAE or autoregressive model) that maximizes the log-likelihood of the observation $x_t$ given the representation $Z_t$:

$$ \mathcal{L}_{\text{gen}} = \mathbb{E}_{x_t} [\log p(x_t | Z_t)]. $$

Using the decomposition $x_t = (s_t, n_t)$, the likelihood factorizes as (using chain rule of probability):

$$ \log p(s_t, n_t | Z_t) = \log p(s_t | Z_t) + \log p(n_t | s_t, Z_t). $$

To maximize this objective (globally), the representation $Z_t$ must maximize both terms.
* Maximizing $\log p(s_t | Z_t)$ requires $Z_t$ to encode the signal $s_t$.

* Maximizing $\log p(n_t | s_t, Z_t)$ requires $Z_t$ to encode the nuisance $n_t$ (so that the posterior entropy $H(n_t | Z_t)$ is minimized).

**Conclusion:** A solution $Z_t$ that discards $n_t$ (i.e. $Z_t$ depends only on $s_t$) is *suboptimal* for the generative objective because it fails to reconstruct the nuisance component $n_t$. Thus, generative models *force* the representation to be maximal (retaining $N$).

## H.2 Case 2: VJEPA Prediction Objective

Consider the VJEPA objective, which maximizes the likelihood of the *future target* $Z_{t+\Delta}$ given the current representation $Z_t$:

$$ \mathcal{L}_{\text{VJEPA}} = \mathbb{E}[\log p(Z_{t+\Delta} | Z_t)]. $$

By the definition of nuisance variables, $Z_{t+\Delta}$ is conditionally independent of $n_t$ given $s_t$. Therefore, the true predictive distribution satisfies:

$$ p(Z_{t+\Delta} | s_t, n_t) = p(Z_{t+\Delta} | s_t). $$

Now compare two candidate representations:

1. **Maximal representation:** $Z_t^{\max} = (s_t, n_t)$.

2. **Minimal representation:** $Z_t^{\min} = s_t$ (discards $n_t$).

Substituting these into the objective:

$$ \mathcal{L}_{\text{VJEPA}}(Z_t^{\max}) = \mathbb{E}[\log p(Z_{t+\Delta} | s_t, n_t)] = \mathbb{E}[\log p(Z_{t+\Delta} | s_t)] = \mathcal{L}_{\text{VJEPA}}(Z_t^{\min}). $$

**Conclusion:** The minimal representation $Z_t^{\min}$ achieves the exact same optimal loss value as the maximal representation. Since the objective does not penalize the absence of $n_t$, VJEPA *admits* (is invariant to) minimal solutions that discard nuisance variability.

# I Derivation of the BJEPA Predictive Factorization

In this section, we derive the relationship between the joint posterior $p(Z_T | Z_C, \eta)$ and the individual predictive factors used in the BJEPA architecture (Eq. 30). We aim to express the posterior probability of the target latent state $Z_T$, given the historical context $Z_C$ and auxiliary information $\eta$, in terms of the individual conditional probabilities $p(Z_T | Z_C)$ and $p(Z_T | \eta)$.

**Assumption: Conditional Independence.** We assume that the history encoder and the auxiliary encoder provide independent information about the target state. Formally, we assume that $Z_C$ and $\eta$ are conditionally independent given the true target $Z_T$:

$$ p(Z_C, \eta | Z_T) = p(Z_C | Z_T) p(\eta | Z_T). \qquad (56) $$

This implies that if the true future state is known, the past history and the auxiliary constraints (e.g. goals) do not provide additional information about each other.