Skip to content
Bud's Wiki

Flow Matching in Feature Space for Stochastic World Modeling

TL;DR

Section titled “TL;DR”

FlowWM (FAIR + Inria, Porcher/Carion/Alahari/Chen) trains a stochastic video world model by running flow matching directly in a frozen DINOv3 encoder’s high-dimensional feature space, instead of first compressing to a VAE latent. The paper argues that deterministic feature-space predictors (DINO-WM, DINO-Foresight) collapse multimodal futures into blurry means, while VAE-based stochastic predictors have latents that are poor for downstream perception; FlowWM keeps DINOv3’s semantic richness and proper multi-modality. Key ingredients are a shallow-wide projection head sized above per-patch latent dim, a resolution-aware noise-schedule shift, and a differentiable one-step projection that lets temporal-consistency and task-driven losses supervise the endpoint without integrating the ODE. On a controlled Bouncing-Shapes benchmark FlowWM cuts F1 error from 14.4 (best prior stochastic baseline) to 4.31; on a new Waymo-based FuturePerception benchmark it lifts object-detection AP(3) from 17.5 (VAE-WM) / 16.5 (prior feature-space stochastic WM) to 21.7 and improves depth RMSE from 0.10 to 0.078.

Key claims

Section titled “Key claims”
  • Deterministic predictors trained with regression losses provably collapse multimodal futures into predictions that may correspond to no valid future [§1, Appendix 7].
  • Standard DiT recipes tuned for low-dimensional VAE latents are unstable or suboptimal when applied directly to high-dimensional pretrained features; three architectural/schedule choices are load-bearing to make it work [§1, §3.2, §3.4].
  • The projection head width should exceed the per-patch latent dimension (paper uses d_head = 4096 > d_latent = 384 * H * W / patch); narrower heads severely underperform, and gains from added depth saturate at 2 DiT layers [§3.2, Fig. 15b, §5.4].
  • A differentiable one-step projection — approximating the denoised endpoint as z_1 ≈ z_t + (1 − t) · v_θ(z_t, t) — lets temporal-consistency and task-driven losses be applied at the (projected) endpoint without backpropagating through the full ODE sampling trajectory [§3.3, Eq. 3].
  • The temporal-consistency loss (matching predicted-endpoint temporal derivatives to ground-truth DINOv3 feature derivatives) improves object-detection AP(3) by ~5.2% relative over the standard flow-matching baseline at no additional training FLOPs [§3.3, Eq. 5, Table 3].
  • Task-driven supervision — backpropagating a frozen detector’s loss through the one-step-projected endpoint — provides further gains when the downstream task is known at training time, with timestep weighting mattering [Table 4].
  • Timestep-schedule shifting toward noisier steps (Eq. 6, originally from SD3 for high-resolution images) is critical when increasing latent dimension or prediction horizon in video world modeling — the same principle transfers [§3.4, Table 5].
  • On Bouncing Shapes (controlled synthetic with all valid futures enumerated), FlowWM reduces F1 error from 14.4 (best prior stochastic baseline, ref. 48) to 4.31, closing most of the gap to the Oracle at 1.00 [Table 1].
  • On FuturePerception (a new Waymo-based benchmark), FlowWM beats every baseline on object detection and depth: AP(3) 21.7 vs 18.2 (VAE-WM), 17.1 (prior stochastic feature-space WM), 15.2 (deterministic-DINO); depth RMSE 0.078 vs 0.10, δ₁ 0.723 vs 0.57 [Table 2].
  • Re-encoding VAE-WM predictions through DINOv3 (VAE-decode → DINOv3-encode) does not recover DINOv3-space performance and slightly hurts (17.5 → 16.5 AP(3)) — the information is already lost at the prediction stage, so no post-hoc re-encoding compensates for a VAE-latent bottleneck [§5.3].
  • Best-of-N detection scaling: performance improves consistently with more sampled futures, whereas deterministic predictors receive no benefit — a concrete practical payoff of the stochastic formulation [Fig. 16].
  • Deterministic and stochastic predictors degrade at different rates over horizon: stochastic degrades slower because it can maintain uncertainty rather than averaging incompatible futures [§5.5, Fig. 6].
  • Pixel FVD (via a lightweight decoder trained from DINOv3 latents) improves ~2× over the deterministic predictor, even though pixel generation is not the training objective [§5.3 Pixel generation].

Method

Section titled “Method”

FlowWM’s design is a targeted port of the Diffusion Transformers with Representation Autoencoders (RAE) recipe from image generation to stochastic video world modeling. Each frame is encoded independently by a frozen DINOv3 ViT-S with token features averaged over layers 3/6/9/12. A DiT-style transformer with cross-attention to context frames takes a noisy target-latent sequence and predicts velocity; the flow-matching probability path is the standard linear interpolation x_t = (1 − t) · z_0 + t · z_1, and the ground-truth velocity v* = z_1 − z_0 is regressed with MSE. Positional embeddings decompose into temporal and spatial components; RoPE applies to both z_t and the context. QK-norm is used for stability. The head is deliberately shallow-and-wide (following RAE): projection dimension d_head = 4096 exceeds the per-patch latent dimension, which the ablations confirm is the load-bearing choice — going deeper does not help, but shrinking width kills performance.

Two auxiliary losses use the same trick — a one-step projection ẑ_1 = z_t + (1 − t) · v_θ(z_t, t) — to avoid ODE integration during training. The temporal-consistency loss L_TC = MSE(Δt(ẑ_1), Δt(z*_1)) enforces that the endpoint’s temporal derivative matches ground truth’s; the task-driven loss runs a frozen object-detector over ẑ_1 and backpropagates its detection loss through the projection into the world model. The timestep sampling distribution is shifted per Eq. 6 to redistribute mass toward noisier steps when increasing latent dimension or prediction horizon. Inference uses a 50-step Euler solver from t=0 to t=1.

Two benchmarks are introduced. Bouncing Shapes is a synthetic 2D box with two objects and probabilistic bounces, so all valid futures can be enumerated and precision/recall over modes can be measured against ground truth. FuturePerception is derived from Waymo Open Dataset (1024² resolution, 4 context → 12 future frames at 10 FPS) and evaluated via a linear-probe detector (DINO-DETR trained on raw DINOv3 features) and a lightweight depth head (Depth Anything V2 pseudo-labels); predictions on future features are directly consumed by the frozen probes, so the metric measures preservation of task-relevant information, not pixel reconstruction.

Results

Section titled “Results”

Synthetic Bouncing Shapes (F1 error, lower is better): FlowWM 4.31, FlowWM w/o TC 4.53, prior stochastic feature-space WM (ref. 48) 14.4, deterministic predictor 17.1, DINO-WM 17.8, Oracle 1.00 [Table 1]. The gap between deterministic (17.1) and FlowWM (4.31) is the empirical price of collapsing multimodal futures; the gap between the prior stochastic baseline (14.4) and FlowWM (4.31) is the empirical price of using an ill-tuned standard-DiT recipe in high-dimensional feature space.

FuturePerception object detection: AP(1) / AP(3) — FlowWM 20.9 / 21.7, WAN-2.2 VAE-WM 17.5 / 18.2, ref. 48 16.5 / 17.1, deterministic-DINO 15.2 / 15.2, DINO-WM 14.5 / 14.5, Oracle 65.1 / 65.1 [Table 2]. Depth prediction: FlowWM RMSE 0.078 / δ₁ 0.723 / δ₂ 0.858 / δ₃ 0.913 vs VAE-WM 0.10 / 0.57 / 0.80 / 0.90 [Table 2]. Pixel FVD via a lightweight DINOv3-latent decoder improves ~2× over the deterministic predictor [§5.3].

Why it’s interesting

Section titled “Why it’s interesting”

This is the first filed paper to instantiate the Representation Autoencoders recipe for a stochastic world model — the concept page’s explicit open question of “where does the ‘compress vs preserve’ tradeoff land at video scale?” now has a data point that answers it for the world-modeling downstream task specifically, and answers it in favor of preserving DINOv3’s high-dimensional features rather than compressing to a VAE latent. Sister paper to VideoRAE: Taming Video Foundation Models for Generative Modeling via Representation Autoencoders on the generative video side: both freeze a strong video/image encoder and adapt the RAE architectural machinery to video, but VideoRAE targets decoding to pixels while FlowWM targets predicting future features for perception — the two together cover both halves of the “frozen visual foundation model as latent space” thesis. It also closes the loop with Back to the Features: DINO as a Foundation for Video World Models (DINO-world) which had established latent-predictive world modeling on frozen DINOv2 but stayed deterministic; FlowWM makes the same architectural bet but adds the stochasticity that DINO-world’s own paper flagged as a limitation for real-world driving. On the RAE recipe side, FlowWM confirms the dimension-aware noise shift survives outside image generation (matching the surviving-ingredient story from Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders) and adds a new element the RAE line has not previously needed — a one-step projection that lets temporal-consistency and task-driven losses supervise the endpoint without ODE-differentiation, a technique that generalizes cleanly to any flow-matching model that wants downstream-task backprop without paying Autoregressive Adversarial Post-Training for Real-Time Interactive Video Generation-style computational costs.

See also

Section titled “See also”