PaperlayerView exact revision on arXiv

Source-linked AI summary

NowcastDiT: Diffusion Transformers are Effective Precipitation Nowcasters

Haoran Xu, Xingzhuo Guo, Yuchen Zhang, Jincheng Zhong, Jianmin Wang, Mingsheng Long

arXiv:2609.37038v1cs.LGcs.AIcs.CV

TL;DR

Precipitation nowcasting needs to capture complex, uncertain rainfall evolution, but the need for specialized diffusion architectures remains underexplored. NowcastDiT adapts a standard Diffusion Transformer with scheduled noise correlation and timestep-aware reinforcement learning, achieving strong benchmark performance across perceptual quality and meteorological skill.

  • Problem

    Existing precipitation generators use task-specific designs, leaving insufficiently explored how necessary specialization is beyond a standard diffusion model.

    Source: §1 · p. 1

  • Method

    NowcastDiT keeps a standard DiT backbone and adds a dynamics-aware noise prior and timestep-aware reinforcement learning for precipitation-specific requirements.

    Source: §1 · p. 2

  • Results

    NowcastDiT achieves state-of-the-art performance in perceptual quality and meteorological skill on SEVIR and MRMS; aggregate CSI improves 8.7% on SEVIR and 15.1% on MRMS relative to the strongest baselines.

    Source: §1 · p. 2, §5.2 · p. 7

  • Takeaways & Limitations

    The results support standard DiT as an effective foundation for precipitation nowcasting, with targeted adaptations for domain-specific requirements.

    Source: §1 · p. 2

    The physical-plausibility rationale relies on the suggestion that data-driven DiTs can learn coherent motion and local evolution from radar sequences.

    Source: §3 · p. 3

Abstract

from arXiv · show

Precipitation nowcasting demands accurate short-term forecasts under strong spatiotemporal variability. Diffusion models are well suited to modeling complex precipitation distributions, yet existing approaches often introduce increasingly specialized designs, leaving the capability of a standard diffusion architecture underexplored. We show that a standard Diffusion Transformer already provides a simple and scalable foundation for precipitation nowcasting, with domain-specific requirements accommodated naturally within its design space. Based on this principle, we develop NowcastDiT and instantiate this flexibility through two complementary adaptations: a dynamics-aware noise prior for temporally coherent forecasts, and end-to-end reinforcement learning with timestep-aware rewards for meteorological skill. Experiments on SEVIR and MRMS benchmarks show that NowcastDiT achieves state-of-the-art performance in both perceptual quality and meteorological skill. These results suggest that standard DiT can serve as an effective foundation for precipitation nowcasting.

1 INTRODUCTION

Precipitation nowcasting needs to model uncertain, complex rainfall evolution, while prior generative approaches often build in task-specific mechanisms. This work asks whether a standard DiT, with targeted adaptations, can provide an effective foundation instead.

1 INTRODUCTION

  • The paper questions whether precipitation nowcasting needs specialized generative models, given that such designs can narrow exploration of foundational model designs.Prior approaches include physics-based motion constraints, intensity continuity regularization, and decomposition-based frameworks.Source: §1 · p. 1
  • The authors argue that standard DiTs can model the shared spatiotemporal dynamics and uncertainty of precipitation forecasting and conditional video generation.They also argue that precipitation-specific requirements can fit within the broader design space of diffusion models.Source: §1 · pp. 1–2
  • Figure 1 contrasts specialized diffusion nowcasters with standard DiT and NowcastDiT’s minimal adaptation without changing the backbone.The figure frames the paper’s approach as adapting a standard architecture rather than replacing it with a specialized one.Source: Figure 1 · p. 2
  • NowcastDiT retains a standard DiT backbone and adds a dynamics-aware noise prior and timestep-aware reinforcement learning for temporal coherence and meteorological skill.Experiments on SEVIR and MRMS report state-of-the-art performance across perceptual quality and meteorological skill.Source: §1 · p. 2

2 PRELIMINARIES

The preliminaries formulate precipitation nowcasting as conditional generation of future latent frames from past observations using diffusion and flow matching.

2 PRELIMINARIES

  • A VAE encodes future radar frames as target latents and past frames as the conditioning sequence.The model generates future latents conditioned on encoded history.Source: §2 · p. 2
  • The diffusion model learns a conditional velocity field by denoising linear interpolations between Gaussian noise and target latents.The training objective is flow matching.Source: §2 · p. 2
  • At inference, the model solves the learned velocity field from Gaussian noise to a predicted latent and decodes it into future precipitation frames.The update starts from z0 = ϵ and produces ˆx1:S = D(ˆz1).Source: §2 · p. 2

3 REVISITING PRECIPITATION NOWCASTING REQUIREMENTS

The section argues that precipitation nowcasting shares core spatiotemporal modeling demands with video generation, while its data and objectives motivate targeted adaptations rather than specialized architectures.

3 REVISITING PRECIPITATION NOWCASTING REQUIREMENTS

  • Precipitation fields span multiple spatial and temporal scales and require models that represent hierarchical dependencies and uncertainty in future evolution.The section presents DiTs’ scalable spatiotemporal representations as a suitable foundation.Source: §3 · p. 3
  • Data-driven DiTs may capture coherent motion and local evolution from radar sequences without explicit physical constraints.The section supports this possibility by pointing to generative models’ demonstrated ability to learn intuitive physics from video.Source: §3 · p. 3
  • Sparse, long-tailed precipitation intensities and meteorological objectives such as CSI and HSS call for precipitation-aware adaptations, not necessarily a specialized model.Precipitation evolution combines large-scale advection with localized growth and decay.Source: §3 · p. 3
  • The proposed separation retains a standard DiT for general spatiotemporal modeling and reserves specialization for requirements that differ from generic generation.The section’s takeaway is targeted adaptation rather than model redesign.Source: §3 · p. 3

4 NOWCASTDIT

NowcastDiT keeps a standard Diffusion Transformer backbone and targets precipitation-specific needs through scheduled cross-frame noise correlation and timestep-aware reward optimization.

4.1 TRANSFORMER BACKBONE

  • NowcastDiT uses a standard DiT conditioned on encoded observations, with QK-Norm and 3D RoPE among its backbone choices.It avoids task-specialized forecasting modules, explicit physical constraints, and cascaded architectures.Source: §4 · p. 3, §4.1 · pp. 3–4, Figure 2 · p. 4

4.2 TEMPORAL COHERENCE WITH SCHEDULED NOISE CORRELATION

  • DyPro correlates noise across forecast frames, with coupling that relaxes as denoising proceeds to retain frame-specific precipitation variation.The prior preserves each frame’s standard Gaussian marginal while changing cross-frame covariance; its sampler updates correlation at each step with one network evaluation per step, matching standard Euler sampling in styles.Source: §4.2 · pp. 4–5

4.3 METEOROLOGICAL SKILL WITH TIMESTEP-AWARE REWARDS

  • Post-training uses GRPO to optimize meteorological skills after flow-matching pretraining, using sampled forecasts and selected stochastic transitions.The strategy separates general generative modeling from skill optimization.Source: §4.3 · p. 5
  • The timestep-aware reward shifts from precipitation structure near the noise endpoint to local detail near the clean endpoint.Rnoise combines pooled low-threshold CSI and SSIM, while Rclean combines high-threshold CSI and negative LPIPS; the supplied table caption identifies a performance comparison on SEVIR and MRMS but provides no cell values.Source: §4.3 · p. 6, Table 1 · p. 6

5 EXPERIMENTS

The experiments evaluate NowcastDiT on two radar benchmarks and compare it with deterministic and generative baselines. They report benchmark gains and analyze how model scale, guidance, noise correlation, and reward weighting affect forecasting.

5.1 EXPERIMENTAL SETUP

  • Experiments compare deterministic and generative baselines on SEVIR and MRMS, assessing meteorological skill and perceptual quality across defined radar forecast settings.SEVIR uses 5 observed VIL frames to predict 20 future frames at 128 × 128 resolution; MRMS uses 4 observed frames to predict 20 at 256 × 256 resolution.Source: §5 · p. 6, §5.1 · pp. 6–7

5.2 MAIN RESULTS

  • 8.7% and 15.1% aggregate CSI gains and 2.9% and 13.8% HSS gains over strongest baselines on SEVIR and MRMS, respectively.LPIPS falls by 4.8% and 4.1%, while high-threshold CSI gains range from 19.9%–31.4%; NowcastDiT also achieves the best SSIM among generative methods on SEVIR.Source: §5.2 · p. 7
  • Selected SEVIR and MRMS forecasts show NowcastDiT retaining localized intense regions at longer leads and better preserving MRMS's main rainband shape.Figure 3 shows forecasts at 10-minute intervals through 100 minutes; Figure 4 shows 20-minute intervals through 200 minutes.Source: Figure 3 · p. 7, §5.2 · p. 8, Figure 4 · p. 8

5.3 ANALYSES

  • Increasing model size consistently improves aggregate CSI, supporting scaling as a way to strengthen the standard DiT backbone.Figure 5(a) reports CSI gains relative to the first setting in that panel and labels absolute CSI values.Source: §5.3 · p. 8, Figure 5 · p. 9
  • CFG and RL each improve aggregate CSI over the base model, and their combination achieves the best performance.The authors interpret this as meteorological alignment providing an optimization axis beyond standard diffusion guidance.Source: §5.3 · p. 8, Figure 5 · p. 9
  • All tested nonzero DyPro correlation strengths improve aggregate CSI over the i.i.d. noise baseline, with performance peaking at a moderate strength.Performance declines as correlation increases further; the default correlation strength is α = 0.5.Source: §5.3 · p. 8, Figure 5 · p. 9
  • Timestep-aware reward weighting outperforms static weighting across mixed, CSI-only, and perceptual-only objectives.Mixed metrics provide a better balance between meteorological skill and perceptual quality.Source: §5.3 · pp. 8–9, Figure 5 · p. 9

6 RELATED WORK

Prior precipitation diffusion methods use task-specific mechanisms, leaving open whether specialized designs are necessary for effective nowcasting. Advances in video diffusion and post-training motivate NowcastDiT's exploration of precipitation requirements within a standard DiT design space.

Diffusion models for precipitation nowcasting

  • Prior precipitation diffusion methods pair domain-specific mechanisms with architectural or framework designs, leaving unresolved whether specialization is necessary for effective nowcasting.Examples include intensity-continuity guidance, decomposed training, and cascaded latent-space modeling.Source: §6 · p. 9

Advances in generative modeling techniques

  • Video-diffusion noise priors and reinforcement-learning post-training show how temporal coherence and task-specific alignment can fit within a standard DiT's design space.NowcastDiT develops these strategies for precipitation dynamics and meteorological skill.Source: §6 · p. 9

7 CONCLUSION

The conclusion presents a standard Diffusion Transformer with targeted diffusion and training adaptations as sufficient for effective precipitation nowcasting. The appendix describes a frozen, spatial-only VAE used in pretraining and post-training.

7 CONCLUSION

  • A standard Diffusion Transformer with targeted diffusion and training adaptations is presented as sufficient for effective precipitation nowcasting, with strong meteorological skill and perceptual quality on SEVIR and MRMS.Future work will evaluate generalization across observation systems, resolutions, and longer forecast horizons.Source: §7 · p. 10

A.1 VAE ARCHITECTURE AND TRAINING

  • The VAE compresses spatial resolution only, operates independently across frames, and remains frozen during flow-matching pretraining and RL post-training.Table 2 covers encoder, decoder, and latent-projection parameters.Source: §A.1 · p. 14, Table 2 · p. 14

A.2 TRANSFORMER ARCHITECTURE, TRAINING, AND INFERENCE

The conditioner adds encoded history to noisy patch embeddings while preserving standard DiT blocks; inference also uses guidance to adjust historical conditioning.

A.2 TRANSFORMER ARCHITECTURE, TRAINING, AND INFERENCE

  • Observed latent frames occupy the first L conditioning positions, while future positions are zero-masked before frame-wise patch encoding.The conditioner then adds spatial-topology and frame-index embeddings and applies a SiLU activation and zero-initialized 1 × 1 convolution.Source: §A.2 · p. 14
  • Conditioned embeddings are flattened into spatiotemporal tokens for standard DiT blocks, without changing the blocks or their input–output interface.The conditioner requires neither cross-attention nor a task-specific fusion pathway.Source: §A.2 · p. 14
  • Table 3 collects the architecture, flow-matching pretraining, and pretrained-model inference configurations for MRMS and SEVIR.Its parameter count includes the conditioner, time embedding, and output layer.Source: §A.2 · p. 14, Table 3 · p. 15
  • Classifier-free guidance combines conditional and zero-conditioned velocity predictions from the same pre-RL checkpoint to adjust historical conditioning.A guidance weight of 1 recovers the conditional prediction without guidance; pretraining uses w = 1.6.Source: §A.2 · p. 15

B FOUNDATION, TRAINING AND SAMPLING WITH DYPRO

DyPro changes the noise process over forecast frames while retaining a flow-matching formulation, and supplies deterministic and stochastic sampling dynamics.

B FOUNDATION, TRAINING AND SAMPLING WITH DYPRO

  • DyPro uses timestep-dependent shared strength α(1 − t) to correlate Gaussian innovations recursively across frames, with independent noise recovered at α = 0 or t = 1.The recursion can be applied as correlate and inverted by decorrelate.Source: B FOUNDATION, TRAINING AND SAMPLING… · p. 16

B.1 CORRELATED NOISE

  • The framewise recursion defines a lower-triangular transform; its correlated noise remains Gaussian with standard Gaussian frame marginals.The transform is invertible, and at α = 0 or t = 1 the noise is independent across frames.Source: §B.1 · p. 16

B.2 LINEAR INTERPOLATION UNDER CORRELATED NOISE

  • Linear interpolation under the time-varying noise transform defines a probability path from a correlated Gaussian prior to the conditional clean-data distribution.The path’s marginal velocity includes both the clean-data/noise prediction and a term for how the noise transform changes with diffusion time.Source: §B.2 · p. 17

B.3 TRAINING TARGET AND DETERMINISTIC SAMPLING

  • The training target regresses clean latent minus correlated noise, and deterministic sampling holds the clean estimate fixed within each step.The update transports estimated innovations between timesteps with one network evaluation per step and reduces to the standard Euler update when α = 0.Source: §B.3 · pp. 17–18

B.4 STOCHASTIC SAMPLING

  • The stochastic extension refreshes independent-innovation coordinates with an Ornstein–Uhlenbeck process while preserving their standard Gaussian distribution.Setting κt = 0 recovers deterministic dynamics; the step update mixes innovations with independent Gaussian noise before applying the transform at the next timestep.Source: §B.4 · p. 18

C TIMESTEP-AWARE POST-TRAINING

Post-training uses selected stochastic transitions to optimize terminal forecast rewards, with reward weighting tied to diffusion time and guidance applied at inference.

C.1 MIXED SAMPLING AND WINDOW SCHEDULING

  • Rollouts use a contiguous active window of stochastic transitions and deterministic updates elsewhere, with the final transition to data always deterministic.The window start cycles during training; the default schedule uses N = 10, K = 4, stride 2, and shifts every 20 iterations.Source: §C.1 · p. 19

C.2 TIMESTEP-AWARE REWARD EVALUATION

  • Active transitions receive rewards from the same final forecast, with metric weights set by each transition’s diffusion time rather than by intermediate noisy latents.The reward combines pooled low-threshold CSI and SSIM with high-threshold CSI and negative LPIPS components.Source: §C.2 · p. 19

C.3 MIXGRPO OBJECTIVE AND POLICY SCORES

  • MixGRPO compares group-relative forecast quality separately at each active transition and optimizes a clipped local surrogate only over stochastic transitions.Deterministic transitions still affect the sampled final forecast, and equal group rewards yield zero advantages.Source: §C.3 · pp. 19–20

C.4 TRAINING ALGORITHM

  • The training loop decodes each trajectory’s final forecast, computes its noise and clean rewards, weights them by timestep, then group-normalizes advantages for the policy update.The helper uses a shared DyPro initial latent across the group.Source: §C.4 · pp. 20–21

C.5 DOMAIN GUIDANCE

  • Domain Guidance combines the RL-adapted conditional branch with the frozen pre-RL model as the zero-conditioned reference at inference.Guidance is disabled during RL rollouts; using the pre-RL model for both branches recovers classifier-free guidance.Source: §C.5 · p. 21

C.6 POST-TRAINING AND REWARD CONFIGURATION

  • The post-training configuration keeps the pretrained VAE frozen and uses each rollout batch for one gradient update before synchronizing the rollout policy.Table 4 summarizes post-training, reward, and evaluation settings.Source: §C.6 · p. 21

D.1 DATASETS AND PREPROCESSING

The section specifies the SEVIR and MRMS forecasting data, evaluation metrics, lead-time comparisons, and a cumulative component ablation.

D.1 DATASETS AND PREPROCESSING

  • SEVIR forecasts use five observed VIL frames to predict 20 future frames at 128 × 128 resolution, while MRMS uses 24-frame sequences sampled at 10-minute intervals.SEVIR splits are event-time based; MRMS uses 2016–2020 for development and 2021 for testing.Source: §D.1 · pp. 22–23, Table 5 · p. 22

D.2 EVALUATION METRICS AND AGGREGATION

  • Meteorological skill is measured with CSI and HSS, while perceptual quality uses LPIPS and SSIM; skill scores are averaged uniformly across lead times and thresholds.Higher CSI, HSS, and SSIM and lower LPIPS indicate better performance.Source: §D.2 · p. 23

E.1 FORECAST SKILL ACROSS LEAD TIMES

  • Figure 6 compares CSI and HSS across lead times on MRMS and SEVIR, including NowcastDiT with and without RL post-training.The plotted ranges are 10–200 minutes for MRMS and 5–100 minutes for SEVIR.Source: §E.1 · p. 23

E.2 CUMULATIVE COMPONENT ABLATION

  • Table 6 presents a cumulative SEVIR ablation that successively adds 3D RoPE and QK-Norm, DyPro, and timestep-aware RL to the standard DiT baseline.Each row retains all components added in earlier rows.Source: §E.2 · p. 23, Table 6 · p. 24

E.3 CFG AND SKILL-AWARE POST-TRAINING

The ablations examine how classifier-free guidance (CFG), reinforcement learning (RL), DyPro correlation strength, and timestep-aware reward variants affect forecast skill and perceptual metrics. Additional visualizations compare selected forecasts across matched lead times.

E.3 CFG AND SKILL-AWARE POST-TRAINING

  • The combined CFG and RL setting has the highest aggregate CSI, HSS, and SSIM among the four settings, while CFG alone has the lowest LPIPS.The metrics do not necessarily improve together.Source: Table 7 · p. 24, §E.3 · p. 24

E.4 DYPRO CORRELATION STRENGTH

  • Among the tested DyPro settings, α = 0.5 gives the highest aggregate CSI on SEVIR; all ablation runs omit RL post-training to isolate the noise prior.The sweep varies cross-frame noise correlation strength.Source: Table 8 · p. 24, §E.4 · p. 24

E.5 TIMESTEP-AWARE REWARD VARIANTS

  • The six reward variants combine low- and high-threshold CSI with SSIM and LPIPS, including a timestep-aware perceptual reward that shifts from SSIM at the noise endpoint to negative LPIPS at the clean endpoint.All variants share a cyclic window schedule, whose contribution is not separately ablated.Source: §E.5 · pp. 24–25, Table 9 · p. 25

E.6 ADDITIONAL VISUALIZATIONS

  • Four selected cases compare NowcastDiT with the standard DiT baseline and a task-specific baseline, showing inputs, ground truth, and forecasts at matched lead times.The cases cover SEVIR and MRMS and supplement the aggregate evaluation.Source: §E.6 · p. 25, Figure 8 · p. 27, Figure 9 · p. 27
  • Figure 6 plots CSI and HSS across forecast lead times, while Figure 7 plots LPIPS and RMSE during SEVIR reward-ablation training.Lower LPIPS and RMSE indicate better performance in Figure 7.Source: Figure 6 · p. 26, Figure 7 · p. 26

Related papers

Loading 2609.37038v1…