Adaptive Flow Matching for Resolving Small-Scale Physics

Stathi Fotiadis, Noah D Brenowitz, Tomas Geffner, Yair Cohen, Michael Pritchard, Arash Vahdat, Morteza Mardani

Abstract: Conditional diffusion and flow models are effective for super-resolving small-scale details in natural images. However, in physical sciences such as weather, three major challenges arise: (i) spatially misaligned input-output distributions (PDEs at different resolutions lead to divergent trajectories), (ii) misaligned and distinct input-output channels (channel synthesis), (iii) several channels with diverse stochasticity scales (multiscale). To address these, we propose to first encode inputs into a latent base distribution that is closer to the target, then apply Flow Matching to generate small-scale physics. The encoder captures deterministic components, while Flow Matching adds stochastic details. To handle uncertainty in the deterministic part, we inject noise via an adaptive noise scaling mechanism, dynamically adjusted by maximum-likelihood estimates of the encoder’s predictions. Experiments on real-world weather data (including super-resolution from 25 km to 2 km scales in Taiwan) and in synthetic Kolmogorov flow datasets show that our proposed Adaptive Flow Matching (AFM) framework outperforms existing methods and produces better-calibrated ensembles.

Lay Summary: Weather forecasts often miss important small-scale details, such as the exact path of a thunderstorm, because they rely on low-resolution simulations that smooth out finer structures. This limits our ability to prepare for extreme weather events like floods or heatwaves. Our research presents a machine learning method that enhances these coarse forecasts by recovering realistic small-scale features, even when the input and output data are not perfectly aligned or come from different sources. The model first captures large-scale, predictable patterns and then adds the missing fine details using controlled randomness. This combination improves both the accuracy and reliability of the predictions. Our approach could support earlier and more precise detection of severe weather, helping communities respond more effectively. Beyond weather, the technique may be useful in other scientific fields such as climate modeling or medical imaging where data resolution and alignment pose similar challenges.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Applications

Keywords: flow matching, diffusion models, weather downscaling, distribution misalignment

Submission Number: 7837