Flow Matching in Feature Space for Stochastic World Modeling
For world modeling in robotics and autonomous systems, this work provides a method to generate diverse, temporally consistent futures while maintaining high-quality features for downstream perception.
FlowWM introduces flow matching directly in pretrained feature space (e.g., DINOv3) for stochastic world modeling, addressing the trade-off between forecasting uncertainty and preserving perception-relevant information. It improves perception performance, mode coverage, and horizon robustness on synthetic and real-world benchmarks.
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.