Flow Matching in Feature Space for Stochastic World Modeling

2026-06-27

AI Summary

Existing visual world models struggle to preserve information useful for downstream perception tasks while modeling future uncertainty. This work introduces stochastic flow matching in high-dimensional pretrained visual feature spaces—such as DINOv3—for the first time, constructing a stochastic world model equipped with a differentiable single-step projection mechanism tailored to this space to enable efficient training. By integrating temporal consistency constraints with task-driven optimization objectives, the proposed method substantially improves performance on perception tasks, enhances mode coverage in multimodal future prediction, and increases robustness in long-horizon forecasting across both synthetic and real-world benchmarks, demonstrating its effectiveness and generalizability.

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 diffusion recipes suboptimal. To address this, we investigate the design choices needed for feature-space flow matching and introduce a differentiable one-step projection mechanism that enables efficient 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.

Problem

• world modeling
• stochastic prediction
• feature space
• multimodal futures
• perception performance

Innovation

• flow matching
• stochastic world modeling
• pretrained feature space
• temporal consistency
• multimodal forecasting