Iterative Refinement Diffusion for Super-Resolved Data Assimilation of Multiscale Physical Systems

Mrigank Dhingra; Ramchandran Muthukumar; Rebecca Willett; Omer San

cs.LGAug 13, 2026

Iterative Refinement Diffusion for Super-Resolved Data Assimilation of Multiscale Physical Systems

Authors: Mrigank Dhingra, Ramchandran Muthukumar, Rebecca Willett, Omer San

Organizations: Department of Mechanical and Aerospace Engineering, University of Tennessee, Knoxville, TN, 37996, USA · UChicago Data Science Institute, University of Chicago, Chicago, IL, 60637, USA

Abstract

Recovering high-resolution states from sparse, low-resolution observations is a central challenge in scientific machine learning and data assimilation. Classical data assimilation exploits temporal information through forecast-analysis cycles, but often requires repeated access to expensive high-resolution forecast models. Generative super-resolution can recover unresolved structure from coarse observations, but is commonly used as a one-shot mapping that does not fully exploit constraints from past states. We introduce Iterative Refinement (IR), a learned data assimilation framework that combines these perspectives. Instead of performing a single coarse-to-fine reconstruction, IR decomposes the task into resolution-wise forecast-analysis operations across a multiresolution hierarchy. At each stage, a shared neural operator with resolution-dependent spectral mode slicing provides a dynamical prior, while a shared conditional diffusion corrector uses the current coarser-resolution state to produce a refined posterior at the next finer resolution. We evaluate IR on one-dimensional stochastically forced Burgers dynamics and two-dimensional Kraichnan turbulence. On the challenging 256x256 Kraichnan benchmark, IR achieves an RMSE of 0.184 and an SSIM of 0.836, outperforming spectral upsampling, one-shot diffusion super-resolution, enhanced deep super-resolution, and an autoregressive forecaster. On the more constrained Burgers testbed, IR remains competitive with one-shot diffusion, which achieves the lowest RMSE. These results show that one-shot generative reconstruction can be effective for simpler settings, while hierarchical forecast-analysis refinement becomes advantageous in strongly multiscale and underdetermined regimes. Overall, IR combines temporal priors, generative correction, and multiresolution reconstruction for learned data assimilation in complex physical systems.

Figures & tables

Figure 1: Preview of the main spectral recovery result on the 2D Kraichnan turbulence benchmark. The shaded blue region denotes the observed low-wavenumber modes resolved by the coarse $32\times 32$ input, while the shaded red region denotes unresolved modes with $k>16$ that must be reconstructed. Spectral upsampling loses energy beyond the coarse Nyquist limit, EDSR is overly dissipative, and one-shot diffusion recovers only part of the high-wavenumber tail. Iterative refinement most closely follows the ground-truth spectrum in the unresolved range, demonstrating improved recovery of fine-scale turbulent structure.

Figure 2: Preview of the accuracy–cost trade-off for super-resolved data assimilation on the 2D Kraichnan benchmark. Each point reports the overall RMSE at $256\times 256$ against the per-step wall time on a logarithmic scale. Lower and further left is better. The solver EnKF, labeled as “Solver EnKF (oracle dynamics knowledge),” achieves the lowest RMSE because it advances an ensemble using the high-resolution pseudospectral solver. The learned EnKF is much faster but less accurate. Iterative refinement provides an intermediate learned assimilation strategy, substantially improving over the learned EnKF while avoiding repeated online full-solver ensemble forecasts.

Figure 3: Illustration of the sequential super-resolution inverse problem. The finest-resolution trajectory contains coherent small-scale structures that are only partially visible after coarse observation. Given only the low-resolution sequence, many high-resolution trajectories may be compatible with the same observations, making the reconstruction problem ill-posed. Temporal information from previous high-resolution estimates is therefore essential for narrowing the space of admissible fine-scale states.

Figure 4: Multiresolution state hierarchy used in the proposed formulation. The coarsest level $r=0$ corresponds to the observed state, while the finest level $r=R$ is the target reconstruction. Downsampling operators $\mathcal{D}_{r+1\to r}$ define the coarse representations, and upsampling operators $\mathcal{U}_{r\to r+1}$ transfer information upward through the hierarchy. This decomposition converts a single large coarse-to-fine inverse problem into a sequence of smaller refinement tasks.

Figure 5: High-level view of iterative refinement across time and resolution. At each physical time step, an IR update block receives the previous posterior at the same resolution and the current conditioning signal from the coarser level, then produces the current posterior analysis at that resolution. Horizontal arrows denote temporal autoregressive propagation, while vertical arrows denote coarse-to-fine conditioning within the same time step. The right panel expands a single IR update block into its three components: a shared mode-sliced FNO forecaster, a conditioned upsampler, and a shared diffusion corrector.

Figure 6: Detailed one-step forecast–analysis cascade for iterative refinement. At each target resolution level $r$ , the shared mode-sliced FNO forecaster advances the previous posterior $x_{t-1}^{(r)}$ to produce the forecast prior $\widetilde{x}_{t}^{(r)}$ . The lifted conditioning signal $c_{t}^{(r-1)\uparrow}$ bypasses the forecaster and enters only the shared diffusion corrector, which produces the posterior analysis $x_{t}^{(r)}$ . The posterior from each level is then propagated upward as conditioning for the next finer refinement stage.

Figure 7: A single FNO is initialized at the finest spectral bandwidth; coarser-resolution inference uses progressively smaller low-frequency slices of the same learned weights.

Figure 8: Shared conditional diffusion corrector used to implement the analysis step at each refinement stage. The denoiser receives three aligned spatial channels: the noisy target sample, the FNO forecast prior, and the upsampled conditioning field from the previous coarser level. These channels are concatenated and passed through a conditional U-Net that predicts the diffusion noise. The diffusion timestep and resolution index are embedded and injected into the network through FiLM conditioning, allowing the same corrector $G_{\theta}$ to adapt across diffusion steps and resolution levels.

Figure 9: Two-stage training strategy for the forecast–analysis pipeline. In Stage 1, a single shared FNO forecaster is trained across resolution levels using mode slicing. In Stage 2, the FNO weights are frozen and the trained forecasters are run in teacher-forced mode to generate one-step forecast priors. These forecasts are time-aligned with the corresponding current coarse states and target fine states to form diffusion training tuples, which are then used to train the shared corrector $G_{\theta}$ .

Method	Temporal prior	Obs. correction	Hierarchy	Stochastic
Spectral upsampling	–	–	–	–
EDSR	–	✓	–	–
One-shot diffusion	–	✓	–	✓
FNO-only forecaster	✓	–	–	–
Iterative refinement	✓	✓	✓	✓

Table 1: Comparison of baseline methods.

Method	RMSE
Spectral Upsample	0.084944 (0.017270)
EDSR	0.021447 (0.004078)
One-Shot SR	0.002776 (0.000245)
Iterative Refinement	0.004355 (0.000254)
FNO-only forecaster	–

Table 2: Benchmark results on the 1D stochastic Burgers testbed, evaluated on an aligned subset of

5

trajectories over

700

time steps at resolution

512

. Values are reported as mean RMSE with the standard deviation across trajectories shown in parentheses. Lower RMSE is better.

Method	RMSE	SSIM
Spectral Upsample	0.335469 (0.008368)	0.669100 (0.110947)
One-Shot SR	0.349025 (0.127302)	0.748400 (0.078439)
EDSR	0.238840 (0.004079)	0.786113 (0.003680)
Iterative Refinement	0.18401 (0.013212)	0.835870 (0.004520)
FNO-only forecaster	2.568060 (0.198535)	–

Table 3: Benchmark results on the 2D Kraichnan turbulence testbed, evaluated on an aligned subset of

3

trajectories over

200

time steps at resolution

256\times 256

. Values are reported as mean with the standard deviation across trajectories shown in parentheses. Lower RMSE is better; higher SSIM is better.

Figure 10: Temporal rollout comparison on the 2D Kraichnan turbulence benchmark at resolution $256\times 256$ for a representative trajectory. Rows show the ground truth, EDSR reconstruction and error, one-shot diffusion reconstruction and error, and iterative refinement reconstruction and error. While both one-shot baselines recover the large-scale vortex organization, iterative refinement more accurately preserves fine filaments and local vortex geometry, yielding consistently smaller residual error throughout the rollout.

Figure 11: Local zoom comparison on the 2D Kraichnan benchmark for a representative final-time snapshot. The highlighted region shows that iterative refinement best reproduces the shape and intensity of the vortex core, whereas EDSR yields a smoother reconstruction and the one-shot diffusion model introduces visible local distortions. This close-up illustrates the advantage of hierarchical analysis for recovering fine-scale turbulent structure.

Figure 12: Trajectory-averaged energy spectrum $E(k)=|\hat{u}_{k}|^{2}$ for the 1D stochastic Burgers benchmark at four representative times ( $t=T/4,\,T/2,\,3T/4,$ and $T$ ). spectral upsampling is limited by the coarse observation bandwidth, while EDSR exhibits an earlier high-wavenumber roll-off. The one-shot diffusion and iterative refinement models both closely track the ground-truth spectrum across the resolved wavenumber range, indicating accurate recovery of fine-scale spectral content.

Figure 13: Radial energy spectrum for the 2D Kraichnan turbulence benchmark at resolution $256\times 256$ , shown at four evenly spaced trajectory times $t=T/4$ , $T/2$ , $3T/4$ , and $T-1$ , with spectra averaged over the test trajectories. The vertical dotted line marks the coarse-grid Nyquist limit ( $k=16$ ). spectral upsampling loses energy rapidly beyond this limit, while EDSR remains overly dissipative at moderate and high wavenumbers. One-shot diffusion improves the spectral reconstruction but still underestimates the high- $k$ tail. Iterative refinement provides the closest match to the ground-truth spectrum across all four time points, demonstrating more accurate and temporally persistent recovery of fine-scale turbulent structure.

Panel of Figure 14 — Figure 14: Temporal stability diagnostics on the 2D Kraichnan benchmark at resolution $256\times 256$ . (a) RMSE over time. Iterative refinement maintains the lowest and flattest error profile throughout the rollout, outperforming one-shot diffusion, EDSR, spectral upsampling, and the autoregressive FNO-only baseline, which diverges rapidly. (b) Temporal consistency measured by the frame-to-frame displacement norm $\|\omega_{t}-\omega_{t-1}\|_{2}$ . Iterative refinement most closely tracks the ground-truth temporal evolution, whereas spectral upsampling is overly smooth and the one-shot and EDSR baselines exhibit larger deviations.

Figure 15: Snapshot comparison between data-assimilation baselines and iterative refinement on the 2D Kraichnan benchmark at resolution $256\times 256$ . The top row shows the ground truth, spectral upsampling, solver EnKF, learned EnKF using an FNO forecast, and iterative refinement. The bottom row shows the absolute error for each reconstruction method. The solver EnKF and iterative refinement both recover the main vortex and filament structures well, while the learned EnKF exhibits larger localized artifacts and spectral upsampling misses substantial fine-scale content.

Figure 16: Radial energy spectrum comparison between EnKF variants and iterative refinement on the 2D Kraichnan benchmark. Spectra are shown at four representative times and averaged over test trajectories. The solver EnKF, which uses the pseudospectral solver as the ensemble forecast model, most closely follows the ground-truth spectrum. Iterative refinement also recovers the spectral decay well without requiring online solver access. The learned EnKF, which replaces the solver with an FNO forecast, is less reliable at high wavenumbers and may introduce nonphysical spectral content.

Panel of Figure 17 — Figure 17: Cascade-depth ablation on the 1D Burgers benchmark. All variants use the same trained forecasters and shared diffusion corrector; only the inference cascade is changed. The full 3-stage cascade, $64\rightarrow 128\rightarrow 256\rightarrow 512$ , achieves the lowest aggregate RMSE, maintains the most stable temporal error profile, and yields the smallest spectral error beyond the coarse-observation Nyquist limit.

Panel of Figure 18 — Figure 18: Cascade-depth ablation on the 2D Kraichnan turbulence benchmark. The one-shot 1-stage variant attempts a direct $32\times 32\rightarrow 256\times 256$ reconstruction, the 2-stage variant inserts an intermediate $128\times 128$ refinement, and the full 3-stage cascade further includes a $64\times 64$ stage. The deeper cascade consistently improves aggregate accuracy, lowers the temporal error floor, and most faithfully reconstructs the high-wavenumber portion of the energy spectrum beyond the coarse-grid Nyquist limit.

Figure 19: Schematic of the three propagation-signal variants considered in the ablation study. After each refinement stage, the next finer stage is conditioned on either (i) the corrected posterior from the previous stage, (ii) the FNO forecast prior at that stage, or (iii) the original coarse observation upsampled to the required intermediate resolution. The default iterative refinement method corresponds to posterior propagation.

Panel of Figure 20 — Figure 20: Propagation-signal ablation on the 1D Burgers benchmark. All variants use the same trained FNO forecasters and diffusion corrector; only the signal passed from one refinement stage to the next is changed. Propagating the corrected posterior yields the lowest aggregate RMSE, the most stable temporal behavior, and the smallest spectral error. Using the FNO forecast or repeatedly reusing the upsampled coarse observation removes the benefit of progressive posterior refinement and substantially degrades the final high-resolution reconstruction.

Panel of Figure 21 — Figure 21: Propagation-signal ablation on the 2D Kraichnan turbulence benchmark. The posterior-propagation variant remains most accurate at the finer resolutions and achieves the lowest final RMSE. The forecast and posterior variants exhibit similar radial energy spectra, indicating that both retain comparable spectral amplitudes across wavenumbers; however, the posterior still yields substantially lower RMSE because the diffusion correction improves the spatial alignment and structural fidelity of vortical features. Reusing only the upsampled coarse observation performs worst at fine scales, showing that progressive posterior refinement is essential for accurate multiscale reconstruction.

Panel of Figure 22 — Figure 22: Summary of reconstruction metrics across all methods. Top: 1D Burgers benchmark, where one-shot diffusion achieves the lowest RMSE and spectral RMSE, with iterative refinement close behind. Bottom: 2D Kraichnan turbulence benchmark, where iterative refinement achieves the best RMSE, spectral RMSE, and SSIM, demonstrating the advantage of hierarchical correction in the more challenging multiscale setting.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 23: Autoregressive FNO-only diagnostic on the 1D Burgers benchmark at resolution $512$ . The FNO-only rollout initially follows the large-scale structure but becomes unstable and blows up at $t=54$ , after which the predicted trajectory loses physical consistency. The diffusion posterior remains closely aligned with the ground-truth shock evolution over the full time horizon.

Figure 24: Autoregressive FNO-only diagnostic on the 2D Kraichnan benchmark at resolution $256\times 256$ for a representative trajectory. Columns show increasing time, while rows compare ground truth, the FNO-only baseline, its absolute error, the diffusion posterior, and its absolute error. The FNO-only rollout accumulates substantial error and develops spurious small-scale structures, whereas the diffusion posterior preserves the correct vortex and filament organization with much smaller error throughout the sequence.

Figure 25: Snapshot comparison on the 1D stochastic Burgers benchmark at resolution $512$ for a representative trajectory at four times ( $t=0$ , $T/4$ , $T/2$ , and $T-1$ ). For each method, the reconstructed signal is shown together with the pointwise absolute error relative to the ground truth. Both the one-shot diffusion and iterative refinement models closely track the Burgers shock structure, while EDSR improves over spectral upsampling but exhibits larger localized deviations near sharp transitions.

Figure 26: Space–time comparison on the 1D stochastic Burgers benchmark for a representative trajectory. The Hovmöller plots show the full temporal evolution for the ground truth and the reconstructions from EDSR, one-shot diffusion, iterative refinement, and spectral upsampling. The learned methods recover the dominant space–time structure well, with one-shot diffusion and iterative refinement remaining closest to the ground truth over the full rollout.

Panel of Figure 27 — Figure 27: Temporal stability diagnostics on the 1D Burgers benchmark at resolution $512$ , averaged over the test trajectories. (a) RMSE over time. One-shot diffusion achieves the lowest temporal RMSE, while iterative refinement remains close and stable throughout the rollout. (b) Temporal consistency measured by the frame-to-frame displacement norm $\|u_{t}-u_{t-1}\|_{2}$ . One-shot diffusion and iterative refinement closely follow the ground-truth temporal variation, whereas EDSR and spectral upsampling are more dissipative.

Figure 28: Forecast-prior diagnostic on the 1D Burgers benchmark. At a representative time step, the one-step FNO forecast already captures the global solution profile across resolutions, but deviates near the shock and steep-gradient regions. The diffusion posterior corrects these local errors and aligns more closely with the ground truth, showing that the corrector acts primarily as a targeted refinement of the forecast prior.

Figure 29: Forecast-prior diagnostic on the 2D Kraichnan benchmark at resolution $256\times 256$ . Rows show the ground truth, the one-step FNO prior, the diffusion posterior, and the magnitude of the posterior correction relative to the prior. The forecast already captures the large-scale flow organization, while the diffusion corrector applies localized adjustments concentrated around filaments and vortex interfaces, indicating residual analysis rather than full re-synthesis.

Figure 30: Sensitivity of iterative refinement to the DDIM stochasticity parameter $\eta$ , reported as mean $\pm$ standard deviation across test trajectories. Top row: reconstruction RMSE. Bottom row: temporal smoothness measured by frame-to-frame $\ell_{2}$ displacement. In the 1D Burgers case, performance changes only mildly with $\eta$ . In the 2D Kraichnan case, deterministic or weakly stochastic sampling performs poorly, while moderate stochasticity $(\eta\geq 0.5)$ yields both lower RMSE and substantially smoother temporal behavior.

Figure 31: Sensitivity of iterative refinement to the number of DDIM reverse steps, shown as mean $\pm$ standard deviation across test trajectories together with wall-clock inference time. In the 1D Burgers benchmark, RMSE improves sharply from $10$ to $25$ steps and then plateaus, indicating diminishing returns beyond moderate step counts. In the 2D Kraichnan benchmark, RMSE is relatively insensitive to the number of steps over the tested range, while computational cost increases steadily, highlighting a clear quality–cost trade-off.

PDE & Domain
Parameter	1D Burgers	2D Kraichnan
Domain	$[0,2\pi]$ , periodic	$[0,2\pi]^{2}$ , periodic
Solver resolution	$N=512$	$256\times 256$
Solver timestep	$\Delta t=10^{-4}$	$\Delta t=5\times 10^{-3}$
Viscosity $\nu$	$5\times 10^{-3}$	$1\times 10^{-3}$
Large-scale damping	Linear friction $\alpha=0.1$	Ekman drag $\mu=0.05$

Appendix

Table 4: PDE parameters and dataset configuration for the two testbeds.

FNO Forecaster
Parameter	1D Burgers	2D Kraichnan
Fourier layers	3	4
Channel width	32	32
Mode truncation $k_{\text{max}}$	$N_{r}/4$	$N_{r}/4$ per dim
Activation	GELU	GELU
Diffusion Corrector U-Net (shared across all stages)

Appendix

Table 5: Model architecture details. The diffusion corrector

G

is shared across all resolution stages; a shared FNO forecaster

F_{r}

is trained per stage.

FNO Forecaster Training
Parameter	1D Burgers	2D Kraichnan
Epochs	100	80
Batch size	64	32
Optimizer	AdamW	AdamW
Learning rate	$10^{-3}$	$10^{-3}$
Weight decay	$10^{-4}$	$10^{-4}$

Appendix

Table 6: Training hyperparameters. All experiments use a single NVIDIA RTX 4090 (24 GB).

Model / Component	1D Burgers	2D Kraichnan
Spectral Upsample	0	0
EDSR [ 28 ]	445,121	1,662,977
One-Shot Diffusion SR U-Net	2,280,641	6,970,689
IR Diffusion Corrector $G$ (ours)	2,412,353	4,171,537

Appendix

Table 7: Learnable parameter counts for the diffusion corrector and competing baselines. For the Iterative Refinement (IR) method only the shared U-Net corrector

G

is listed; FNO forecasters are additional but are also required by ablation variants and are not part of the generative corrector comparison. Spectral Upsample has no learnable parameters (analytical zero-padding).

Data Assimilation

Data assimilation combines observations with dynamical models to estimate hidden system states and their uncertainty over time.

See all 58 papers

Super-Resolving Coarse-Resolution Weather Forecasts With Flow Matching
Apr 1, 2026Aymeric Delefosse, Anastase Charantonis, Dominique BéréziatWeather Forecasting Probabilistic Forecasting
Inria, ARCHES, Paris, France · Sorbonne Université, CNRS, LIP6, Paris, France
Physics-Consistent Diffusion for Efficient Fluid Super-Resolution via Multiscale Residual Correction
Feb 25, 2026Zhihao Li, Shengwei Dong, Chuang Yi +5Physics-Informed Diffusion Models Image Super-Resolution
The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · SandGold AI Research, Guangzhou, China · The Hong Kong University of Science and Technology, Hong Kong SAR, China +1
ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing
May 14, 2026Yixuan Jia, Siyi Chen, Yida Pan +9Data Assimilation
University of Michigan, Ann Arbor · Michigan State University · Los Alamos National Laboratory

Apr 1, 2026J/K move · Enter open · S save

Super-Resolving Coarse-Resolution Weather Forecasts With Flow Matching

Aymeric Delefosse, Anastase Charantonis, Dominique Béréziat

Inria, ARCHES, Paris, France · Sorbonne Université, CNRS, LIP6, Paris, France

Weather Forecasting Probabilistic Forecasting

Open paper AskPDF arXiv

Machine learning-based weather forecasting models now surpass state-of-the-art numerical weather prediction systems, but training and operating these models at high spatial resolution remains computationally expensive. We present a modular framework that decouples forecasting from spatial resolution by applying learned generative super-resolution as a post-processing step to coarse-resolution forecast trajectories. We formulate super-resolution as a stochastic inverse problem, using a residual formulation to preserve large-scale structure while reconstructing unresolved variability. The model is trained with flow matching exclusively on reanalysis data and is applied to global medium-range forecasts. We evaluate (i) design consistency by re-coarsening super-resolved forecasts and comparing them to the original coarse trajectories, and (ii) high-resolution forecast quality using standard ensemble verification metrics and spectral diagnostics. Results show that super-resolution preserves large-scale structure and variance after re-coarsening, introduces physically consistent small-scale variability, and achieves competitive probabilistic forecast skill at 0.25° resolution relative to an operational ensemble baseline, while requiring only a modest additional training cost compared with end-to-end high-resolution forecasting.

Abstract

Data Assimilation

Explore similar work

Super-Resolving Coarse-Resolution Weather Forecasts With Flow Matching

Physics-Consistent Diffusion for Efficient Fluid Super-Resolution via Multiscale Residual Correction

ForcingDAS: Unified and Robust Data Assimilation via Diffusion Forcing