Probabilistic JEPA Predictor Using Flow Matching

Idea in one sentence: replace JEPA’s point predictor with conditional flow matching over future latent trajectories, so the model can represent several distinct plausible future regimes instead of averaging them.

The Core Idea

A deterministic JEPA predictor produces one future latent trajectory. When several futures are plausible, a point predictor tends to select or average them.

The proposed model instead samples multiple future JEPA latent trajectories conditioned on:

• observed history;
• system context;
• candidate actions or control inputs;
• known future exogenous variables.

A decoder maps each sampled latent trajectory back to the original multivariate time-series space when exact values, constraints, or downstream costs are needed.

flowchart LR
  H["History + context + actions"]
  E["Frozen JEPA encoder"]
  F["Conditional flow matching"]
  Z["Multiple future latent trajectories"]
  D["Decoder"]
  X["Multiple future numeric trajectories"]

  H --> E
  E --> F
  H --> F
  F --> Z
  Z --> D
  H --> D
  D --> X

The goal is not merely to produce diverse-looking samples. The samples should represent the actual plausible future regimes and their probabilities.

What We Want From The Model

A useful probabilistic predictor should preserve:

Property	Meaning
Distinct future regimes	Recovery, persistent degradation, and failure remain separate modes rather than becoming one averaged trajectory.
Coherent trajectories	Time order, regime duration, cross-channel relationships, and system constraints remain valid.
Calibrated probabilities	A regime predicted with probability $p$ should occur roughly that often in comparable held-out cases.
Action sensitivity	Changing a real action or control input should move probability mass between the correct future regimes.
Rare but important outcomes	Failures, threshold crossings, extreme values, and long recovery tails should not disappear in average metrics.
Usable rollout cost	Sampling should be cheap enough for forecasting, risk estimation, or candidate-action evaluation.

Here multi-modal means several modes of a probability distribution, not several data modalities.

Why Flow Matching?

There is one primary reason:

Flow matching may let us model and sample from a complex, non-Gaussian distribution containing several separated future regimes.

This is the central hypothesis, not an established result for this setting. Flow matching does not guarantee that the model will discover all regimes, assign them calibrated probability mass, or avoid invalid trajectories between modes. The experiment must test exactly those properties.

Everything else—training behavior, sampling speed, and implementation convenience—is secondary to this reason.

What MSE-Optimal Forecasting Changes

Expectations vs. Realities provides direct numeric time-series evidence for why the idea is needed. If the same observed history permits several futures, the conditional mean is the unique MSE-optimal deterministic prediction, but its variance omits the irreducible conditional variance. Whenever that missing fraction is nonzero, the mean cannot also match the marginal distribution of realized futures.

The paper reports that, across nine forecasting benchmarks, keeping candidates within 5% of the best MSE still exposes a median 17.3% improvement in horizon-wise marginal realism. This turns the opening intuition into a concrete minimum evaluation gate: the proposed model should report an MSE–realism Pareto front and compare alternatives inside an application-dependent MSE band instead of claiming success from one point metric.

This result motivates a distributional predictor, not flow matching or JEPA specifically. It also uses marginal $W_1$ across contexts, which can look good while assigning a plausible future to the wrong context or breaking temporal order. The main experiment must therefore keep the stronger requirements already listed here: separated-mode coverage, calibrated mass, rare tails, multivariate constraints, action-sensitive probability shifts, and a deployable selector or decision test.

What FlowWM Changes

FlowWM is the closest current empirical precedent for this idea. It applies conditional flow matching to structured future trajectories in frozen DINOv3 feature space and tests whether samples cover enumerated visual futures.

It is not a complete implementation of the idea:

Dimension	FlowWM	This idea
Representation	Frozen DINOv3 features	Future latent trajectories from a deterministic JEPA-style encoder
Domain	Passive video	Multivariate time series and eventually action-conditioned world models
Conditioning	Past frames	History, context, candidate actions/control inputs, known exogenous variables
Grounding	Frozen detection/depth heads and a lightweight pixel decoder	Numeric decoder, oracle-decoder ceiling, and decoder-use probes
Main comparison	Deterministic and alternative visual latent models	Matched raw-trajectory flow, raw flow with JEPA features, and latent flow
Distribution test	Mode precision/recall plus best-of-$N$ perception	Mode coverage, calibrated mass, tails, invalid regions, and action-sensitive shifts
Decision test	None	Action ranking, planner regret, and closed-loop utility

FlowWM provides three useful positive updates:

1. The core mechanism is plausible. On Bouncing Shapes, its F1 mode error is $4.31$ versus $14.4$ for the prior stochastic feature-space baseline and $17.1$ for a matched deterministic predictor.
2. High-dimensional feature flow needs specific engineering. A wide velocity head, timestep shifting toward noisier regions, and full-rank temporal noise matter more than simply making the DiT backbone deeper.
3. One-step projection is a useful training interface. Endpoint-level temporal, constraint, decoder, or task losses can be applied without differentiating through the full ODE trajectory.

The main research claim remains open because FlowWM does not test calibrated mode probabilities, unequal or rare branches, action conditioning, raw-versus-latent flow under matched budgets, numeric time-series fidelity, or closed-loop decisions. Its best-of-$N$ Waymo metric also uses the realized future to choose among samples; that is proposal-coverage evidence, not a deployable selector or risk estimate.

What Beckmann Transport Models Change

Beckmann Transport Models offers a possible speed-oriented branch for this idea. Instead of integrating a time-dependent conditional velocity field, learn either a time-independent autonomous field or a direct transport map from noise to the target future-latent distribution. A one/few-pass map could materially reduce planner rollout cost when many candidate futures are needed.

BTM also adds a crucial distributional warning: reaching the right support or covering every visible mode does not guarantee the right probability mass. Its weighted-atom experiment should motivate unequal and rare future-regime tests, regional mass error, and basin leakage in addition to best-of-$N$ coverage.

This is not yet evidence for the proposed time-series model. The BTM paper does not test conditional future trajectories, action or control-input conditioning, calibrated regime probabilities, multivariate numeric constraints, or closed-loop planning. Its singular-support assumption can be made formal by zero-padding, but that does not prove that JEPA future latents occupy a meaningful manifold. The minimal experiment SHOULD therefore compare ordinary conditional flow with a conditional BTM map at matched training compute and at 1, 2, and 5 network evaluations, while keeping calibration, rare-mode coverage, invalid-between-mode mass, action sensitivity, and wall-clock latency as separate endpoints.

What IMLE And Explorative Modeling Change

Original IMLE introduced data-to-nearest-generated-sample training in 2018, and peer-reviewed conditional IMLE already used the exact per-condition hard best-of-many objective needed for a direct future-trajectory generator. Explorative Modeling is the modern scaling and cross-modal evidence source, not the origin of the estimator.

This IMLE/XM family provides the strongest simple baseline added after this idea was written. It samples $K$ raw or latent future candidates and backpropagates only through the candidate closest to the observed future. This directly tests the idea’s core failure mode: can candidate width stop a point predictor from averaging separated future regimes?

IMLE/XM does not remove the main reason to study flow matching. The hard minimum is primarily a coverage objective and does not determine how much probability mass each future regime receives. Flow matching remains the primary candidate when the target is a calibrated predictive belief rather than only a proposal set. The minimal experiment should therefore name the baseline historically as conditional IMLE/Forward XM, not replace flow with it.

Why Predict In JEPA Latent Space?

The future can also be modeled directly in the raw time-series space. The JEPA latent route is interesting because it may:

• focus generation on regimes and reusable system dynamics;
• remove high-dimensional nuisance variation;
• improve long-horizon coherence;
• reduce rollout cost when many candidate futures are needed;
• transfer better across systems or channel configurations.

But latent prediction can also discard exact values, spikes, threshold crossings, or action-sensitive details. Therefore JEPA latent flow must be compared against a direct raw-trajectory flow. It should not be preferred merely because it sounds more structured.

Normalization Must Not Delete A Future Mode

Noise or Signal? does not test JEPA, flow matching, probabilistic forecasting, or control, but it exposes an upstream failure mode for this experiment. Mean/std RevIN can let one extreme observation dominate a window and the inverse transform, while median/MAD normalization can intentionally suppress the same observation. For telemetry, that observation may be corruption, or it may be the action transient, threshold crossing, workload burst, or failure precursor that separates future regimes.

The first FM-JEPA experiment SHOULD therefore make normalization an explicit ablation axis rather than an invisible preprocessing constant:

• no instance normalization beyond fixed training-set scaling;
• causal mean/std RevIN;
• causal median/MAD normalization;
• stats-as-context, where raw units, window location/scale, missingness, and outlier flags remain available to the encoder and decoder.

The oracle JEPA decode must recover raw-unit spikes, thresholds, action effects, and rare-regime identity under each choice. Flow matching cannot reconstruct a branch cue that the input transform or frozen encoder has already erased, and calibrated latent samples are not enough when inverse normalization distorts the raw operational distribution.

Proposed Model

1. Deterministic JEPA Encoder

Train a standard deterministic JEPA-style encoder and freeze it. It provides:

• a representation of the observed history;
• target latent trajectories for real futures;
• optional intermediate features for the decoder.

2. Flow-Matching JEPA Predictor

Train conditional flow matching to transform noise into a structured future latent trajectory.

The target should be a sequence or multi-scale latent grid, not one pooled vector for the entire future. Local latent features can preserve spikes and short motifs; global features can represent regime identity and long-range dynamics.

FlowWM suggests three starting hypotheses rather than fixed requirements: make the velocity projection head at least as wide as the generated latent feature dimension, shift timestep sampling as latent dimension and horizon grow, and keep enough independent noise across time to represent branch-local uncertainty. Each choice still needs a matched ablation for multivariate time series.

3. Deterministic Decoder

Train a separate decoder that reconstructs the future numeric trajectory from:

• the sampled future latent trajectory;
• recent raw history;
• selected intermediate JEPA features;
• actions, context, and known exogenous variables.

The first decoder should be deterministic. Otherwise both the flow predictor and decoder can model the same uncertainty, making it unclear where the future modes are represented.

FlowWM’s one-step projection can be tested as an efficient way to apply temporal-difference, constraint, decoder, or downstream-cost losses to a predicted clean endpoint. The full sampled endpoint must still be checked separately: a useful one-step auxiliary loss does not guarantee that the numerically integrated trajectory has the same properties.

The Crucial Comparison

Model	Purpose
Raw trajectory flow	Tests whether flow matching can model the required future distribution directly, without a JEPA bottleneck.
Raw flow + JEPA features	Separates the benefit of JEPA conditioning from the benefit of generating in latent space.
Conditional BTM direct map over raw and JEPA-latent trajectories	Tests whether autonomous one/few-pass transport preserves calibrated future mass and constraints while reducing rollout latency.
Conditional IMLE / Forward XM over raw and JEPA-latent trajectories	Tests whether hard best-of-$K$ candidate matching is sufficient to preserve separated regimes without learning a continuous flow density.
Fast-LeWM-style deterministic action-prefix predictor	Tests whether probabilistic latent flow improves calibrated mode coverage, action ranking, and planner regret beyond a strong parallel multi-horizon point predictor rather than only beyond recursive one-step rollout.
Oracle JEPA decode	Decodes the real future JEPA latent trajectory and measures how much important information the latent representation already lost.
JEPA latent flow + decoder	Tests the proposed probabilistic JEPA predictor end to end.

The oracle decode should be run before training the latent flow. If real future latents cannot recover critical values, rare regimes, or action effects, generated latents will not fix the problem.

FlowWM does not include the matched raw-trajectory flow or raw-flow-plus-feature controls above. Its DINOv3-versus-VAE comparison supports semantic feature choice for perception, but it cannot establish that latent generation is better than direct raw generation for this target.

Fast-LeWM supplies the deterministic action-conditioned comparison that FlowWM lacks. Its dense action-prefix objective reduces serial candidate-rollout cost, but it predicts one latent future and does not model calibrated separated modes. A latent-flow claim should therefore beat both recursive one-step rollout and this stronger parallel point-prediction baseline under matched action prefixes, planner budget, horizon, and end-to-end latency.

Relation To VJEPA

VJEPA provides a useful probabilistic JEPA interface, but the demonstrated predictor uses an independent diagonal Gaussian.

That is a weak option for this task because one Gaussian is unimodal. It can represent uncertainty around one center, but it can still average separated futures or place probability between valid regimes.

VJEPA could use a richer flow-based predictor. In that case, flow matching—not the variational Gaussian formulation—is the mechanism expected to model the complex future distribution. VJEPA is therefore relevant context, but not the main focus or a required baseline.

What Must Be Proven

The idea succeeds only if the experiments show that:

1. flow matching preserves distinct future regimes rather than producing one broad cloud;
2. sampled regime frequencies are calibrated;
3. trajectories remain temporally and cross-channel coherent;
4. rare and decision-relevant futures are retained;
5. candidate actions move probability mass between the correct regimes;
6. the decoder actually uses the sampled future latent;
7. the JEPA latent route provides a measurable advantage over direct raw flow.

Best-of-$N$ oracle scores should be reported only as proposal-coverage diagnostics. Deployment credit requires calibrated sample frequencies or a selector that ranks futures without observing the realized future.

Useful decoder checks are simple:

• remove the future latent;
• shuffle future latents between examples;
• replace a recovery latent with a failure latent;
• decode real future latents and inspect what remains missing.

If the output barely changes, the decoder is ignoring the probabilistic JEPA prediction.

Minimal Experiment

Stage 1: Controlled Branching Dynamics

Build a synthetic multivariate system with:

• two or more known future regimes;
• known regime probabilities;
• action-dependent changes in those probabilities;
• invalid trajectories between modes;
• rare failure and recovery tails.

This gives a ground-truth test of whether flow matching learned the distribution rather than just sample diversity.

FlowWM’s Bouncing Shapes benchmark is a strong implementation template, but the first time-series benchmark should go beyond it by adding unequal branch probabilities, rare branches, typed actions or control inputs, explicit invalid between-mode regions, and multivariate constraints. Otherwise the experiment can test geometric mode coverage without testing calibration or decision relevance.

Stage 2: Action-Conditioned Simulator

Move to a simulator with explicit state, actions or control inputs, exogenous variables, and replayable futures. CityLearn and Grid2Op are possible existing candidates. The proposed Kubernetes OTEL Control Gym is the operations-specific target: evaluate whether sampled futures preserve distinct recovery/failure regimes, move probability mass correctly under typed interventions, and improve constrained Pareto planning rather than only passive prediction.

Stage 3: Real Multivariate Data

Use operational data only when action timing, target, parameters, status, and outcomes are explicit enough. Otherwise differences between sampled futures should be described as conditional predictions, not counterfactual effects.

How To Judge Raw Versus Latent Flow

Prefer JEPA latent flow if it matches the raw model on regime coverage, calibration, tails, and action sensitivity while improving at least one of:

• data efficiency;
• long-horizon coherence;
• transfer;
• rollout speed or memory;
• candidate-action evaluation cost.

Prefer raw trajectory flow if it preserves important numeric detail better and the latent route offers no clear efficiency or transfer advantage.

Reject both if the samples look plausible but fail calibration, tail coverage, or action-sensitive evaluation.

Main Risks

• Flow matching may still drop rare modes or connect modes through invalid trajectories.
• JEPA latents may preserve semantic regimes while losing exact numeric detail.
• A powerful decoder may ignore the sampled future latent and predict from history alone.
• Observational action logs may make action effects look causal when they are confounded.
• Better sample quality may not improve action ranking or planning.
• Conditional IMLE / Forward XM may cover every regime while assigning the wrong sample frequencies, and Reverse XM may collapse without explicit entropy or coverage control.

Expected Outcome

The working hypothesis is:

• raw flow will be strongest when exact values and short-horizon detail dominate;
• JEPA latent flow will be strongest for long horizons, high-dimensional observations, nuisance-heavy data, and repeated planning rollouts;
• multi-scale JEPA latents may provide the best compromise between semantic future regimes and dense numeric fidelity.

Relation To Foundation TSFM Agenda

Agenda slot	Intended contribution
Multi-modal future distributions	Test whether flow matching can preserve separated, calibrated future regimes.
Latent-state prediction	Replace one JEPA future point with a distribution over future latent trajectories.
Semantic state vs numeric detail	Compare direct raw generation, oracle latent decoding, and generated latent decoding.
Control and counterfactuals	Measure whether candidate actions move probability mass between future regimes correctly.

Open Questions

• Does flow matching actually preserve separated future regimes in multivariate time series? FlowWM partially answers the mechanism question outside numeric time series by covering enumerable visual futures, but calibration, dense numeric recovery, and action-sensitive regimes remain open.
• When does calibrated raw/latent flow beat the much simpler conditional-IMLE / Forward-XM candidate-coverage baseline under matched training and serving cost?
• Can a conditional BTM direct map preserve calibrated rare future-regime mass under matched training compute while beating ordinary conditional flow on end-to-end latency?
• When does JEPA latent generation outperform direct raw-trajectory generation?
• Which JEPA layers preserve both regime semantics and exact numeric detail?
• How structured should the future latent trajectory be across time and scale?
• Can the model remain calibrated under rare regimes and distribution shift?
• Does better probabilistic modeling improve action ranking and closed-loop decisions?

Collaboration

If this direction resonates with you, I would be happy to talk with like-minded people, collaborate on research, and work on use-cases together.

Ideas are not the bottleneck. Hands are. Time-series modeling should be moving at least as fast as vision, audio, and robotics.

• Email: [email protected]
• X: @chemeris
• Telegram: @alexanderchemeris

Related Pages

• Foundation Time-Series Model Research Agenda
• Latent-State Time-Series Modeling
• Latent-Space Predictive Learning
• World Models
• Joint Embedding Predictive Architecture
• Time-Series Benchmark Hygiene
• Action-Conditioned Time-Series Datasets
• VJEPA
• Sundial
• World Models
• Reconstruction Or Semantics?
• RAEv2
• FlowWM
• Beckmann Transport Models
• Explorative Modeling
• Fast-LeWM
• Expectations vs. Realities
• Kubernetes OTEL Control Gym