Stochastic Flow Matching for Resolving Small-Scale Physics
This work addresses super-resolution for weather and physics applications, offering a novel approach to handle deterministic and stochastic components, but it is incremental as it builds on existing flow and diffusion models.
The paper tackled the problem of super-resolving small-scale details in physical sciences like weather, where challenges include misaligned distributions, multi-scale dynamics, and limited data, by proposing a stochastic flow matching framework that encodes inputs to a latent base distribution and adds stochastic details, resulting in significant outperformance over existing methods on real-world and PDE-based datasets.
Conditioning diffusion and flow models have proven effective for super-resolving small-scale details in natural images.However, in physical sciences such as weather, super-resolving small-scale details poses significant challenges due to: (i) misalignment between input and output distributions (i.e., solutions to distinct partial differential equations (PDEs) follow different trajectories), (ii) multi-scale dynamics, deterministic dynamics at large scales vs. stochastic at small scales, and (iii) limited data, increasing the risk of overfitting. To address these challenges, we propose encoding the inputs to a latent base distribution that is closer to the target distribution, followed by flow matching to generate small-scale physics. The encoder captures the deterministic components, while flow matching adds stochastic small-scale details. To account for uncertainty in the deterministic part, we inject noise into the encoder output using an adaptive noise scaling mechanism, which is dynamically adjusted based on maximum-likelihood estimates of the encoder predictions. We conduct extensive experiments on both the real-world CWA weather dataset and the PDE-based Kolmogorov dataset, with the CWA task involving super-resolving the weather variables for the region of Taiwan from 25 km to 2 km scales. Our results show that the proposed stochastic flow matching (SFM) framework significantly outperforms existing methods such as conditional diffusion and flows.