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 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.
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.
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.
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.
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. |
| 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.
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:
- flow matching preserves distinct future regimes rather than producing one broad cloud;
- sampled regime frequencies are calibrated;
- trajectories remain temporally and cross-channel coherent;
- rare and decision-relevant futures are retained;
- candidate actions move probability mass between the correct regimes;
- the decoder actually uses the sampled future latent;
- the JEPA latent route provides a measurable advantage over direct raw flow.
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.
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 candidates.
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.
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?
- 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