# FLEX: A Backbone for Diffusion-Based Modeling of Spatio-temporal Physical Systems

arXiv: https://arxiv.org/abs/2505.17351 · HTML: https://mlsanity.com/paper/2505.17351/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 56 of its 62 references that are in the map.
- First posted on arXiv 23 May 2025.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 30 of 56)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets](https://mlsanity.com/paper/2311.15127/) (2023)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer](https://mlsanity.com/paper/2408.06072/) (2024)
- [DPM-Solver: A Fast ODE Solver for Diffusion Probabilistic Model Sampling in Around 10 Steps](https://mlsanity.com/paper/2206.00927/) (2022)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Consistency Models](https://mlsanity.com/paper/2303.01469/) (2023)
- [Image Super-Resolution via Iterative Refinement](https://mlsanity.com/paper/2104.07636/) (2021)
- [Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2304.08818/) (2023)
- [Variational Diffusion Models](https://mlsanity.com/paper/2107.00630/) (2021)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [FourCastNet: A Global Data-driven High-resolution Weather Model using Adaptive Fourier Neural Operators](https://mlsanity.com/paper/2202.11214/) (2022)
- [Latte: Latent Diffusion Transformer for Video Generation](https://mlsanity.com/paper/2401.03048/) (2024)
- [All are Worth Words: A ViT Backbone for Diffusion Models](https://mlsanity.com/paper/2209.12152/) (2022)
- [MagicVideo: Efficient Video Generation With Latent Diffusion Models](https://mlsanity.com/paper/2211.11018/) (2022)
- [Activating More Pixels in Image Super-Resolution Transformer](https://mlsanity.com/paper/2205.04437/) (2022)
- [Simple diffusion: End-to-end diffusion for high resolution images](https://mlsanity.com/paper/2301.11093/) (2023)
- [SRDiff: Single Image Super-Resolution with Diffusion Probabilistic Models](https://mlsanity.com/paper/2104.14951/) (2021)

## Built on by (most cited first, 8 of 8)

- [Physics-Constrained Fine-Tuning of Flow-Matching Models for Generation and Inverse Problems](https://mlsanity.com/paper/2508.09156/) (2025)
- [HydroDiffusion: Diffusion-Based Probabilistic Streamflow Forecasting with a State Space Backbone](https://mlsanity.com/paper/2512.12183/) (2025)
- [On The Hidden Biases of Flow Matching Samplers](https://mlsanity.com/paper/2512.16768/) (2025)
- [Flow Matching and Diffusion Models via PointNet for Generating Fluid Fields on Irregular Geometries](https://mlsanity.com/paper/2601.03030/) (2026)
- [Particle-Guided Diffusion for Gas-Phase Reaction Kinetics](https://mlsanity.com/paper/2603.05139/) (2026)
- [UQ-SHRED: uncertainty quantification of shallow recurrent decoder networks for sparse sensing via engression](https://mlsanity.com/paper/2604.01305/) (2026)
- [Modelling Gas-Phase Reaction Kinetics with Guided Particle Diffusion Sampling](https://mlsanity.com/paper/2604.16461/) (2026)
- [Variational Streaming Flow: Probabilistic Forecasting in Physical Time](https://mlsanity.com/paper/2610.00976/) (2026)

## Abstract (arXiv)

We introduce FLEX (FLow EXpert), a backbone architecture for generative modeling of spatio-temporal physical systems using diffusion models. FLEX operates in the residual space rather than on raw data, a modeling choice that we motivate theoretically, showing that it reduces the variance of the velocity field in the diffusion model, which helps stabilize training. FLEX integrates a latent Transformer into a U-Net with standard convolutional ResNet layers and incorporates a redesigned skip connection scheme. This hybrid design enables the model to capture both local spatial detail and long-range dependencies in latent space. To improve spatio-temporal conditioning, FLEX uses a task-specific encoder that processes auxiliary inputs such as coarse or past snapshots. Weak conditioning is applied to the shared encoder via skip connections to promote generalization, while strong conditioning is applied to the decoder through both skip and bottleneck features to ensure reconstruction fidelity. FLEX achieves accurate predictions for super-resolution and forecasting tasks using as few as two reverse diffusion steps. It also produces calibrated uncertainty estimates through sampling. Evaluations on high-resolution 2D turbulence data show that FLEX outperforms strong baselines and generalizes to out-of-distribution settings, including unseen Reynolds numbers, physical observables (e.g., fluid flow velocity fields), and boundary conditions.