PaperlayerView exact revision on arXiv

Source-linked AI summary

Just Weather Scoring: Efficient End-to-end Nowcasting with Distributional Diffusion

Jannik Wiese, Johannes Schusterbauer, Tommaso Martorella, Björn Ommer

arXiv:2610.12189v1cs.LGcs.CV

TL;DR

Probabilistic precipitation nowcasting needs efficient forecasts that represent uncertainty, while existing diffusion approaches often depend on separate components and costly iterative denoising. JWS forecasts directly in radar space using Masked Asynchronous Diffusion and a simple distributional objective, achieving state-of-the-art probabilistic performance on SEVIR and MeteoNet.

  • Problem

    Existing diffusion-based nowcasting approaches often use separately trained compression or deterministic forecasting components, and iterative denoising remains costly.

    Source: §1 · p. 1

  • Method

    JWS is an end-to-end radar-space diffusion model combining clean-context Masked Asynchronous Diffusion with a hyperparameter-free distributional objective.

    Source: §1 · p. 2, §3.2 · pp. 4–5

  • Results

    JWS achieves state-of-the-art probabilistic nowcasting performance on SEVIR and MeteoNet; its smallest model remains competitive with approximately 17× fewer parameters than FlowCast.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    JWS quantifies uncertainty by ensembling directly in radar space, without relying on compression models.

    Source: §5 · p. 10

Abstract

from arXiv · show

Generative diffusion models are well-suited for probabilistic precipitation nowcasting, but existing approaches often rely on separately trained compression or deterministic forecasting components and remain costly at inference due to iterative denoising. We introduce Just Weather Scoring (JWS), a single-stage, end-to-end diffusion model which addresses both issues by forecasting directly in radar space and enabling few-step generation. Radar-space modeling greatly simplifies training and inference and eliminates uncertainty arising from lossy compression. JWS combines Masked Asynchronous Diffusion, a timestep-sampling scheme that preserves clean context while adapting diffusion training to high-dimensional spatio-temporal data, with a simple scoring-rule objective that aligns training with probabilistic forecasting and unlocks few-step generation. On the SEVIR and MeteoNet benchmarks, JWS achieves state-of-the-art probabilistic forecasting performance at reduced training and inference cost. Even our smallest model remains competitive using substantially fewer parameters and more than 17x faster inference.

1 INTRODUCTION

Probabilistic precipitation nowcasting needs detailed, rapidly updated forecasts that represent uncertainty, but prior diffusion approaches often depend on multiple trained stages. JWS forecasts directly in radar space, combining Masked Asynchronous Diffusion with a distributional objective to support efficient probabilistic generation.

1 INTRODUCTION

  • Generative forecasts can represent precipitation uncertainty, whereas regression averages plausible outcomes and becomes blurrier as uncertainty grows.The need to resolve spatial detail and update rapidly makes uncertainty-aware nowcasting important for applications such as renewable energy planning and emergency response.Source: §1 · pp. 1–2
  • Prior diffusion nowcasting often relies on compression, deterministic forecasts, or residual tasks, adding stages that can introduce bias and uncertainty.Small errors in precipitation intensity or location can have consequential impacts.Source: §1 · p. 2
  • JWS operates directly on uncompressed radar data with a single transformer, while Masked Asynchronous Diffusion preserves clean context and adapts timestep sampling to forecasting.The model predicts clean radar maps and downsamples only through patchification.Source: §1 · p. 2
  • Simple Distributional Diffusion Models use a hyperparameter-free objective that interpolates distributional CRPS and squared error, supporting few-step generation.The objective combines scoring-rule supervision with diffusion model coverage.Source: §1 · p. 2
  • Combining MAD and sDDM generates radar sequences in pixel space, avoids additional pipeline stages, and directs computation to generative forecasting.The paper presents the combination as enabling simple and efficient training and inference without lossy compression or deterministic priors.Source: §1 · p. 2
  • Approximately 17× fewer parameters than FlowCast, JWS’s smallest model remains competitive with substantially lower inference cost than prior methods.The paper attributes these results to its combined components.Source: §1 · p. 2

2 RELATED WORK

Related work addresses the cost and coverage challenges of probabilistic nowcasting through generative models, pixel-space methods, and few-step sampling. The paper positions its approach against limitations in compression-based, prior-conditioned, and tuned diffusion methods.

Precipitation nowcasting

  • Regression objectives yield blurry forecasts, motivating generative nowcasting; diffusion offers mode coverage, while latent diffusion and deterministic priors are common ways to reduce denoising cost.The cited discussion also notes concerns about reduced coverage from priors and additional uncertainty from lossy compression.Source: §2 · p. 3

Pixel-space modeling

  • Pixel-space nowcasting is challenging at high dimensionality; prior methods use specialized architectures or model deterministic-forecast residuals, whereas this work uses standard transformers end to end.The passage contrasts these approaches with pixel-space modeling using standard transformers.Source: §2 · p. 3

Few-step sampling

  • Distillation reduces latency but adds training stages, while flow-based approaches incur training overhead; distributional diffusion enables few-step inference, and this work replaces tuned loss schedules with CRPS–squared-error interpolation.The cited discussion describes iDDM as relying on tuned interdependent parameters.Source: §2 · p. 3

3 METHOD

JWS performs conditional spatiotemporal diffusion directly on radar sequences, using clean-context conditioning and asynchronous, resolution-aware timestep sampling. It combines clean-sequence prediction with a particle-based scoring objective for distributional forecasts.

3.1 PRELIMINARIES

  • The task is to model the distribution of future precipitation frames given past radar observations, using diffusion to transform noise into radar-sequence samples and form forecast ensembles.Flow Matching trains a velocity predictor on interpolated noisy sequences; independent conditional samples support ensemble-based uncertainty estimation.Source: §3.1 · pp. 3–4

3.2 PIXEL-SPACE NOWCASTING

  • JWS generates radar forecasts directly in pixel space, while randomly selected clean context frames support forecasting, infilling, or conditioning on corrupted data.Unlike FREUD’s multiple reconstructions per forecast, JWS directly samples radar forecasts with one model and removes the decoding stage.Source: §3.2 · p. 4, Figure 2 · p. 4
  • MAD samples framewise timesteps around a shared sequence-level center, preserving clean context; σ=0 recovers RaMViD, while σ>0 with pU<1 retains asynchronous denoising.Draws outside [0, 1] are replaced with uniform draws on the corresponding side of the center.Source: §3.2 · pp. 4–5
  • For sequence dimensionality D=THW, timestep shifting increases intermediate corruption when D>D0, while preserving the τ=0 and τ=1 endpoints.The method leaves clean context unchanged; the figure caption notes that higher-resolution radar fields retain more structure at fixed corruption, motivating stronger noise.Source: §3.2 · p. 5
  • The model predicts clean sequences rather than velocity directly, using an HDiT-style hierarchy with spatial attention at high resolution and factorized spatiotemporal attention in the lower-resolution bottleneck.The clean prediction is converted to velocity for the velocity-space loss.Source: §3.2 · p. 5

3.3 SIMPLE DISTRIBUTIONAL DIFFUSION

  • sDDM combines CRPS and particlewise MSE with timestep-dependent weighting; auxiliary-noise particles provide ensemble predictions, and shared early layers reduce repeated computation.The loss transitions from marginal distributional supervision toward regression as τ approaches 1.Source: §3.3 · pp. 6–7

4 EXPERIMENTS

Experiments evaluate JWS on SEVIR and MeteoNet, progressively testing timestep sampling, architecture, and the distributional objective before system-level comparisons. Results cover benchmark skill, inference efficiency, deterministic conditioning, extrapolation, and calibration.

4 EXPERIMENTS

  • The evaluation uses CRPS, SSIM, HSS, and CSI over six thresholds, alongside rank histograms and reliability index for calibration.Ablations use 50 sampling steps and 512 samples; experiments begin with 24M-parameter models unless noted otherwise.Source: §4 · p. 7, §4.1 · p. 7

4.1 BUILDING JWS

  • Rescaled uniform timestep sampling gives superior CRPS and RI, while MAD with σ = 0.1 preserves clean conditioning and outperforms naive Diffusion Forcing.Logit-normal sampling with noise scaling achieves the best CSI; MAD with pU < 1 preserves clean conditioning.Source: §4.1 · p. 7
  • Hierarchical models improve pixel-space nowcasting, while factorized spatiotemporal attention reduces computational cost.MAD supports effective nowcasting with standard transformers without specialized architectural modifications.Source: §4.1 · p. 7
  • sDDM performs strongly at low and high NFE, and JWS outperforms alternative methods with a single function evaluation.Squared-error Flow Matching requires many iterations, while training only with CRPS deteriorates at high NFE.Source: Figure 6 · p. 8, §4.1 · p. 8

4.2 SYSTEM-LEVEL WEATHER NOWCASTING

  • JWS achieves better CRPS, SSIM, and RI on SEVIR, and the best reported CRPS and SSIM on MeteoNet; guidance improves HSS and CSI on both benchmarks.JWS-T remains competitive with larger models, and the 9M-parameter model outperforms prior methods in CRPS on MeteoNet.Source: Table 5 · p. 8, §4.2 · p. 8, Table 6 · p. 9
  • Deterministic priors improve CSI and HSS up to τ = 0.2, but any deterministic conditioning consistently degrades calibration.Framewise confidence also matters: cleaner early-frame priors and noisier later-frame priors improve CRPS, reliability, and CSI over uniform prior strength.Source: §4.2 · p. 9, Figure 9 · p. 9
  • Without additional training, JWS extrapolates from 12 to 36 predicted frames, achieving superior CRPS and calibration while remaining competitive in CSI.It also substantially improves over the deterministic EarthFormer prior.Source: Figure 10 · p. 10, §4.2 · p. 10
  • JWS rank histograms are flatter than FREUD's and CasCast's, with calibration remaining superior over lead time.At +60 min, ensembles retain global structure with local variability, and uncertainty heatmaps highlight high-precipitation regions.Source: Figure 12 · p. 10, Figure 11 · p. 10, §4.2 · p. 10

5 CONCLUSION

JWS is an end-to-end generative nowcasting method designed for rapid updates and calibrated uncertainty. It operates without compression, supports ensemble uncertainty estimation in radar space, and achieves state-of-the-art probabilistic nowcasting on SEVIR and MeteoNet.

5 CONCLUSION

  • JWS avoids compression models, so ensembling quantifies uncertainty directly in radar space.Training with a distributional objective enables few-step sampling and reduces inference latency.Source: §5 · p. 10
  • JWS achieves state-of-the-art probabilistic nowcasting performance on the SEVIR and MeteoNet benchmarks.The method scales with model size and inference compute, while efficient settings remain competitive.Source: §5 · p. 10
  • The simplified training and inference pipeline is intended to support adoption by meteorological practitioners.The conclusion presents ease of adoption as an aim.Source: §5 · p. 10

AI USE STATEMENT

The paper reports uses of generative AI in methodology, experiments, implementation, and manuscript preparation, with human review and verification procedures. JWS uses publicly available datasets and standard transformers, and the authors release code, checkpoints, and evaluation pipelines.

AI USE STATEMENT

  • Generative AI tools assisted with methodology and experiments, implementation, literature review, figure creation, and paper editing.The authors state that AI tools were not used for theoretical or conceptual frameworks, mathematical claims, hypothesis formulation, dataset cleaning, qualitative analysis, or result interpretation.Source: AI USE STATEMENT · p. 11
  • Researchers reviewed AI-assisted work, discussed methodology and experiments, and manually verified AI-generated figures.LLM-generated code was tested, and literature summaries were checked by at least two researchers after manual review.Source: AI USE STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

  • JWS uses publicly available datasets and standard transformers with minimal modifications, and the authors release code, checkpoints, and evaluation pipelines.The statement says these materials and choices make JWS easy to reproduce from scratch.Source: REPRODUCIBILITY STATEMENT · p. 11

A IMPLEMENTATION DETAILS

JWS uses standard Llama-style transformers across four model sizes, with hierarchical attention and patchification. Training expands an ensemble and applies timestep sampling, diffusion noise, and a distributional loss; inference uses Euler updates and sampled noise for ensembles.

A IMPLEMENTATION DETAILS

  • JWS variants range from 9M to 340M parameters and differ in layer count and hidden dimensionality.Hierarchical models use 1×16×16 patchification, neighborhood attention over 5×5 regions, and factorized spatial and temporal attention at the bottleneck.Source: A IMPLEMENTATION DETAILS · p. 18
  • The training step interpolates noise and clean radar data, predicts K pixel-space sequences, and converts predictions and targets to velocity space.The simple loss combines a CRPS term, which attracts ensemble predictions to the target, with an interaction term that encourages ensemble diversity.Source: A IMPLEMENTATION DETAILS · p. 19
  • The simple loss mixes CRPS and mean squared error, with the mixing weight determined by diffusion time.The appendix identifies this as L_simple = (1 − t) × CRPS + t × MSE.Source: A IMPLEMENTATION DETAILS · p. 19
  • Training samples framewise diffusion timesteps, applies dimensionality-aware shifting, and uses up to 13 clean conditioning frames.Unconditional training occurs with probability pU = 0.25; σ = 0.1 caps the timestep-sampling standard deviation.Source: A IMPLEMENTATION DETAILS · p. 19
  • JWS-T and JWS-S train for 300k iterations, while JWS-B trains for 200k and JWS-L for 100k after checkpoint initialization.Training uses SEVIR and MeteoNet sequences; MeteoNet frames are resized to 384×384 for training and restored to 400×400 for metric calculation.Source: A IMPLEMENTATION DETAILS · pp. 19–20
  • Inference integrates the ODE with Euler updates, keeps auxiliary noise fixed along each trajectory, and samples different initial noises for ensemble forecasts.The default setting uses 50 function evaluations; model variants use noise conditioning in bottleneck and high-resolution layers.Source: A IMPLEMENTATION DETAILS · pp. 19–20

B ADDITIONAL ABLATIONS

The additional ablations report JWS’s lead-time, calibration, and compute comparisons alongside analyses of model scaling, solver choice, and training design. Several design variants, including filtering and loss reweighting, do not improve overall performance.

FLOP comparison

  • 6.6 vs 59.5 TFLOPs: the smallest JWS model uses about ninefold fewer inference FLOPs than CasCast while achieving superior CRPS.JWS-L uses more inference FLOPs than CasCast but is faster in wall-clock time, highlighting reduced pipeline overhead.Source: B ADDITIONAL ABLATIONS · p. 20

Performance over lead time

FlowCast calibration comparison

  • At long lead times, JWS-B is better calibrated and outperforms FlowCast on CRPS, whose overconfidence and reliability worsen with horizon.Figure 16 describes FlowCast as showing substantially stronger overconfidence, especially at long lead times.Source: B ADDITIONAL ABLATIONS · p. 20, Figure 16 · p. 21

Patch size, model, and data scaling

  • More tokens and FLOPs improve CRPS, CSI, and HSS; the smallest JWS-T outperforms FREUD LSM-L, and larger models perform better.Large models overfit when trained directly at full resolution; downsampled pretraining mitigates this, while reduced training data causes JWS-L to overfit and greater data diversity stabilizes scaling.Source: Figure 17 · p. 21, Figure 18 · p. 21, Figure 19 · p. 22

Confidence-aware priors and solvers

  • Confidence-aware priors improve CRPS, calibration, and CSI over uniform prior strength; Euler sampling performs best for CRPS, while DDM leads on HSS and SSIM.The prior design uses cleaner deterministic priors for early frames and noisier priors for later frames; Euler outperforms Heun and SDE integration.Source: Figure 21 · p. 22, Figure 22 · p. 23

Training particles and expansion layer

  • Increasing training particles or expanding the ensemble earlier does not improve performance, and earlier expansion increases training cost.The ablation favors expansion at the last bottleneck layer; the small-scale verification likewise finds no consistent benefit from more particles or earlier expansion.Source: B ADDITIONAL ABLATIONS · p. 23, Figure 23 · p. 23, Figure 24 · p. 23

Auxiliary conditioning and timestep spread

  • Residual-stream concatenation is the best tested ξ-injection mechanism, while keeping final layers clean deteriorates performance; σ = 0.1 performs best for MAD timestep spread.The timestep comparison includes uniform and uncorrelated alternatives.Source: B ADDITIONAL ABLATIONS · p. 23, Table 8 · p. 24, Table 9 · p. 24

Loss reweighting and data filtering

  • Neither filtering nor loss reweighting improves overall performance, and all tested interventions worsen CRPS, supporting training on the full dataset with equal weights.The MeteoNet table reports increased CRPS with reweighting and resampling; the SEVIR table reports minimal impact from filtering and loss weighting.Source: B ADDITIONAL ABLATIONS · p. 23, Table 11 · p. 24, Table 10 · p. 24

C METRICS AND EVALUATION PROTOCOL

This section names the evaluation protocol and defines the metrics used to assess probabilistic precipitation forecasts. The metrics cover forecast error, structural similarity, event skill, and ensemble calibration.

C METRICS AND EVALUATION PROTOCOL

C.1 METRIC DEFINITIONS

  • CRPS is a proper scoring rule for univariate probabilistic regression that generalizes Mean Absolute Error.Its ensemble approximation balances distance from the observation against ensemble spread.Source: §C.1 · p. 24
  • SSIM measures structural similarity between ground truth and mean forecasts, with sliding-window results averaged using Gaussian weighting.For single-channel data, the formula uses image means, variances, and covariance.Source: §C.1 · p. 25
  • HSS measures pixelwise positional skill against chance, while CSI measures correctly predicted events excluding true negatives; both average over six thresholds.HSS ranges from 0 for no improvement over chance to 1 for perfect prediction.Source: §C.1 · p. 25
  • RI measures rank-histogram deviation from uniformity, while SSR compares ensemble spread with forecast error; SSR near 1 indicates calibration.The rank histogram has N + 1 possible ranks for an N-member ensemble.Source: §C.1 · p. 25

C.2 EVALUATION PROTOCOL

The evaluation protocol uses SEVIR for most experiments and establishes benchmark results on its full test set. Plots and ablations use a 512-sample subset because of computational constraints.

C.2 EVALUATION PROTOCOL

  • Unless stated otherwise, experiments use 50 network evaluations and compare methods with 10 ensemble members.The main SEVIR results use the full test set; plots and ablations use 512 samples due to computational constraints.Source: §C.2 · p. 25
  • The 512-sample evaluation closely matches larger evaluation runs and has low inter-seed variance.Table 12 reports mean ± standard deviation across eight evaluation seeds.Source: §C.2 · p. 25, Table 12 · p. 26
  • Experiments run on four Nvidia GH200 GPUs with sample parallelization, while latency is measured on one GPU; literature baselines are verified where possible.Previously unreported metrics are recomputed from available checkpoints when authors provide trained weights.Source: §C.2 · p. 26

C.3 DATASETS

The section specifies the SEVIR and MeteoNet evaluation data and protocols, then surveys related work on machine-learning nowcasting, pixel-space diffusion, and few-step sampling.

C.3 DATASETS

  • SEVIR and MeteoNet use established evaluation protocols, with 60-minute forecasts from 65-minute context on SEVIR and 12-frame forecasts from 12-frame context on MeteoNet.SEVIR contains 20,393 events from the continental United States; MeteoNet uses the southeastern France region, resized for model processing while metrics remain at native resolution.Source: §C.3 · p. 26

D.1 MACHINE LEARNING FOR NOWCASTING

  • JWS uses a single end-to-end pixel-space model, addressing compression uncertainty and the multi-stage designs used by prior nowcasting approaches.Prior work uses latent compression, deterministic forecasts, or residual prediction; JWS instead applies a standard scalable architecture directly to pixel-space nowcasting.Source: §D.1 · p. 27

D.2 PIXEL-SPACE DIFFUSION

  • Pixel-space generation has used specialized denoisers and token-efficient architectures, while JWS applies standard transformers and denoising procedures directly to weather nowcasting.Prior approaches include latent diffusion, resolution scheduling, neural-field prediction, and local or latent-token attention.Source: §D.2 · pp. 27–28

D.3 FEW-STEP SAMPLING

  • JWS builds on distributional diffusion for few-step weather generation, while improved Distributional Diffusion Models combine time-dependent distributional and MSE objectives.Other few-step approaches can require distillation, Jacobian-vector products, or multiple evaluations to construct training targets; distributional diffusion trains ensembles with scoring-rule supervision.Source: §D.3 · p. 29

E QUALITATIVE RESULTS

The qualitative results show ensemble variability and localized uncertainty, compare JWS forecasts with prior methods, and illustrate how forecast detail changes with inference steps.

E.1 QUALITATIVE FORECAST ENSEMBLES

  • Ensemble examples show three forecast members across the 60-minute lead time, with high-precipitation regions exhibiting variability.Additional SEVIR examples show three members and forecast steps at 10-minute intervals.Source: Figure 26 · p. 30, Figure 27 · p. 30, Figure 28 · p. 31

E.2 QUALITATIVE COMPARISONS

  • Across examples, JWS better captures local and global patterns, low-precipitation dynamics, and sharper, more realistic forecasts than prior methods.FlowCast and JWS both capture localized high precipitation accurately.Source: Figure 30 · p. 32, Figure 31 · p. 33, Figure 32 · p. 34

E.3 QUALITATIVE UNCERTAINTY

  • At +60 minutes, variance maps highlight localized uncertainty in complex precipitation regions across five ensemble members.The examples are from the SEVIR dataset.Source: Figure 35 · p. 37

E.4 QUALITATIVE INFERENCE-TIME SCALING

  • At +60 minutes, fewer inference evaluations preserve global structure but can blur forecasts, while more steps recover details and variability.The examples show 1 NFE capturing global structure, 2 NFE resolving patch artifacts, and 10 NFE approaching the 50 NFE setting.Source: Figure 36 · p. 38, Figure 37 · p. 39, Figure 38 · p. 40
Loading 2610.12189v1…