Source-linked AI summary

MDN-VJEPA: Multimodal Probabilistic Joint-Embedding World Models

Yongchao Huang

arXiv:2610.13956v1cs.LG

TL;DR

When a context admits branching futures, point predictions or unimodal beliefs may fail to preserve the distinct alternatives needed for prediction and planning. MDN-VJEPA models those alternatives with a Gaussian-mixture JEPA transition, coupling modes across trajectories and encouraging consistency across temporal compositions. Controlled experiments report faithful branching-law recovery and improvements in planning, trajectory coherence, and multi-horizon agreement.

  • Problem

    A point prediction can collapse branching futures to their conditional mean, while a Gaussian predictor remains unimodal and cannot represent separated alternatives.

    Source: §1 · p. 4, §6 · p. 23
  • Method

    MDN-VJEPA uses a multimodal probabilistic transition, trajectory-level modes, and distributional consistency between direct and recursively composed predictions.

    Source: §3.2 · p. 10, §9 · p. 41
  • Results

    The controlled experiments report faithful recovery of branching laws, better navigation decisions with the full mixture than its conditional mean, more coherent trajectories, and reduced multi-horizon disagreement.

    Source: §7.8 · p. 38
  • Takeaways & Limitations

    The results support preserving which futures make up predictive uncertainty, how they form trajectories, and how beliefs compose through time.

    Source: §9 · p. 41

    The controlled tests do not demonstrate the complete PIB–MDN-VJEPA system with a learned high-dimensional encoder on large-scale natural video or real robotic interaction.

    Source: §8.6 · p. 40

Abstract

from arXiv · show

Joint-Embedding Predictive Architectures (JEPAs) learn world models by predicting future representations rather than reconstructing observations. Variational JEPA (VJEPA) gives this prediction probabilistic semantics and permits flexible conditional predictive families. We propose MDN-VJEPA, a multimodal probabilistic JEPA that instantiates the latent predictor as a mixture density network (MDN), representing future states with a Gaussian mixture whose components encode alternative futures, their probabilities, and within-mode uncertainty. For sequential prediction, we address two additional structural requirements: multi-modal futures should remain temporally coherent across a trajectory, and predictive beliefs should remain consistent when propagated to the same horizon through different temporal path compositions, as motivated by the Chapman--Kolmogorov equation. Together, these considerations define a three-level hierarchy of multimodal prediction, coherent multimodal trajectories, and compositionally consistent multimodal dynamics. Controlled experiments show that MDN-VJEPA closely recovers known branching predictive laws, that retaining the complete multimodal belief rather than its conditional mean can lead to materially different and substantially better planning decisions under branching futures, and that trajectory-level coupling and distributional consistency improve temporal coherence and multi-horizon agreement. These results support a simple principle: a probabilistic JEPA should preserve not only how uncertain the future is, but which distinct futures constitute that uncertainty, how they form coherent trajectories, and how those beliefs compose through time.

1 Introduction

Point prediction can average valid branching futures into an unrealizable mean, while a single Gaussian remains unable to represent distinct modes. MDN-VJEPA uses Gaussian mixtures and adds trajectory-level coherence and distributional consistency across horizons.

1 Introduction

  • When histories admit multiple futures, a point predictor may average them into an unrealizable route, while a single Gaussian remains unimodal.A Gaussian mixture can retain separate alternatives and quantify their probabilities and within-mode uncertainty.Source: §1 · p. 4, Figure 1 · p. 5
  • MDN-VJEPA instantiates the probabilistic JEPA predictor as a neural Gaussian mixture, with component means, covariances, and weights encoding future locations, local uncertainty, and predictive probabilities.Its mixture head preserves distinct futures rather than averaging them together.Source: §1 · pp. 4, 6
  • Retaining the complete predictive mixture matters for planning because evaluating cost at the mean can differ from averaging cost over futures, potentially obscuring a rare failure branch.The contribution is not the Gaussian-mixture head alone, but coherent, uncertainty-aware, compositionally consistent JEPA dynamics.Source: §1 · pp. 4, 6
  • A single categorical mode governs the future block, so each mixture component represents a candidate trajectory rather than independently chosen frame-level alternatives.Component labels may permute, split, or change meaning across contexts; semantic trajectory identity is not guaranteed.Source: §1 · p. 5
  • Direct multi-step predictions are encouraged to agree with recursively composed predictions at the same horizon, comparing full distributions under a Chapman–Kolmogorov-inspired principle.This adds compositional consistency beyond temporal coherence alone.Source: §1 · p. 5
  • The framework combines PIB’s compact predictive state, JEPA’s retention of predictive information, and the MDN’s capacity to preserve alternative futures.Temporal coherence and distributional path consistency structure how multimodal beliefs evolve.Source: §1 · p. 6

2 Related Work

Related work situates MDN-VJEPA within probabilistic JEPA, multimodal latent dynamics, and latent planning. The paper distinguishes its continuous mixture transition from prior mixture and trajectory approaches, and frames direct-versus-composed distribution agreement as distinct from ordinary multi-step supervision.

2.1 JEPA, VJEPA, and predictive information

  • MDN-VJEPA extends probabilistic JEPA from choosing a predictive family to coupling multimodal beliefs across time and composing them across horizons (Huang, 2026d; Huang, 2026a).Related work also distinguishes its sequential transition law from joint context–target mixture models (Huang, 2026b).Source: §2.1 · p. 7

2.2 Mixture density networks and multimodal latent world models

  • Although MDNs and MDN-RNNs predate this work (Bishop, 1994; Ha & Schmidhuber, 2018), MDN-VJEPA applies mixtures to JEPA transitions with trajectory and horizon constraints.Unlike finite atomic branches, Gaussian components also parameterize uncertainty within each branch (Song et al., 2026).Source: §2.2 · p. 7

2.3 Latent dynamics and planning

  • Unlike reconstruction-oriented latent dynamics, JEPA predicts in representation space; MDN-VJEPA compares direct and recursively composed distributions rather than merely supervising longer rollouts.Retaining alternatives supports expected-cost, chance-constrained, or tail-sensitive planning.Source: §2.3 · p. 8

3 From Probabilistic JEPA to MDN-VJEPA

MDN-VJEPA uses a Gaussian-mixture predictor within a time-indexed probabilistic JEPA, preserving alternative future hypotheses and their probabilities and within-component uncertainty. The section also describes scalable covariance choices, context-dependent mixture utilization, and how the model relates to Gaussian and deterministic JEPA predictors.

3.1 Time-indexed probabilistic JEPA

  • The time-indexed formulation encodes observed history as a predictive state and predicts future target representations conditioned on that state and side information.An EMA target encoder supplies stop-gradient target representations during predictive training.Source: §3.1 · pp. 8–9, Figure 3 · p. 9

3.2 The mixture-density transition

  • MDN-VJEPA’s Gaussian-mixture transition defines a normalized conditional law whose component means, covariances, and weights represent candidate futures, within-mode uncertainty, and predictive support.As a probability kernel, it can be sampled, integrated, recursively composed, and evaluated with probabilistic scoring and calibration criteria.Source: §3.2 · p. 10

3.3 Covariance parameterization

  • Diagonal covariance is the scalable default, while low-rank-plus-diagonal covariance is an option when within-component uncertainty has cross-dimensional correlations.The low-rank extension has O(dr) parameter and storage complexity, rather than O(d^2) for dense covariance.Source: §3.3 · p. 11

3.4 Effective mixture complexity

  • The fixed maximum K allows context-dependent effective mixture complexity, but the effective component count measures mixture utilization rather than the number of semantically distinct futures.Its value is approximately 1 when mass concentrates on one component and equals K when weights are uniform.Source: §3.4 · pp. 11–12

3.5 Nested model hierarchy

  • With K = 1 the model recovers a Gaussian VJEPA predictor; with a deterministic target and fixed isotropic covariance, its objective is equivalent to squared-error prediction.MDN-VJEPA is a finite Gaussian-mixture choice within general VJEPA; the deterministic relation is objective-level, not literal architectural inclusion.Source: §3.5 · p. 12

4 Coherent Multimodal Latent Dynamics

The section develops a hierarchy for sequential multimodal prediction: horizon-wise mixtures can yield incoherent trajectories, shared trajectory modes couple futures across horizons, and distributional consistency relates direct predictions to recursively composed dynamics. It also distinguishes local transition modes from trajectory modes and separates within-mode from between-mode predictive uncertainty.

4.1 Level I: Why horizon-wise mixtures are insufficient

  • Accurate multimodal marginals do not determine a coherent trajectory: independently chosen horizon modes can splice incompatible futures, such as alternating left and right turns.The missing structure is a coupling of alternatives across horizons.Source: §4.1 · p. 13

4.2 Level II: Trajectory-level modes

  • Level II replaces a product of horizon-wise mixtures with a mixture of trajectory products, using one shared mode to couple predictions across the directly predicted future block.Each horizon remains multimodal; conditional independence within a mode is a simple realization, with joint covariance available for stronger dependencies.Source: §4.2 · pp. 13–14

4.3 Level III: Distributional consistency across prediction horizons

  • Level III regularizes direct horizon-h predictions toward recursively composed one-step predictive laws using a distributional discrepancy that is invariant to component-label permutations.Under the stated true-kernel assumptions, the direct and composed distributions coincide; agreement alone does not ensure predictive correctness.Source: §4.3 · pp. 15–16

4.4 Latent transition-mode interpretation

  • The one-step transition uses a local categorical mode to select a mode-specific continuous transition, but recursive rollout does not require that label to persist across time.The paper distinguishes this local indicator from the trajectory-level mode and does not posit a persistent discrete Markov chain.Source: §4.4 · p. 16

4.5 Decomposing predictive uncertainty

  • Predictive covariance separates expected uncertainty within a future mode from uncertainty due to separation among competing future hypotheses.This learned-model decomposition is not by itself an aleatoric–epistemic decomposition; parameter or representation uncertainty requires additional treatment.Source: §4.5 · p. 17

5 A Compositional Training Objective

This section presents a three-level training objective that combines level-specific predictive likelihoods with representation and mixture safeguards, progressing from horizon-wise multimodal prediction to coherent trajectories and consistency across predictive paths. At Level III, it adds supervised recursive dynamics and consistency between direct and composed predictive laws.

5.1 Level-wise predictive likelihoods

  • Level I fits horizon-wise mixture likelihoods, whereas Levels II–III use a shared mixture component to explain the full trajectory; Level III also supervises its recursive one-step kernel.The Level-I likelihood does not require the same component to explain every horizon, unlike the shared-mode trajectory likelihood.Source: §5.1 · pp. 18–19

5.2 Predictive-information regularization

  • Target-side KL regularizes the stochastic future representation through a trainable target-view path, while context-side KL compresses the predictive state to remove unnecessary history information while retaining what predicts the future.The EMA target encoder supplies stop-gradient predictive targets and receives changes through EMA updates rather than direct gradient descent.Source: §5.2 · p. 20

5.3 Preventing degenerate solutions

  • Distinct safeguards address representation collapse with predictive and representation supervision, dead mixture components with a weak responsibility-based penalty, and Gaussian variance collapse with a variance floor.The component penalty discourages persistently negligible responsibility without forcing uniform mixture weights; the variance floor is the default protection against pathological likelihood spikes.Source: §5.3 · p. 21

5.4 Level-III consistency during optimization

  • Level III's CK loss aligns direct and recursively composed laws by updating recursive dynamics toward a stop-gradient direct reference, alongside direct trajectory and one-step supervision.Agreement alone cannot establish predictive correctness, so the consistency loss is used with direct trajectory and one-step likelihood supervision.Source: §5.4 · p. 22

5.5 Overall level-wise objective

  • The level-aware objective combines a level-specific direct likelihood with representation regularization and a component-use safeguard, adding supervised one-step dynamics and CK consistency for Level III.A one-step transition may also be trained for Level II when recursive rollout is needed; the CK term is activated only for Level III.Source: §5.5 · pp. 22–23
  • Level I learns multimodal horizon-wise predictions, Level II couples them into coherent trajectory-level alternatives, and Level III encourages consistency between directly predicted and recursively composed predictive laws.This level-wise organization mirrors the model construction described in Section 4.Source: §5.5 · p. 23

6 Decision-Theoretic Consequences of Multimodal Prediction

Retaining a multimodal predictive belief can change planning decisions because its mean may not represent a realizable future, and its branches carry distinct probabilities and consequences. Planners can query this belief at different levels of computational fidelity.

6 Decision-Theoretic Consequences of Multimodal Prediction

  • Replacing a multimodal predictive distribution with its conditional mean can change planning decisions, especially when the mean lies between realizable futures.The complete distribution separates uncertainty over which future occurs from variability within each future.Source: §6 · p. 23, §6.1 · p. 24

6.1 From mean-only to distributional planning

  • The learned mixture exposes each branch's probability and mode-conditioned consequence to the planner.Source: §6.1 · pp. 23–24
  • A low-probability branch can contribute explicitly to predicted failure probability even when it has little influence on the conditional mean.Source: §6.1 · p. 24
  • For nonlinear costs or constraints, evaluating the cost at the mean prediction generally differs from averaging cost over possible futures.Candidate action sequences with similar mean predictions can therefore have substantially different expected costs or failure probabilities.Source: §6.1 · p. 24

6.2 Branch-wise approximations and sampling

  • Branch-wise component trajectories retain major alternatives and their predictive probabilities while approximating each mode by its conditional mean.This can be useful when branches matter to decisions but extensive stochastic rollouts are unnecessarily expensive.Source: §6.2 · p. 25
  • A planner can trade computational cost against fidelity by using the mean, component trajectories, or full hierarchical samples.Component trajectories preserve branch structure but ignore within-mode variation; full hierarchical sampling retains both between- and within-mode uncertainty.Source: §6.2 · p. 25

7 Experiments

The experiments evaluate MDN-VJEPA’s multimodal predictions across planning, trajectory coherence, and temporal-composition consistency. They assess the effects of preserving full predictive beliefs, coupling trajectories into coherent alternatives, and regularizing direct and composed predictions, including predictive-metric tradeoffs.

7.3 Baselines and ablations

  • The controlled Branch-JEPA-style reference is a finite-support trajectory predictor, not a reproduction of the complete external training stack; experiments selectively isolate predictive-family and structural effects.The comparisons include an atomic weighted support over trajectory hypotheses, which lacks continuous within-branch uncertainty.Source: §7.1 · p. 26, §7.3 · p. 28, §7.5 · p. 31

7.1 Common experimental protocol and evaluation criteria

  • Evaluation uses independent data splits and matched sampling budgets where specified, and multimodal branches are never selected using ground-truth futures at deployment.Predictive evaluation includes density or sample-based scores, calibration and branch coverage where available, and direct–composed distribution discrepancies for compositional consistency.Source: §7.1 · p. 26

7.2 Three-level evaluation hierarchy

  • The three-level hierarchy tests multimodal marginals, shared-mode trajectory coherence, and direct-versus-recursively-composed distribution consistency as distinct predictive requirements.A shared latent mode specifies a complete candidate trajectory; the final level evaluates both compositional agreement and fidelity to observed targets.Source: §7.2 · p. 27

7.4 Experiment 1: controlled branching dynamics

  • MDN-VJEPA reached energy distance 0.0022 ± 0.0007, approximately 87× smaller than the parameter-matched Gaussian’s 0.1913 ± 0.0024.Its energy score, 0.5554 ± 0.0087, was nearly the oracle’s 0.5541 ± 0.0078. The advantage was specific to separated futures: at δ ∈ {0.05, 0.10}, the Gaussian was as accurate as or marginally better than the mixture.Source: Table 3 · p. 29, §7.4 · p. 29, Figure 7 · p. 30
  • The learned mixture recovered branch probabilities with MAE 0.0115 ± 0.0025 and calibrated probabilities with expected calibration error 0.0271 ± 0.0058.Relative errors for within-mode and between-mode covariance were 5.71% ± 4.12% and 3.69% ± 1.04%, respectively; with λuse = 0, the effective component count was 1.787 ± 0.008, with no evidence of collapse in the balanced base regime.Source: §7.4 · p. 30

7.5 Experiment 2: stochastic robot navigation

  • MDN-VJEPA without CK achieved energy distance 0.0026 ± 0.0006, approximately 8× smaller than the atomic predictor’s 0.0212 ± 0.0004.Its energy score was 0.3049 ± 0.0023, close to the oracle’s 0.3037 ± 0.0016. The atomic predictor and both MDN variants achieved full branch coverage; without CK, branch-probability MAE was 0.0162 ± 0.0048 and branch ECE was 0.0241 ± 0.0102.Source: Table 4 · p. 31, §7.5 · p. 31
  • Using the same MDN-VJEPA+CK model, full-distribution planning reduced collision rate from 0.5563 ± 0.0095 to 0.0025 ± 0.0003 and raised success rate from 0.4437 ± 0.0095 to 0.9975 ± 0.0003.The true expected objective fell from 5.6473 ± 0.0820 to 1.1828 ± 0.0018, close to the oracle value 1.1769 ± 0.0017; the planners differed only in whether they used the conditional mean or complete mixture.Source: §7.5 · pp. 31–32, Table 5 · p. 32
  • CK reduced direct–composed energy distance from 0.1585 ± 0.0097 to 0.1236 ± 0.0119, approximately 22%, across all five seeds.The composed prediction’s energy score improved from 0.3806 ± 0.0119 to 0.3741 ± 0.0136, while direct predictive energy score changed from 0.3052 ± 0.0055 to 0.3076 ± 0.0089, trading slight direct accuracy for better compositional agreement.Source: §7.5 · pp. 32–33
  • In stochastic navigation, CK preserved multimodality with a minimum effective component count of 1.849, but direct energy distance worsened from 0.0026 ± 0.0006 to 0.0121 ± 0.0113 and branch-probability MAE from 0.0162 ± 0.0048 to 0.0535 ± 0.0281.The authors interpret CK as a structural consistency regularizer rather than an unconditional improvement to one-step or direct predictive metrics.Source: §7.5 · p. 33

7.6 Experiment 3: stochastic video prediction in representation space

  • Level II preserved perfect marginal branch coverage while reducing semantic switch rate from 0.6379 ± 0.0023 to zero and raising coherent branch coverage from 0.0759 ± 0.0069 to 1.0000 ± 0.0000.Both levels had marginal branch coverage 1.0 and similarly small branch-probability errors, indicating that coherence improved without Level I simply failing to learn the branches.Source: Table 6 · p. 34, Figure 10 · p. 34, §7.6 · pp. 34–35
  • Level II reduced trajectory energy distance from 0.6246 ± 0.0075 to 0.0151 ± 0.0076 and best-of-64 trajectory RMSE from 0.2063 ± 0.0028 to 0.0430 ± 0.0003, near the 0.0412 ± 0.0001 oracle.The trajectory-level model produced three coherent future families, while Level I frequently spliced branches despite recovering horizon-wise support.Source: Table 7 · p. 35, Figure 11 · p. 35, §7.6 · p. 35

7.7 Experiment 4: direct vs. composed multi-horizon prediction

  • Adding CK reduced normalized direct–composed energy-distance AUC from 0.02040 ± 0.00227 to 0.01716 ± 0.00162, approximately 15.9%.The reduction occurred across the evaluated multi-horizon range, although the models were effectively tied at h = 16.Source: §7.7 · pp. 36–37, Table 8 · p. 37, Figure 12 · p. 37
  • In Experiment 4, CK largely retained fidelity but had a small composed-path tradeoff: terminal energy distance rose from 0.00407 ± 0.00136 to 0.00644 ± 0.00157, while effective component counts remained similar.At h = 16, direct and composed Energy Scores remained close to the oracle; the authors characterize CK as a structural consistency regularizer, not a universal predictive-metric improvement.Source: §7.7 · pp. 37–38

7.8 Summary of empirical findings

  • Across Experiments 1–2, MDN-VJEPA’s energy distance was 0.0022 ± 0.0007 versus Gaussian 0.1913 ± 0.0024, and full-mixture planning cut collisions from 0.5563 ± 0.0095 to 0.0025 ± 0.0003.The summary frames these findings alongside coherent trajectory prediction and reduced multi-horizon disagreement as evidence for preserving distinct futures and their composition through time.Source: §7.8 · p. 38

8 Discussion

The discussion frames MDN-VJEPA’s contribution as structuring continuous multimodal beliefs across futures and time, not simply replacing a Gaussian with a GMM. It distinguishes branch uncertainty, trajectory coherence, and compositional consistency while noting limits in the experiments and model assumptions.

8.1 Our contribution is not “GMM instead of Gaussian”

  • The predictive contribution combines continuous mixtures, shared trajectory modes, and distributional path consistency, rather than claiming GMM heads alone as novel [Huang, 2026d; Song et al., 2026].The discussion also notes that controlled experiments isolate predictive structure rather than empirically ablating PIB representation learning.Source: §8.1 · pp. 38–39

8.2 Atomic branches vs. Gaussian components

  • Both atomic branches and Gaussian components represent weighted alternatives, but Gaussian components additionally model within-branch variation; neither representation is uniformly preferable.In the controlled benchmark, atomic and MDN predictors recover both route families, while MDN has lower energy distance and collision probability; this does not establish general dominance over Branch-JEPA.Source: §8.2 · p. 39

8.3 Mode coherence is not global mode identity

  • A shared trajectory-level mode eliminates semantic switching while preserving horizon-wise branch coverage in the controlled generator, but does not guarantee global component identity.Components may permute, split, or represent the same semantic future; effective component count is not necessarily the number of semantic futures.Source: §8.3 · pp. 39–40

8.4 Consistency is not predictive accuracy

  • CK reduces direct–composed discrepancy, but consistency alone does not ensure predictive accuracy and can trade off against fidelity.In Experiment 4, discrepancy falls across all five seeds while direct likelihood and Energy Score remain essentially unchanged; recursive rollout has a small terminal fidelity tradeoff.Source: §8.4 · p. 40

8.5 Long-horizon mixture propagation

  • Nonlinear MDN composition may leave the finite-GMM family, motivating sample-based distribution comparison without analytic closure or component alignment.Propagating and comparing ancestral rollout populations adds computational cost at long horizons; projecting them back to a compact mixture is left outside the core formulation.Source: §8.5 · p. 40

8.6 Limitations and scope

  • The controlled studies do not demonstrate the complete PIB–MDN-VJEPA system with a learned high-dimensional encoder on natural video or real robotic interaction.Other stated boundaries include fixed component count, possible failure of Gaussian components on strongly non-elliptical branches, conditional horizon independence under Level II, and CK’s cost and fidelity tradeoffs; the model is not a full switching state-space model.Source: §8.6 · pp. 40–41

9 Conclusion

MDN-VJEPA preserves distinct weighted future hypotheses and develops their prediction through multimodality, trajectory coherence, and compositional consistency. The appendix clarifies the notation and uncertainty distinctions used to describe that hierarchy.

9 Conclusion

  • MDN-VJEPA preserves distinct future latent states with probabilities and within-mode uncertainty, and controlled experiments support its branching-law, planning, coherence, and compositional-consistency results.The paper’s conclusion is that probabilistic world models should preserve which futures constitute uncertainty, how they form trajectories, and how beliefs compose through time.Source: §9 · p. 41

A Symbols, Objectives, and Notation

A.2 Encoder and predictor notation

  • Table 10 summarizes encoder, predictor, and distribution notation.The encoder and its EMA copy are roles of a shared trainable encoder, not separately parameterized encoders.Source: §A.2 · p. 43, Table 10 · p. 44

A.3 Mixture and latent-mode notation

  • The trajectory-level mode M is shared across a directly predicted future block, whereas M_t+1 is a local recursive-transition choice; labels need not persist across rollout.The model does not introduce a separate discrete transition law p(M_t+1 | M_t).Source: §A.3 · p. 45, Table 11 · p. 45

A.4 Three predictive levels

  • The hierarchy progresses from horizon-wise multimodal marginals to a shared-mode trajectory law, then to compositional consistency between direct and recursive predictions.Level II replaces a product of mixtures with a mixture of trajectory products.Source: §A.4 · pp. 45–46

A.5 Predictive uncertainty

  • Predictive covariance separates within-mode uncertainty from between-mode branching uncertainty, but this is not an aleatoric–epistemic decomposition.The appendix also identifies the Level-II mean and conditional covariance as summaries of the predictive distribution.Source: §A.5 · p. 46

A.6 Training objectives

The appendix organizes training objectives, direct–composed consistency notation, and planning notation. It distinguishes predictive-hierarchy losses from representation terms and describes sample-based comparison for recursively composed distributions.

A.6 Training objectives

  • The training objective combines predictive and auxiliary terms, while representation terms and the utilization safeguard are orthogonal to the Level-I–III predictive hierarchy.The appendix presents the full objective and summarizes objective terms and coefficients in Tables 12 and 13.Source: §A.6 · pp. 46–47, Table 12 · p. 46, Table 13 · p. 47

A.7 Direct–composed consistency notation

  • Because nonlinear MDN composition is generally not a finite closed-form K-component GMM, the consistency comparison uses ancestral samples and permutation-invariant energy distance.In Experiment 4, direct samples are a stop-gradient reference while recursive samples remain differentiable; the direct predictor retains its ordinary predictive-likelihood gradients.Source: §A.7 · p. 47

A.8 Planning notation

  • Table 14 summarizes notation for distribution-aware planning, and the branch-separation parameter is locally defined for Experiment 1.Source: Table 14 · p. 47, §A.8 · p. 48

A.9 Evaluation terminology

The evaluation terminology distinguishes scoring a predictive distribution against an observed outcome from comparing two complete predictive distributions.

A.9 Evaluation terminology

  • Energy Score compares a predictive distribution with an observed outcome, whereas energy distance compares two complete distributions.Table 15 is identified as summarizing the principal evaluation quantities.Source: §A.9 · p. 48, Table 15 · p. 48

B Marginalization of the trajectory-level mixture

Sharing one mode across a future trajectory changes how horizon-wise predictions are coupled without removing their Gaussian-mixture form. Thus, matching horizon-wise marginals does not imply matching joint trajectory distributions.

B Marginalization of the trajectory-level mixture

  • Each horizon-wise marginal retains the K-term Gaussian-mixture form when a trajectory-level mode is shared across the future block.The coupling changes the joint distribution across time, not the individual horizon’s capacity to represent separated modes.Source: B Marginalization of the… · p. 49
  • Levels I and II can have identical marginals at every horizon yet different joint trajectories because only Level II shares one mode across horizons.Level I can select mixture identities independently at each horizon; Level II uses one shared trajectory-level mode.Source: B Marginalization of the… · p. 50
  • The horizon-wise mixture result also holds when each mode defines a joint Gaussian over the concatenated trajectory, not only under conditional factorization.Marginalizing the joint Gaussian at a horizon and then the shared mode again yields a Gaussian mixture.Source: B Marginalization of the… · p. 50

C Additional Experimental Details and Results

The experiments ran in Google Colab on a virtualized Intel Xeon host, with scripts using CUDA when available and otherwise falling back to CPU.

C Additional Experimental Details and Results

C.1 Experiment 1: controlled branching dynamics

In a controlled latent-space branching test, MDN-VJEPA closely recovers the known predictive law, substantially outperforming Gaussian predictors in the base condition. The experiments also examine mixture recovery, identifiability, robustness, and calibration.

C.1 Experiment 1: controlled branching dynamics

  • The controlled latent-space protocol generates train, validation, and test contexts independently and uses known branch prototypes only for post-prediction semantic matching.The predictor is an MLP; deterministic prediction uses squared error, while Gaussian VJEPA and MDN-VJEPA use negative log-likelihood with diagonal covariance and a variance floor.Source: §C.1 · p. 51, Table 16 · p. 51
  • 0.0022 ± 0.0007: MDN-VJEPA’s energy distance, versus 0.1922 ± 0.0034 for the standard Gaussian and 0.1913 ± 0.0024 for the parameter-matched Gaussian.Across 5 paired seeds, the matched-Gaussian-minus-MDN energy-distance gap was 0.1891 ± 0.0024; the corresponding NLL gap was 1.6192 ± 0.0090.Source: §C.1 · p. 51, Table 17 · p. 51
  • Mixture-recovery diagnostics show small probability and calibration errors, covariance-decomposition terms recovered within a few percent on average, and effective component counts close to true complexity.Near-balanced responsibilities arise in the base regime with zero global branch-logit bias, despite disabling the component-usage regularizer.Source: §C.1 · p. 51, Table 18 · p. 52
  • When branch separation is small, density recovery can remain accurate even though identifying individually meaningful components is weakly identifiable.This limitation applies at δ ∈ {0.05, 0.10}, where the components overlap strongly.Source: §C.1 · p. 52, Table 19 · p. 52
  • As within-mode noise rises, the matched Gaussian improves from 0.2205 ± 0.0049 at σz = 0.08 to 0.0144 ± 0.0015 at σz = 0.80, while MDN-VJEPA remains close to the true law.The Gaussian approximation is also most challenged near balanced branch probabilities and improves as one branch dominates.Source: §C.1 · p. 52, Table 20 · p. 52, Figure 13 · p. 53
  • The aggregate 5-seed branch-probability ECE is 0.0271 ± 0.0058; the representative reliability plot is diagnostic, with synthetic labels and probabilities excluded from training.Source: §C.1 · p. 53, Figure 14 · p. 53

C.2 Experiment 2: stochastic robot navigation

In the navigation experiment, full-distribution planning is more useful than planning from the same model’s mean, while CK training improves direct-versus-composed distributional agreement. The results also expose trade-offs in direct calibration and planning behavior.

C.2 Experiment 2: stochastic robot navigation

  • CK fine-tuning jointly optimizes direct trajectory and one-step likelihoods, adding a gradually increased consistency coefficient and hard straight-through Gumbel–Softmax samples for CK training.Both CK and no-CK models begin from the same warm-start and cloned parameters.Source: §C.2 · p. 55
  • Both MDN-VJEPA and the controlled atomic predictor identify both route families, with weighted support precision and branch coverage of 1.0 for every multimodal model.MDN-VJEPA also reproduces the continuous law around the branches, while matched-support trajectory errors are 0.0047 ± 0.0005 for MDN and 0.0047 ± 0.0006 for MDN+CK.Source: §C.2 · p. 56, Table 23 · p. 56, Figure 15 · p. 56
  • The MDN-mixture expected-cost and CVaR planners approach the oracle objective; the Gaussian achieves essentially zero collisions only with a substantially longer route.The atomic predictor retains short trajectories but underestimates some local collision uncertainty; every method has goal rate 1.0.Source: §C.2 · pp. 56–57, Table 24 · p. 57
  • Across collision penalties λcol ∈ {2, 4, 8, 16}, mean-only collision rate stays approximately 0.56 while full-distribution planning stays near 0.0025.The world model is not retrained for this sweep.Source: §C.2 · p. 57, Figure 16 · p. 57
  • CK training reduces direct–composed energy distance from 0.1585 ± 0.0097 to 0.1236 ± 0.0119, while recursively composed target Energy Score improves from 0.3806 ± 0.0119 to 0.3741 ± 0.0136.The reduction in direct–composed energy distance occurs in every seed; direct-target Energy Score changes only slightly.Source: §C.2 · p. 57, Table 25 · p. 57
  • The navigation results show a trade-off: unconstrained MDN gives the most accurate direct conditional law, while CK sacrifices a small amount of direct calibration for improved recursive consistency.Retaining the complete resulting predictive distribution remains more useful for planning than collapsing it to its mean.Source: §C.2 · p. 58

C.3 Experiment 3: stochastic video prediction in representation space

Experiment 3 isolates whether multimodal predictions represent coherent futures, not just accurate horizon-wise branches. Sharing one mixture choice across the trajectory recovers a joint predictive law closer to the oracle than independent horizon-wise choices.

  • Level I chooses mixture components independently across horizons, whereas Level II shares one mixture choice across the complete trajectory.The models use the same horizon-specific Gaussian component parameterization, isolating temporal coupling as the tested distinction.Source: §C.3 · p. 59
  • Both levels retain all three branches at each horizon, but Level I frequently switches semantic branches while Level II preserves coherent trajectories.Level I achieves full marginal branch coverage; Level II maintains that coverage while recovering coherent complete trajectories.Source: §C.3 · p. 61
  • Level-II component-mean sequences have zero semantic and any-switch rates across all five seeds, with learned means following each branch and narrow local uncertainty.Figure 18 shows one component describing one coherent branch together with within-branch uncertainty.Source: §C.3 · pp. 61–62, Figure 18 · p. 62
  • The nominal 90% sample interval covers about 0.87–0.88 for all probabilistic models, including the oracle, under the finite-sample protocol.Coverage uses empirical quantiles from 64 predictive samples, so the paper interprets it relative to the oracle rather than as analytic calibration.Source: §C.3 · p. 61
  • Level II substantially improves aggregate Energy Score over Level I and the unimodal Gaussian reference, and approaches the generator oracle.The comparison is visualized in Figure 17.Source: §C.3 · p. 61, Figure 17 · p. 61
  • Level II recovers a coherent joint predictive distribution close to the oracle trajectory law, while Level I remains close to a unimodal Gaussian.Level I models individual horizon-wise branch supports accurately, but its independent choices can combine into trajectories that are not realizable semantic futures.Source: Figure 17 · p. 61, §C.3 · p. 62

C.4 Experiment 4: direct versus composed multi-horizon prediction

Experiment 4 tests whether distributional consistency regularization makes direct and recursively composed predictions agree across horizons. It reduces their integrated disagreement across the tested range, with a small long-horizon fidelity tradeoff for the composed path.

C.4 Experiment 4: direct versus composed multi-horizon prediction

  • The two arms share the same architecture and checkpoint before fine-tuning; Level III adds the CK discrepancy without adding model parameters.Both receive the same minibatch order and optimization budget; Level II sets λCK = 0.Source: §C.4 · pp. 63–64
  • Level III reduces direct–composed energy-distance discrepancy through h = 12 and is comparable to Level II at h = 16; raw and normalized energy-distance AUC improve in all five seeds.At h = 16, the difference is negligible and slightly favors Level II; the main conclusion reflects the complete horizon range.Source: Figure 19 · p. 65, §C.4 · pp. 65, 68, 69
  • Fixed-bandwidth MMD2-AUC decreases from 0.00992 ± 0.00206 to 0.00866 ± 0.00111, but improves in only 2 of 5 seeds.The paper treats MMD as a less stable complementary diagnostic and retains energy distance as the primary evidence.Source: §C.4 · p. 65, Figure 20 · p. 66
  • The direct path retains essentially unchanged target fidelity, while the recursively composed path incurs a small long-horizon fidelity tradeoff.Both models' Energy Scores remain close to the simulator oracle, despite the composed-path energy-distance increase at h = 16.Source: §C.4 · pp. 66–67
  • The consistency gain does not collapse modes: both models retain multimodal endpoint structure, with similar event-probability and component-utilization diagnostics.The representative horizon-16 distributions show that CK changes agreement between computational paths rather than collapsing predictions to one mode.Source: §C.4 · pp. 67–68, Figure 24 · p. 68

Disclaimer

  • The author reports AI assistance in development and implementation, while taking responsibility for the final manuscript and its conclusions.The author estimates the contribution split as approximately 70%:30% between author and ChatGPT.Source: Disclaimer · p. 69
Loading 2610.13956v1…