HyperAIHyperAI

Command Palette

Search for a command to run...

Flow Matching in Feature Space for Stochastic World Modeling

François Porcher Nicolas Carion Karteek Alahari Shizhe Chen

Abstract

World modeling requires forecasting uncertain futures while preserving information useful for downstream perception. Existing visual world models often struggle to satisfy both goals: VAE-based stochastic models operate in low-dimensional reconstruction latents, which can limit perception performance, while deterministic predictors using strong pretrained features collapse multimodal futures into a single blurry mean. In this work, we propose FlowWM, a stochastic world model that performs flow matching directly within pretrained feature space (e.g., DINOv3). This is challenging because pretrained features are substantially high-dimensional, making standard difusion recipes suboptimal. To address this, we investigate the design choices needed for feature-space flow matching and introduce a diferentiable one-step projection mechanism that enables eficient training with temporal consistency and task-driven objectives. We evaluate FlowWM on two benchmarks: a synthetic benchmark for systematic evaluation of accuracy and diversity, and a real-world benchmark FuturePerception. FlowWM improves perception performance, mode coverage, and horizon robustness, validating our proposed design for stochastic world modeling in high-dimensional feature spaces.

One-sentence Summary

FAIR at Meta and Inria propose FlowWM, a stochastic world model that performs flow matching directly in high-dimensional pretrained feature space (e.g., DINOv3) and uses a differentiable one-step projection to enforce temporal consistency and task-driven objectives, yielding improved perception, mode coverage, and horizon robustness on both a synthetic benchmark and the real-world FuturePerception benchmark.

Key Contributions

  • FlowWM is a stochastic world model that performs flow matching directly in high-dimensional pretrained visual feature spaces (e.g., DINOv3), enabling diverse future predictions while retaining semantic structure for downstream tasks.
  • A differentiable one-step projection mechanism incorporates temporal consistency and task-driven objectives during training, avoiding backpropagation through the full diffusion trajectory and improving efficiency.
  • The paper contributes a synthetic benchmark for systematic accuracy and diversity evaluation, the real-world FuturePerception benchmark, and demonstrates that FlowWM improves perception performance, mode coverage, and horizon robustness over prior methods.

Introduction

In autonomous driving and robotics, visual world models forecast future states from past context, but the choice of representation is critical: predictions must preserve semantic and geometric information for downstream perception and planning, not just pixel fidelity. Prior work either uses VAE latents optimized for reconstruction rather than perception, or relies on deterministic predictors in pretrained feature spaces that collapse multimodal futures into a single average, yielding unrealistic outcomes. The authors propose FlowWM, a stochastic world model that applies flow matching directly in high-dimensional pretrained feature spaces such as DINOv3. They introduce architectural and scheduling choices tailored to high-dimensional latents, along with a differentiable one-step projection that enables efficient training with temporal coherence and task-driven objectives without backpropagating through the full sampling trajectory. FlowWM improves perception performance, mode coverage, and horizon robustness over deterministic and stochastic baselines in both a synthetic bouncing-shapes environment and a real-world autonomous driving benchmark.

Dataset

The authors evaluate latent world models using two complementary benchmarks, one synthetic and one real-world, designed to measure perception quality relevant for planning and control.

  • Synthetic benchmark

  • Source: Custom-generated controlled stochastic environments.

  • Composition: Simple objects in highly structured settings.

  • Purpose: Enables rigorous, reproducible evaluation with full ground truth.

  • Usage: Models are tested on future frame prediction, and perception metrics are computed directly from the generated sequences.

  • Real-world benchmark: FuturePerception

  • Source: Built entirely from the Waymo Open Dataset, originally created for object tracking.

  • Composition & filtering: No data filtering is applied; the dataset retains the natural distribution of real-world autonomous driving scenes.

  • Video properties: High-resolution frames (1920×1280) at 10 FPS.

  • Task definition: 4 context frames are used to predict the next 12 frames, corresponding to 1.2 seconds into the future.

  • Processing: No additional cropping or metadata generation is mentioned; the raw frames are used directly.

  • Usage: The benchmark is used exclusively for evaluation. Latent world models generate future frame predictions, and the quality of the latent predictions is assessed through downstream perception tasks:

  • Object detection: Evaluated on large vehicles, pedestrians, and cyclists with AP_L(N) (best-of-N AP to capture multi-modal futures).

  • Depth estimation: Assessed with standard RMSE and threshold accuracy δ_i.

  • Split: No explicit training split is described; the benchmark serves as a fixed evaluation protocol measuring how well predicted frames support perception modules.

Method

The authors introduce FlowWM, a stochastic latent space world model that operates directly in the high-dimensional latent spaces of pretrained vision encoders. Given a video sequence, the initial context frames are encoded frame-by-frame using a frozen pretrained encoder (e.g., DINOv3) to obtain context latents. The model then predicts the subsequent latent frames. Evaluation is performed on downstream perception tasks using frozen decoders applied to the predicted features.

Refer to the framework diagram:

To generate future latents, the authors employ flow matching with a standard linear probability path. Let x0x_0x0 denote a noise latent and x1x_1x1 a data latent corresponding to future frames conditioned on a context latent xctxx_{\text{ctx}}xctx. A continuous interpolation variable τ(0,1)\tau \in (0, 1)τ(0,1) is used along the probability path, computing a linear interpolation between noise and data as: xτ=(1τ)x0+τx1x_{\tau} = (1 - \tau) x_0 + \tau x_1xτ=(1τ)x0+τx1 For this linear path, the ground-truth velocity field is u(xτ,τ)=x1x0u^{\star}(x_{\tau}, \tau) = x_1 - x_0u(xτ,τ)=x1x0. A neural network uθ(xτ,xctx,τ)u_{\theta}(x_{\tau}, x_{\text{ctx}}, \tau)uθ(xτ,xctx,τ) is trained to approximate this velocity, learning the conditional distribution p(xfuturexctx)p(x_{\text{future}} | x_{\text{ctx}})p(xfuturexctx). Sampling is performed by integrating the ODE dxτdτ=uθ(xτ,xctx,τ)\frac{dx_{\tau}}{d\tau} = u_{\theta}(x_{\tau}, x_{\text{ctx}}, \tau)dτdxτ=uθ(xτ,xctx,τ) from τ=0\tau=0τ=0 to τ=1\tau=1τ=1, starting from x0pnoisex_0 \sim p_{\text{noise}}x0pnoise.

As shown in the figure below:

The model adopts a Transformer-based latent generative architecture. It takes a noisy target latent xτx_{\tau}xτ as input and is conditioned on encoded context frames xctxx_{\text{ctx}}xctx via cross-attention. Positional embeddings are decomposed into temporal and spatial components, applying rotary positional encodings (RoPE) to both inputs. The architecture consists of a backbone for multimodal fusion and a projection head for velocity prediction. The backbone contains 2 DiT-style Transformer blocks (dimension 256) with query-key normalization. The resulting features are fed through AdaLN conditioning into a shallow but wide projection head composed of 2 wide projection layers (dimension 1024). This wide projection dimension ensures the head's hidden dimension exceeds the per-patch latent dimension, which is critical for accurately predicting velocity fields in high-dimensional latent spaces.

The training process involves multiple objectives. The standard flow matching objective optimizes the conditional expected squared error: LFM(θ)=Ex1pdata,x0pnoise,τU(0,1)[uθ(xτ,τ)(x1x0)22]\mathcal{L}_{\text{FM}}(\theta) = \mathbb{E}_{x_1 \sim p_{\text{data}}, x_0 \sim p_{\text{noise}}, \tau \sim \mathcal{U}(0,1)} \left[ \| u_{\theta}(x_{\tau}, \tau) - (x_1 - x_0) \|_2^2 \right]LFM(θ)=Ex1pdata,x0pnoise,τU(0,1)[uθ(xτ,τ)(x1x0)22] To ensure temporal consistency and avoid artifacts like flickering, the authors introduce a regularization term inspired by Physics-Informed Neural Networks. Instead of integrating the ODE, which is computationally expensive, they use a one-step projection to obtain a differentiable estimate of the predicted endpoint: x~1(τ)=xτ+(1τ)uθ(xτ,τ)\tilde{x}_1(\tau) = x_{\tau} + (1 - \tau) u_{\theta}(x_{\tau}, \tau)x~1(τ)=xτ+(1τ)uθ(xτ,τ) They then enforce alignment between the ground-truth temporal differences and those of the predicted latents obtained via this projection, yielding an auxiliary temporal loss: Ltemporal=t=1T1Δx~tΔxtGT2\mathcal{L}_{\text{temporal}} = \sum_{t=1}^{T-1} \| \Delta \tilde{x}_t - \Delta x_t^{\text{GT}} \|^2Ltemporal=t=1T1∥Δx~tΔxtGT2 Furthermore, if a downstream task is known, task-driven objectives can be incorporated. Using the one-step projection to obtain a differentiable estimate x~1\tilde{x}_1x~1, a frozen downstream detector (e.g., for object detection) is applied, and its detection loss is backpropagated through the projection and the world model to steer representations towards semantic utility.

Finally, to maintain a stable signal-to-noise ratio when increasing the dimension of the generated tensor, the authors apply a resolution-dependent timestep shift. The original timestep τ\tauτ is transformed as: τ=ατ1+(α1)τ\tau' = \frac{\alpha \tau}{1 + (\alpha - 1) \tau}τ=1+(α1)τατ This rescales diffusion time by a factor α\alphaα toward noisier timesteps, preserving endpoints while redistributing probability mass toward larger noise levels.

Experiment

The evaluation uses a synthetic bouncing shapes benchmark and a real-world Waymo-based FuturePerception benchmark, which assess latent world model predictions through downstream perception tasks such as object detection and depth estimation. FlowWM's stochastic modeling in high-dimensional semantic latent spaces consistently outperforms deterministic and low-dimensional VAE-based alternatives, better capturing multi-modal futures and maintaining accuracy over longer prediction horizons. Ablations confirm that wide projection heads, temporal consistency loss, and a shifted timestep schedule are crucial design choices, while stochastic sampling steadily improves performance as the number of sampled futures increases.

FlowWM achieves dramatically lower prediction errors than deterministic models and a prior stochastic approach on the Bouncing Shapes benchmark. The F1 error falls from over 14 to under 5, with temporal consistency loss providing further improvement. Gains stem from both more accurate samples and better coverage of plausible futures. FlowWM yields substantially lower precision, recall, and F1 errors compared to deterministic predictors and the earlier stochastic method. Adding temporal consistency loss further improves all metrics, reducing the F1 error even closer to the oracle. Deterministic models and the prior stochastic baseline all have F1 errors above 14, while FlowWM cuts this to below 5, approaching the oracle's near-zero error.

FlowWM consistently outperforms all baselines on both object detection and depth estimation, with particularly large gains over deterministic and VAE-based world models. Using semantic latents provides a clear advantage over VAE latents, and re-encoding VAE outputs into semantic space fails to close the gap, confirming that the representation quality, not just the prediction method, drives downstream performance. Stochastic modeling proves essential in real-world settings, as FlowWM degrades more slowly over time and benefits from sampling multiple futures. FlowWM achieves the highest detection and depth scores, substantially surpassing the deterministic predictor and VAE-based world models. Re-encoding VAE predictions into DINOv3 space slightly degrades detection performance, indicating that VAE latent spaces lack the semantic fidelity needed for perception tasks. Compared to the deterministic predictor, FlowWM reduces the Fréchet Video Distance by 43%, producing future frames that are distributionally closer to real driving videos. Deterministic world models accumulate error more rapidly over time, while stochastic models preserve multiple plausible futures, leading to more stable long-horizon predictions. Additional sampling consistently improves detection performance for FlowWM, whereas deterministic models gain no benefit from extra samples.

Evaluated on the Bouncing Shapes benchmark and real-world driving perception tasks, FlowWM achieves dramatic error reductions compared to deterministic and prior stochastic models, cutting F1 error from over 14 to below 5 on Bouncing Shapes and approaching the oracle. The use of semantic latents is crucial for downstream performance, and stochastic modeling yields slower error accumulation and benefits from sampling multiple futures, leading to a 43% reduction in Fréchet Video Distance. These results confirm that flow matching with semantic representations produces more accurate and distributionally realistic future predictions.


Build AI with AI

From idea to launch — accelerate your AI development with free AI co-coding, out-of-the-box environment and best price of GPUs.

AI Co-coding
Ready-to-use GPUs
Best Pricing

HyperAI Newsletters

Subscribe to our latest updates
We will deliver the latest updates of the week to your inbox at nine o'clock every Monday morning
Powered by MailChimp