VJEPA: Variational Joint Embedding Predictive Architectures as Probabilistic World Models

Background

• JEPA provides a third self-supervised paradigm that predicts missing or future representations, avoiding observation-level likelihoods while emphasizing predictable, task-relevant structure.I-JEPA and V-JEPA demonstrate strong vision and video performance, while H-JEPA motivates latent dynamics for control and long-horizon reasoning.Source: §1 · p. 5

Problem

• Existing JEPA world models remain largely deterministic and regression-trained, leaving predictive uncertainty, probabilistic semantics, and the conditions for control-sufficient representations unformalized.This gap is especially important because current approaches implicitly assume predictive sufficiency without characterizing when or why it holds.Source: §1 · p. 5

Contributions

• BJEPA factorizes predictive belief into learned dynamics and structural prior experts, enabling constraint and physical-law injection alongside zero-shot task transfer through a Product of Experts.The modular factorization connects JEPA to Bayesian filtering and supports structural priors in latent planning.Source: §1 · pp. 5–6
• In a noisy environment, VJEPA and BJEPA filter high-variance distractors, demonstrating uncertainty-aware world modeling without autoregressive observation reconstruction.This empirical result addresses whether temporal structure requires modeling nuisance observations directly.Source: §1 · pp. 5–6
• VJEPA formulates JEPA as a probabilistic predictive state-space model, showing that sequential prediction can remain likelihood-free and that representations can serve as sufficient information states for planning and optimal control.The formulation also provides formal guarantees for collapse avoidance.Source: §1 · p. 6

2.1 Predictive State Representations

• PSRs represent state through sufficient statistics of future observations, generalize Markov and hidden-state models, and motivate VJEPA as a neural amortized PSR.PSRs were introduced by Sutton et al. and formalized independently by Singh et al..Source: §2.1 · p. 6

2.2 World Models and Latent Dynamics

• PlaNet and Dreamer learn probabilistic latent dynamics with observation reconstruction, whereas JEPA-family models prioritize predictive sufficiency over pixel fidelity; MuZero omits reconstruction but uses value and policy heads.VJEPA differs by learning latent dynamics through representation prediction without rewards, policies, or observation likelihoods.Source: §2.2 · p. 6

2.3 Active Inference

• VJEPA shares Active Inference’s latent belief propagation but omits sensory likelihoods, preferences, and policy evaluation, remaining a self-supervised model of predictive belief dynamics.Active Inference frames perception, action, and learning as variational inference, with perception minimizing free energy and action minimizing expected future free energy.Source: §2.3 · p. 7

2.4 Joint Embedding Predictive Architectures (JEPAs)

• VJEPA extends JEPA with an explicit variational predictive distribution over future latent states, enabling uncertainty estimation and probabilistic control beyond deterministic regression objectives.Prior JEPA methods do not provide a principled likelihood-based account of predictive uncertainty, whereas VJEPA adds likelihood and KL regularization. [5]Source: §2.4 · pp. 7–8

2.5 JEPA World Models and Latent-Space Planning

• VJEPA complements deterministic JEPA world models by enabling belief propagation, distributional planning, and formal connections to PSRs and Bayesian filtering without observation reconstruction.Existing JEPA planning methods support latent-space control but leave uncertainty, belief propagation, and sufficient-information-state guarantees implicit. [76]Source: §2.5 · pp. 8–9

• JEPA predicts embeddings of missing or future data, encouraging predictable task-relevant structure while discarding nuisance variability rather than reconstructing observations or contrasting samples.Source: §3 · p. 9
• The framework partitions inputs into context and target subsets, encodes them separately, and predicts the target representation from context plus structural side information such as spatial or temporal location.The formulation supports spatial masking, as in I-JEPA, and temporal masking, as in video JEPA, while remaining modality-agnostic.Source: §3 · p. 9
• An exponential-moving-average target encoder changes more slowly than the predictor, stabilizing training and preventing representational collapse [55, 54, 45, 26].Source: §3 · p. 9
• JEPA minimizes regression error between predicted and target embeddings without reconstructing observations or modeling pixel-level likelihoods.This objective organizes the latent space around predictable relationships in the data.Source: §3 · p. 10
• The usual mean-squared-error loss is equivalent to maximizing the log-likelihood of the target embedding under an isotropic Gaussian centered at the prediction with fixed variance.This implicit simple predictive distribution motivates explicitly modeling uncertainty in the variational extension.Source: §3 · p. 10

4 VJEPA: Variational Formulation of JEPA

• VJEPA replaces JEPA’s point estimate with an explicit predictive distribution over target embeddings, enabling uncertainty-aware prediction and multimodal futures conditioned on context and target structure.This is a minimal probabilistic extension that preserves JEPA’s context-target structure, conditioning mechanism, and training asymmetry.Source: §4 · p. 10, Table 1 · p. 11

4.1 Probabilistic Predictive Model

• The probabilistic predictor pϕ(ZT | ZC, ξT ) conditions on deterministic context encoding ZC and target side information ξT, while observation reconstruction remains unoptimized.VJEPA optionally factorizes an observation model, but its learning signal comes solely from predictive distribution matching in representation space.Source: §4.1 · p. 11

4.2 Inference Model from the Target Encoder

• The target encoder’s EMA-based inference distribution qθ′(ZT | xT ) supplies a slowly evolving representation-space posterior, stabilizing training and preventing degenerate fixed points.The inference distribution may be a diagonal Gaussian whose mean and variance are parameterized by the target encoder.Source: §4.2 · p. 11

4.3 Variational Objective

• VJEPA minimizes predictive negative log-likelihood plus β-weighted KL regularization to a fixed prior, introducing uncertainty, allowing multimodal futures, and yielding collapse-avoidance guarantees under mild conditions.With fixed variance, the objective reduces to JEPA’s squared-loss objective.Source: §4.3 · p. 12

4.4 Schematic Overview

• Training samples target representations from qθ′, computes their predictive negative log-likelihood, adds KL regularization, and supports Gaussian, mixture, or normalizing-flow predictors.The reparameterization trick may be used for continuous inference distributions, and the deterministic JEPA procedure is recovered as a special case.Source: Figure 1 · p. 13, §4.5 · pp. 13–14

4.6 Prediction and Uncertainty Propagation

• VJEPA represents future uncertainty as a latent belief, supporting point, MAP, mode, moment, or Monte Carlo predictions and multi-step distributional rollouts for planning and control.Sampling yields plausible future representations and can approximate downstream expectations even for complex predictive distributions; sequential propagation is analogous to Bayesian filtering.Source: §4.6 · pp. 14–15
• Unlike ensembles that represent uncertainty through parameter disagreement, VJEPA captures it within one probabilistic representation-space model, using Monte Carlo sampling without maintaining multiple models.This provides a distributional generalization of ensemble world models while avoiding their computational overhead and calibration issues.Source: §4.6 · p. 15

4.7 Collapse Avoidance

• Under target diversity and nontrivial conditioning, every global optimum of the VJEPA objective is non-collapsed because collapsed contexts incur irreducible excess loss.The theorem attributes collapse avoidance primarily to the predictive objective’s structure, with KL regularization further discouraging degenerate target encodings.Source: §4.7 · pp. 15–16

5 Time-indexed VJEPA as a Latent Space Dynamical System

• Time-indexed JEPA/VJEPA defines latent states, predictive transitions, and belief propagation, enabling filtering and control directly in representation space.This interpretation remains consistent with the original JEPA training objective.Source: §5 · pp. 16–17

5.1 Background: State-Space Models

• A state-space model supports prediction and control by propagating beliefs over latent states through a transition model.This belief is conditioned on available observations and prior controls.Source: §5.1 · p. 17

5.2 From JEPA Masking to Temporal Dynamics

• Temporal JEPA predicts arbitrary future latent states at chosen horizons, including skip-time predictions, without requiring an autoregressive observation likelihood.The target can correspond to a future time index, and multi-step prediction need not be reduced to iterated one-step transitions.Source: §5.2 · p. 17

5.3 Latent Transition Model Induced by JEPA/VJEPA

• VJEPA induces a stochastic latent transition model, optionally conditioned on intervening controls and horizon information, while training without an observation-level likelihood.When the horizon is one, the formulation reduces to a standard one-step state-space model.Source: §5.3 · p. 18

5.4 Sequential but Not Autoregressive

• JEPA/VJEPA is sequential but not autoregressive because it propagates beliefs and predicts latent states without specifying or factorizing p(x_t+1 | x_≤t).Observation-level autoregressive objectives force models to represent full observation entropy, including textures and noise irrelevant to downstream control.Source: §5.4 · p. 19

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

• Clip-level V-JEPA captures temporal cues through one-shot masked-region prediction but does not enforce recursively compositional, Markovian latent dynamics across horizons.Time-indexed recursive prediction adds the structural assumptions needed for temporal composition, uncertainty propagation, and control.Source: §5.5 · pp. 19–20

6.1 Background: POMDPs and Belief-State Control

• Classical POMDP control uses a Bayesian belief as a sufficient information state, reducing history-dependent planning to a fully observable belief-space MDP.The belief summarizes observation-action history and supports policies that map the current belief directly to actions.Source: §6.1 · pp. 20–21

6 JEPA/VJEPA for Control

• VJEPA enables optimal control from a predictively sufficient latent state, so observation reconstruction and explicit Bayesian belief estimation are unnecessary under realizability and optimization assumptions.The learned state preserves the control-relevant predictive information of a true belief or predictive state representation, up to approximation error.Source: §6.2 · pp. 21–22

6.3 Predictive Sufficiency for Control (Formalization)

• If future cost-relevant variables depend on history only through Z_t, latent predictive sufficiency implies cost sufficiency and an optimal policy can be chosen as π⋆(u_t | Z_t).This consequence holds when stage costs are measurable with respect to the latent trajectory; VJEPA’s guarantee remains conditional on realizability, expressive models, and near-global optimization.Source: §6.3 · pp. 22–23

6.4 Latent-Space Planning Objective

• VJEPA replaces intractable Bayesian belief propagation with a learned predictive latent state and transition model, making the H-step objective formally equivalent to stochastic MPC in representation space.The predictive distribution enables risk-neutral planning and risk-sensitive criteria such as variance penalties or CVaR without an explicit observation likelihood.Source: §6.4 · p. 23, Latent-space planning objective · p. 24

6.5 Sampling vs. Point Prediction Planning

• Point-prediction rollouts are efficient when predictive uncertainty is weak, unimodal, and approximately action-independent, whereas sampling-based rollouts handle multimodality, partial observability, and rare costly outcomes.Monte Carlo aggregation also supports robust and risk-sensitive planning by changing how sampled trajectories are evaluated.Source: §6.5 · p. 24, §6.7 · p. 26

6.6 Algorithm: VJEPA Sampling-Based Latent MPC

• VJEPA-MPC estimates the current predictive state, samples candidate action sequences and latent dynamics trajectories, and selects actions using their cumulative predicted costs.The procedure is presented as random-shooting or MPPI-style sampling-based latent MPC.Source: §6.6 · pp. 24–25

6.7 When is MAP Control Optimal?

• Under Gaussian predictive dynamics with quadratic costs and action-independent covariance, minimizing expected cost is equivalent to deterministic planning on the predictive mean, hence MAP control.When covariance depends on actions, costs are non-quadratic, or dynamics are multimodal, point planning can fail to account for uncertainty and tail risk.Source: §6.7 · pp. 25–27

7.1 Variational Maximization of Predictive Information

• VJEPA loss minimization is equivalent to maximizing a lower bound on mutual information between current and future latent states.The bound follows from the Barber–Agakov variational lower bound, while the target encoder’s slowly evolving entropy is treated as fixed during optimization.Source: §7.1 · pp. 27–28

7.2 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.Under the signal-plus-nuisance formulation, it admits minimal sufficient statistics that discard nuisance variability while preserving predictive mutual information, unlike reconstruction objectives that retain noise.Source: §7.2 · pp. 28–29

7.3 Contrast with Generative Objectives

• Generative world models maximize mutual information between latent states and raw pixels, forcing representations to encode high-entropy, task-irrelevant noise.The decomposition separates predictable signal information from conditional nuisance information; minimizing the generative loss penalizes failure to capture the latter.Source: §7.3 · pp. 29–30
• VJEPA predicts in a regularized latent space whose target encoder discards pixel noise, so its objective contains no penalty for ignoring nuisance variability.This lets the model filter high-frequency observation noise before prediction and allocate capacity to causal dynamics relevant for planning and control.Source: §7.3 · p. 30

8 Bayesian JEPA (BJEPA)

• BJEPA factorizes prediction into a learned dynamics expert and a constraint-based expert, enabling principled incorporation of goals, constraints, and task structure.The dynamics captures typical latent evolution, while the prior injects auxiliary information such as goals, safety constraints, or physical manifolds.Source: §8 · p. 30, §8.1 · p. 31

8.1 Bayesian Factorization via Product of Experts

• The Product of Experts yields a Bayesian update whose posterior concentrates on states that are both physically reachable and task-compliant.Under conditional independence of context and auxiliary information given the target, the formulation is exact up to normalization and marginal scaling, while restricting planning to the intersected latent space.Source: §8.1 · p. 31

8.2 The BJEPA Architecture

• BJEPA adds a latent-space prior encoder to the JEPA/VJEPA context-target architecture, and VJEPA is recovered exactly with a uniform prior.Gaussian likelihood and prior experts admit a closed-form Kalman-like fusion, while all components operate in representation space without reconstruction, reward supervision, or autoregressive likelihood factorization.Source: §8.2 · pp. 32–33

8.3 Training and Inference

• BJEPA trains data-driven dynamics with optional structural regularization, then introduces task-specific priors at inference to enable zero-shot latent-space planning.Training uses a variational negative log-likelihood with KL-based regularization, while inference fuses dynamics and priors so plans remain dynamically feasible and task-compliant; sampling preserves multimodal uncertainty when needed.Source: §8.3 · pp. 34–36

8.4 Sequential BJEPA as Latent Bayesian Filtering

• Sequential BJEPA performs Bayesian filtering in representation space by predicting with dynamics and correcting with task or constraint priors, thereby unifying filtering and planning without explicit observation likelihoods.The resulting posterior supports a belief-space MPC interpretation in which physically feasible beliefs are intersected with task-relevant constraints.Source: §8.4 · pp. 36–37

8.5 MAP Optimality under BJEPA

• BJEPA’s Gaussian Product of Experts makes MAP planning equivalent to minimizing expected quadratic cost when expert covariances are action-independent and the prior is action-independent.The posterior covariance is constant across actions, so the trace term in expected cost does not affect optimization; the Gaussian mean equals its MAP estimate.Source: §8.5 · pp. 37–38

8.6 Energy-Based Priors and Latent Optimization

• Energy-based priors allow non-convex constraints and task potentials, while the learned Gaussian dynamics anchors latent solutions to physically plausible futures.Because VJEPA/BJEPA components are differentiable, planning reduces to gradient-based energy minimization directly in representation space.Source: §8.6 · p. 38

8.7 Algorithm: Gradient-Based Planning with BJEPA

• BJEPA implements model predictive control by optimizing action sequences through differentiable latent rollouts, using the learned likelihood as dynamics constraints and the energy prior as the cost.At each step, the first action is executed and the latent state is updated before replanning.Source: §8.7 · pp. 38–39

8.8 Modularity, Transfer, and Continual Learning

• BJEPA separates objective-agnostic dynamics from task-specific priors, enabling zero-shot task transfer by swapping priors and mitigating catastrophic forgetting as new goals are introduced.The shared dynamics model preserves general physics knowledge, while lightweight priors provide task specificity at inference time.Source: §8.8 · p. 39
• BJEPA combines likelihood-free representation learning with explicit uncertainty and modular task specification, unlike frameworks that require observation reconstruction or monolithic reward functions.This combination preserves JEPA’s avoidance of observation reconstruction while adding probabilistic planning and structured priors.Source: §8.8 · p. 40, Table 3 · p. 40

9.1 Environment Setup

• The environment uses Dx = 20 observations and Ds = 4 latent state dimensions, with a stable rotating signal and a distractor scaled up to σ = 8.0, where SNR reaches −2.2 dB.The distractor variance can grow to approximately 64× the signal variance, creating the intended high-variance nuisance challenge.Source: §9.1 · pp. 40–41

9.2 Model Architectures and Evaluation Alignment

• The comparison uses five linear models with latent dimension Dz = Ds = 4, evaluating each model according to its temporal prediction target with a linear-probe R2 metric on held-out generalization data.VAE is evaluated against st, JEPA and VJEPA predictor outputs against st+1, and BJEPA’s fused mean µpost against st+1.Source: §9.2 · p. 41, §9.3 · p. 41

9.3 Results and Analysis

• At σ = 8.0, JEPA, VJEPA, and BJEPA maintain R2 > 0.84 for signal recovery, while VAE falls to R2 ≈0.50 and AR to R2 = 0.578; deterministic JEPA reaches R2 = 0.93.Table 4 and Fig.3 report the cross-model performance comparison; probabilistic VJEPA and BJEPA also show superior training stability across seeds.Source: §9.3 · p. 41, Table 4 · p. 42, Figure 3 · p. 42
• At σ = 8.0, VAE and AR track high-frequency noise, whereas VJEPA and BJEPA filter the distractor and track the underlying true signal st+1.The qualitative filtering comparison is shown in Fig.4 using dashed generative-model reconstructions and solid joint-embedding reconstructions.Source: §9.3 · p. 41, Figure 4 · p. 43

10 Discussion

• VJEPA and BJEPA formalize JEPA as probabilistic latent dynamical modeling, bridging deterministic representation learning with probabilistic world models.The framework treats JEPA as a principled approach for learning latent dynamical systems rather than merely a self-supervised training technique.Source: §10 · p. 43

10.1 Unifying JEPA with Probabilistic State-Space Models

• VJEPA learns non-autoregressive latent transitions pϕ(Zt+1|Zt), enabling belief propagation and filtering without modeling pixel densities.This decouples sequential modeling from autoregressive observation likelihoods used by standard world models such as Dreamer and generative sequence models such as Transformers.Source: §10.1 · p. 43

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

• VJEPA prioritizes predictive mutual information rather than input variance, ignoring high-variance nuisance noise and acting as a semantic noise filter for planning.Theorem 4 links its loss to a lower bound on I(Zt; Zt+∆), while Proposition 1 establishes nuisance invariance; reconstruction objectives instead behave like PCA.Source: §10.2 · pp. 43–44

10.3 Applications and Broader Impact

• VJEPA supports latent-space control across robotics, games, and language, while BJEPA enables swapping goal or constraint priors without retraining dynamics.Robotics uses VJEPA-MPC for visual control and BJEPA for zero-shot task transfer; games use reward-agnostic latent simulation with MCTS, while language can use token planning toward target semantic states.Source: §10.3 · pp. 44–45

10.4 Limitations and Future Directions

• Unimodal Gaussian predictive heads can underperform deterministic JEPA on multimodal dynamics, motivating GMM or latent-diffusion heads.Averaging distinct outcomes can produce blurry or incorrect beliefs, and stronger KL regularization can collapse representations while weaker regularization can miscalibrate variance estimates.Source: §10.4 · p. 45

11 Conclusion

• VJEPA and BJEPA establish a probabilistic JEPA framework by modeling predictive latent distributions and factorizing dynamics from constraint priors.VJEPA resolves deterministic JEPA’s ambiguity in stochastic futures, while BJEPA enables zero-shot task transfer and constraint satisfaction through a Product of Experts.Source: §11 · pp. 45–46
• VJEPA maximizes a variational lower bound on predictive mutual information without autoregressive observation likelihoods, filtering nuisance variability while enabling credible intervals from multiple samples.The predictive distribution supports principled uncertainty estimation unavailable to deterministic JEPA.Source: §11 · p. 46
• Predictive latent states can suffice for optimal control, making pixel reconstruction unnecessary while unifying JEPA with Predictive State Representations and Bayesian filtering.The framework learns temporal dynamics p(Zt+1|Zt) without factorizing the high-dimensional observation density p(xt+1|x≤t).Source: §11 · p. 46

A Equivalence Between Squared Loss and Gaussian Likelihood

• Deterministic JEPA regression is equivalent to maximum likelihood under an isotropic Gaussian predictive model with fixed variance, motivating VJEPA’s learned uncertainty structure.The equivalence follows because the Gaussian negative log-likelihood differs from squared error only by parameter-independent terms.Source: A Equivalence Between Squared Loss and… · pp. 52–53

B Relationships Between Information-Theoretic Quantities

• The appendices establish the entropy, cross-entropy, KL, mutual-information, chain-rule, symmetry, ELBO, and non-negativity identities underpinning VJEPA and BJEPA’s theoretical analysis.These relationships support the variational objective and the proof of a non-collapsed global optimum.Source: B Relationships Between… · p. 53, §B.2 · p. 53, §B.3 · p. 54

C Compare VJEPA with Hidden Markov Models (HMMs)

• VJEPA resembles an HMM through latent states, transitions, and belief-style filtering, but omits observation likelihoods, permits direct multi-step prediction, and learns representation geometry.Its encoder-defined target distribution is not an HMM posterior, and the model retains uncertainty propagation without reconstructing irrelevant observations.Source: C Compare VJEPA with Hidden Markov… · p. 56, §C.3 · p. 57, §C.4 · p. 57

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

• Sequential VJEPA can be interpreted as an HMM-like latent transition process without emissions, preserving probabilistic semantics and uncertainty propagation while avoiding sensory reconstruction.This design is closer to continuous Predictive State Representations than classical HMM fitting.Source: §C.5 · p. 57, §C.4 · p. 57

D Kalman-Filter-Like Interpretation of VJEPA

• VJEPA mirrors Bayesian-filter prediction in latent space and becomes exact for linear-Gaussian dynamics, generalizing Kalman-style belief updates through nonlinear amortized inference.Interpreting Zt as a sufficient statistic for belief history lets pϕ implement a generalized transition update.Source: D Kalman-Filter-Like Interpretation of… · p. 58, §D.2 · p. 58

D.2 VJEPA as Generalized Latent Prediction

• VJEPA captures Kalman-style prediction by modeling nonlinear transitions, heteroscedastic aleatoric uncertainty, and propagated belief uncertainty in latent space.The mean predictor models state-dependent dynamics, while sampling from latent input beliefs propagates epistemic uncertainty alongside covariance output by the predictor.Source: §D.2 · pp. 58–59

D.3 Amortized Correction

• VJEPA amortizes the measurement correction by mapping updated observation histories directly to posterior latent states through its context encoder.Training aligns the predictive belief with the encoder’s inferred belief, making the predictor consistent with the encoder’s implicit Bayesian updates without explicit likelihood or Kalman-gain equations.Source: §D.3 · p. 59

E A Particle-Filter Interpretation of VJEPA

• The particle filter predicts particles with pϕ, reweights them using either an optional decoder likelihood or a qθ′-based latent pseudo-likelihood, then resamples to prevent degeneracy.The latent update compares encoder-informed probability with a fixed reference prior, while resampling draws particles according to normalized weights.Source: E A Particle-Filter Interpretation of… · pp. 59–60
• Sequential VJEPA forms a full latent belief-propagation mechanism: pϕ supplies dynamics, while qθ′ injects measurement information directly in latent space.This particle-filter interpretation supports multimodal beliefs without Gaussian restrictions and is illustrated by the three-stage procedure in Fig. 5.Source: E A Particle-Filter Interpretation of… · pp. 59–60, Figure 5 · p. 61

F Detailed Proof of Theorem 1

• Under the stated assumptions, VJEPA admits no collapsed global optimum because collapsed context representations incur an irreducible prediction gap.The gap arises from fitting one unconditional distribution to multiple distinct target distributions; the KL regularizer does not create this strict separation.Source: F Detailed Proof of Theorem 1 · p. 61, Theorem 1: No Collapsed Global Optimum… · p. 63

Theorem 1: No Collapsed Global Optimum under Target Diversity

• The proof proceeds by showing that a collapsed context encoder removes context dependence, yielding an optimal unconditional predictor whose loss exceeds that of a non-collapsed solution.The excess loss is characterized through conditional mutual information, while the KL term is independent of ZC and therefore does not eliminate the gap.Source: Theorem 1: No Collapsed Global Optimum… · p. 62
• Target diversity and nontrivial conditioning ensure that non-collapsed representations can model context-dependent target distributions and achieve strictly lower loss than collapse.In the realizable limit, the non-collapsed predictor reaches conditional target entropy, satisfying pred = EξT ExT |ξT H(ZT | XT = xT , ξT ) < Lcoll.Source: Theorem 1: No Collapsed Global Optimum… · pp. 62–63

Theorem 3: Optimality of MAP (Mean) Control under Quadratic Costs

• Under Gaussian dynamics with action-independent covariance and quadratic next-state costs, minimizing expected cost is equivalent to minimizing cost at the predictive mean and MAP point.A Gaussian’s unique MAP point is its mean, so mean planning and MAP control coincide.Source: Theorem 3: Optimality of MAP (Mean)… · p. 64
• The proof decomposes the expected quadratic cost into a mean-dependent term and the trace term tr(QcΣ), with cross terms vanishing under zero-mean deviations.Because Σ does not depend on the action, tr(QcΣ) is constant and does not change the argmin.Source: Theorem 3: Optimality of MAP (Mean)… · p. 65
• If covariance depends on the action, or costs are non-quadratic or predictive distributions are non-Gaussian or multimodal, expected-cost minimization generally cannot use a single point estimate.In these cases, the optimal action may depend on predictive uncertainty rather than only the predictive mean or MAP point.Source: Theorem 3: Optimality of MAP (Mean)… · p. 66
• The post-action cost c(Zt+1, ut) matches prediction-based stochastic MPC and belief-space control, while Markov expectations can absorb it into an equivalent cost on (Zt, ut).The paper adopts post-action indexing because costs are evaluated after applying the action and observing its predicted effect.Source: Theorem 3: Optimality of MAP (Mean)… · p. 66

H Proof of Proposition 1 (Invariance to Nuisance)

• The formal derivation shows that generative objectives force representations to encode nuisance information, while VJEPA objectives do not.Source: H Proof of Proposition 1 (Invariance… · p. 66
• The proof models each observation as x_t = (s_t, n_t), where s_t predicts future target Z_t+∆ and n_t satisfies I(n_t; Z_t+∆|s_t) = 0.Z_t denotes the representation of x_t.Source: H Proof of Proposition 1 (Invariance… · p. 66

H.1 Case 1: Generative / Reconstruction Objective

• Generative reconstruction maximizes log p(s_t, n_t | Z_t) = log p(s_t | Z_t) + log p(n_t | s_t, Z_t), requiring representation of both components.The signal term requires encoding s_t, while the nuisance term requires encoding n_t to minimize H(n_t | Z_t).Source: §H.1 · pp. 66–67
• A generative solution that discards n_t is suboptimal because it cannot reconstruct the nuisance, forcing a maximal representation that retains N.Source: §H.1 · p. 67

H.2 Case 2: VJEPA Prediction Objective

• VJEPA instead maximizes likelihood of the future target representation Z_t+∆ given the current representation Z_t.Source: §H.2 · p. 67
• Because n_t is conditionally irrelevant given s_t, the maximal and minimal representations—respectively retaining or discarding n_t—achieve the same predictive objective value.The minimal representation therefore achieves the exact same optimal loss as the maximal representation, since VJEPA does not penalize absence of n_t.Source: §H.2 · p. 67
• VJEPA consequently admits nuisance-invariant minimal solutions that discard nuisance variability without sacrificing optimal predictive performance.Source: §H.2 · p. 67

I Derivation of the BJEPA Predictive Factorization

• BJEPA factorizes the posterior p(ZT | ZC, η) into history- and auxiliary-conditioned predictive experts after absorbing observed-input terms into a normalization constant.The derivation assumes ZC and η are conditionally independent given ZT, and approximates the marginal target prior as uniform or absorbed into expert biases.Source: I Derivation of the BJEPA Predictive… · pp. 67–68

J.1 Environment and Data Generation

• The Noisy TV experiment uses a Linear-Gaussian environment with 20-dimensional observations, 4-dimensional signals and distractors, 6000 training steps, 2000 testing steps, and nine noise scales from 0.0 to 8.0.Observations combine the signal, scaled distractors, and sensor noise.Source: §J.1 · pp. 68–69

J Model Architectures and Training Strategy

• All five models use bias-free linear maps from observations to four-dimensional latents, comparing pixel-space VAE and AR reconstruction with latent-space JEPA, VJEPA, and BJEPA prediction.The latent dimension equals the true signal dimension, enforcing a capacity trade-off between signal and distractor encoding.Source: §J.2 · p. 69

J.2.1 Generative Baselines (Pixel Reconstruction)

• Pixel-space baselines must reconstruct high-variance distractors, with VAE modeling each observation and AR predicting the next observation through a bottleneck.The stated hypotheses are that VAE prioritizes high-variance distractors and AR tracks predictable distractor texture, potentially neglecting the signal.Source: §J.2.1 · p. 69

J.2.2 Joint-Embedding Architectures (Latent Autoregression)

• JEPA predicts future latents, VJEPA parameterizes a probabilistic predictive distribution, and BJEPA fuses its dynamics expert with a learnable static prior through a Product of Experts.VJEPA and BJEPA sample only from the target encoder and analytically evaluate the predictor’s Gaussian negative log-likelihood; target sampling prevents predictor variance collapse, while the diagonal-Gaussian assumption cannot represent multimodal futures.Source: §J.2.2 · pp. 70–71

J.2.3 Temporal Alignment in Training and Evaluation

• Training uses current-next observation pairs for predictive models, whose latent targets are zt+1, while VAE uses single-observation reconstruction; evaluation correspondingly uses st for VAE and st+1 for predictive models.BJEPA evaluates its hard-fused posterior mean as the constrained next-state prediction.Source: §J.2.3 · pp. 71–72

J.2.4 Linear Probe Evaluation Protocol

• The linear-probe protocol fits ordinary-least-squares regressions from each model’s architecture-specific latent or predictor output to the temporally aligned ground-truth signal and scores held-out R2.An R2 near 1 indicates signal capture, whereas an R2 near 0 indicates failure to disentangle signal from distractor.Source: §J.2.4 · pp. 72–73

J.3 Optimization and Hyperparameters

• Optimization uses Adam with full-batch training; BJEPA’s structural prior is static but learnable, whereas its target-regularization reference prior is the fixed unit Gaussian N(0, I).The optimization and loss hyperparameters are summarized in Table 6.Source: §J.3 · pp. 73–74, Table 6 · p. 74

J.4 Results

• 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 results indicate that likelihood-based objectives encode high-variance distractors, causing representation collapse on the target task.Source: §J.4 · p. 74, Table 7 · p. 76
• Training losses are not compared across methods because generative models optimize pixel reconstruction, whereas JEPA-based models optimize differently scaled latent-space objectives.Raw loss magnitudes therefore do not define a common metric for cross-method comparison.Source: §J.4 · p. 74
• At Noise Scale 3.0, JEPA temporarily drops to 0.841 Test while VJEPA and BJEPA remain above 0.90, indicating greater probabilistic training stability.The deterministic predictor later recovers at Scale 4.0, but the instability is attributed to optimization difficulty under jittery stochastic targets.Source: §J.4 · p. 75
• Across auxiliary random seeds, deterministic JEPA occasionally collapses to R2 < 0.5, whereas VJEPA and BJEPA consistently remain more stable through aleatoric-uncertainty modeling.The variance outputs of the probabilistic methods help account for stochastic targets during optimization.Source: §J.4 · p. 75
• At σ = 8.0, JEPA, VJEPA, and BJEPA maintain R2 > 0.85 signal recovery, with JEPA reaching R2 = 0.93 in the reported run.Table 7 summarizes performance across all noise scales, while the reported run identifies JEPA as the highest-accuracy model despite greater training variance.Source: §J.4 · p. 75, Table 7 · p. 76

J.5 Computational Resources

• All experiments were conducted using Google Colab’s standard CPU runtime, with an Intel(R) Xeon(R) CPU at 2.20GHz, 2 logical CPUs, 13.61 GB RAM, and 242.49 GB disk.Source: §J.5 · p. 77
• For the latent-dimension Dz = 4 linear toy experiment, training was CPU-only and individual model runs completed in approximately 10–35 seconds.The reported runtime is associated with the lightweight experimental setup.Source: §J.5 · p. 77