---
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.
author:
- Yongchao Huang [^1]
bibliography:
- reference.bib
date: 2026-01-20
title: "**VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models**"
---
```{=latex}
\newcommand{\cmark}{\ding{51}}
```
```{=latex}
\newcommand{\xmark}{\ding{55}}
```
```{=latex}
\maketitle
```
```{=latex}
\tableofcontents
```
```{=latex}
\newpage
```
# Introduction {#sec:intro}
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 [@kingma2014vae; @rezende2014stochastic; @oord2016pixelrnn; @raffel2020t5], and (ii) contrastive learning, which discriminates between positive and negative pairs [@hadsell2006dimensionality; @oord2018cpc; @chen2020simclr; @he2020moco]. Both approaches suffer from fundamental limitations: generative models must model nuisance variability and high-entropy details irrelevant for downstream tasks [@alemi2017ib; @higgins2017beta; @locatello2019challenging], while contrastive methods require carefully engineered negative sampling schemes and can introduce representational bias [@tschannen2019mutual; @saunshi2019theoretical; @chen2020simclr].
Joint Embedding Predictive Architectures (JEPA) offer a third paradigm [@lejepa2022ijepa; @bardes2024vjepa]. Rather than reconstructing observations or contrasting samples, JEPA learns by predicting *representations* of missing or future data [@lejepa2022ijepa]. This removes the need for observation-level likelihoods and explicitly biases learning toward predictable, task-relevant structure [@bardes2024vjepa; @larsen2016autoencoding; @Singh2004]. Recent instantiations such as I-JEPA [@lejepa2022ijepa] and V-JEPA [@bardes2024vjepa; @bardes2024vjepa_openReview] have demonstrated strong empirical performance on vision and video tasks. Early on, LeCun [@LeCun2022HJepa] 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. [@Bardes2025VJEPA2] 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. [@terver2025jepaworldmodels] 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. [@Destrade2026ValueGuidedJEPA] 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 [@lejepa2022ijepa; @bardes2024vjepa]. 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.
# Related Work {#sec:related_work}
## Predictive State Representations
Predictive State Representations (PSRs) were introduced by Sutton et al. [@Sutton2001] and independently formalized by Singh et al. [@Singh2004]. They replace latent states with sufficient statistics of future observations, representing the state as $$s_t = p(o_{t+1:t+H} \mid o_{\le t}),$$ where $H$ denotes the prediction horizon. Sutton et al. [@Sutton2001] 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. [@Singh2004] 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$.
## World Models and Latent Dynamics
Model-based latent dynamics have been widely explored in reinforcement learning. PlaNet [@hafner2018planet] 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 [@hafner2019dreamer] 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} \mid z_t, a_t), \quad p(x_t \mid z_t),$$ where the observation likelihood $p(x_t \mid z_t)$ is explicitly modeled.
In contrast, JEPA [@lejepa2022ijepa], V-JEPA (video-JEPA) [@bardes2024vjepa] 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 [@Schrittwieser2020muzero] 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.
## 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 [@friston2010free; @friston2012perceptions; @friston2017active]. 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)} \left[\log q(z) - \log p(x, z)\right],$$ 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 \mid x)$, thereby implementing perception as approximate Bayesian inference. This formulation was introduced and developed in early work by Friston [@friston2010free] and extended to perception-action loops in subsequent studies [@friston2012perceptions].
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 [@friston2017active].
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 \mid 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.
## Joint Embedding Predictive Architectures (JEPAs)
JEPA was introduced to self-supervised learning by Assran et al. [@lejepa2022ijepa] 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 [@lejepa2022ijepa].
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 [@bardes2024vjepa]. 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 [@bardes2023mcJEPA] incorporates motion and content prediction objectives for richer video representation [@bardes2023mcJEPA], S-JEPA applies the JEPA principle to skeletal action recognition [@abdelfattah2024sjepa], and Point-JEPA adapts JEPA to point cloud self-supervised learning [@saito2025pointjepa]. Other variants explore JEPA in graph domains [@skenderi2025graphJEPA] and in joint vision-language models [@chen2025vlJEPA]. Additional work (e.g. C-JEPA [@mo2024cJEPA], DSeq-JEPA [@he2025dseqJEPA]) 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 [@bardes2022VICReg] in C-JEPA [@mo2024cJEPA]) 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.
## 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 [@LeCun2022HJepa]. 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. [@terver2025jepaworldmodels] 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. [@Destrade2026ValueGuidedJEPA] 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. [@Bardes2025VJEPA2] 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.
# JEPA: Background and Deterministic Formulation {#sec:jepa}
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 [@lejepa2022ijepa; @bardes2024vjepa]. 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 [@lejepa2022ijepa]) or temporal masking (as in video JEPA [@bardes2024vjepa]), but the formulation itself is agnostic to the specific modality. The context is mapped to a latent representation using a trainable encoder $$\label{eq:context_encoder}
Z_C = f_\theta(x_C),$$ which produces a compact embedding summarizing the available information. The target is encoded using a separate encoder $$\label{eq:target_encoder}
Z_T = f_{\theta'}(x_T),$$ whose parameters $\theta'$ are updated as an exponential moving average (EMA) of the context encoder parameters[@mnih2015human; @mnih2016asynchronousmethodsdeepreinforcement; @lillicrap2019continuouscontroldeepreinforcement; @grill2020byol]: $$\label{eq:target_encoder_weights_EMA}
\theta' \leftarrow \tau \theta' + (1-\tau)\theta.$$ This asymmetric design, inspired by prior work [@grill2020byol], 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 $$\label{eq:context_predictor}
\hat{Z}_T = g_\phi(Z_C, \xi_T),$$ 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: $$\label{eq:JEPA_squared_loss}
\begin{aligned}
\mathcal{L}_{\text{JEPA}}
&= \left\| \hat{Z}_T - Z_T \right\|^2 \\
&= \left\| g_\phi \left(Z_C, \xi_T \right) - Z_T \right\|^2 \\
&= \left\| g_\phi \left(f_\theta(x_C), \xi_T \right)
- f_{\theta'}(x_T) \right\|^2 \\
&= \left\| g_\phi \left(f_\theta(x_C), \xi_T \right)
- f_{\tau \theta' + (1-\tau)\theta}(x_T) \right\|^2 ,
\end{aligned}$$ where the target encoder parameters are updated as per Eq. `\eqref{eq:target_encoder_weights_EMA}`{=latex}. 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.`\ref{app:lsq_equivalence_with_MLE}`{=latex} a brief derivation): $$\label{eq:JEPA_loss_prob_explain}
\mathcal{L}_{\text{JEPA}}
\propto
- \log p(Z_T \mid Z_C),
\quad
p(Z_T \mid Z_C)
=
\mathcal{N} \left(
Z_T \middle|
g_\phi(Z_C, \xi_T), \sigma^2 I
\right),$$ 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.
# VJEPA: Variational Formulation of JEPA {#sec:vjepa}
Eq. `\eqref{eq:JEPA_squared_loss}`{=latex} shows that the deterministic JEPA objective is equivalent to maximum likelihood under an isotropic Gaussian predictive model with fixed variance (Eq. `\eqref{eq:JEPA_loss_prob_explain}`{=latex}). 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.`\ref{tab:jepa_vs_vjepa}`{=latex}; these features will be progressively introduced in later sections.
```{=latex}
\footnotesize
```
::: {#tab:jepa_vs_vjepa}
**Feature** **JEPA** **VJEPA**
---------------------------------------- ----------------------------------- ------------------------------------------------
Target prediction $\hat Z_T = g_\phi(Z_C,\xi_T)$ $p_\phi(Z_T \mid Z_C,\xi_T)$
Prediction type Point estimate Predictive distribution
Training objective MSE (equally Gaussian likelihood) Negative log-likelihood $+$ KL
Uncertainty modeling Implicit (fixed variance) Explicit (learned variance / distribution)
Multi-modal futures Not supported Supported
Collapse avoidance Architectural (EMA, asymmetry) Objective-level (predictive mismatch + KL reg)
Probabilistic interpretation Implicit Gaussian Explicit variational model
Compatibility with filtering / control Limited Natural (belief propagation)
: Comparison between deterministic JEPA and its variational extension VJEPA. VJEPA preserves the structural design of JEPA while endowing it with an explicit probabilistic semantics.
:::
## 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. `\eqref{eq:context_encoder}`{=latex}. 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. `\eqref{eq:context_predictor}`{=latex}. VJEPA introduces a probabilistic predictive model over the target representation: $$\label{eq:vjepa_predictive_model}
p_\phi(Z_T \mid Z_C, \xi_T),$$ 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 \mid Z_C, \xi_T)=\mathcal{N}\!\big(\mu_\phi(Z_C, \xi_T),\,\Sigma_\phi(Z_C)\big)$; if $p_\phi(Z_T \mid Z_C, \xi_T)=g_\phi(Z_C,\xi_T)$ is chosen, VJEPA reduces to the deterministic JEPA. The probabilistic, predictive model Eq. `\eqref{eq:vjepa_predictive_model}`{=latex} can be further decomposed, e.g. formulated as a Bayesian posterior (see later Section.`\ref{sec:bjepa}`{=latex}).
*Optionally*, one may also posit an observation model $$\label{eq:vjepa_obs_model_optional}
p_\psi(x_T \mid Z_T),$$ yielding the joint factorization $$\label{eq:vjepa_joint}
p_{\phi,\psi}(Z_T, x_T \mid x_C, \xi_T)
=
p_\phi(Z_T \mid Z_C, \xi_T)\, p_\psi(x_T \mid Z_T).$$ In this text, in the spirit of JEPA, we *do not* optimize an observation-level reconstruction objective over $p_\psi(x_T \mid Z_T)$; the learning signal is provided solely through the predictive distribution in representation space, $p_\phi(Z_T \mid Z_C,\xi_T)$.
## 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: $$\label{eq:vjepa_inference_target}
q_{\theta'}(Z_T \mid x_T),$$ which may be instantiated as a diagonal Gaussian whose mean and variance are parameterized by $f_{\theta'}(x_T)$. The EMA update in Eq. `\eqref{eq:target_encoder_weights_EMA}`{=latex} ensures that the inference target evolves slowly, which stabilizes training and prevents degenerate fixed points.
## Variational Objective
Given $(x_C, x_T, \xi_T)$, VJEPA trains the predictive model $p_\phi(Z_T\mid Z_C,\xi_T)$ to match the target-encoder distribution $q_{\theta'}(Z_T\mid x_T)$. A simple and effective objective is the *regularized negative log-likelihood*: $$\label{eq:vjepa_objective}
\mathcal{L}_{\text{VJEPA}}
=
\mathbb{E}_{x}
\mathbb{E}_{Z_T \sim q_{\theta'}(\cdot \mid x_T)}
\Big[-\log p_\phi(Z_T \mid Z_C, \xi_T)\Big]
+
\beta\,
\mathbb{E}_{x}
\mathrm{KL} \left(q_{\theta'}(Z_T \mid x_T)\,\|\,p(Z_T)\right),$$ 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\mid Z_C,\xi_T)$, and (iii) yields collapse-avoidance guarantees under mild conditions (see later Section `\ref{sec:properties_of_VJEPA}`{=latex}).
#### Deterministic JEPA as a special case.
If we set $$q_{\theta'}(Z_T\mid x_T)=\delta(Z_T-f_{\theta'}(x_T))$$ and $$p_\phi(Z_T\mid Z_C,\xi_T)
=
\mathcal{N}(Z_T\mid g_\phi(Z_C,\xi_T), \sigma^2 I)$$ with fixed $\sigma^2$, then minimizing Eq. `\eqref{eq:vjepa_objective}`{=latex} reduces (up to additive constants) to the squared-loss JEPA objective in Eq. `\eqref{eq:JEPA_squared_loss}`{=latex}.
## Schematic Overview
The complete VJEPA architecture, illustrating the probabilistic predictor, target sampling, and the variational loss components, is schematically shown in Fig.`\ref{fig:vjepa_architecture}`{=latex}.
Variational JEPA (VJEPA) architecture. The context xC is encoded into ZC, conditioning a probabilistic predictor pϕ(ZT ∣ ZC, ξT). The target xT is encoded into a distribution qθ′ (typically via EMA), from which a sample ZT is drawn. Training minimizes the sum of a negative log-likelihood term (red box) and a KL regularization term (blue box). fθ′ can output sufficient statistics (e.g. mean and covariance for Gaussians) for parameterizing qθ′.
## Training Procedure
Algorithm `\ref{alg:vjepa}`{=latex} 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.`\ref{sec:toy_experiment}`{=latex}.
```{=latex}
\begin{algorithm}[H]\caption{VJEPA Training}
\label{alg:vjepa}
\begin{algorithmic}[1]
\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$
\FOR{each minibatch $\{x\}$}
\STATE Sample a context-target partition $(x_C, x_T)$ and corresponding target structure $\xi_T$
\STATE Encode context:
$
Z_C \leftarrow f_\theta(x_C)
$
\STATE Compute inference distribution from target:
$
q_{\theta'}(Z_T \mid x_T)
$
\STATE Sample target representation (e.g. using reparameterization):
$
Z_T \sim q_{\theta'}(Z_T \mid x_T)
$
\STATE Evaluate predictive distribution:
$
p_\phi(Z_T \mid Z_C, \xi_T)
$
\STATE Compute variational loss:
$
\mathcal{L}
=
-\log p_\phi(Z_T \mid Z_C, \xi_T)
+ \beta \,\mathrm{KL} \left(q_{\theta'}(Z_T \mid x_T)\,\|\,p(Z_T)\right)
$
\STATE Update $\theta$ and $\phi$ via gradient descent on $\mathcal{L}$
\STATE Update target encoder via EMA:
$
\theta' \leftarrow \tau \theta' + (1-\tau)\theta
$
\ENDFOR
\end{algorithmic}
\end{algorithm}
```
#### Remark.
Algorithm `\ref{alg:vjepa}`{=latex} reduces to standard JEPA training when the inference distribution $q_{\theta'}(Z_T \mid x_T)$ collapses to a Dirac delta, $q_{\theta'}(Z_T \mid x_T)=\delta \left(Z_T-f_{\theta'}(x_T)\right),$ and the predictive distribution is chosen as an isotropic Gaussian with fixed variance, $p_\phi(Z_T \mid Z_C,\xi_T)=\mathcal{N} \left(Z_T \mid g_\phi(Z_C,\xi_T), \sigma^2 I\right).$ In this case, minimizing the VJEPA objective in Eq. `\eqref{eq:vjepa_objective}`{=latex} is equivalent (up to constants) to minimizing the squared-loss JEPA objective in Eq. `\eqref{eq:JEPA_squared_loss}`{=latex}.
## 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. `\eqref{eq:context_encoder}`{=latex}) $$Z_C = f_\theta(x_C),$$ and then form the predictive distribution (i.e. the predictor Eq. `\eqref{eq:vjepa_predictive_model}`{=latex}) $$p_\phi(Z_T \mid Z_C, \xi_T).$$
#### Point prediction.
A point estimate of the target representation may be obtained from the predictive distribution $p_\phi(Z_T \mid 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 \mid 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 \mid 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}}
=
\operatorname*{mode} \left[p_\phi(Z_T \mid Z_C, \xi_T)\right],$$ 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 \mid 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 \mid Z_C, \xi_T),
\qquad 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 \left(Z_T^{(k)}\right),$$ 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 $$\label{eq:VJEPA_multi_step_prediction}
p(Z_{t+\Delta})
=
\int p_\phi(Z_{t+\Delta} \mid Z_t, \xi_{t+\Delta})\, p(Z_t)\, dZ_t,$$ 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} \mid Z_t^{(m)}, \xi_{t+1}),
\qquad 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 \mid 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.
## Collapse Avoidance {#sec:properties_of_VJEPA}
Here we formalize the claim that VJEPA discourages representational collapse.
::: definition
**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 \qquad \forall x_C.$$*
:::
::: theorem
**Theorem 1** (No Collapsed Global Optimum under Target Diversity). *`\label{thm:collapse_avoidance}`{=latex} Consider the VJEPA objective in Eq. `\eqref{eq:vjepa_objective}`{=latex}. Assume:*
1. *(**Target diversity**) There exist target inputs $x_T$ and $x_T'$ such that $q_{\theta'}(\cdot\mid x_T) \neq q_{\theta'}(\cdot\mid x_T')$ (as distributions).*
2. *(**Nontrivial conditioning**) The predictive family $\{p_\phi(\cdot\mid 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\mid Z_C,\xi_T) \neq p_\phi(\cdot\mid Z_C',\xi_T)$ for some $\phi$.*
*Then any globally optimal solution of the VJEPA training objective Eq. `\eqref{eq:vjepa_objective}`{=latex} 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\mid 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\mid \xi_T):=\mathbb{E}_{x_T\mid \xi_T}[q_{\theta'}(Z_T\mid x_T)]$, hence achieves $$\inf_{p} \ \mathbb{E}_{x_T\mid \xi_T}\,\mathbb{E}_{Z_T\sim q_{\theta'}(\cdot\mid x_T)}
[-\log p(Z_T)]
=
H_{q}(Z_T\mid \xi_T),$$ which equals $H_q(Z_T\mid X_T,\xi_T) + I_q(Z_T;X_T\mid \xi_T)$, where $I(\cdot;\cdot)$ denotes mutual information (see Appendix.`\ref{app:info_theory_identities}`{=latex}). By Assumption (i), we have $I_q(Z_T;X_T\mid \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. `\eqref{eq:vjepa_objective}`{=latex}. The KL regularizer further discourages degenerate target encodings by penalizing pathological $q_{\theta'}(\cdot\mid x_T)$, but is not the main source of the strict gap. A more detailed proof can be found in Appendix.`\ref{app:collapse_avoidance_proof}`{=latex}. `\hfill `{=latex}$\square$
::: remark
*Remark 1*. Theorem `\ref{thm:collapse_avoidance}`{=latex} shows that collapse is not merely discouraged but is fundamentally incompatible with global optimality of the VJEPA objective (Eq. `\eqref{eq:vjepa_objective}`{=latex}) 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.
:::
# Time-indexed VJEPA as a Latent Space Dynamical System {#sec:JEPA_as_dynamics}
VJEPA always learns a conditional predictive relation in representation space of the form (Eq. `\eqref{eq:vjepa_predictive_model}`{=latex}): $$Z_T \sim p_\phi(Z_T \mid 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 [@bardes2024vjepa; @bardes2024vjepa_openReview].
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.`\ref{tab:jepa_dynamics_cases}`{=latex} summarizes the different regimes in which JEPA-style objectives [^2] are commonly used.
```{=latex}
\footnotesize
```
::: {#tab:jepa_dynamics_cases}
**What is Predicted** **Structure Learned** **Dynamical System?**
---------------------------------------- --------------------------- ---------------------------------------- ------------------------------------------
Image JEPA (I-JEPA [@lejepa2022ijepa]) Masked image regions Spatial / semantic relations `\texttimes `{=latex}
Clip-level V-JEPA [@bardes2024vjepa] Masked video blocks Intra-clip spatiotemporal correlations `\texttimes`{=latex} (not compositional)
**Time-indexed JEPA** $Z_{t+\Delta}$ from $Z_t$ Latent state transitions across time `\checkmark `{=latex}
: 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.
## 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: $$\label{eq:ssm_factorization}
p(s_{1:T},x_{1:T}\mid u_{1:T-1})
=
p(s_1)\prod_{t=1}^{T-1} p(s_{t+1}\mid s_t,u_t)\prod_{t=1}^T p(x_t\mid s_t),$$ where $u_t$ is a control/action input. Classical HMMs (see Appendix.`\ref{app:hmmlink}`{=latex}) 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 `\eqref{eq:ssm_factorization}`{=latex} is that prediction and control can be performed by propagating a *belief* over latent states $p(s_t\mid x_{\le t},u_{ 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 `\eqref{eq:vjepa_dynamics}`{=latex} is $$\label{eq:vjepa_controlled_dynamics}
Z_{t+\Delta} \sim p_\phi \big(Z_{t+\Delta}\mid Z_t,\xi_{t+\Delta},u_{t:t+\Delta-1}\big),$$ 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. `\eqref{eq:vjepa_obs_model_optional}`{=latex}) $$\label{eq:ssm_obs_optional}
x_{t+\Delta} \sim p_\psi(x_{t+\Delta}\mid Z_{t+\Delta}),$$ which would turn the latent process in Eq. `\eqref{eq:vjepa_dynamics}`{=latex} into a fully generative state-space model. In this case, the implied generative structure is $$Z_t \longrightarrow Z_{t+\Delta} \longrightarrow 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 \longrightarrow x_t \longrightarrow x_{t+1},$$ and therefore require modeling the full conditional distribution $p(x_{t+1}\mid 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}\mid Z_{t+\Delta})$. The decoder is indexed at $t+\Delta$ to mirror the skip-time prediction structure of JEPA (Eq. `\eqref{eq:vjepa_dynamics}`{=latex}); when $\Delta=1$, this formulation reduces to the standard one-step state-space model.
## Sequential but Not Autoregressive {#sec:non_autoregressive}
We clarify here a common confusion of equating \`\`sequential modeling" with \`\`autoregressive modeling". Autoregression refers specifically to a likelihood factorization over observations: $$\label{eq:autoregressive_factorization}
p(x_{1:T})
=
\prod_{t=1}^T p(x_t\mid x_{Bayesian JEPA (BJEPA) architecture. A predictive likelihood models dynamics while a prior encoder maps auxiliary input η 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.
#### 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 \mid \eta) \propto 1.$$ In this case, the Bayesian posterior predictor becomes $$p(Z_T \mid Z_C, \eta) = p_{\text{like}}(Z_T \mid Z_C),$$ which reduces exactly to the VJEPA predictive distribution. Thus, BJEPA generalizes VJEPA without altering its training objective.
## Training and Inference {#sec:bjepa_training_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 \mid Z_C)$. While we can assume an uninformative prior[^17] 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. `\eqref{eq:vjepa_objective}`{=latex}, the general BJEPA training objective is: `\small`{=latex} $$\label{eq:bjepa_training_objective}
\mathcal{L}_{\text{BJEPA}}
=
\mathbb{E}_{x} \left[
\underbrace{
\mathbb{E}_{Z_T \sim q_{\theta'}(\cdot \mid x_T)}
\Big[ -\log p_{\text{like}}(Z_T \mid Z_C) \Big]
}_{\text{Maximize Predictive Likelihood}}
+
\beta
\underbrace{
\mathrm{KL} \left(q_{\theta'}(Z_T \mid x_T)\,\|\,p_{\text{ref}}(Z)\right)
}_{\text{Regularize Information State}}
+
\gamma
\underbrace{
\mathrm{KL} \left(p_{\text{like}}(Z_T \mid Z_C)\,\|\,p_{\text{struct}}(Z_T)\right)
}_{\text{Enforce Structural Constraints}}
\right].$$
- **Likelihood ($\mathcal{L}_{\text{VJEPA}}$):** The first two terms (controlled by $\beta$) constitute the standard VJEPA objective (Eq.`\ref{eq:vjepa_objective}`{=latex}). 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_{like}$) deviates too far from the \"task compliant\" distribution ($p_{prior}$), therefore forces the predicted dynamics to adhere to the structural prior. In our \"Noisy TV\" experiments (Section `\ref{sec:toy_experiment}`{=latex}), 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. `\eqref{eq:target_encoder_weights_EMA}`{=latex}, while the variance parameters $\sigma^2_{\theta'}$ are learned directly via gradient descent on the KL regularization term (the second term in Eq. `\eqref{eq:bjepa_training_objective}`{=latex}). 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 \mid Z_C)$ rather than the full posterior $p(Z_T \mid Z_C, \eta)$ (Eq. `\eqref{eq:bjepa_posterior}`{=latex}). This factorization is deliberate: the likelihood 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 \mid \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 \mid \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 \mid \eta_{\text{img}}) \approx \mathcal{N}(Z_T \mid f_{\theta'}(\eta_{\text{img}}), \sigma^2_{\text{goal}} 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 \mid \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[^18] (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 \mid Z_C)$ predicts where the system *will* go based on history $Z_C$.
2. *Prior Step:* the prior expert $p_{\text{prior}}(Z_T \mid \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 \mid Z_C, \eta) \propto p_{\text{like}}(Z_T \mid Z_C) \times p_{\text{prior}}(Z_T \mid \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.
The planner then samples or optimizes latent states $Z_T$ that maximize the joint posterior: $$\log p(Z_T \mid Z_C, \eta) \cong \log p_{\text{like}}(Z_T \mid Z_C) + \log p_{\text{prior}}(Z_T \mid \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 \mid 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. `\eqref{eq:bjepa_posterior}`{=latex}, prediction and planning may formally proceed via:
- **MAP prediction:** $$\hat Z_T^{\text{MAP}}
=
\arg\max_{Z_T}
\big[
\log p_{\text{like}}(Z_T \mid Z_C)
+
\log p_{\text{prior}}(Z_T \mid \eta)
\big];$$
- **Sampling-based prediction:** $$Z_T^{(m)} \sim p(Z_T \mid Z_C, \eta),
\qquad 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.
## Sequential BJEPA as Latent Bayesian Filtering {#sec:bjepa_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} \mid Z_t)$ propagates the state forward based on system dynamics.
2. *Correction (Constraint Update):* The latent prior $p_{\text{prior}}(Z_{t+1} \mid \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: $$\label{eq:bjepa_filter}
p(Z_{t+1} \mid Z_t, \eta_{t+1})
\propto
\underbrace{p_{\text{like}}(Z_{t+1} \mid Z_t)}_{\text{Dynamics (Prediction)}}
\times
\underbrace{p_{\text{prior}}(Z_{t+1} \mid \eta_{t+1})}_{\text{Constraint (Correction)}}.$$ 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} \mid 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
*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} \mid 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.
:::
## MAP Optimality under BJEPA
We now show that the sufficiency of MAP planning (Theorem `\ref{thm:map_optimal}`{=latex} of VJEPA) extends to the BJEPA architecture, provided the \`\`experts" are well-behaved.
::: theorem
**Theorem 5** (Optimality of MAP Planning under BJEPA). *`\label{thm:bjepa_map_optimal}`{=latex} Consider a one-step planning problem where the belief is formed by a BJEPA posterior: $$p(Z_{t+1} \mid Z_t, \eta, u_t) \propto p_{\text{like}}(Z_{t+1} \mid Z_t, u_t) \times p_{\text{prior}}(Z_{t+1} \mid \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^\star = \arg\min_{u_t} c(\hat{Z}_{t+1}^{\text{MAP}}, u_t).$$*
:::
::: proof
*Proof.* Since the posterior is formed by the product of two Gaussians, it is itself a Gaussian: $$p(Z_{t+1} \mid \cdot) = \mathcal{N}(Z_{t+1} \mid \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[^19] into a mean-dependent term and a trace term $\mathrm{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. ◻
:::
## Energy-Based Priors and Latent Optimization {#sec:bjepa_ebm}
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 \mid \eta) \propto \exp\!\big(-E_{\phi'}(Z_T;\eta)\big),$$ 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^{\mathrm{MAP}}
=
\arg\min_{Z_T}
\Big[
\underbrace{
\tfrac{1}{2}\|Z_T-\mu_\phi(Z_C)\|_{\Sigma^{-1}}^2
}_{\text{Dynamics Consistency}}
+
\underbrace{
E_{\phi'}(Z_T;\eta)
}_{\text{Task Constraint}}
\Big].$$ 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.
## Algorithm: Gradient-Based Planning with BJEPA {#sec:bjepa_algo}
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 `\ref{alg:bjepa_mpc}`{=latex}. 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.
```{=latex}
\begin{algorithm}[H]\caption{BJEPA Gradient-Based MPC}
\label{alg:bjepa_mpc}
\begin{algorithmic}[1]
\REQUIRE Current history $x_{\le t}$, Task $\eta$, Horizon $H$
\REQUIRE Differentiable Dynamics $p_{\text{like}}$, Energy Prior $E_{\phi'}$
\STATE \textbf{Initialize:} Estimate current state $Z_t = f_\theta(x_{\le t})$
\STATE \textbf{Initialize:} Action sequence $u_{0:H-1}$ (e.g. random or warm-start)
\WHILE{optimization budget not exceeded}
\STATE \textit{1. Forward Rollout (Dynamics Expert)}
\STATE $\hat{Z}_0 \leftarrow Z_t$
\FOR{$k=0$ \TO $H-1$}
\STATE $\hat{Z}_{k+1} \leftarrow \mu_{\text{dyn}}(\hat{Z}_k, u_k)$ \COMMENT{Predict future mean}
\ENDFOR
\STATE \textit{2. Evaluate BJEPA Posterior (Energy)}
\STATE $J \leftarrow 0$
\FOR{$k=1$ \TO $H$}
\STATE $J_{\text{prior}} \leftarrow E_{\phi'}(\hat{Z}_k; \eta)$ \COMMENT{Task Constraint Energy}
\STATE $J \leftarrow J + J_{\text{prior}} + \lambda \|u_{k-1}\|^2$ \COMMENT{Add control cost}
\ENDFOR
\STATE \textit{3. Backward Pass (Optimization)}
\STATE $\nabla_{u} J \leftarrow \text{Backpropagate}(J)$
\STATE $u_{0:H-1} \leftarrow u_{0:H-1} - \alpha \nabla_{u} J$ \COMMENT{Gradient update}
\ENDWHILE
\STATE \textbf{Execute:} Apply $u_0$ to environment
\STATE \textbf{Recurse:} Receive $x_{t+1}$, update $Z_{t+1}$, repeat.
\end{algorithmic}
\end{algorithm}
```
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}}$.
## Modularity, Transfer, and Continual Learning {#sec:bjepa_transfer}
A fundamental advantage of the BJEPA factorization (Eq. `\ref{eq:bjepa_posterior}`{=latex}) 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 \mid \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.`\ref{tab:world_model_comparison}`{=latex} 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.
```{=latex}
\resizebox{\linewidth}{!}{
\begin{tabular}{lccccc}
\toprule
Method & Obs.\ Reconstruction & Latent Uncertainty & Planning Space & Task Specification & Reward/Cost \\
\midrule
Dreamer & \checkmark & \checkmark & Latent & Monolithic & Learned \\
JEPA-WM & \xmark & \xmark & Latent & Monolithic & Learned \\
VJEPA (Ours) & \xmark & \checkmark & Latent & Implicit & None \\
\textbf{BJEPA (Ours)} & \xmark & \checkmark & \textbf{Latent} & \textbf{Modular Prior} & \textbf{Energy / Prior} \\
\bottomrule
\end{tabular}}
```
```{=latex}
\begin{flushleft}
\scriptsize \textbf{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 \mid 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.
\end{flushleft}
```
# Toy Experiment: a \`\`Noisy TV" Linear System {#sec:toy_experiment}
We present an analytically tractable experiment designed to stress-test the *nuisance invariance* property (Proposition `\ref{prop:nuisance_invariance}`{=latex}) 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[^20], 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.`\ref{app:toy_experiment_details}`{=latex}.
## 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 = C s_t + D (\sigma d_t) + \epsilon_t.$$ The latent processes evolve as follows:
- *The Signal ($s_t$):* evolves according to a stable rotation, $s_{t+1} = A_{\text{rot}} s_t + w_t$, where $w_t \sim \mathcal{N}(0, 0.1^2 I)$. The orthogonal transition ensures the signal remains bounded and maintains constant variance, fixing the theoretical recoverability throughout the sequence.
- *The Distractor ($d_t$):* evolves as a \`\`sticky" random walk, $d_{t+1} = 0.9 d_t + v_t$. We scale the distractor by $\sigma \in [0, 8]$, allowing its variance to grow up to $\approx 64\times$ the signal variance. At $\sigma=8.0$, the SNR drops to $-2.2$ dB, effectively burying the signal in noise.
## Model Architectures and Evaluation Alignment {#subsec:architectures}
We compare five models using a latent dimension $D_z = D_s = 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_t$ represents the current state. It is evaluated against the current signal $s_t$.
- *AR (Pixel-Predictive):* the bottleneck $z_t$ is optimized to predict future pixels $x_{t+1}$. Consequently, $z_t$ is evaluated on its ability to recover the *future* signal $s_{t+1}$.
- *JEPA/VJEPA (Latent-Predictive):* we extract the *predictor output* (the explicit prediction of $Z_{t+1}$) for the linear probe. These are evaluated against $s_{t+1}$.
- *BJEPA (Bayesian Fusion):* at inference, BJEPA uses *hard fusion* (Product of Experts) to combine dynamics with a task prior. The resulting fused mean $\mu_{post}$ is evaluated against $s_{t+1}$.
## Results and Analysis
Performance is measured via the coefficient of determination ($R^2$) using a linear probe fit on training data and evaluated on generalization test data.
#### Failure of Generative Models.
As shown in Table `\ref{tab:toyExperimentResults}`{=latex} and Fig.`\ref{fig:results_grid}`{=latex}, generative models (VAE, AR) suffer catastrophic degradation as noise increases. At scale $\sigma=8.0$, VAE test signal recovery drops to $R^2 \approx 0.50$ while its noise recovery remains high ($R^2 = 0.62$), confirming it has prioritized high-variance distractors. Similarly, the AR model drops to $R^2 = 0.578$.
#### Robustness of Joint-Embedding Architectures.
All joint-embedding architectures (JEPA, VJEPA, BJEPA) demonstrate robustness to the \`\`Noisy TV" distractor, maintaining $R^2 > 0.84$ even at scale 8.0. Deterministic JEPA achieved the highest accuracy ($R^2=0.93$), though it exhibited momentary instability at scale 3.0 ($R^2=0.841$). Probabilistic VJEPA and BJEPA showed superior training stability across seeds.
#### Qualitative Filtering.
Fig.`\ref{fig:reconstructions}`{=latex} 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_{t+1}$ with high fidelity.
```{=latex}
\footnotesize
```
```{=latex}
\setlength{\tabcolsep}{5pt}
```
```{=latex}
\renewcommand{\arraystretch}{1.05}
```
::: {#tab:toyExperimentResults}
Scale (SNR) Model $\uparrow$Signal $R^2$ (Tr/Te) $\downarrow$Noise $R^2$ (Tr/Te) $\downarrow$Time (Tr/Te)
--------------- ------- --------------------------------------------------- ---------------------------------------------------- --------------------------------
0.0 (inf dB) VAE [**1.000**]{.underline} / [**1.000**]{.underline} NA / NA 12.6s / 0.01s
AR 0.999 / 0.999 NA / NA [**6.4s**]{.underline} / 0.01s
JEPA 0.947 / 0.930 NA / NA 16.9s / 0.01s
VJEPA 0.999 / 0.999 NA / NA 13.9s / 0.01s
BJEPA 0.987 / 0.981 NA / NA 23.5s / 0.01s
4.0 (3.8 dB) VAE 0.822 / 0.730 0.512 / 0.458 12.5s / 0.01s
AR 0.839 / 0.756 0.394 / 0.338 [**6.5s**]{.underline} / 0.01s
JEPA [**0.999**]{.underline} / [**0.999**]{.underline} [**0.004**]{.underline} / [**-0.010**]{.underline} 16.6s / 0.02s
VJEPA 0.994 / 0.993 0.025 / -0.007 13.8s / 0.01s
BJEPA 0.920 / 0.899 0.213 / 0.156 23.2s / 0.01s
8.0 (-2.2 dB) VAE 0.613 / 0.499 0.656 / 0.620 12.3s / 0.01s
AR 0.680 / 0.578 0.491 / 0.449 [**7.1s**]{.underline} / 0.01s
JEPA [**0.947**]{.underline} / [**0.930**]{.underline} [**0.226**]{.underline} / [**0.183**]{.underline} 16.1s / 0.01s
VJEPA 0.905 / 0.870 0.299 / 0.251 13.4s / 0.01s
BJEPA 0.896 / 0.841 0.292 / 0.238 23.0s / 0.01s
: Performance metrics ($R^2$) across models at representative noise scales. $Tr$ denotes training set, $Te$ denotes test set. Full table is listed in Appendix.`\ref{app:toy_experiment_details}`{=latex}.
:::
{#fig:results_grid width="0.65\\linewidth"}
{#fig:reconstructions width="1.0\\linewidth"}
# Discussion {#sec: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.
## 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 [@hafner2019dreamer]) and generative sequence models (e.g. Transformers [@vaswani2023attentionneed]) 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.`\ref{sec:JEPA_as_dynamics}`{=latex} and Section.`\ref{sec:jepa_control}`{=latex}), 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.`\ref{eq:VJEPA_multi_step_prediction}`{=latex}) 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.
## Nuisance Invariance: The \`\`PCA *vs.* CCA" Distinction
Our \"Noisy TV\" experiment (Section `\ref{sec:toy_experiment}`{=latex}) provides empirical validation for the Information-Theoretic analysis in Section `\ref{sec:information_theoretic}`{=latex}. 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 generalization[^21] of *Principal Component Analysis (PCA)* [@Michal2019vaePCA]. 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[^22] (CCA [@weenink2003cca])* or the *Information Bottleneck*. As established in Theorem `\ref{thm:mi_bound}`{=latex} (Section `\ref{sec:information_theoretic}`{=latex}), 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 `\ref{prop:nuisance_invariance}`{=latex}) makes VJEPA as a *semantic noise filter* for downstream planning.
## Applications and Broader Impact {#sec:applications}
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 [@hafner2019dreamer]) 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} \mid Z_t, u_t)$ predicts the consequences of motor commands. Planning can be performed via the *VJEPA-MPC* algorithm (Algo.`\ref{alg:jmpc}`{=latex}) 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* [@Schrittwieser2020muzero], 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} \mid Z_t, u_t)$ captures the transition rules and physics of the game world. This enables planning via Monte Carlo Tree Search (MCTS [@MCTS_wiechowski_2022]) 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.
- **Control Inputs:** The tokens (or chunks of text) act as discrete actions driving the system ($u_t \in \mathcal{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} \mid Z_t, u_t).$$
In this framework, a response is not evaluated by its likelihood $\prod p(u_t \mid u_{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ϕ 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θ′ (or an optional decoder pψ). Finally, Resampling is performed to avoid particle degeneracy, resulting in an updated particle set representing the posterior belief for the next time step.
# Detailed Proof of Theorem `\ref{thm:collapse_avoidance}`{=latex} {#app:collapse_avoidance_proof}
We provide a detailed proof that the VJEPA objective admits no collapsed global optimum under the stated assumptions in Theorem `\ref{thm:collapse_avoidance}`{=latex}.
`\noindent`{=latex}`\fbox{%
\textbf{Theorem~\ref{thm:collapse_avoidance}: 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\mid x_T)}
\big[-\log p_\phi(Z_T \mid Z_C, \xi_T)\big]
+
\beta\,\mathbb{E}_x
\mathrm{KL}\!\left(q_{\theta'}(Z_T\mid x_T)\,\|\,p(Z_T)\right).
\]
Assume:
\begin{enumerate}
\item (\textbf{Target diversity}) There exist $x_T,x_T'$ such that
$q_{\theta'}(\cdot\mid x_T)\neq q_{\theta'}(\cdot\mid x_T')$.
\item (\textbf{Nontrivial conditioning}) The predictive family
$\{p_\phi(\cdot\mid Z_C,\xi_T)\}$ can represent different distributions for different
$Z_C$ (for fixed $\xi_T$).
\end{enumerate}
Then no global minimizer of $\mathcal{L}_{\text{VJEPA}}$ satisfies
$f_\theta(x_C)\equiv c$ for all $x_C$.
%
}`{=latex}
::: proof
*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 \qquad \forall x_C.$$ Then the predictive model cannot depend on the context content and reduces to $$p_\phi(Z_T \mid Z_C,\xi_T) = p_\phi(Z_T \mid c,\xi_T),$$ i.e. an *unconditional* distribution for each $\xi_T$. Under this restriction, the first term of the VJEPA objective becomes $$\label{eq:collapsed_loss}
\mathcal{L}_{\text{pred}}^{\text{coll}}
=
\mathbb{E}_{\xi_T}
\mathbb{E}_{x_T\mid \xi_T}
\mathbb{E}_{Z_T\sim q_{\theta'}(\cdot\mid x_T)}
\big[-\log p_\phi(Z_T\mid c,\xi_T)\big].$$
#### Step 2: Optimal unconditional predictor.
For a fixed $\xi_T$, define the *aggregated target distribution* $$\bar q(Z_T\mid \xi_T)
:=
\mathbb{E}_{x_T\mid \xi_T}\big[q_{\theta'}(Z_T\mid x_T)\big].$$ It is a standard result that the distribution minimizing cross-entropy with respect to $\bar q$ is $\bar q$ itself. Hence, $$\inf_{p(\cdot\mid c,\xi_T)}
\mathbb{E}_{x_T\mid\xi_T}
\mathbb{E}_{Z_T\sim q_{\theta'}(\cdot\mid x_T)}
[-\log p(Z_T)]
=
H(\bar q(\cdot\mid \xi_T)).$$ Therefore, the best achievable collapsed prediction loss equals $$\label{eq:collapsed_entropy}
\mathcal{L}_{\text{pred}}^{\text{coll}}
=
\mathbb{E}_{\xi_T}\, H(Z_T\mid \xi_T).$$
#### Step 3: Decomposition via conditional mutual information.
By standard entropy identities (see Appendix.`\ref{app:info_theory_identities}`{=latex}), $$H(Z_T\mid \xi_T)
=
H(Z_T\mid X_T,\xi_T)
+
I(Z_T;X_T\mid \xi_T).$$ Assumption (i) (*target diversity*) implies that $$I(Z_T;X_T\mid \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\mid Z_C,\xi_T)$ and therefore approximate $q_{\theta'}(Z_T\mid x_T)$ conditionally. In the realizable limit, $$\mathcal{L}_{\text{pred}}^{\text{non-coll}}
=
\mathbb{E}_{\xi_T}
\mathbb{E}_{x_T\mid \xi_T}
H(Z_T\mid X_T=x_T,\xi_T)
<
\mathcal{L}_{\text{pred}}^{\text{coll}}.$$
#### Step 5: Role of the KL regularizer.
The KL term $$\mathrm{KL}(q_{\theta'}(Z_T\mid 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 {#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 `\ref{thm:collapse_avoidance}`{=latex} 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. ◻
:::
# Detailed Proof of Theorem `\ref{thm:map_optimal}`{=latex} {#app:map_optimal_proof}
We restate Theorem `\ref{thm:map_optimal}`{=latex} in Section.`\ref{sec:map_optimal}`{=latex} for convenience.
`\noindent`{=latex}`\fbox{%
\textbf{Theorem \ref{thm:map_optimal}: Optimality of MAP (Mean) Control under Quadratic Costs}
Consider a one-step control problem with latent predictive model
$p_\phi(Z_{t+1}\mid 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}\mid Z_t,u_t) = \mathcal{N} \big(Z_{t+1}\mid \mu_\phi(Z_t,u_t),\,\Sigma\big),
\]
where the covariance $\Sigma$ does \emph{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^\star)^\top Q_c (Z_{t+1}-z^\star) + r(u_t),
\]
with $Q_c \succeq 0$ and arbitrary $r(\cdot)$.
Then the action minimizing expected cost,
\[
u_t^\star \in \arg\min_{u_t} \mathbb{E} \left[c(Z_{t+1},u_t)\mid Z_t,u_t\right],
\]
is equivalently obtained by minimizing the cost at the predictive mean:
\[
u_t^\star \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 \emph{MAP control}.
%
}`{=latex}
::: proof
*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),
\qquad
\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 $$\label{eq:quadratic_cost_objective}
J(u_t)
:=
\mathbb{E} \left[c(Z,u_t)\mid Z_t,u_t\right]
=
\mathbb{E} \left[(Z-z^\star)^\top Q_c (Z-z^\star)\right] + r(u_t).$$ in which the only nontrivial term is the expectation of the quadratic form.
#### Step 1: Expand the quadratic form.
Let $d := Z - z^\star$. Then $$d^\top Q_c d
=
(Z-z^\star)^\top Q_c (Z-z^\star).$$ Insert and subtract the mean $\mu$: $$Z - z^\star
=
(Z-\mu) + (\mu - z^\star).$$ Therefore, $$(Z-z^\star)^\top Q_c (Z-z^\star)
=
\big((Z-\mu)+(\mu-z^\star)\big)^\top Q_c \big((Z-\mu)+(\mu-z^\star)\big).$$ Expand the product into four terms: $$\begin{aligned}
&\big((Z-\mu)+(\mu-z^\star)\big)^\top Q_c \big((Z-\mu)+(\mu-z^\star)\big) \\
&\quad=
\boldsymbol{(Z-\mu)^\top Q_c (Z-\mu)}
+ (Z-\mu)^\top Q_c (\mu-z^\star)
+ (\mu-z^\star)^\top Q_c (Z-\mu)
+ \boldsymbol{(\mu-z^\star)^\top Q_c (\mu-z^\star)}.
\end{aligned}$$
#### 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}\big[(Z-\mu)^\top Q_c (\mu-z^\star)\big]
=
\mathbb{E}[Z-\mu]^\top Q_c (\mu-z^\star)
=
0.$$ Similarly, $$\mathbb{E}\big[(\mu-z^\star)^\top Q_c (Z-\mu)\big]
=
(\mu-z^\star)^\top Q_c \mathbb{E}[Z-\mu]
=
0.$$
*(ii) The constant term remains.* Since $(\mu-z^\star)^\top Q_c (\mu-z^\star)$ is deterministic given $u_t$, $$\mathbb{E}\big[(\mu-z^\star)^\top Q_c (\mu-z^\star)\big]
=
(\mu-z^\star)^\top Q_c (\mu-z^\star).$$
*(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}\big[\varepsilon^\top Q_c \varepsilon\big] = \mathrm{tr}(Q_c\Sigma).$$ To see this, use the identity for any vector $\varepsilon$ and matrix $Q_c$: $$\varepsilon^\top Q_c \varepsilon
=
\mathrm{tr} \left(\varepsilon^\top Q_c \varepsilon\right)
=
\mathrm{tr} \left(Q_c \varepsilon\varepsilon^\top\right),$$ where we used $\mathrm{tr}(a)=a$ for a scalar $a$, and cyclicity of trace [^27]. Taking expectations, $$\mathbb{E}\big[\varepsilon^\top Q_c \varepsilon\big]
=
\mathbb{E} \left[\mathrm{tr} \left(Q_c \varepsilon\varepsilon^\top\right)\right]
=
\mathrm{tr} \left(Q_c \,\mathbb{E}[\varepsilon\varepsilon^\top]\right)
=
\mathrm{tr}(Q_c\Sigma),$$ where linearity of expectation and linearity of trace justify exchanging $\mathbb{E}$ and $\mathrm{tr}$.
#### Step 3: Assemble the expected quadratic cost.
Combining the above, $$\mathbb{E} \left[(Z-z^\star)^\top Q_c (Z-z^\star)\right]
=
(\mu-z^\star)^\top Q_c (\mu-z^\star) + \mathrm{tr}(Q_c\Sigma).$$ Hence the full objective in Eq. `\eqref{eq:quadratic_cost_objective}`{=latex} becomes $$J(u_t)
=
(\mu_\phi(Z_t,u_t)-z^\star)^\top Q_c (\mu_\phi(Z_t,u_t)-z^\star)
+
\mathrm{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 $\mathrm{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^\star)^\top Q_c (\mu_\phi(Z_t,u_t)-z^\star)
+ 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.
#### Step 5: Mean equals MAP for a Gaussian.
For $Z\sim\mathcal{N}(\mu,\Sigma)$ with $\Sigma\succ 0$, the log-density is $$\log p(Z) = -\tfrac{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. ◻
:::
::: remark
*Remark 3* (Assumptions matter). If $\Sigma=\Sigma(u_t)$ depends on the action, then the term $\mathrm{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
*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 [@mesbah2016stochasticmpc; @rawlings2017MPC].
Some reinforcement learning and dynamic programming formulations instead write the stage cost as $c(Z_t,u_t)$ [@sutton2018rl; @puterman2014mdp]. The two conventions are equivalent under expectation for Markov dynamics, since $$\mathbb{E}\!\left[c(Z_{t+1},u_t)\mid Z_t,u_t\right]$$ 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.
:::
# Proof of Proposition `\ref{prop:nuisance_invariance}`{=latex} (Invariance to Nuisance) {#app:proof_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} \mid s_t) = 0$.
Let $Z_t$ be the representation of $x_t$.
## 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} \left[ \log p(x_t \mid Z_t) \right].$$ Using the decomposition $x_t = (s_t, n_t)$, the likelihood factorizes as (using chain rule of probability): $$\log p(s_t, n_t \mid Z_t) = \log p(s_t \mid Z_t) + \log p(n_t \mid s_t, Z_t).$$ To maximize this objective (globally), the representation $Z_t$ must maximize both terms.
- Maximizing $\log p(s_t \mid Z_t)$ requires $Z_t$ to encode the signal $s_t$.
- Maximizing $\log p(n_t \mid s_t, Z_t)$ requires $Z_t$ to encode the nuisance $n_t$ (so that the posterior entropy $H(n_t \mid 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$).
## 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} \left[ \log p(Z_{t+\Delta} \mid Z_t) \right].$$ 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} \mid s_t, n_t) = p(Z_{t+\Delta} \mid s_t).$$ Now compare two candidate representations:
1. **Maximal representation:** $Z_t^{\text{max}} = (s_t, n_t)$.
2. **Minimal representation:** $Z_t^{\text{min}} = s_t$ (discards $n_t$).
Substituting these into the objective: $$\mathcal{L}_{\text{VJEPA}}(Z_t^{\text{max}})
= \mathbb{E} [\log p(Z_{t+\Delta} \mid s_t, n_t)]
= \mathbb{E} [\log p(Z_{t+\Delta} \mid s_t)]
= \mathcal{L}_{\text{VJEPA}}(Z_t^{\text{min}}).$$ **Conclusion:** The minimal representation $Z_t^{\text{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.
# Derivation of the BJEPA Predictive Factorization {#app:poe_derivation}
In this section, we derive the relationship between the joint posterior $p(Z_T \mid Z_C, \eta)$ and the individual predictive factors used in the BJEPA architecture (Eq. `\ref{eq:bjepa_posterior}`{=latex}). 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 \mid Z_C)$ and $p(Z_T \mid \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$: $$\label{eq:cond_indep}
p(Z_C, \eta \mid Z_T) = p(Z_C \mid Z_T) \, p(\eta \mid Z_T).$$ 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.
#### Derivation.
Starting with Bayes' theorem for the posterior: $$p(Z_T \mid Z_C, \eta) = \frac{p(Z_C, \eta \mid Z_T)\, p(Z_T)}{p(Z_C, \eta)}.$$ Substituting the conditional independence assumption from Eq. `\eqref{eq:cond_indep}`{=latex}: $$p(Z_T \mid Z_C, \eta) = \frac{p(Z_C \mid Z_T)\, p(\eta \mid Z_T)\, p(Z_T)}{p(Z_C, \eta)}.$$ Next, we invert the individual likelihood terms using Bayes' theorem to express them in terms of the predictive distributions (the \"experts\"): $$p(Z_C \mid Z_T) = \frac{p(Z_T \mid Z_C)\, p(Z_C)}{p(Z_T)}
\quad \text{and} \quad
p(\eta \mid Z_T) = \frac{p(Z_T \mid \eta)\, p(\eta)}{p(Z_T)}.$$ Substituting these back into the posterior equation: $$p(Z_T \mid Z_C, \eta) = \frac{\left( \frac{p(Z_T \mid Z_C)\, p(Z_C)}{p(Z_T)} \right) \left( \frac{p(Z_T \mid \eta)\, p(\eta)}{p(Z_T)} \right) p(Z_T)}{p(Z_C, \eta)}.$$ Simplifying the terms: $$p(Z_T \mid Z_C, \eta) = \frac{p(Z_T \mid Z_C)\, p(Z_T \mid \eta) \, p(Z_C) \, p(\eta)}{p(Z_T) \, p(Z_C, \eta)}.$$ Since $Z_C$ and $\eta$ are observed inputs (constants with respect to the optimization of $Z_T$), the terms $p(Z_C)$, $p(\eta)$, and $p(Z_C, \eta)$ can be absorbed into a normalization constant $\mathcal{Z}$. The relation reduces to: $$p(Z_T \mid Z_C, \eta) \propto \frac{p(Z_T \mid Z_C)\, p(Z_T \mid \eta)}{p(Z_T)}.$$ In the BJEPA formulation, we approximate the marginal prior over targets $p(Z_T)$ as uniform (or absorbed into the learned bias of the experts), yielding the standard Product of Experts form: $$p(Z_T \mid Z_C, \eta) \propto p(Z_T \mid Z_C)\, p(Z_T \mid \eta).$$
# Toy Experiment Details {#app:toy_experiment_details}
Here we provide the specific implementation details, hyperparameters, and architectural and training specifications used to generate the results in Section `\ref{sec:toy_experiment}`{=latex}.
## Environment and Data Generation
The synthetic environment is a Linear-Gaussian system designed to simulate the \`\`Noisy TV" scenario.
- `\text{Dimensions:}`{=latex} Observation dimension $D_x = 20$, Signal dimension $D_s = 4$, Distractor dimension $D_d = 4$.
- **Dynamics:**
- *Signal:* $s_{t+1} = A_{\text{rot}} s_t + w_t$, where $A_{\text{rot}}$ is an orthogonal matrix generated via QR decomposition of a random Gaussian matrix. Process noise $w_t \sim \mathcal{N}(0, 0.1^2 I)$.
- *Distractor:* $d_{t+1} = 0.9 d_t + v_t$. Process noise $v_t \sim \mathcal{N}(0, 0.3^2 I)$.
- **Observations:** $x_t = C s_t + D (\sigma d_t) + \epsilon_t$. Matrices $C$ and $D$ are random matrices with column-wise unit norm normalization. Sensor noise $\epsilon_t \sim \mathcal{N}(0, 0.01^2 I)$.
- **Dataset:** We generate a continuous training trajectory of $T_{\text{train}}=6000$ steps and a testing trajectory of $T_{\text{test}}=2000$ steps.
- **Noise Scales:** The scaling factor $\sigma$ is swept from $0.0$ to $8.0$ in 9 equidistant steps.
## Model Architectures and Training Strategy {#subsec:architectures_app}
We implement five models to map the current observation $x_t \in \mathbb{R}^{D_x}$ to a latent representation $z_t \in \mathbb{R}^{D_z}$. To impose a hard trade-off between signal and distractor encoding, we set the latent dimension equal to the true signal dimension ($D_z = D_s = 4$). All models utilize strictly linear transformations ($f(x) = Wx$) without bias terms (\`bias=False\`), matching the linear nature of the data generation process.
We compare two modeling paradigms: **Generative Reconstruction** (VAE, AR) operating in pixel space, and **Joint-Embedding Prediction** (JEPA, VJEPA, BJEPA) operating via latent-space autoregression.
### Generative Baselines (Pixel Reconstruction)
These models must reconstruct the full observation $x$, including high-variance distractors.
#### Linear Variational Autoencoder (VAE).
The VAE treats the problem as a static modeling task. It maps $x_t$ to a latent distribution $q(z_t|x_t)$ and reconstructs $x_t$.
- *Architecture:* A probabilistic encoder with two linear heads $\mu_\phi, \log\sigma^2_\phi$ mapping $D_x \to D_z$, and a deterministic decoder $\hat{x}_t = W_{dec} z_t$.
- *Objective:* The Evidence Lower Bound (ELBO)[^28]: $$\mathcal{L}_{\text{VAE}} = \| x_t - \hat{x}_t \|^2 + \beta D_{KL}(q(z_t|x_t) \| \mathcal{N}(0, I)).$$
- *Hypothesis:* To minimize reconstruction error, the VAE must prioritize features with the highest variance. We predict it will encode distractors and ignore the signal.
#### Linear Pixel-Autoregressive (AR).
The AR model is a predictive autoencoder operating in pixel space. It compresses $x_t$ into a bottleneck to predict the *next* observation $x_{t+1}$.
- *Architecture:* A deterministic encoder $z_t = W_{enc} x_t$ mapping $D_x \to D_z$ and a predictive decoder $\hat{x}_{t+1} = W_{pred} z_t$ mapping $D_z \to D_x$.
- *Objective:* Mean Squared Error in pixel space: $\mathcal{L}_{\text{AR}} = \| x_{t+1} - \hat{x}_{t+1} \|^2$.
- *Hypothesis:* Since distractors are predictable ($d_{t+1} \approx 0.9 d_t$), the AR model must track them to minimize prediction error, wasting capacity on distractor texture.
### Joint-Embedding Architectures (Latent Autoregression)
These models are time-indexed, order-1 Markovian in latent space, trained to predict $z_{t+1}$ without reconstructing pixels.
#### Linear JEPA (Deterministic).
JEPA aligns the predicted next state with the encoded future state.
- *Architecture:*
1. Context Encoder: $z_t = W_{enc} x_t$.
2. Predictor: A linear layer $\hat{z}_{t+1} = W_{pred} z_t$ mapping $D_z \to D_z$.
3. Target Encoder (EMA): Same architecture as Context Encoder.
- *Objective:* VICReg loss, minimizing prediction error $\| \hat{z}_{t+1} - z'_{t+1} \|^2$ with variance/covariance regularization to prevent collapse.
#### Linear Probabilistic VJEPA.
VJEPA extends JEPA by treating the prediction as a probabilistic distribution.
- *Architecture:*
1. Deterministic Encoder: $z_t = W_{enc} x_t$.
2. Probabilistic Predictor: Two linear heads outputting $\mu_{pred}, \log\sigma^2_{pred} = f_\phi(z_t)$ mapping $D_z \to D_z$.
3. Target Encoder (EMA): Two heads outputting $\mu_{\theta'}, \log\sigma^2_{\theta'}$ mapping $D_x \to D_z$.
- *Objective:* The regularized negative log-likelihood (Eq. `\ref{eq:vjepa_objective}`{=latex}): $$\mathcal{L}_{\text{VJEPA}} = \mathbb{E}_{\tilde{z}_{t+1} \sim q_{\theta'}} [-\log p_\phi(\tilde{z}_{t+1} | z_t)] + \beta D_{KL}(q_{\theta'}(z'_{t+1}|x_{t+1}) \| \mathcal{N}(0, I)).$$
#### Linear Bayesian JEPA (BJEPA).
BJEPA incorporates a static Bayesian prior to filter out unpredictable dynamics by fusing a *Predictive Likelihood* (Dynamics Expert) with a *Latent-Space Prior* (Static Expert).
- *Architecture:* Identical to VJEPA, but adds learnable static parameters $\mu_{prior}, \log\sigma^2_{prior}$ (initialized to zero) to define a preferred latent manifold.
- *Inference:* BJEPA performs *hard fusion* via a Product of Experts (PoE): $$\mu_{post}, \log\sigma^2_{post} = \text{PoE}((\mu_{dyn}, \sigma^2_{dyn}), (\mu_{prior}, \sigma^2_{prior})).$$
- *Objective:* $\mathcal{L}_{\text{BJEPA}} = \mathcal{L}_{\text{VJEPA}} + \gamma D_{KL}(p_{\text{like}}(z_{t+1}|z_t) \| p_{\text{prior}}(z_{t+1}))$.
#### Stochastic Training Strategy.
In our implementation, we parameterize the context encoder $f_\theta(x_C)$ as a deterministic map, delegating uncertainty modeling entirely to the predictive transition $p_\phi(Z_T \mid Z_C)$ and the target inference $q_{\theta'}(Z_T \mid x_T)$. For both VJEPA and BJEPA, optimization of the variational objective (Eq. `\ref{eq:vjepa_objective}`{=latex} and Eq. `\ref{eq:bjepa_training_objective}`{=latex}) requires estimating an expectation over the target distribution $\mathbb{E}_{Z_T \sim q_{\theta'}}$. Following standard practice in stochastic gradient variational inference (SGVI), we approximate this expectation using a single Monte Carlo sample ($K=1$) drawn from the target encoder's distribution $q_{\theta'}(Z_T \mid x_T)$ via the reparameterization trick at each training step. While this single-sample estimate is noisy for any individual data point, the stochasticity of mini-batch gradient descent (SGD) smooths the objective over the course of training, providing an unbiased gradient estimator.
Note that the sampling operation is applied *only* to the target encoder $q_{\theta'}$ to obtain a concrete regression target $Z_T$. We do *not* sample from the predictor $p_\phi$. Instead, the predictor outputs the parameters of its belief distribution (e.g. $\mu_\phi, \sigma^2_\phi$), and we analytically evaluate the negative log-likelihood density (Gaussian in our implementation) of the sampled target $Z_T$ under these parameters. Mathematically, this evaluates how well the predictor's belief explains the specific realization drawn from the target encoder's belief.
#### Necessity of Target Sampling.
Sampling $Z_T$ from the target distribution $q_{\theta'}$ (rather than simply using its mean $\mu_{\theta'}$) is strictly necessary to prevent *variance collapse* in the predictor. Consider the Gaussian negative log-likelihood term minimized by the predictor: $$\mathcal{L}_{\text{pred}} \approx \frac{(Z_T - \mu_\phi)^2}{2\sigma^2_\phi} + \frac{1}{2}\log(\sigma^2_\phi).$$ If we were to use the target mean as the training target (i.e. set $Z_T = \mu_{\theta'}$), the predictor could learn to output $\mu_\phi \approx \mu_{\theta'}$ perfectly, driving the squared error term to zero. This would encourage the variance $\sigma^2_\phi$ to collapse to zero (or $\log\sigma^2_\phi \to -\infty$) to minimize the second term, resulting in a deterministic, overconfident model. By sampling $Z_T \sim \mathcal{N}(\mu_{\theta'}, \sigma^2_{\theta'})$, the target becomes a \`\`moving target" that fluctuates around the mean according to the target's uncertainty. The squared error term $(Z_T - \mu_\phi)^2$ remains non-zero, forcing the predictor to maintain a non-zero variance $\sigma^2_\phi$ to cover the distribution of potential futures.
#### Implementation and Limitations.
In our experimental implementation, we parameterize both the target encoder $q_{\theta'}$ and the predictor $p_\phi$ as independent diagonal Gaussian distributions: $$q_{\theta'}(Z \mid x_T) = \mathcal{N}(Z; \mu_{\theta'}(x_T), \text{diag}(\sigma^2_{\theta'}(x_T))),
\quad
p_\phi(Z \mid Z_C) = \mathcal{N}(Z; \mu_\phi(Z_C), \text{diag}(\sigma^2_\phi(Z_C))).$$ While this choice simplifies optimization and admits closed-form KL terms, the unimodal Gaussian assumption represents a limitation. In environments with complex, multi-modal bifurcations (e.g. distinct \`\`go left" vs. \`\`go right" futures), a single Gaussian may average the modes, resulting in a high-variance belief centered on an invalid state. Future work could address this by implementing more expressive heads, such as Gaussian Mixture Models or normalizing flows.
### Temporal Alignment in Training and Evaluation {#subsubsec:temporal_alignment}
All models receive observation pairs $(x_t, x_{t+1})$ during training, where $x_t$ denotes the current observation and $x_{t+1}$ denotes the next-step observation. The temporal alignment differs between model families based on their predictive structure:
#### Training Alignment.
- *VAE:* trained on single observations $x_t$ with reconstruction target $x_t$ (no temporal prediction).
- *AR:* trained on pairs $(x_t, x_{t+1})$ with pixel-space prediction target $x_{t+1}$.
- *JEPA/VJEPA/BJEPA:* trained on pairs $(x_t, x_{t+1})$ with latent-space prediction target $z_{t+1} = f_{\theta'}(x_{t+1})$.
#### Evaluation Alignment.
As predictive models (AR, JEPA, VJEPA, BJEPA) are trained to predict future states, their latent representations should be evaluated against the *future* ground-truth signal $s_{t+1}$, not the current signal $s_t$. Concretely:
- *VAE:* The encoder output $z_t = \mu_\phi(x_t)$ represents the current state and is evaluated against $s_t$.
- *AR:* The encoder output $z_t = W_{enc} x_t$ is designed to predict $x_{t+1}$, hence is evaluated against $s_{t+1}$.
- *JEPA:* The predictor output $\hat{z}_{t+1} = W_{pred} z_t$ explicitly predicts the next latent state and is evaluated against $s_{t+1}$.
- *VJEPA:* The predictor mean $\mu_{pred}(z_t)$ represents the expected next latent state and is evaluated against $s_{t+1}$.
- *BJEPA:* At inference, the hard-fused posterior mean $\mu_{post} = \text{PoE}(\mu_{dyn}, \mu_{prior})$ represents the constrained prediction of the next latent state and is evaluated against $s_{t+1}$.
This temporal alignment ensures that each model's representation is evaluated against the appropriate ground-truth signal consistent with its predictive objective.
### Linear Probe Evaluation Protocol {#subsubsec:linear_probe}
To assess how well each model's latent representation captures the true underlying signal $s_t$ (while ignoring the distractor $d_t$), we employ a linear probe methodology. For each model, we extract latent representations and fit a linear regression to predict the ground-truth signal, then measure the coefficient of determination ($R^2$) on held-out test data.
#### Latent Extraction.
The latent representation used for probing differs by model, reflecting their distinct architectures:
- *VAE:* $z_t^{\text{probe}} = \mu_\phi(x_t)$ --- the encoder mean, representing the current state.
- *AR:* $z_t^{\text{probe}} = W_{enc} x_t$ --- the encoder output, which implicitly predicts the next state.
- *JEPA:* $z_t^{\text{probe}} = W_{pred} (W_{enc} x_t) = \hat{z}_{t+1}$ --- the predictor output, explicitly predicting the next latent state.
- *VJEPA:* $z_t^{\text{probe}} = \mu_{pred}(W_{enc} x_t)$ --- the predictor mean, representing the expected next latent state.
- *BJEPA:* $z_t^{\text{probe}} = \mu_{post}$ --- the hard-fused posterior mean computed via Product of Experts: $$\mu_{post} = \frac{\sigma^2_{prior} \cdot \mu_{dyn} + \sigma^2_{dyn} \cdot \mu_{prior}}{\sigma^2_{dyn} + \sigma^2_{prior}},$$ where $(\mu_{dyn}, \sigma^2_{dyn})$ are the dynamics predictor parameters and $(\mu_{prior}, \sigma^2_{prior})$ are the learned static prior parameters.
#### Probe Fitting and Evaluation.
Given extracted latents $\{z_i^{\text{probe}}\}_{i=1}^{N_{\text{train}}}$ from the training set and corresponding ground-truth signals $\{s_i^{\text{target}}\}_{i=1}^{N_{\text{train}}}$, we fit a linear regression: $$\hat{s} = W_{\text{probe}} z^{\text{probe}} + b_{\text{probe}},$$ where $(W_{\text{probe}}, b_{\text{probe}})$ are learned via ordinary least squares. The target signal $s^{\text{target}}$ is chosen according to the temporal alignment described above: $s_t$ for VAE, and $s_{t+1}$ for all predictive models (AR, JEPA, VJEPA, BJEPA).
We then evaluate on the test set by computing the $R^2$ score: $$R^2 = 1 - \frac{\sum_i \| s_i^{\text{target}} - \hat{s}_i \|^2}{\sum_i \| s_i^{\text{target}} - \bar{s} \|^2},$$ where $\bar{s}$ is the mean of the test targets. An $R^2$ close to 1 indicates that the latent representation successfully captures the true signal, while an $R^2$ close to 0 indicates that the representation has failed to disentangle signal from distractor.
#### Interpretation.
This evaluation protocol directly tests the core hypothesis: predictive models operating in latent space (JEPA, VJEPA, BJEPA) should learn representations that capture the predictable signal while discarding unpredictable distractors, whereas generative models (VAE, AR) operating in pixel space are forced to allocate capacity to high-variance distractors regardless of their task relevance. The use of predictor outputs (rather than encoder outputs) for JEPA-family models ensures we evaluate the representation that the model is actually trained to produce, i.e. the predicted future state.
## Optimization and Hyperparameters
All models are trained using the *Adam* optimizer [@kingma2017haha_dont_ask_me_who_is_Adam_my_name_is_Yong]. We use full-batch training (processing the continuous trajectory as a single batch) for stability in this linear setting.
::: {#tab:hyperparams}
**Parameter** **Value** **Description**
----------------------------------------------- ---------------------------------------------------------------- -----------------------------------------
*Global Optimization*
Random Seed 111 For reproducibility
Optimizer Adam [@kingma2017haha_dont_ask_me_who_is_Adam_my_name_is_Yong] Standard implementation
Learning Rate $1 \times 10^{-3}$ Constant throughout training
Training Steps 6000 Sufficient for convergence
*JEPA Family Common Settings*
EMA Decay ($\tau$) 0.99 Target encoder update rate
*Model-Specific Coefficients*
**JEPA** (VICReg Loss)
`\quad `{=latex}Invariance Coeff 25.0 MSE term weight
`\quad `{=latex}Variance Coeff 25.0 Hinge loss on standard deviation
`\quad `{=latex}Covariance Coeff 1.0 Off-diagonal decorrelation weight
**VJEPA**
`\quad `{=latex}$\beta$ (Eq. `\ref`{=latex}) 0.01 Target encoder KL regularization weight
**BJEPA**
`\quad `{=latex}$\beta$ (Eq. `\ref`{=latex}) 0.01 Target encoder KL regularization weight
`\quad `{=latex}$\gamma$ (Eq. `\ref`{=latex}) 0.1 **Structural Prior** KL weight
: Optimization settings and loss function hyperparameters.
:::
#### Note on BJEPA Priors.
In the Python code implementation for the \`\`Noisy TV" experiment, the Structural Prior used for the $\gamma$ term is *static* (time-invariant) but *learnable*. It is parameterized as a Gaussian with parameters $\mu_{\text{prior}}, \sigma^2_{\text{prior}}$ that are optimized alongside the model weights to find the stable manifold of the signal. The reference prior used for the $\beta$ term (target regularization) is a fixed unit Gaussian $\mathcal{N}(0, I)$.
## Results
We measure performance using the *coefficient of determination* ($R^2$) between the learned latent $Z_t$ and the true signal $s_t$ via a linear probe. The quantitative results [^29] are summarized in Table.`\ref{tab:toyExperimentResults}`{=latex} and visualized across all 9 noise scales in Fig.`\ref{fig:results_grid}`{=latex}. Additionally, Fig.`\ref{fig:reconstructions}`{=latex} provides a qualitative view of the signal reconstruction at low, medium, and high noise levels.
#### Failure of Generative Models.
As shown in Table.`\ref{tab:toyExperimentResults_full}`{=latex} and Fig.`\ref{fig:results_grid}`{=latex}, the reconstructive baselines (VAE, AR) suffer catastrophic performance degradation as the distractor noise scale increases. At the highest noise scale ($\sigma=8.0$), VAE and AR signal recovery drops to $R^2 \approx 0.50$ and $0.58$ respectively, while their noise recovery remains high ($R^2 = 0.62$ and $R^2 = 0.45$). This confirms that likelihood-based objectives force the models to encode high-variance distractors, leading to representation collapse on the target task.
#### Robustness of Joint-Embedding Architectures.
In contrast, all joint-embedding architectures (JEPA, VJEPA, BJEPA) demonstrate robustness to the \`\`Noisy TV" distractor. They maintain high signal recovery ($R^2 > 0.85$) even at $\sigma=8.0$, effectively filtering out the distractor dynamics. In this specific trial, the deterministic JEPA achieved the highest final accuracy ($R^2=0.93$), though it exhibited higher variance in training stability compared to the probabilistic variants[^30].
#### Stability and Optimization.
The high-variance noise makes the target embeddings ($z_{t+1}$) jittery and unpredictable. As standard JEPA is deterministic, it struggles to average out this jitter, leading to potential training instability. Even in this successful run, JEPA showed a momentary instability at Noise Scale 3.0, dropping to $0.841$ (Test) while VJEPA/BJEPA remained at $> 0.90$. This collapse at Scale 3.0 likely indicates optimization difficulties where the variance regularization term fought against the predictive loss. This hints at the optimization difficulty deterministic models face when the noise magnitude creates complex gradients, even if the model recovered at Scale 4.0.
```{=latex}
\footnotesize
```
```{=latex}
\setlength{\tabcolsep}{5pt}
```
```{=latex}
\renewcommand{\arraystretch}{1.05}
```
::: {#tab:toyExperimentResults_full}
Scale (SNR) Model $\uparrow$Signal $R^2$ (Tr/Te) $\downarrow$Noise $R^2$ (Tr/Te) $\downarrow$Time (Tr/Te)
--------------- ------- --------------------------------------------------- ---------------------------------------------------- --------------------------------
0.0 (inf dB) VAE [**1.000**]{.underline} / [**1.000**]{.underline} NA / NA 12.6s / 0.01s
0.0 (inf dB) AR 0.999 / 0.999 NA / NA [**6.4s**]{.underline} / 0.01s
0.0 (inf dB) JEPA 0.947 / 0.930 NA / NA 16.9s / 0.01s
0.0 (inf dB) VJEPA 0.999 / 0.999 NA / NA 13.9s / 0.01s
0.0 (inf dB) BJEPA 0.987 / 0.981 NA / NA 23.5s / 0.01s
1.0 (15.8 dB) VAE 0.978 / 0.973 0.065 / 0.002 12.6s / 0.01s
1.0 (15.8 dB) AR 0.983 / 0.980 0.041 / -0.007 [**6.6s**]{.underline} / 0.01s
1.0 (15.8 dB) JEPA 0.947 / 0.930 0.230 / 0.178 16.9s / 0.01s
1.0 (15.8 dB) VJEPA [**0.999**]{.underline} / [**0.998**]{.underline} [**0.006**]{.underline} / [**-0.012**]{.underline} 13.8s / 0.01s
1.0 (15.8 dB) BJEPA 0.975 / 0.976 0.029 / -0.003 23.5s / 0.01s
2.0 (9.8 dB) VAE 0.903 / 0.853 0.283 / 0.194 12.8s / 0.01s
2.0 (9.8 dB) AR 0.923 / 0.889 0.179 / 0.094 [**6.4s**]{.underline} / 0.01s
2.0 (9.8 dB) JEPA 0.947 / 0.930 0.213 / 0.172 17.0s / 0.01s
2.0 (9.8 dB) VJEPA [**0.996**]{.underline} / [**0.995**]{.underline} [**0.020**]{.underline} / -0.012 14.2s / 0.01s
2.0 (9.8 dB) BJEPA 0.966 / 0.967 0.042 / [**-0.023**]{.underline} 23.3s / 0.01s
3.0 (6.3 dB) VAE 0.852 / 0.770 0.458 / 0.393 12.5s / 0.01s
3.0 (6.3 dB) AR 0.870 / 0.801 0.328 / 0.256 [**6.5s**]{.underline} / 0.01s
3.0 (6.3 dB) JEPA 0.916 / 0.841 [**0.000**]{.underline} / [**-0.003**]{.underline} 16.8s / 0.01s
3.0 (6.3 dB) VJEPA [**0.947**]{.underline} / [**0.930**]{.underline} 0.224 / 0.176 13.5s / 0.01s
3.0 (6.3 dB) BJEPA 0.946 / 0.919 0.106 / 0.078 23.0s / 0.01s
4.0 (3.8 dB) VAE 0.822 / 0.730 0.512 / 0.458 12.5s / 0.01s
4.0 (3.8 dB) AR 0.839 / 0.756 0.394 / 0.338 [**6.5s**]{.underline} / 0.01s
4.0 (3.8 dB) JEPA [**0.999**]{.underline} / [**0.999**]{.underline} [**0.004**]{.underline} / [**-0.010**]{.underline} 16.6s / 0.02s
4.0 (3.8 dB) VJEPA 0.994 / 0.993 0.025 / -0.007 13.8s / 0.01s
4.0 (3.8 dB) BJEPA 0.920 / 0.899 0.213 / 0.156 23.2s / 0.01s
5.0 (1.8 dB) VAE 0.777 / 0.681 0.543 / 0.492 12.5s / 0.01s
5.0 (1.8 dB) AR 0.806 / 0.717 0.418 / 0.364 [**6.4s**]{.underline} / 0.01s
5.0 (1.8 dB) JEPA [**0.947**]{.underline} / [**0.930**]{.underline} [**0.200**]{.underline} / [**0.183**]{.underline} 16.5s / 0.01s
5.0 (1.8 dB) VJEPA 0.945 / 0.928 0.227 / 0.190 14.1s / 0.01s
5.0 (1.8 dB) BJEPA 0.884 / 0.841 0.317 / 0.240 22.3s / 0.01s
6.0 (0.3 dB) VAE 0.729 / 0.627 0.577 / 0.531 12.3s / 0.01s
6.0 (0.3 dB) AR 0.770 / 0.677 0.437 / 0.386 [**7.1s**]{.underline} / 0.01s
6.0 (0.3 dB) JEPA [**0.999**]{.underline} / [**0.999**]{.underline} [**0.007**]{.underline} / [**-0.011**]{.underline} 16.0s / 0.01s
6.0 (0.3 dB) VJEPA 0.947 / 0.930 0.197 / 0.202 13.7s / 0.01s
6.0 (0.3 dB) BJEPA 0.935 / 0.914 0.205 / 0.171 22.8s / 0.01s
7.0 (-1.1 dB) VAE 0.675 / 0.566 0.615 / 0.574 12.7s / 0.01s
7.0 (-1.1 dB) AR 0.729 / 0.631 0.462 / 0.416 [**7.1s**]{.underline} / 0.01s
7.0 (-1.1 dB) JEPA [**0.946**]{.underline} / [**0.929**]{.underline} [**0.209**]{.underline} / 0.191 16.1s / 0.01s
7.0 (-1.1 dB) VJEPA 0.939 / 0.918 0.264 / [**0.182**]{.underline} 13.4s / 0.01s
7.0 (-1.1 dB) BJEPA 0.894 / 0.851 0.291 / 0.237 23.1s / 0.01s
8.0 (-2.2 dB) VAE 0.613 / 0.499 0.656 / 0.620 12.3s / 0.01s
8.0 (-2.2 dB) AR 0.680 / 0.578 0.491 / 0.449 [**7.1s**]{.underline} / 0.01s
8.0 (-2.2 dB) JEPA [**0.947**]{.underline} / [**0.930**]{.underline} [**0.226**]{.underline} / [**0.183**]{.underline} 16.1s / 0.01s
8.0 (-2.2 dB) VJEPA 0.905 / 0.870 0.299 / 0.251 13.4s / 0.01s
8.0 (-2.2 dB) BJEPA 0.896 / 0.841 0.292 / 0.238 23.0s / 0.01s
: Performance metrics ($R^2$) across models at all noise scales. $Tr$ denotes training set, $Te$ denotes test set. Note the robustness of JEPA family models at high noise levels (e.g. Scale 8.0) compared to generative baselines.
:::
## Computational Resources
All experiments were conducted using the standard CPU runtime provided by Google Colab [@Edwards2024Colab]. The key hardware specifications are summarized below:
- **Platform:** Google Colab (Linux x86_64)
- **CPU:** Intel(R) Xeon(R) CPU @ 2.20GHz
- **CPU Topology:** 2 Logical CPUs (1 Physical Core, 2 Threads per core)
- **Memory (RAM):** 13.61 GB
- **Disk Space:** 242.49 GB
- **Instruction Set:** x86_64 with AVX2 and FMA support
Given the lightweight nature of the linear toy experiment (latent dimension $D_z=4$), training was performed entirely on the CPU, with individual model runs completing in approximately 10-35 seconds (see Table `\ref{tab:toyExperimentResults_full}`{=latex}).
## Code Availability
The complete source code and data generation procedures used for the experiments in this paper are available at: .
[^1]: Author email: yongchao.huang@abdn.ac.uk
[^2]: The temporal formulation in this section also applied to JEPA.
[^3]: 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.
[^4]: 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\mid s_{t-1})\prod_{t=1}^T p(x_t\mid s_t),$ which implies, after marginalizing latent states $\{s_t\}$, $p(x_{1:T})=\prod_{t=1}^T p(x_t\mid x_{