VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang

Source-linked AI summary

VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Yongchao Huang

arXiv:2601.14354v1Jan 20, 2026cs.LG

TL;DR

Existing JEPA and related world models leave predictive uncertainty and sufficient-state guarantees underformalized, limiting their probabilistic interpretation for stochastic control. The paper introduces VJEPA as a variational predictive model and BJEPA as a modular extension, finding robust distractor filtering while enabling uncertainty estimation without observation likelihoods.

Problem
Existing JEPA world models rely on deterministic regression, leaving uncertainty over future latent states and predictive sufficiency for planning and control implicit.
Method
VJEPA learns a predictive distribution over future latent states with a variational objective, while BJEPA combines learned dynamics and structural priors through a Product of Experts.
Results
VJEPA and BJEPA maintain R2 > 0.84 under the Noisy TV distractor at scale 8.0, while deterministic JEPA reaches R2 = 0.93.
Takeaways & Limitations
VJEPA supports predictive uncertainty estimation, belief propagation, and sufficient-information-state control without autoregressive observation reconstruction.
Takeaways & Limitations
The unimodal Gaussian predictive head can underperform or average distinct modes in complex multimodal environments, motivating more expressive heads.

Abstract

from arXiv · show

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 \textit{deterministic} regression objectives, which mask probabilistic semantics and limit its applicability in stochastic control. In this work, we introduce \emph{Variational JEPA (VJEPA)}, a \textit{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 \emph{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.

1 Introduction

JEPA avoids observation reconstruction by predicting latent representations, but existing world-model formulations remain deterministic and leave uncertainty and predictive sufficiency unformalized. This work introduces probabilistic JEPA formulations that connect latent prediction to control, Bayesian filtering, and structural constraints.

JEPA motivation: JEPA predicts representations of missing or future data instead of reconstructing observations or contrasting samples.This removes observation-level likelihoods and emphasizes predictable, task-relevant structure.
Open gap: Existing JEPA world models use deterministic regression objectives, leaving uncertainty over future latent states and predictive sufficiency for control implicit.They do not formalize when or why the learned representation is sufficient for planning and control.
Core direction: The paper frames JEPA as a probabilistic predictive state-space model whose representations can serve as sufficient information states for control and planning.The formulation does not require autoregressive likelihood factorization.
Contributions: VJEPA introduces a probabilistic JEPA formulation, while BJEPA factorizes predictive belief into learned dynamics and structural priors for constraints and zero-shot task transfer.BJEPA uses a Product of Experts to combine these components.
Empirical validation: In a noisy environment, VJEPA and BJEPA filter high-variance distractors while providing uncertainty-aware world modeling without autoregressive observation reconstruction.The experiment supports the feasibility of uncertainty-aware latent prediction in noisy settings.

2 Related Work

Related approaches connect predictive representations to PSRs, probabilistic latent dynamics, Active Inference, and JEPA world models, but differ in scalability, observation modeling, or uncertainty treatment. VJEPA adds an explicit predictive distribution while retaining likelihood-free representation prediction.

Predictive State Representations: Predictive State Representations represent state through sufficient statistics of future observations and can generalize Markov and hidden-state models.The paper positions VJEPA as a neural, amortized PSR that compresses predictive distributions into latent embeddings.
Latent world models: PlaNet and Dreamer learn probabilistic latent dynamics alongside an explicit observation likelihood, including reconstruction of observations from latent states.Their latent dynamics are trained through an evidence lower bound over pixels.
Likelihood-free prediction: JEPA, V-JEPA, and VJEPA discard observation reconstruction and instead learn predictive latent representations that emphasize future information and predictive sufficiency.This contrasts with latent world models that explicitly model p(x_t | z_t).
Comparison with MuZero: MuZero also avoids observation reconstruction, whereas VJEPA learns latent dynamics purely through representation prediction without rewards, policies, or observation likelihoods.This makes VJEPA a self-supervised model for both control and general representation learning.
Active Inference: Active Inference uses variational inference and Bayesian filtering in latent world models, while VJEPA omits sensory likelihoods, preferences, and policy evaluation.The paper characterizes VJEPA as a model of predictive latent belief dynamics rather than a full Active Inference system.
Position of VJEPA: Prior JEPA methods and world models provide empirical planning evidence but leave uncertainty, belief propagation, and sufficient-information-state guarantees implicit.VJEPA addresses this gap through an explicit predictive distribution over future latent states and connections to Bayesian filtering and PSRs.

3 JEPA: Background and Deterministic Formulation

JEPA learns predictable, task-relevant structure by predicting target embeddings from context rather than reconstructing pixels. Its deterministic formulation uses an asymmetric encoder design and regression objective, which motivates VJEPA’s explicit probabilistic extension.

Core principle: JEPA predicts embeddings of missing or future data, encouraging predictable task-relevant structure while discarding nuisance variability.The framework avoids both observation reconstruction and positive-negative contrastive discrimination.
Context-target setup: The input is partitioned into context and target subsets, with the context encoding available information and the target representing held-out content.The partition can use spatial or temporal masking and is modality-agnostic.
Asymmetric encoders: A separate target encoder is updated by exponential moving average, making target embeddings evolve more slowly than the predictor.This asymmetric design stabilizes training and helps prevent representational collapse.
Prediction: The predictor maps the context representation to the target embedding using side information such as spatial location or temporal index.It does not directly access the target observations.
Deterministic objective: Training minimizes a regression loss between predicted and target embeddings rather than reconstructing observations or modeling pixel-level likelihoods.The objective enforces consistency between representations and organizes the latent space around predictable relationships.
Probabilistic interpretation: The mean-squared-error JEPA loss has a natural interpretation as maximizing the log-likelihood of target embeddings under a fixed-variance isotropic Gaussian.This implicit predictive-distribution assumption motivates explicitly modeling uncertainty in VJEPA.

4 VJEPA: Variational Formulation of JEPA

VJEPA extends deterministic JEPA by replacing point prediction with an explicit predictive distribution over target representations while preserving JEPA’s context-target structure and training asymmetry. Its variational objective supports uncertainty-aware prediction, collapse avoidance, and representation-space modeling without observation reconstruction.

4 VJEPA: Variational Formulation of JEPA: VJEPA replaces deterministic point prediction with a distribution over possible target embeddings while preserving JEPA’s context-target structure, conditioning, and training asymmetry.The extension gives JEPA explicit probabilistic semantics without altering its core architectural principles.
4.1 Probabilistic Predictive Model: The context encoder maps xC to ZC, while pϕ(ZT | ZC, ξT ) predicts a stochastic target representation conditioned on context and target structure.The target structure ξT can specify spatial location, temporal index, or masking pattern.
4.1 Probabilistic Predictive Model: VJEPA does not optimize an observation-level reconstruction objective; its learning signal comes solely from predicting the target distribution in representation space.An optional observation model may be written, but it is not optimized in the JEPA-style objective.
4.2 Inference Model from the Target Encoder: The EMA target encoder defines an amortized inference distribution qθ′(ZT | xT ), whose slowly evolving target stabilizes training and prevents degenerate fixed points.The inference distribution may be instantiated as a diagonal Gaussian with mean and variance parameterized by fθ′(xT ).
4.3 Variational Objective: VJEPA trains pϕ(ZT | ZC, ξT ) to match qθ′(ZT | xT ) using a regularized negative log-likelihood with a prior-based KL term.The fixed prior can be N(0, I), with β > 0 controlling regularization strength.
4.3 Variational Objective: The objective introduces explicit predictive uncertainty, permits multi-modal futures, and provides collapse-avoidance guarantees under mild conditions.The predictor can use expressive distributions such as learned Gaussians, mixtures, or normalizing flows.
4.6 Prediction and Uncertainty Propagation: VJEPA predictions can be represented as beliefs over future embeddings, propagated through moments or samples rather than reduced to a single point forecast.Within one probabilistic model, Monte Carlo sampling represents uncertainty without maintaining multiple ensemble models.
4.7 Collapse Avoidance: Under target diversity and nontrivial conditioning, every global optimum of the VJEPA objective is non-collapsed because constant context representations incur irreducible excess loss.The strict gap arises from the informative conditioning available through non-collapsed context representations, while KL regularization further discourages pathological target encodings.

5 Time-indexed VJEPA as a Latent Space Dynamical System

Time-indexing JEPA contexts and targets turns predictive representation learning into a latent state-space model with transitions, belief propagation, and control in representation space. This sequential interpretation remains likelihood-free over observations, but requires temporal indexing and reuse of predictions across steps.

Latent dynamical-system formulation: Time-indexed VJEPA defines a latent state, predictive transition model, and belief-propagation mechanism in representation space.The encoder supplies the state representation, while the predictor models future latent states and supports filtering-style updates.
Latent dynamical-system formulation: JEPA induces a latent dynamical system only when contexts and targets are explicitly time-indexed and predictions are reused across steps.Clip-level V-JEPA otherwise captures short-range correlations without defining a compositional state transition model.
Predictive transitions: VJEPA predicts arbitrary future or masked latent states and can incorporate intervening control inputs between the context and target times.The formulation supports skip-time prediction rather than requiring only one-step transitions.
Likelihood-free sequential modeling: JEPA/VJEPA performs sequential latent prediction and belief propagation without specifying or factorizing an observation likelihood.This separates sequential modeling in latent space from autoregressive modeling of observations and avoids modeling irrelevant observation entropy.
Sequential composition: The time-indexed predictor supports recursive belief-state evolution analogous to classical state-space models, enabling filtering, multi-step prediction, and planning in representation space.The extension adds structural assumptions needed for temporal composition, uncertainty propagation, and control beyond one-shot clip-level prediction.

6 JEPA/VJEPA for Control

JEPA/VJEPA can support control without reconstructing observations when its latent representation is predictively sufficient for future task consequences. Under this condition, planning and stochastic MPC can operate directly over latent trajectories and incorporate uncertainty-sensitive objectives.

Predictive information states: A control-relevant latent representation need only predict future task consequences under candidate actions, not reconstruct observations.An observation decoder is optional when objectives can be expressed through predictive latent features or a cost head.
Predictive information states: Under realizability and sufficient model expressivity, VJEPA representations can inherit the predictive sufficiency of a true information state for optimal control.The learned representation is an information-preserving transformation of the underlying predictive state up to approximation error.
Predictive information states: If latent predictive sufficiency holds and costs are measurable from the latent trajectory, cumulative-cost distributions depend on history only through the latent representation.This is the control-relevant predictive-sufficiency condition used to derive optimal policies based only on the representation.
Optimal control: When predictive sufficiency holds over the planning horizon, an optimal policy exists that depends on history only through the latent representation.Thus, planning can be reduced from observation histories to a latent information state.
Distributional planning: Stochastic VJEPA enables distributional planning by propagating latent trajectories and applying risk-neutral or risk-sensitive objectives to their distribution.Variance penalties and CVaR can be incorporated when uncertainty should affect action selection.
MAP control: Mean or MAP planning is optimal under the stated quadratic-cost Gaussian assumptions because covariance terms are action-independent constants.Outside those assumptions, action-dependent covariance, non-quadratic costs, or multimodal predictions make sampling-based rollouts important.

7 Information-Theoretic Analysis

VJEPA frames latent prediction as an information-theoretic objective: minimizing predictive negative log-likelihood maximizes a lower bound on mutual information between present and future representations. Unlike reconstruction objectives, this predictive bottleneck can discard nuisance variability while retaining task-relevant information.

Predictive mutual information: VJEPA loss minimization maximizes a lower bound on I(Zt; Zt+∆) through the Barber-Agakov variational bound.The target entropy is treated as fixed by the slowly evolving target encoder.
Predictive Information Bottleneck: VJEPA acts as a Predictive Information Bottleneck, retaining information about the future while remaining invariant to non-predictive details of the past.This contrasts with autoencoders that preserve information needed to reconstruct the full input history.
Nuisance invariance: VJEPA can admit minimal sufficient representations that discard nuisance variability while preserving predictive information about future latent states.The paper contrasts this with reconstruction objectives, which are driven to retain high-variance nuisance details.
Generative versus predictive objectives: Generative objectives maximize information between latent states and raw future observations, including nuisance-noise information when it contributes to reconstruction.The mutual-information decomposition separates task-relevant signal from conditional information about noise.
Nuisance invariance: Autoregressive observation objectives penalize representations that fail to predict high-entropy future noise, whereas VJEPA’s latent target removes that penalty.Prediction is therefore focused on latent dynamics relevant to planning and control rather than pixel-level noise.

8 Bayesian JEPA (BJEPA)

BJEPA extends VJEPA with a modular Bayesian factorization: a learned dynamics expert is combined with a constraint-based prior through a Product of Experts. This yields uncertainty-aware latent filtering and planning while supporting task transfer through prior substitution.

Architecture: BJEPA factorizes predictive belief into a learned dynamics expert and a constraint-based expert.The factorization preserves JEPA’s representation-learning design while adding explicit Bayesian conditioning in latent space.
Architecture: The predictive likelihood encodes typical latent dynamics, while the prior injects goals, safety constraints, or physical manifolds through auxiliary input η.The two factors are treated as conditionally independent information sources about the future state.
Bayesian fusion: The Product of Experts restricts predictions to states that are both dynamically feasible and task-compliant.In the Gaussian case, the combination admits a closed-form latent-space update analogous to a Kalman update.
Design properties: BJEPA retains representation-space operation without requiring observation reconstruction, reward supervision, or autoregressive likelihood factorization.Its explicit posterior supports uncertainty-aware planning while preserving JEPA’s likelihood-free observation modeling.
Bayesian fusion: BJEPA reduces exactly to the VJEPA predictive distribution when the auxiliary prior is absent, so it generalizes VJEPA without changing its training objective.This makes the prior an additional conditioning mechanism rather than a replacement for the predictive model.
Modularity and transfer: Swapping the prior expert enables zero-shot transfer to new tasks without retraining the learned dynamics model.Task-specific information is injected at inference time, preserving the shared model’s environmental knowledge.
Filtering and planning: Sequential BJEPA updates act as Bayesian filtering in representation space, with auxiliary constraints functioning as virtual observations for stochastic latent MPC.The posterior belief combines physically feasible futures with goal- or constraint-consistent futures.

9 Toy Experiment: a “Noisy TV” Linear System

The Noisy-TV experiment tests whether latent-predictive models ignore high-variance distractors when the true signal is low-dimensional. Joint-embedding models remain robust at the highest noise level, while generative baselines degrade and track the distractor.

Experimental setup: The experiment uses a 20-dimensional observation with a 4-dimensional signal and a scaled sticky-random-walk distractor.At σ = 8.0, the distractor variance reaches approximately 64× the signal variance and the SNR falls to −2.2 dB.
Experimental setup: Models are evaluated by linear-probe R2 for signal recovery on training and generalization test data.JEPA and VJEPA use predicted future latents, while BJEPA evaluates its fused posterior mean.
Results: At σ = 8.0, VAE test signal recovery falls to R2 ≈0.50 while noise recovery remains R2 = 0.62, and AR signal recovery falls to R2 = 0.578.These results indicate that the generative baselines prioritize high-variance distractors.
Results: JEPA, VJEPA, and BJEPA maintain R2 > 0.84 at σ = 8.0, with deterministic JEPA reaching R2 = 0.93.The probabilistic variants showed superior training stability across seeds, while deterministic JEPA exhibited momentary instability at σ = 3.0.
Qualitative filtering: At σ = 8.0, VAE and AR reconstructions track high-frequency noise, whereas VJEPA and BJEPA track the underlying true signal.The figure compares dashed generative reconstructions with solid joint-embedding reconstructions against the black ground-truth signal.

10 Discussion

The discussion interprets JEPA as probabilistic latent state-space modeling rather than observation generation. It connects predictive mutual information to nuisance filtering and highlights probabilistic expressivity and regularization balance as remaining limitations.

Core interpretation: The paper presents VJEPA and BJEPA as probabilistic latent dynamical frameworks that formalize JEPA beyond a self-supervised training trick.The discussion positions this as a bridge between deterministic representation learning and probabilistic world models.
Core interpretation: VJEPA decouples sequential modeling from autoregressive pixel likelihoods by learning a time-indexed latent transition model.The transition kernel supports belief propagation and filtering without requiring autoregressive observation factorization.
Information-theoretic interpretation: Generative reconstruction prioritizes high-variance input directions, whereas VJEPA prioritizes dimensions with high predictive mutual information and can ignore noise regardless of magnitude.The discussion characterizes these objectives as nonlinear analogues of PCA and CCA or the Information Bottleneck, respectively.
Applications: The framework is presented as applicable to robotics, games, and language planning because latent dynamics can support control or trajectory optimization without pixel-level reconstruction.These application discussions describe sensory, game-mechanics, and semantic states as latent planning spaces.
Limitations: Probabilistic VJEPA can slightly underperform deterministic JEPA or require more tuning because a unimodal Gaussian predictive head may average distinct future modes.The paper identifies multimodal bifurcations, such as choosing left or right around an obstacle, as a motivating case for richer predictive heads.
Limitations: The balance between predictive and KL-regularization terms is important: excessive regularization can collapse representations, while insufficient regularization can leave variance estimates uncalibrated.The paper proposes adaptive balancing mechanisms as future work.

11 Conclusion

The conclusion presents VJEPA as a probabilistic JEPA framework and BJEPA as its modular extension for uncertainty-aware latent planning. It connects predictive representations to control, Bayesian filtering, and likelihood-free sequential modeling.

VJEPA recasts JEPA as variational inference over future latent states, clarifying the probabilistic semantics of predictive embeddings.
VJEPA maximizes a variational lower bound on predictive mutual information without requiring an autoregressive observation likelihood.
The framework enables principled uncertainty estimation by sampling predictive distributions while avoiding explicit modeling of high-entropy nuisance variables.
BJEPA factorizes predictive belief into learned dynamics and structural prior experts, supporting constraint injection and zero-shot task transfer.
Under standard conditions, predictive latent states suffice to derive optimal control, making pixel reconstruction redundant for planning.
Sequential VJEPA admits Bayesian-filter and particle-filter interpretations, propagating latent beliefs through dynamics and surrogate measurement updates.

F Detailed Proof of Theorem 1

Theorem 1 shows that, under target diversity and nontrivial conditioning assumptions, collapsed context representations cannot be globally optimal for the VJEPA objective. The proof attributes collapse avoidance to an irreducible prediction-loss gap rather than architectural heuristics.

A collapsed context encoder forces the predictor to use a single unconditional distribution for each conditioning variable.
The proof assumes target diversity and a predictive family capable of representing different distributions for different context representations.
The collapsed predictor incurs a strictly positive irreducible excess loss when target embeddings vary across inputs.
A non-collapsed representation can condition predictions on context and achieve strictly lower loss than the collapsed solution.
The KL regularizer prevents pathological target-encoder collapse but does not create the strict separation between collapsed and non-collapsed optima.
Therefore, no collapsed representation is globally optimal under the theorem’s stated assumptions.

G Detailed Proof of Theorem 3

Theorem 3 establishes when stochastic control can optimize at the predictive mean, and therefore at the MAP point, instead of integrating over the full predictive distribution. This reduction depends on Gaussian predictions, quadratic costs, and action-independent covariance.

The theorem considers one-step control with a latent probabilistic dynamics model and a stage cost evaluated after the action.
For quadratic costs and Gaussian predictions with action-independent covariance, expected cost separates into mean-dependent and covariance terms.
The covariance contribution is constant across actions, so minimizing expected cost is equivalent to minimizing cost at the predictive mean.
Because the Gaussian predictive mean is its unique MAP point, mean planning is equivalently MAP control.
If covariance depends on actions, or costs or predictive distributions are non-quadratic or non-Gaussian, uncertainty generally affects the optimal action.

H Proof of Proposition 1 (Invariance to Nuisance)

The proposition contrasts generative objectives with VJEPA under nuisance variables that are irrelevant to future targets. Generative likelihood maximization favors retaining nuisance information, whereas VJEPA permits equally optimal representations that discard it.

The setup decomposes observations into a future-relevant signal and a nuisance variable conditionally independent of the future target.
Generative likelihood maximization requires representations to encode both the signal and nuisance information to reconstruct observations.
Consequently, a generative representation that discards nuisance information is suboptimal because it cannot reconstruct the nuisance component.
VJEPA predicts the future target from the current representation rather than reconstructing the observation itself.
Because the nuisance is conditionally irrelevant to the future target, minimal representations that discard nuisance variability can achieve the same optimal VJEPA loss as maximal representations.

I Derivation of the BJEPA Predictive Factorization

The section derives BJEPA’s predictive factorization by combining a history-conditioned dynamics expert with an auxiliary prior under conditional independence. It also specifies the noisy linear experiment and contrasts deterministic, probabilistic, and Bayesian joint-embedding models.

Predictive factorization: BJEPA factorizes the predictive belief into a history-conditioned expert p(ZT | ZC) and an auxiliary expert p(ZT | η), up to normalization.The derivation treats the target prior as uniform or absorbed into learned expert bias, yielding a Product of Experts form.
Predictive factorization: The factorization assumes that history ZC and auxiliary information η are conditionally independent given the true target state ZT.Under this assumption, their likelihood terms multiply when forming the posterior over ZT.
Experimental setup: The experiment compares pixel-space generative reconstruction with latent-space JEPA prediction in a linear Noisy TV environment containing signal and distractor dimensions.All models use a four-dimensional latent bottleneck, while generative models reconstruct observations that include high-variance distractors.
Model variants: VJEPA models future latent states probabilistically, using Gaussian target and predictive distributions rather than deterministic latent regression.The target encoder supplies a distribution over future latents, while the predictor outputs belief parameters for those targets.
Model variants: BJEPA extends VJEPA with a static latent-space prior that fuses with the predictive likelihood to filter unpredictable dynamics.The added prior is represented by learnable static parameters defining a preferred latent manifold.
Stochastic training: Training samples targets from qθ′ instead of using their mean, preventing the predictor from collapsing to an overconfident zero-variance belief.Sampling keeps the regression target variable, forcing the predictor to maintain non-zero predictive variance.

J.4 Results

In the Noisy TV experiment, generative baselines degrade as distractor noise increases, whereas JEPA-family models retain strong signal recovery at the highest noise level. Probabilistic variants also show greater training stability across seeds, although deterministic JEPA reaches the best accuracy in the reported run.

Evaluation: The evaluation reports R2 between learned latents and true signal across nine noise scales, using both training and test metrics.Table 7 summarizes performance across models, with Tr denoting training and Te denoting test.
Generative baselines: At σ = 8.0, VAE and AR signal recovery falls to R2 ≈0.50 and 0.58, while noise recovery remains R2 = 0.62 and 0.45.The result is attributed to likelihood-based objectives encoding high-variance distractors and collapsing target-task representations.
Joint-embedding models: At σ = 8.0, JEPA, VJEPA, and BJEPA maintain signal recovery above R2 > 0.85 while filtering distractor dynamics.In the reported trial, deterministic JEPA achieves R2 = 0.93, the highest final accuracy.
Stability: At noise scale 3.0, JEPA drops to 0.841 test performance, whereas VJEPA and BJEPA remain above 0.90.The passage links this momentary instability to optimization difficulties under high-variance target embeddings.
Stability: Across different random seeds, deterministic JEPA occasionally collapses below R2 < 0.5 at intermediate noise scales, while VJEPA and BJEPA are consistently more stable.The probabilistic methods account for aleatoric uncertainty through their variance outputs.