Stochastic Flow Matching for Resolving Small-Scale Physics

TL;DR

The proposed stochastic flow matching (SFM) framework significantly outperforms existing methods such as conditional diffusion and flows and injects noise into the encoder output using an adaptive noise scaling mechanism, to account for uncertainty in the deterministic part.

Abstract

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.

Figures (15)

• Figure 1: Overview of the Stochastic Flow Matching (SFM) Method. The encoder transforms (coarse-res.) inputs into a latent distribution more aligned with the (fine-res.) target. It generates channels absent in the input and corrects both spatial and channel misalignments, such as repositioning the typhoon's eye to its more accurate location, and generating radar data. From the latent space, FM generates small-scale physics by transporting samples from $p({\mathbf{z}})$ to $p({\mathbf{x}})$ via the velocity field $\boldsymbol{\nu}(\mathbf{x}, t)$.
• Figure 2: SFM vs. baselines for different weather variables. SFM generates more physically consistent outputs, while UNet output appears blurred, and CDM struggles to accurately reconstruct radar reflectivity. Note that radar reflectivity is not present in the input data and is entirely generated as a new channel.
• Figure 3: SFM power spectra vs. baselines for CWA downscaling. SFM exhibits superior spectral fidelity, closely aligning with the ground truth across all variables, with a particularly strong fidelity for the purely generated radar reflectivity. It consistently outperforms CorrDiff, especially in capturing high-frequency details across all variables.
• Figure 4: SFM vs. Baselines for Kolmogorov Flow and $\tau=10$: when the figures are zoomed in, it is apparent that SFM aligns closer to the ground truth, and the presence of high-frequency artifacts in the baseline models becomes more noticeable.
• Figure 5: Visualization of ERA5-CWA Dataset Variables. The top row shows input variables such as temperature and wind at coarse resolution, while the bottom row presents the corresponding fine-resolution target variables. The maximum radar reflectivity is absent from the input variables and must be constructed by the model. This key misalignment between the low- and high-resolution data increases the complexity of the problem beyond standard super-resolution tasks.