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
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×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×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 Dr+1→r define the coarse representations, and upsampling operators Ur→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 xt−1(r) to produce the forecast prior xt(r) . The lifted conditioning signal ct(r−1)↑ bypasses the forecaster and enters only the shared diffusion corrector, which produces the posterior analysis xt(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θ 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θ .
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×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×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)=∣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×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.
Figure 14: Temporal stability diagnostics on the 2D Kraichnan benchmark at resolution 256×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 ∥ωt−ω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×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.
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→128→256→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.
Figure 18: Cascade-depth ablation on the 2D Kraichnan turbulence benchmark. The one-shot 1-stage variant attempts a direct 32×32→256×256 reconstruction, the 2-stage variant inserts an intermediate 128×128 refinement, and the full 3-stage cascade further includes a 64×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.
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.
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.
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×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.
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 ∥ut−ut−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×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 η , reported as mean ± standard deviation across test trajectories. Top row: reconstruction RMSE. Bottom row: temporal smoothness measured by frame-to-frame ℓ2 displacement. In the 1D Burgers case, performance changes only mildly with η . In the 2D Kraichnan case, deterministic or weakly stochastic sampling performs poorly, while moderate stochasticity (η≥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 ± 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.
Parameter
1D Burgers
2D Kraichnan
PDE & Domain
Domain
[0,2π] , periodic
[0,2π]2 , periodic
Solver resolution
N=512
256×256
Solver timestep
Δt=10−4
Δt=5×10−3
Viscosity ν
5×10−3
1×10−3
Large-scale damping
Linear friction α=0.1
Ekman drag μ=0.05
Appendix
Table 4: PDE parameters and dataset configuration for the two testbeds.
Parameter
1D Burgers
2D Kraichnan
FNO Forecaster
Fourier layers
3
4
Channel width
32
32
Mode truncation kmax
Nr/4
Nr/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 Fr is trained per stage.
Parameter
1D Burgers
2D Kraichnan
FNO Forecaster Training
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).
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.
Existing image SR and generic diffusion models transfer poorly to fluid SR: they are sampling-intensive, ignore physical constraints, and often yield spectral mismatch and spurious divergence. We address fluid super-resolution (SR) with ReMD (Residual-Multigrid Diffusion), a physics-consistent diffusion framework. At each reverse step, ReMD performs a multigrid residual correction: the update direction is obtained by coupling data consistency with lightweight physics cues and then correcting the residual across scales; the multiscale hierarchy is instantiated with a multi-wavelet basis to capture both large structures and fine vortical details. This coarse-to-fine design accelerates convergence and preserves fine structures while remaining equation-free. Across atmospheric and oceanic benchmarks, ReMD improves accuracy and spectral fidelity, reduces divergence, and reaches comparable quality with markedly fewer sampling steps than diffusion baselines. Our results show that enforcing physics consistency inside the diffusion process via multigrid residual correction and multi-wavelet multiscale modeling is an effective route to efficient fluid SR. Our code are available on https://github.com/lizhihao2022/ReMD.
Zhihao Li, Shengwei Dong, Chuang Yi +5
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
Data assimilation (DA) estimates the state of an evolving dynamical system from noisy, partial observations, and is widely used in scientific simulation as well as weather and climate science. In practice, filtering methods rely on frame-to-frame transition models. However, these models are fragile when observations are non-Markovian (when they form only a partial slice of a higher-dimensional latent state as in real-world weather data): they tend to accumulate errors over long horizons. At the same time, learned DA methods typically commit to a single regime, either filtering (nowcasting, real-time forecasting) or smoothing (retrospective reanalysis), which splits what should be a shared prior across application-specific pipelines. To address both issues, we introduce ForcingDAS, a unified and robust DA framework. Built on Diffusion Forcing with an independent noise level assigned to each frame, ForcingDAS learns a joint-trajectory prior instead of frame-to-frame transitions. This allows it to capture long-horizon temporal dependencies and reduce error accumulation. In addition, the same trained model spans the full filtering to smoothing spectrum at inference time. Specifically, nowcasting, fixed-lag smoothing, and batch reanalysis are selected through the inference schedule alone, without retraining. We evaluate ForcingDAS on 2D Navier-Stokes vorticity, precipitation nowcasting, and global atmospheric state estimation. Across all settings, a single model is competitive with or outperforms both learned and classical baselines that are specialized for individual regimes, with the largest gains observed on real-world weather benchmarks.
Yixuan Jia, Siyi Chen, Yida Pan +9
University of Michigan, Ann Arbor · Michigan State University · Los Alamos National Laboratory