PreDiff: Precipitation Nowcasting with Latent Diffusion Models
Authors: Zhihan Gao
Xingjian Shi
Boran Han
Hao Wang
Xiaoyong Jin
Danielle C. Maddix
Yi Zhu
Mu Li
Bernie Wang
Venue: NeurIPS 2023 poster (2023)
Keywords: Machine Learning for Earth Science
Spatiotemporal Forecasting
Generative Models
Diffusion Models
TL;DR: We introduce a two-stage pipeline for probabilistic precipitation nowcasting: 1) We develop a latent diffusion model PreDiff. (2) We incorporate an auxiliary knowledge alignment mechanism, allowing the integration of domain expertise.

Abstract: Earth system forecasting has traditionally relied on complex physical models that are computationally expensive and require significant domain expertise.
In the past decade, the unprecedented increase in spatiotemporal Earth observation data has enabled data-driven forecasting models using deep learning techniques.
These models have shown promise for diverse Earth system forecasting tasks but either struggle with handling uncertainty or neglect domain-specific prior knowledge, resulting in averaging possible futures to blurred forecasts or generating physically implausible predictions.
To address these limitations, we propose a two-stage pipeline for probabilistic spatiotemporal forecasting: 1) We develop *PreDiff*, a conditional latent diffusion model capable of probabilistic forecasts. 2) We incorporate an explicit knowledge alignment mechanism to align forecasts with domain-specific physical constraints. 
This is achieved by estimating the deviation from imposed constraints at each denoising step and adjusting the transition distribution accordingly.
We conduct empirical studies on two datasets: N-body MNIST, a synthetic dataset with chaotic behavior, and SEVIR, a real-world precipitation nowcasting dataset. 
Specifically, we impose the law of conservation of energy in N-body MNIST and anticipated precipitation intensity in SEVIR. 
Experiments demonstrate the effectiveness of PreDiff in handling uncertainty, incorporating domain-specific prior knowledge, and generating forecasts that exhibit high operational utility.

PreDiff: Precipitation Nowcasting with Latent Diffusion Models
Zhihan Gao"
Xingjian Shif Boson Al
Hong Kong University of Science and Technology zhihan, gao@connect...... ust.hk
xshiabeconnect.. ust.hk
Boran Han
Hao Wang AWS AI Labs
Xiaoyong Jin Amazon jxiaoyon@amazon.coom com
AWS boranhan@amazon.comm com
howngz@amazon.co com
Danielle Maddix
Yi Zhu
Mu Li
Yuyang Wang
AWS AI Labs
Boson AI
Boson AI
AWS Al Labs
dmmaddix0@mazon.co com
yi@boson. ai
mu@boson.cm ai
yuyavang@amazon.coo com
Abstract
Earth system forecasting has traditionally relied on complex physical models that are computationally expensive and require significant domain expertise. In the past decade, the unprecedented increase in spatiotemporal Earth observation data has enabled data-driven forecasting models using deep learning techniques. These models have shown promise for diverse Earth system forecasting tasks. However, they either struggle with handling uncertainty knowledge; as a result, they tend to suffer from averaging possible futures to blurred forecasts or generating physically implausible predictions. To address these limitations, we propose a two-stage pipeline for probabilistic spatiotemporal forecasting: 1) We develop Difff conditional latent diffusion model capable of probabilistic forecasts. 2) We incorporate an explicit knowledge alignment mechanism to align forecasts with domain-specific physical constraints. This is achieved by estimating the deviation from imposed constraints at each denoising step and adjusting the transition distribution accordingly. We conduct empirical studies on two datasets: N N-body MNIST, synthetic dataset with chaotic behavior, and SEVIR, a real-world precipitation nowcasting dataset. Specifically, we impose the law of conservation of energy in N-body MNIST and anticipated precipitation intensity in SEVIR. Experiments demonstrate the effectiveness of PreDiff han- dling uncertainty, incorporating domain-speciff prior knowledge, and generating forecasts that exhibit high operational utility.
neglect domain-speeific prior
Introduction
Earth's intricate climate system significantly influences daily life. Precipitation nowcasting, tasked with delivering accurate rainfall forecasts for the near future (c.g., 0-6 hours), is vital for decision- making across numerous industries and services. Recent advancements in data-driven deep learning (DL) techniques have demonstrated promising potential this field, rivaling conventional numerical methods g (5) with their advantages of being more skillful (5) efficient [37) and scalable B.
However, accurately predicting the future rainfall remains challenging for data-driven algorithms.
WWrk conducted during internship Amazon. Work conducted while at Amazon.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

The generates blurry predictions. This is caused by the high variability and complexity inherent Earth's ' climatic system. Even minor differences in initial conditions can lead to vastly divergent outcomes that are difficult to predict. Most methods adopt point estimation of the future rainfall and are trained by minimizing pixel-wise loss functions (e.g., mean-squared error). These methods lack the capability of capturing multiple plausible futures and will generate blurry forecasts which lose important operational details. Therefore, what are needed instead are probabilistic models that can represent the uncertainty inherent in stochastic systems. The probabilistic models can capture multiple plausible futures, generating diverse high-quality predictions that better align with real-world data.
Earth system forecasting algorithms 47 61 37 69 2 29 3 typically
The emergence of diffusion models (DMs) [22] has enabled powerful probabilisticc frameworks for generative modeling. DMs have shown remarkable capabilities in generating high-quality images A0 and videos 15 231. As a likelihood- model, DMs do not exhibit mode collapse or training instabilities like GANs [IO]. Compared to autoregressive (AR) models 53
45 4
46
[3
39 65
that generate images pixel-by-pixce
DMs can produce higher resolution images faster and with
higher quality. They are also better at handling uncertainty 662 34
57 59 without drawbacks like
exposure bias [13 in AR models. Latent diffusion models (LDMs) 42 [2 further improve on DMs
by separating the model into two phases, only applying the costly diffusion in compressed latent space. This alleviates the computational costs of DMs without significantly impairing performance.
Despite DMs' success in image and video generation 42
15
66
36
32
[5]. its application to
precipitation nowcasting and Earth system forecasting is in early stages [6]. One of the major concerns is that this purely data-centric approach lacks constraints and controls from prior knowledge about the dynamic system. Some spatiotemporal forecasting approaches have incorporated domain knowledge by modifying the model architecture or adding extra training losses [1] 1 37 This enables them to be aware of prior knowledge and generate physically plausible forecasts. However, these approaches still face challenges, such requiring design new model architectures or retrain the entire model from scratch when constraints change. More detailed discussions on related works are provided in Appendix A
6], we propose general two-stage pipeline for training data-driven Earth system forecasting model. 1) In the first stage, we focus on capturing the intrinsic semantics in the data by training an LDM. To capture Earth's long-term and complex changes, we instantiate the LDM's core neural network as a UNet-style architecture based on Earthformer R. 2) In the second stage, we inject prior knowledge of the Earth system by training knowledge alignment network that guides the sampling process of the LDM. Specifically the alignment network parameterizes an energy function that adjusts the transition probabilities during each denoising step. This encourages the generation of physically plausible intermediate latent states while suppressing those likely to violate the given domain knowledge. We summarize main contributions as follows:
Inspired by recent success in controllable generative models 68 24
33
We introduce a novel LDM based model PreDiff for precipitation nowcasting.
We propose a general two-stage pipeline for training data-driven Earth system forecasting models. Specifically, we develop knowledge alignment mechanism to guide the sampling process of PreDiff. This mechanism ensures that the generated predictions align with domain-specific prior knowledge better, thereby enhancing the reliability of the forecasts, without requiring any modifications to the trained PreDiff model.
Our method achieves state-of-the-an performance on the N-body MNIST g dataset and attains state-of-the-an perceptual quality on the SEVIR 55
dataset.
2 Method
We follow 47 48 55
to formulate precipitation nowcasting as a spatiotemporal forecasting problem. The -tep observation represented as a spatiotemporal sequence y = [ E XX where H and W denote the spatial resolution, and C' denotes the number of mea- surements at each space-time coordinate. Probabilistic forecasting aims to model the conditional probabilistic distribution p(zly) of the stur-ste-a future xr [r'1f/a E R Hx given the observation y. In what follows, we will present the parameterization of p(aly) by controllable LDM.

---

Earthformer-UNet
Knowledge Alignment
X
Po.$ 22
Pa (Ze zt+1 Zcond)
e A -==l(u((rttt))FF000)
y
denoise
D
zond ZT
Zondd t++
cond Zt
zond Zo
XT
Figure 1: Overview of PreDiff inference with knowledge alignment. An observation sequence y is encoded into a latent context 2cond by the frame-wise encoder E. The latent diffusion model pa(2e Fee1 Acona), which is parameterized by an Earthformer-UNet, then generates the latent future 20 by autoregressively denoising Gaussian noise 27 conditioned on Æcond- It takes the concatenation of the latent context Zcond (in the blue border) and the previous-step noisy latent future žt+1 (in the cyan border) as input, and outputs 2t. The transition distribution each step from 24-1 to 2 can be further refined as Pe.el pe.o(2r (4+1; 14, Fo) via knowledge alignment, according auxiliary prior knowledge. This denoising process iterates from = T to t = O, resulting in a denoised latent future 20. Finally, 2o is decoded back to pixel space by the frame-wise decoder 'D produce the final prediction T. (Best viewed in color).
2.1 Preliminary: Diffusion Models
Diffusion models (DMs) learn the data distribution p(x) by training a model reverse predefined noising process that progressively corrupts the data. Specifically, the noising process is defined as ((ze)2y-11 A (Zx; / -1, (1 az)
T, where zo p(z) is the true data, and ØT ~ (0, I) is random noise. The coefficients a, follow a fixed schedule over the timesteps DMs factorize and parameterize the joint distribution over the data Tp and noisy latents T, as po(our) = p(x))(())))))))))())) I Po (=( where each step of the reverse denoising process is Gaussian distribution po (2, 22) = N (2e t), (o, t)), which is trained to recover #;-1 from zt-
To apply DMs for spatiotemporal forecasting, p(xly) is factorized and parameterized as po(xly) po [To:r dz) = p(æx) I Po ( y dr where ( y) represents the con- ditional denoising transition with the condition y.
2.2 Conditional Diffusion Latent Space
To improve the computational efficiency of DM training and inference, our PreDiff follows LDM to adopt a two-phase training that leverages the benefits lower-dimensional latent representations. The two sequential phases of the PreDiff training are: 1) Training frame-wise variational autoen- coder (VAE) (28 that encodes pixel space into a lower-dimensional latent space, and 2) Training a conditional DM that generates predictions in this acquired latent space.
Frame-wise autoencoder. We follow [7] to train a frame autoencoder using combination of the pixel-wise loss (e.g. L.2 loss) and an adversarial loss. Different from 77. we exclude the perceptual loss since there are standard pretrained models for perception on Earth observation data. Specifically, the encoder & is trained to encode data frame r' E H xC to latent representation 2 E(x') H. Hww.W.XX XCXC XC, The decoder D lears reconstruct the data frame z0 D(2') from

---

the encoded latent. We denote ~ p=(2)x) E RI*H,xW,xX; equivalent to = | [E(x²)] representing encoding a sequence of frames pixel space into latent spatiotemporal sequence. And pD (ar 2) denotes decoding a latent spatiotemporal sequence.
Latent difflsson. latent space as 2cond E R Iinxx H, xWa xC2 as [ The conditional distribution po (2:rr (ccod) of the latent future 2,
With the context y being encoded by the frame-wise encoder into the learned
Rhw xW, xC2 given Zcond is factorized and parameterized as (2):
Zcood PE
(1)
po (20:T)) 2cood) p(2r) ) II 12 Zcood).
(2)
where p(2T) (0, .). As proposed by 22 45 an equivalent parameterization is to have the
DMs learn match the transition noise c0(2), t) of step instead of directly predicting 2t-1. The training objective of PreDiff is simplified as shown in (
LCLDM ((c.)).te~NN100) || c0 L, Zcond) 112.
where (x, y) is a sampled context sequence and target sequence data pair, and given that, 2t. Peee and Acand ~ pe (Zcomdly))
Instantiating po (= Zcoma). Compared to images, modeling spatiotemporal observation data precipitation noweasting poses greater challenges due to their higher dimensionality. We propose replacing the UNet backbone in LDM [42 with Earthformer- UNet. derived from Earthformer's 's en- coder [8], which is known for its ability to model intricate and extensive spatiotemporal dependencies in the Earth system.
Earthformer- UNet adopts hierarchical UNet architecture with self cuboid attention as the building blocks, excluding the bridging cross-attention in the encoder-decoder architecture of Earth- former. More details of the architecture design of Earthformer- UNet are provide in Appendix [B.]] We find Earthformer-UNet to be more stable and effective modeling the transition distribution po( 1 ).omd). takes the concatenation of the encoded latent context Acood and the noisy latent future ,, along the temporal dimension as input, and predicts the one-step-ahead noisy latent future 2t-1 (in practice, the transition noise from 24 to 2--1 is predicted as shown in
2.3 Incorporating Knowledge Alignment
Though DMs hold great promise for diverse and realistic generation, the generated predictions may violate physical constraints, disregard fomain-speccngg prior knowledge, thereby fail give plausible and non-tin non-triivil results [14 44). One possible reason for this is that DMs are not necessar- ily trained on data full compliant with domain knowledge. When trained on such data, there is no guarantee that the generations sampled from the learned distribution will re- main physically realizable. The causes may also stem from the stochastic nature of chaotic systems, the approximation error in denoising steps, etc.
Algorithm 1 One training step of the knowledge alignment network Up
I: (ar, y) sampled from data
2: ~ Uniform (0, T)
2, q(2 20 )e (zolz)
4: Lu + U1 L 1, F(x, w)ll
To address this issue, we propose knowledge alignment to incorporate auxiliary prior knowledge:
F(G.y) = Foly) E Rd
into the diffusion generation process. The knowledge alignment imposes constraint F on the forecast ê. optionally with the observation y. based on domain expertise. E.g., for an isolated physical system, the knowledge E(G,-) -) = Eu km) E R imposes the conservation of energy by enforcing the generation & keep the total energy E( the same as the last observation En( The violation |F y) Fo(u)ll quantifies the deviation of prediction ~ from prior knowledge. The larger violation indicates T diverges further from the constraints. Knowledge alignment hence aims to suppress the probability of generating predictions with large violation. Notice that even the target

---

futures z: from training data may violate the knowledge. i.c. F(z, # Foly), due to noise in data collection or simulation.
Inspired by classifier guidance we achieve knowledge alignment by training a knowledge alignment network (= I, y) estimate F(G. y) from the intermediate latent 2, noising step t. The key idea is to adjust the transition probability distribution pe (+t-11 2t, Rcond) in ( during each latent denoising step to reduce the likelihood of sampling o values expected violate the constraints:
W. Fo) x po(zz|ze+11 Reond) e (1
(5)
where A= is a guidance scale factor. The knowledge alignment network is trained by optimizing the ob-
jective Lu in Alg. [ According to 4. 55 can be approximated by shifting the predicted mean of the
denoising transition ( (= L, #cond) by xx Eg V s, ( L,,,,,,, y) Fo where Eg is the variance
of the original transition distribution pe(2r 2t-1: 1: Rcood) = (2
1.: t, Zeond): Ee Ee(2)+1, t, Rooad)).
Detailed derivation is provided in Appendix C
The training procedure of knowledge alignment is outlined in Alg. 1 The noisy latent 2 for training the knowledge alignment network Us is sampled by encoding the target using the frame-wise encoder & and the forward noising process [=)] eliminating the need for an inference sampling process. This makes the training of the knowledge alignment network Ug independent of the LDM training. At inference time, the knowledge alignment mechanism is applied as a plug-in, without impacting the trained VAE and the LDM. This modular approach allows training lightweight knowledge alignment networks Ug to flexibly explore various constraints and domain knowledge, without the need for retraining the entire model. This stands as key advantage over incorporating constraints into model architectures or training losses.
3 Experiments
We conduct empirical studies and compare PreDiff with other state-of-the-art spatiotemporal fore- casting models on a synthetic dataset N-body MNIST g and a real-world precipitation nowcasting benchmark SEVIR (55] verify the effectiveness of PreDiff in handling the dynamics and uncer- tainty in complex spatiotemporal systems and generating high quality, accurate forecasts. We impose data-specific knowledge alignment: energy conservation on N--odd MNIST and anticipated precipitation intensity on SEVIR. Experiments demonstrate that PreDiff under the guidance of knowledge alignment (PreDifi-KA) is able generate predictions that comply with domain expertise much better, without severely sacrificing fidelity.
3.1
-bodyy MNIST Digits Motion Forecasting
Dataset. The Earth is a chaotic system with complex dynamics. The real-world Earth observation data, such as radar echo maps and satellite imagery, are usually not physically complete. We are unable to directly verify whether certain domain knowledge, like conservation laws of energy and momentum, is satisfied or not. This makes it difficult to verify if a method is really capable of modeling certain dynamics and adhering to the corresponding constraints. To address this, we follow [8 to generate a synthetic dataset named N N-body MNIST which is an extension of MovingMNIST [50). The dataset contains sequences of digits moving subject the gravitational force from other digits.
The governing equation for the motion is d'x. = = = ([a: [[==-=>]-=o ==nnn))))""" where æ; is the spatial dz
Gm,(n; -n;)
coordinates of the i-th digit, G is the gravitational constant, my is the mass of the j-th digit, r is a constant representing the power scale in the gravitational law, and doot is a small softening distance that ensures numerical stability. The motion occurs within 64 x 64 frame. When a digit hits the boundaries of the frame, it bounces back by elastic collision. We use N = 3 for chaotic 3-body motion (35 The forecasting task is to predict 10-step ahead future frames r E IR (00x4466xx given the 1en context y E R 10x64xx64x1 We generate 20,000 sequences for training and 1,000 sequences for testing. Empirical studies on such a synthetic dataset with known dynamics helps provide useful insights for model development and evaluation.
2 Dataset is available at httttss://eviir edu.
1 "Code available at htttps://gitub..r
5

---

Context: y
JN
J
Target:
E.MSE=0.0261
ConvI.STM
E.MSE=0.0329
Earthformer
E.M9E=0.0296
VideoGPT
E.MSE=0.0023
LDM
E.MSE=0.0354
PreDiff
E.MSE=0.0277
PreDiff-KA
N
E.MSE-0.0086
stop sep.
step
step.5
stop step
step
step
Figure 2: A set of example predictions on the N-body MNIST test set. From top to bottom: context sequence target sequence z, predictions by ConvLSTM 47), Earthformer [8 VideoGPT [65], LDM 42 PreDiff, and PreDiff with knowledge alignment (PreDiff-KA). E.MSE denotes the average error between the total energy (kinetic + potential) of the predictions E and the total energy of the last context frame Ely fis). The red dashed line is help the reader to judge the position of the digit "2" in the last frame.
Evaluation. In addition to standard metries MSE, MAE and SSIM, we also report the scores of Fréchet Video Distance (FVD) [51], metric for evaluating the visual quality of generated videos. Similar to Fréchet Inception Distance (FID) [20) for evaluating image generation, FVD estimates the distance between the learned distribution and the true data distribution by comparing the statistics of feature vectors extracted from the generations and the real data. The inception network used in FVD for feature extraction is pre-trained on video classification and is not specifically adapted for processing "unnatural videos" such as spatiotemporal observation data in Earth systems. Consequently. the FVD scores on the N N-body MNIST and SEVIR datasets cannot be directly compared with those on natural video datasets. Nevertheless, the relative ranking of the FVD scores remains meaningful indicator of model ability achieve high visual quality, FVD has shown consistency with expert evaluations across various domains beyond natural images [3] 26). Scores for all involved metrics are calculated using an ensemble of eight samples from each model.
3.1.1 Comparison with the State of the Art
We evaluate seven deterministic spatiotemporal forecasting models: UNet (55] ConyLSTM 47).
PredRNN [T], PhyDNet (TT). E3D-LSTM (60), Rainformer ( and Earthformer 8 well as two probabilistic spatiotemporal forecasting models: VideoGPT [65] and LDM [4). All basclines are trained following the default configurations their officially released code. More implementation details of baselines are provided in Appendix B.2 Results in Table show that PreDiff outperforms these baselines by a large margin in both conventional video prediction metries (i.e., MSE, MAE, SSIM), and a perceptual quality metric, FVD. The example predictions in Fig demonstratee that PreDiff generate predictions with sharp and clear digits in accurate positions. contrast, deterministic baselines resort to generating blurry predictions to accommodate uncertainty. Probabilistic baselines, though producing sharp strokes, either predict incorrect positions or fail to reconstruct the digits. The performance gap between LDM 42 and PreDiff serves as ablation study that highlights the importance of the latent backbone's spatiotemporal modeling capacity. Specifically, the Earthformer- UNet utilized in PreDiff demonstrates superior performance compared to the UNet in LDM 42
3.1.2 Knowledge Alignment: Energy Conservation
In the N-body MNIST simulation, digits move based on Newton's law of gravity, and interact with the boundaries through elastic collisions. Consequently, this system obeys the law of conservation of

---

Table I: Performance comparison on N-hody MNIST. We report conventional frame quality metries (MSE, MAE, SSIM), along with Fréchet Video Distance (FVD) [51] for assessing visual quality, Energy conservation is evaluated via E.MSE and E.MAE between the energy of predictions Eaa (3)
and the initial energy E(y""). Lower values on the energy metrics indicate better compliance with
conservation of energy.
Frame Metrics
Energy Metries
Model
#Param. (M)
MSE +
MAE SSIM
FVD 4
E.MSE L E.MAE 4
Target Persistence
0.000
0.000
1.0000
0.000
0.0132
0.0697
104.9
139.0
0.7270
168.3
UNet [55]
16.6
38.90
94.29
0.8260
142.3
ConVLSTM 47
14.0
32.15
72.64
0,8886
86.31
PredRNN 61
23.8
21.76
54.32
0.9288
20.65
PhyDNet 11]
3.1
28.97
78.66
0.8206
178.0
E3D-LSTM 60
12.9
22.98
62.52
0.9131
22.28
Rainformer
19.2
38.89
96.47
0.8036
163.5
Earthformer 3
7.6
14.82
39.93
0.9538
6.798
VidcoGPT 65
92.2
53.68
77.42
0.8468
39.28
0.0228
0.1092
LDM [42]
410.3
46.29
72.19
0.8773
3.432
0.0243
0.1172
PreDiff PreDifi-KA 129.4
120.7
9,492
25.01
0.9716
0.987
0.0226
0. 0,1083
21.90
43,57
0.9303
4.063
0.0039
0.0443
energy. The total energy of the whole system E(x") at any future time step j during evolution should equal the total energy at the last observation time step E(y"").
We impose the law of conservation of energy for the knowledge alignment on N-body MNIST in the form of
F(a. v) [E(**) ) ..... E( Foly) = [E(y"m)..... E(ytu))²
(6)
(7)
The ground-truth values the total energy Ely (in) and E(x²) are directly accessible since N N-body MNIST is a synthetic dataset from simulation. The total energy can be derived from the velocities (kinetic energy) and positions (potential energy) of the moving digits. A knowledge alignment network Us is trained following Alg. to guide the PreDiff to generate forecasts T that conserve the same energy as the initial step "
To verify the effectiveness of the knowledge alignment on guiding the generations to comply with the law of conservation of energy, we train an energy detector Eder (3 (a that of the forecasts T. We evaluate the energy error between the forecasts and the initial energy using E.MSE(2, y) MSE( MSE(Eade (T), E E(y"")) and E.MAE y) MAE MAe(Edd(()) (d) E(y )). In this evalu- ation, we exclude the methods that generate blurred predictions with ambiguous digit positions. We only focus on the methods that are capable of producing clear digits in precise positions.
that detects the total energy
As illustrated in Table I PreDiff-KA substantially outperforms all baseline methods and PreDiff without knowledge alignment in E.MSE and E.MAR. This demonstrates that the forecasts of PreDiff- KA comply much better with the law of conservation of energy, while still maintaining high visual quality with an FVD score of 4.063.
Furthermore, we detect energy errors in the target data sequences. The first row of Table I indicates that even the target from the training data may not strictly adhere to the prior knowledge. This could be due to discretization errors in the simulation. Table shows that all baseline methods and PreDiff have larger energy errors than the target, meaning purely data-oriented approaches cannot eliminate the impact of noise in the training data. In contrast, PreDiff-KA guided by the law of conservation of energy, overcomes the intrinsic defects in the training data, achieving even lower energy errors compared to the target.
A typical example shown in Fig. 2 demonstrates that while PreDiff precisely reproduces the ground- truth position of digit "2" in the last frame (aligned to the red dashed line), resulting in nearly the same energy error (E.MSE- = 0.0277) as the ground-truth's 's (E.MSE = 0.0261), PreDiff- successfully
"The test MSE of the energy detector is 5.56 x 10 which is much smaller than the scores of E.MSE shown in Table This indicates that the energy detector has high precision and reliability for verifying energy conservation in the model forecasts.
7

---

-50 Min
440 Min
-30 Min
-20 Min
-10 Min
Min
Context: y
31-59
100-133
Target: X
133-160
181-219
219-255
ConvLSTM
CSI - 0.339
CSI-p16-0.47 0.471
Earthformer
CSI = 0.339
CS1-p16=0.455 0.454
VideoGPT
CSI = 0.309
CSI-p16= 0.566
LDM
CSI = 0.266
CSI-p16-0.50 0.568
PreDiff
CSI - 0.329
CSI-p16 0.619
+10 Min
-20 Min
-30 Min
+40 Min
+50 Min
+60 Min
Figure 3: A set of example forecasts from baselines and PreDiff the SEVIR test set. From top
to bottom: context sequence y, target sequence æ, forecasts from ConvLSTM 47 Earthformer [8]
VideoGPT LDM 42). PreDiff.
corrects the motion of digit "2", providing it with physically plausible velocity and position (slightly off the red dashed line). The knowledge alignment ensures that the generation complies better with the law of conservation of energy, resulting in much lower E.MSE 0.0086. On the contrary, none of the evaluated baselines overcome the intrinsic noise from the data, resulting in energy errors comparable to or larger than that of the ground-truth.
Notice that the pixel-wise scores MSE, MAE and SSIM are less meaningful for evaluating PreDiff- KA, since correcting the noise the energy results changing the velocities and positions of the digits. A minor change in the position of digit can cause a large pixel-wise error, even though the digit is still generated sharply and in high quality as shown in Fig. 2
3.2 SEVIR Precipitation Nowcasting
Dataset. The Storm EVent ImageRy (SEVIR) (55] is spatiotemporal Earth observation dataset which consists of 384 km x 384 km image sequences spanning over 4 hours. Images in SEVIR are sampled and aligned across five different data types: three channels (C02, CO9, C13) from the C0 advanced baseline imager, NEXRAD Vertically Integrated Liquid (VIL) mosaics, and GOES-16 Geostationary Lightning Mapper (GLM) flashes. The SEVIR benchmark supports scientific research multiple meteorological applications including precipitation nowcasting, synthetic radar generation, front detection, etc. Due to computational resource limitations, we adopt downsampled version of SEVIR for benchmarking precipitation nowcasting. The task is to predict the future VIL up to 60 minutes (6 frames) given 70 minutes of context VIL (7 frames) at spatial resolution of 128 128, i.e. 16x128x y 7 xxx22xx12xx
Evaluation. Following ation, which is commonly used in precipitation nowcasting and is defined as CSI
55 3
we adopt the Critical Succeas Index (CSI) for evalu-
WHsta #Misses Alarms To count the #Hits (truth=1, pred=l), #Misses (truth=1, pred=0) and #F AFF.Aarmm (truth=0, pred=1). the prediction and the ground-truth are rescaled to the range 0 255 and binarized at thresholds [16, 74, 133, 160, 181, 219]. We also follow 41 to report the CSI at pooling scale 4 x 4 and 16 x 16, which evaluate the performance on neighborhood aggregations multiple spatial scales. These pooled CSI metries assess the models' ability to capture local pattern distributions. Additionally, incorporate FVD [5] and continuous ranked probability score
#Hits

---

-50 Min
40 Min
-30 Min
-20 Min
-10 Min
Min
0:16
Context: y
16:31
31-59
74.100
100-133
Target: X Avgint - 0.137
133-160
181-219
219-255
+40, Avgint = 0.201
+20, g Avgint 0.192
PreDiff
AvgInt 0.126
-20, r Avgint 0.104
-40 Avgint 0.067
-10 Min
-20 Min
-30 Min
+40 Min
+50 Min
+60 Min
Figure 4: A set of example forecasts from PreDiff-KA, i.e., PreDiff under the guidance of anticipated average intensity. From top to bottom: context sequence y, target sequence x, forecasts from PreDiff and PreDiff-KA showcasing different levels of anticipated future intensity (Jrr + no)), where n takes the values of 4, 2, 22 -4.
(CRPS) [ for assessing the visual quality and uncertainty modeling capabilities of the investigated methods. CRPS measures the discrepancy between the predicted distribution and the true distribution. When the predicted distribution collapses into single value, deterministic models, CRPS reduces to Mean Absolute Error (MAF). A lower CRPS value indicates higher forecast accuracy. Scores for all involved metries calculated using ensemble of eight samples from each model.
3.2.1 Comparison to the State of the Art
We adjust the configurations of involved basclines accordingly and tune some of the hyperparameters for adaptation the SEVIR dataset. More implementation details of baselines are provided in Appendix [B.2] The experiment results listed Table 2 show that probabilistic spatiotemporal forecasting methods are not good achieving high CSI scores. However, they are more powerful at capturing the patterns and the true distribution of the data, hence achieving much better FVD scores and CSI-pool16, 16. Qualitative results shown in Fig. demonstrate that CSI is not aligned with human perceptual judgement. For such complex system, deterministic methods give up capturing the real patterns and resort averaging the possible futures, i.e. blurry predictions, to keep the scores from appearing too inaccurate. Probabilistic approaches, of which PreDiff is the best, though are not favored by per-pixel metrics, perform better at capturing the data distribution within local area, resulting in higher CSI-po CSI-pool16 16, lower CRPS, and succeed in keeping the correct local patterns, which can crucial for recognizing weather events. More detailed quantitative results on SEVIR are provided in Appendix D
3.2.2 Knowledge Alignment: Anticipated Average Intensity
Earth system observation data, such as the Vertically Integrated Liquid (VIL) data in SEVIR, are usu- ally not physically complete, posing challenges for directly incorporating physical laws for guidance. However, with highly flexible knowledge alignment mechanism, we can still utilize auxiliary prior knowledge to guide the forecasting effectively. Specifically for precipitation nowcasting on SEVIR, we use anticipated precipitation intensity to align the generations to simulate possible extreme weather events. We denote the average intensity of data sequence as 7(x:) E RR In order to estimate the con- ditional quantiles of future intensity, we train a simple probabilistic time series forecasting model with parametric (Gaussian) distribution p, (I()) )]) = Ni,, (I I[), o, ((I())))) that predict

---

Table 2: Performance comparison on SEVIR. The Critical Success Index, also known as the intersection over union (IoU), is calculated at different precipitation thresholds and denoted as CSI-thresh. CSI reports the mean of CSI-[16, 74, 133, 160, 181, 219). CSI-pools with s 4 and = 16 report the CSI at pooling scales of 4 x 4 and 16 x 16. Besides, we include the continuous ranked probability score (CRPS) for probabilistic forecast assessment, and the scores of Fréchet Video Distance (FVD) for evaluating visual quality.
Metries
Model
#Puram. (M)
FVD 4
CRPS
C8I-pool11
CSI-poo116
Persistence
525.2
0,0526
0.2613
0,3702
0.4690
UNet [5S]
16.6
753.6
0.0353
0.3593
0.4098
0,4805
ConvLSTM 27
14.0
659.7
0,0332
0,4185
0,4452
0.5135
PredRNN [
46.6
663.5
0.0306
0,4080
0.4497
0.5005
PhyDNet
13.7
723.2
0,0319
0.3940
0,4379
0,4854
F3D-LSTM 66
35.6
600.1
0.0297
0.4038
0,4492
0.4961
Rainformer
184.0
760.5
0,0357
0,3661
0,4232
0,4738
Farthformer
690.7
0.4419
0.5005
DGMR [4]
71.5
485.2
0,0435
0,2675
0.3431
0.4832
VideoGPT 65
99.6
261.6
0.0381
0,3653
1143449
0.579%
LDM E
438.6
133.0
0.0280
0.3580
0,4022
0.5522
Preliff
220.5
33.05
0.0246
0,4100
0.4624
0.6244
PreDiff-KA (C 20.)20.))
229.4
34.18
the distribution of the average future intensity T(x) given the average intensity of each context frame (xr))) (abbreviated as [1)")) D). By incorporating F(I, y) = (TT) and Fo(y) = + no, for knowledge alignment, PreDiff-KA gains the capability of generating forecasts for potential extreme cases, e.g., where (()) falls outside the typical range of μ4, = dr.
Fig. shows set of generations from PreDiff and PreDiff-KA with anticipated future intensity 14+ + 107,, n E [-4,- 2, This qualitative example demonstrates that PreDiff is not only capable of capturing the distribution of the future, but also flexible at highlighting possible extreme cases like rainstorms and droughts with the knowledge alignment mechanism, which is crucial for decision-making and precaution.
According to Table 2, the FVD score of PreDiff-KA (34.18) is only slightly worse than the FVD score of PreDiff (33.05). This indicates that knowledge alignment effectively aligns the generations with prior knowledge while maintaining fidelity and adherence to the true data distribution.
Conclusions and Broader Impacts
In this paper, we propose PreDiff, novel latent diffusion model for precipitation nowcasting. We also introduce a general two-stage pipeline for training DL models for Earth system forecasting. Specifically, we develop knowledge alignment mechanism that is capable of guiding PreDiff to generate forecasts in compliance with domain-specific prior knowledge. Experiments demonstrate that our method achieves state-of-the-art performance on N N-body MNIST and SEVIR datasets.
Our work has certain limitations: 1) Benchmark datasets and evaluation metrics for precipitation nowcasting and Earth system forecasting still maturing compared the computer vision domain. While we utilize conventional precipitation forecasting metrics and visual quality evaluation, aligning these assessments with expert judgement remains an open challenge. 2) Effective integration of physical principles and domain knowledge into DL. models for precipitation nowcasting remains an active research area. Close collaboration between DL researchers and domain experts in meteorology and climatology will be key to developing hybrid models that effectively leverage both data-driven learning and scientific theory. While Earth system observation data have grown substantially in recent years, high-quality data remain scarce in many domains. This scarcity can limit PreDiff's s ability accurately capture the true distribution, occasionally resulting in unrealistic forecast hallucinations under the guidance of prior knowledge as it attempts to circumvent the knowledge alignment mechanism. Further research on enhancing the sample efficiency of PreDiff and the knowledge alignment mechanism is needed.
In conclusion, PreDiff represents promising advance in knowledge-aligned DL for Earth system forecasting, but work remains to improve benchmarking, incorporate scientific knowledge, and boost model robustness through collaborative research between AI and domain experts.
10

---

References
[1] Cong Bai, Feng Sun, Jinglin Zhang, Yi Song, and Shengyong Chen. Rainformer: Features extraction balanced network for radar-based precipitation noweasting. IEEE Geoscience and Remote Sensing Letters, 19:1-5, 2022.
[2] Kaifeng Bi, Lingxi Xie, Hengheng Zhang, Xin Chen, Xiaotao Gu. and Qi Tian. Accurate medium-range global weather forecasting with 3d neural networks. Nature, pages 1-6, 2023.
[3] Kang Chen, Tao Han, Junchao Gong, Lei Bai, Fenghua Ling, Jing-Jia Luo, Chen, Leiming Ma, Tianning Zhang, Rui Su, et al. Fengwu: Pushing the skillful global medium-range weather forecast beyond 10 days lead. arXiv preprint arXiv:2304. 02948, 2023.
[4] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8777 2021.
[5] Lasse Espeholt, Shreya Agrawal, Casper Sønderby, Manoj Kumar, Jonathan Heek, Carla Bromberg. Cenk Gazen, Jason Hickey, Aaron Bell, and Nal Kalchbrenner. Skillful twelve hour precipitation forecasts using large context neural networks. arXiv preprint arXiv:2111.07470. 2021.
[6] Patrick Esser, Johnathan Chiu, Parmida Atighehchian, Jonathan Granskog, and Anastasis Germanidis. Structure and content-guided video synthesis with diffusion models. arXiv preprint arXiv:2302.03011 2023.
[7] Patrick Esser, Robin Rombach, and Bjorn Ommer. Taming transformers for high-resolution image synthesis. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12873-12883, 2021.
[8] Zhihan Gao, Xingjian Shi, Hao Wang, Yi Zhu, Yuyang Wang, Mu Li, and Dit- Yan Yeung. Earthformer: Exploring space-time transformers for earth system forecasting. In NeurlPS, 2022.
[9] Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359-378 2007.
[10] Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in neural information processing systems, 27, 2014.
[11] Vincent Le Guen and Nicolas Thome. Disentangling physical dynamics from unknown factors for unsupervised video prediction. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pages 11474-11484 2020.
[12] John Guibas, Morteza Mardani, Zongyi Li, Andrew Tao, Anima Anandkumar, and Bryan Catanzaro. Adaptive fourier neural operators: Efficient token mixers for transformers. arXiv preprint arXiv:2111.133882 2021.
[13] Shantanu Gupta, Hao Wang, Zachary Lipton, and Yuyang Wang. Correcting exposure bias for link recommendation. In ICML, 2021.
[14] Derek Hansen, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, and Michael W. Mahoney. Learning physical models that can respect conservation laws. In Proceedings of the t0th of International Conference on Machine Learning, volume 202, 2023.
[15] William Harvey, Saeid Naderiparizi, Vaden Masrani, Christian Weilbach, and Frank Wood. Flexible diffusion modeling of long videos. arXiv preprint arXiv:2205. 2022.
[16] Yusuke Hatanaka, Yannik Glaser, Geoff Galgon, Giuseppe Torri, and Peter Sadowski. Diffusion models for high-resolution solar forecasts. arXiv preprint arXiv:2302.001 00170, 2023.
[17] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
11

---

[18] Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 1606v 08415, 2016.
[19] Hans Hersbach, Bill Bell, Paul Berrisford, Shoji Hirahara, András Horányi, Joaquín Muñoz- Sabater, Julien Nicolas, Carole Peubey, Raluca Radu, Dinand Schepers. et al. The era5 global reanalysis. Quarterly Journal of the Royal Meteorological Society, 146(730):1999- 1999:999999 2020.
[20] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[21] Geoffrey Hinton, Nitish Srivastava, Alex Krizhevsky, Ilya Sutskever, and Ruslan R Salakhut- dinov. Improving neural networks by preventing co-adaptation of feature detectors. arXiv preprint arXiv: 207.0580, 2012.
[22] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 126840-68511 2020.
[23] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. arXiv preprini 0rX4v:2204.03458, 2022.
[24] Lianghua Huang, Di Chen, Yu Liu, Yujun Shen, Deli Zhao, and Jingren Zhou. Com- poser: Creative and controllable image synthesis with composable conditions. arXiv preprint arXiv:2302.09778, 2023.
[25] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448 pmlr, 2015.
[26] Kevin Kilgour, Mauricio Zuluaga, Dominik Roblek, and Matthew Sharifi. Fry, 'echet audio dis- tance: A metric for evaluating music enhancement algorithms, arXiv preprint arXiv:1812.08466 2018.
[27] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1121v:1412.6980, 2014.
[28] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114, 2013.
[29] Remi Lam, Alvaro Sanchez-Gonzalee Matthew Willson, Peter Wirnsberger, Meire Fortunato, Alexander Pritzel, Suman Ravuri, Timo Ewalds, Ferran Alet, Zach Rasenn et al. Graphcast: Learning skillful medium-ran maangge global weather forecasting. arXiv preprint arXiv:2212. 12794, 2022.
[30] Jussi Leinonen, Ulrich Hamann, Daniele Nerini, Urs Germann, and Gabriele Franch. Latent diffusion models for generative precipitation nowcasting with accurate uncertainty quantification. arXiv preprint arXiv:2304. / 12891, 2023.
[31] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17X1v1171..05101, 2017.
[32] Zhengxiong Luo, Dayou Chen, Yingya Zhang, Yan Huang, Liang Wang, Yujun Shen, Deli Zhao, Jingren Zhou, and Tieniu Tan. Videofusion: Decomposed diffusion models for high-quality video generation. arXiv e-prints, pages arXiv-2303, 2023.
[33] Franccois Mazé and Faez Ahmed. Topodiff: A performance and constraint-guided diffusion model for topology optimization. arXiv preprint 0rX1v.2208.09551, 2022.
[34] Lu Mi, Hao Wang, Yonglong Tian, and Nir Shavit. Training-free uncertainty estimation for neural networks. In AAAI, 2022.
[35] Valtonen MJ, Mauri Valtonen, and Hannu Karttunen. The three-body problem. Cambridge University Press, 2006.
12

---

[36] Haomiao Ni, Changhao Shi, Kai Li, Sharon X. Huang, and Martin Renqiang Min. Conditional image-to-video generation with latent flow diffusion models, 2023.
[37] Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, al. FourCastNet: A global data-driven high-resolution weather model using adaptive fourier neural operators. arXiv preprint arXiv:2202. / 1214, 2022.
[38] Kristina Preuer, Philipp Renz, Thomas Unterthiner, Sepp Hochreiter, and Gunter Klambauer. Fréchet chemnet distance: metric for generative models for molecules in drug discovery. Journal of chemical information and modeling,
2018.
[39] Ruslan Rakhimov, Denis Volkhonskiy, Alexey Artemov, Denis Zorin, and Evgeny Burnaev. Latent video transformer. arXiv preprint arXiv:2006. /0704, 2020.
[40] Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06/25 2022.
[41] Suman Ravuri, Karel Lene, Matthew Willson, Dmitry Kangin, Remi Lam, Piotr Mirowski, Megan Fitzsimons, Maria Athanassiadou, Sheleem Kashem, Sam Madge, et al. Skilful pre- cipitation nowcasting using deep generative models of radar. Nature, 597(7878):6226677, 2021.
[42] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[43] Nataniel Ruiz, Yuanzhen Li, Varun Jampani, Yael Pritch, Michael Rubinstein, and Kfir Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 22500-225 10, 2023.
[44] Nadim Saad, Gaurav Gupta, Shima Alizadeh, and Danielle C. Maddix. Guiding continuous operator learning through physics-based boundary constraints. In Proceedings of the 11 th International Conference Learning Representations, 2023.
[45] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv: 1487, 2022.
[46] Tim Salimans, Andrej Karpathy, Xi Chen, and Diederik P Kingma. Pixelenn++: Improving the pixelcnn with discretized logistic mixture likelihood and other modifications. arXiv preprint arXiv:1701.000 075 2017.
[47] Xingjian Shi, Zhourong Chen, Hao Wang, Dit- Yeung, Wai-Kin Wong, and Wang-chun Woo. Convolutional LSTM network: A machine learning approach for precipitation noweasting. In NeurIPS. volume 28, 2015.
[48] Xingjian Shi, Zhihan Gao, Leonard Lausen, Hao Wang, Dit- Yan Yeung, Wai-kin Wong, and Wang-chun Woo. Deep learning for precipitation nowcasting: A benchmark and new model. In NeurIPS, volume 30, 2017.
[49] Casper Kaae Sønderby, Lasse Espeholt, Jonathan Heek, Mostafa Dehghani, Avital Oliver, Tim Salimans, Shreya Agrawal, Jason Hickey, and Nal Kalchbrenner. Metnet: A neural weather model for precipitation forecasting, arXiv preprint arXiv:2003.12140, 2020.
[50] Nitish Srivastava, Elman Mansimov, and Ruslan Salakhudinov. Unsupervised learning of video representations using LSTMs. In ICML, pages 843-852. PMLR, 2015.
[51] Thomas Unterthiner, Sjoerd van Steenkiste, Karol Kurach, Raphaël Marinier, Marein Michalski, and Sylvain Gelly. Fvd: A new metric for video generation. In DGS@ ICLR, 2019.
13

---

[52] Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling in latent space. In Neural Information Processing Systems (NeurIPS). 2021.
[53] Aaron Van den Oord, Nal Kalchbrenner, Lasse Espeholt, Oriol Vinyals, Alex Graves, et al. Con- ditional image generation with pixelenn decoders, Advances in neural information processing systems, 29, 2016.
[54] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurlPS, volume 30, 2017.
[55] Mark Veillette, Siddharth Samsi, and Chris Mattioli. SEVIR: storm event imagery dataset for deep learning applications in radar and satellite metcorology. Advances in Neural Information Processing Systems, 33:22009-22019, 2020.
[56] Vikram Voleti, Alexia
and Christopher Pal. Masked conditional video diffusion for prediction, generation, and interpolation. arXiv preprint arXiv:2205.09835. 2022.
[57] Hao Wang, SHI Xingjian, and Dit- Yan Yeung. Natural- parametee networks: A class of proba- bilistic neural networks. In NIPS, pages 118-126, 2016.
[58] Hao Wang and Dit- Yan Yeung. Towards bayesian deep learning: A framework and some existing methods. TDKE, 2813333395-33 2016.
[59] Hao Wang and Dit-Yan Yeung. survey on bayesian deep learning, CSUR, 53(5):1-37, 2020.
[60] Yunbo Wang, Lu Jiang. Ming-Hsure Yang. Li-Jia Li, Mingsheng Long, and Li Fei-Fei. Eidetic
3D LSTM: A model for video prediction and beyond. In International conference on learning representations, 2018.
[61] Yunbo Wang, Haixu Wu, Jianjin Zhang, Zhifeng Gao, Jianmin Wang, Philip Yu, and Mingsheng Long. PredRNN: A recurrent neural network for spatiotemporal predictive learning, IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
[62] Ziyan Wang and Hao Wang. Variational imbalanced regression: Fair uncertainty quantification via probabilistic smoothing. In NeurlPS, 2023.
[63] Dirk Weissenborn, Oscar Täckström, and Jakob Uszkoreit. Scaling autoregressive video models. In International Conference on Learning Representations, 2019.
[64] Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pages 3-19, 2018.
[65] Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas. VideoGPT: Video generation using vg-vae and transformers. arXiv preprint arXiv:2104.10157 2021.
[66] Sihyun Yu, Kihyuk Sohn, Subin Kim, and Jinwoo Shin. Video probabilistic diffusion models in projected latent space. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[67] Ye Yuan, Jiaming Song, Umar Iqbal, Arash Vahdat, and Jan Kautz. Physdiff: Physics-guider human motion diffusion model. arXiv preprint arXiv:2212.02500. 2022.
[68] Lvmin Zhang and Maneesh Agrawala. Adding conditional control to text-to-image diffusion models. arXiv preprint arXiv:2302.05543. 2023.
[69] Lu Zhou and Rong-Hur Hua Zhang. A self-attention based neural network for three-dimensional multivariate modeling and its skillful enso predictions. Science Advances, 910)xaa112877 2023.
14

---

A Related Work
Deep learning for precipitation nowcasting In recent years, the field of DL has experienced remarkable advancements, revolutionizing various domains of study. including Earth science. One area where DL has particularly made significant strides is in the field of Earth system forecasting, especially precipitation nowcasting. Precipitation noweasting benefits from the success of DL. architectures such as convolutional neural networks (CNNs), recurrent neural networks (RNNs), and Transformers, which have demonstrated their effectiveness in handling spatiotemporal tensors, the typical formulation for Earth system observation data. ConvLSTM [47). a pioneering approach in DL for precipitation nowcasting, combines the strengths of CNNs and LSTMs processing spatial and temporal data. PredRNN [6] builds upon ConvLSTM by incorporating spatiotemporal memory flow structure. E3D-LSTM [60] integrates 3D CNN to LSTM to enhance long-term high-level relation modeling. PhyDNet [I] incorporated partial differential equation (PDE) constraints in the latent space. MetNet [49] and its successor, MetNet-2 [51. propose architectures based on CSnVLSTM
and dilated CNN, enabling skillful precipitation forecasts up twelve hours ahead. DGMR AT takes an adversarial training approach to generate sharp and accurate nowcasts, addressing the issue of blurry predictions.
In addition to precipitation nowcasting, there has been a surge in the modeling of global weather and medium-range weather forecasting due to the availability of extensive Earth observation data, such as the European Centre for Medium- Range Weather Forecasts (ECMWF)'s ERAS 19 dataset. Several DL-based models have emerged in this area. FourCastNet [37 proposes architecture with Adaptive Fourier Neural Operators (AFNO) [2] as building blocks for autoregressive weather Transformer-
forecasting. Fe ( introduces forecast model that achieves skillful forecasts up to ten days ahead. GraphCast [29] combines graph neural networks with convolutional LSTMs to tackle sub-seasonal forecasting tasks, representing weather phenomena as spatiotemporal graphs. Pangu-Weather Weather (2) proposes a 3D Transformer model with Earth-specific priors and a hierarchical temporal aggregation strategy for medium-range global weather forecasting. While recent years have seen remarkable progress in DL for precipitation nowcasting, existing methods still face some limitations. Some methods are deterministic, failing to capture uncertainty and resulting in blurry generation. Others lack the capability of incorporating prior knowledge, which is crucial for machine learning for science. In contrast, PreDiff' captures the uncertainty in the underlying data distribution via diffusion models, avoiding simply averaging all possibilities into blurry forecasts. Our knowledge alignment mechanism facilitates post-training alignment with physical principles and domain-specing prior knowledge.
global medium- -range weather
Diffusion models Diffusion models (DMs) 22 are class of generative models that have become increasingly popular in recent years. DMs learn the data distribution by constructing a forward process that adds noise to the data, and then approximating the reverse process to remove the noise. Latent diffusion models (LDMs) [42] are variant of DMs that are trained on latent vector outputs from variational autoencoder. LDMs have been shown to be more efficient in both training and inference compared to original DMs. Building on the success of DMs in image generation, DMs have also been adopted for video generation. MCVD 566] trains a DM by randomly masking past and/or future frames in blocks and conditioning on the remaining frames. It generates long videos by autoregressively sampling blocks of frames in a sliding window manner. PVDM (66) projects videos into Iin-eemnniional latent space as 2D vectors, and presents a joint training of unconditional and frame conditional video generations. LFDM [6 employs a flow predictor to estimate latent flows between video frames and learns a DM for temporal latent flow generation. VideoFusion 32 decomposes the transition noise in DMs into per-frame noise and the noise along time axis, and trains two networks jointly to match the noise decomposition. While DMs have demonstrated impressive performance in video synthesis, its applications to precipitation nowcasting and other Earth science tasks have not been well explored. Hatanaka et al. [16] uses DMs to super-resolve coarse numerical predictions for solar forecast. Concurrent to our work, LDCast 30 applies LDMs for precipitation nowcasting. However, LDCast has not studied how to integrate prior knowledge to the DM, which is a unique advantage and novelty of PreDiff.
Conditional controls on diffusion models Another key advantage of DMs is the ability to condition generation on text, class labels, and other modalities for controllable and diverse output. For instance, ControlNet (68) enables fine-tuning pretrained DM by freezing the base model and training a copy
15

---

end-to-ent with conditional inputs. Composer [24] decomposes images into representative factors used as conditions to guide the generation. Beyond text and class labels, conditions in other modalities. including physical constraints, can also be leveraged to provide valuable guidance. TopDiff [33 constrains topology optimization using load, boundary conditions, and volume fraction. Physdiff trains physics-based based motion projection module with reinforcement learning to project denoised motions in diffusion steps into physically plausible ones. Nonetheless, while conditional control has proven to be powerful technique in various domains, its application in DL for precipitation nowcasting remains an unexplored area.
16

---

B Implementation Details
All experiments are conducted on machines with NVIDIA AOGG GPUs (24GB memoery). All models, including PreDiff, knowledge alignment networks and the baselines, can fit in a single GPU without the need for gradient checkpointing or model parallelization.
B.1 PreDiff
Frame-wise autoencoder We follow 7 42 to build frame-wise VAES (not VQVAFs) and train them adversarially from scratch on N-body MNIST and SEVIR frames. As shown in Sec. 2.2 on N-body MNIST dataset, the spatial downsampling ratio is 4 x 4. A frame x E R 664x6411 is encoded to 2² _16x16x3 by parameterizing pe (= (x²) (Ng (x")))((("")) On SEVIR dataset, the spatial downsampling ratio is 8 x 8. A frame xd E IR 128x x28x is encoded to 2 E R 16x 16x4 similarly.
The detailed configurations of the encoder and decoder of the VAE on --body MNIST are shown in Table 3 and Table 4 The detailed configurations of the encoder and decoder of the VAE on SEVIR are shown in Table and T The discriminators for adversarial training on N-body MNIST and SEVIR datasets share the same configurations, which are shown in Table
Latent diffusion model that instantiates pa(ze pa(z111 Zcond) Stemming from Earthformer (8. we build Porthformmr-UNet which is a hierarchical UNet with self cuboid attention [8] layers basic building
not used
2D CNN Upsample
blocks, as shown in Fig. 5
Cuboid Attention Block X.D
On N-body MNIST, it takes the concatenation along the temporal dimension (the sequence length axis) of Acond E R 16x 16x3 and 2, I 1210016x16x 16x3 and outputs 2 E 100016x16xx
res connect
Upsample
input,
Cuboid Attention BD
On SEVIR, it takes the
Downsample
concatenation along the temporal dimension (the sequence length axis) of #cond E 1xx16x16xx and 2, E 1g6x16x16xx input, and outputs 2:-1 R6x16x16x4 Besides, we add the embedding of the denoising step / to the state in front of each cuboid attention block via an embeding layer TEmbed, following (22). The detailed configurations of the Earthformer-UNG is described in Table
XM
Cuboid Attention XD
2D-CNN Downsample
Zeond
2
Figure 5: Earthformer-UNe architec- ture. PreDiff employs an Earthformer- UNet as the backbone for param- eterizing the latent diffusion model (o(2elz44 Zzomd). It takes the concate- nation of the latent context 2cend (in the blue border) and the previous-step noisy latent future 2++1 (in the cyan border) along the temporal dimension (the se- quence length axis) as input, and outputs 2t. (Best viewed in color).
Knowledge alignment networks A knowledge align- ment network parameterizes Up (2e,t, y) to predict F(ū, y) using the noisy latent 24. In practice, we build an Earth- former encoder [8] with a final pooling block as the knowl- edge alignment network to parameterize Us( f, #cona), which takes f,, and the concatenation of Zcond and 2t, in- stead of f, and 2, as the inputs. We find this imple- mentation accurate enough when is small. The detailed configurations of the knowledge alignment network is de- scribed in Table
Optimization We train the frame-wise VAFs using the Adam optimizer [27 following 7 We train the latent Earthformer- and the knowledge alignment network using the Adam optimizer [31] following (8). Detailed configurations are shown in Table 10 Table [I and Table I2 for the frame-wise VAE, the latent Earthformer-
and the knowledge alignment network, respectively. We adopt data parallel and gradient accumulation to use a larger total batch size while the GPU can only afford a smaller micro batch size.
17

---

Table 3: The details of the encoder of the frame-wise VAE on N-body MNIST frames. It encodes an input frame x² R24x64x1 into latent 2 1 Conv3 the 2D convolutional layer with 3 x kernel. GroupNorm32 is the Group Normalization (GN) layer (64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function SiLU(z)
sigmoid(x)) The Attention is the self attention layer [54] that first maps the input to queries Q, keys K and values V by three Linear layers, and then does self attention operation: Attention(z)
Softmax(QK3 /VC)V).
Block
Layer
Resolution
Channels
Input x 2D CNN
64 x 64
1
Conv3 x 3
64 x 64
1 - 128
GroupNorm32
x 64
128
Conv3 x 3
64 64
128
ResNet Block x2
GroupNorm32
64 X 64
128
Conv3 x 3
64 x 64
128
SiLU
64
128
Downsampler
Conv3 X 3
64
64 32 32
128
GroupNorm32
32 x 32
128
Conv3 x 3
32 x 32
128 - 256, 256
ResNet Block x2
GroupNorm32
32 x 32
256
Conv3 3
32 32
256
SiLU
32 x 32
256
Downsampler
Conv3 x 3
32 X 32 -> 16 x 16
256
GroupNorm32
16 X 16
256
Conv3 x 3
X 16
256
> 512, 512
ResNet Block x2
GroupNorm32
16 16
512
Conv3 x 3
16 x 16
512
SiLU
16 x 16
512
GroupNorm32
16 x 16
512
Self Attention Block
Attention
16 16
512
Linear
16 x 16
512
GroupNorm32
16 x 16
512
Conv3 x 3
x 16
512
ResNet Block x2
GroupNorm32
16 16
512
Conv3 x 3
16 x 16
512
SiLU
16 x 16
512
GroupNorm32
16 x 16
512
SiLU
16 16
512
Output Block
Conv3 x 3
16 x 16
512 6
Conv3 x 3
16 x 16
6
18

---

Table 4: The details of the decoder of the frame-wise VAE N-body MNIST frames. It de- codes a latent 2² E 16 16x3 back to a frame in pixel space E R64464x1, Conv3 3 is the 2D convolutional layer with 3 x 3 kernel. GroupNorm32 is the Group Normalization (GN) layer (44) with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function SiLU(x) sigmoid(z). The Attention is the self attention layer [54] that first maps the input to queries Q, keys K and values V by three Linear layers, and then does self attention operation: Attention(z) Sof toof QK² /VC)V).
Block
Layer
Resolution
Channels
Input z
16 16
Conv3 3
16 16
3
2D CNN
Conv3 3
16 x 16
3 + 512
GroupNorm32
16 16
512
Self Attention Block
Attention
16 x 16
512
Linear
16 16
512
GroupNorm32
16 16
512
Conv3
16 16
512
ResNet Block x3
GroupNorm32
16 16
512
Conv3
16 16
512
SiLU
16 16
512
Upsampler
Conv3 x 3
16 x 16 - 32 x 32
512
GroupNorm32
32 32
512
Conv3 3
32 x 32
512
256, 256, 256
ResNet Block
GroupNorm32
32 32
256
Conv3 3
32 x 32
256
SiLU
32 32
256
Upsampler
Conv3
32
32 64 x 64
256
GroupNorm32
64 x 64
256
Conv3
64 64
256 + 128, 128, 128
ResNet Block x3
GroupNorm32
64 x 64
128
Conv3
64 64
128
SiLU
64 64
128
GroupNorm32
64 64
128
Output Block
SiLU
64 64
128
Conv3 3
64 X 64
128 -
19

---

Table The details of the encoder of the frame-wise VAE on SEVIR frames. encodes an input frame x¹ E R 122xxx22xx into latent z E 1616116x4 Conv3 x 3 is the 2D convolutional layer with 3 x kernel. GroupNorm32 is the Group Normalization (GN) layer 64 with 32 groups. SiLU the Sigmoid Linear Unit activation layer [88] with function SiLU(r) = z- gm)id(x).. The Attention the self attention layer [44 that first maps the input to queries Q, keys and values by three
Linear layers, and then does self attention operation: Attention (n)))) Softmax (Qka Boft
Block
Layer
Resolution
Channels
Input x 2D CNN
128 x 128
Conv3
128 128
128
GroupNorm32
128 X 128
128
Conv3 x
128 x 128
128
ResNet Block x2
GroupNorm32
128 128
128
Conv3 x 3
128 x 128
128
SiLU
128 x 128
128
Downsampler
Conv3
128 X 128 64 64
128
GroupNorm32
64 X 64
128
Conv3 x
64 X 64
128 + 256, 256
ResNet Block x2
GroupNorm32
64 64
256
Conv3 x 3
64 X 64
256
SiLU
64 x
256
Downsampler
Conv3
64 X 64 - 32 32
256
GroupNorm32
32 x 32
256
Conv3
32 x 32
256 + 512, 512
ResNet Block x2
GroupNorm32
32 32
512
Conv3 x
32 x 32
512
SiLU
32 x
512
Downsampler
Conv3
32 x 32 - 16 16
512
GroupNorm32
16 X 16
512
Conv3
x
512
ResNet Block x2
GroupNorm32
16 16
512
Conv3 x
16 X 16
512
SiLU
16
512
GroupNorm32
16 x 16
512
Self Attention Block
Attention
16 x 16
512
Linear
16 x 16
512
GroupNorm32
16 16
512
Conv3 x
16 x 16
512
ResNet Block x2
GroupNorm32
x 16
512
Conv3
16 x 16
512
SiLU
16 x 16
512
GroupNorm32
16 16
512
SiLU
16 x 16
512
Output Block
Conv3
16 x 16
512 +
Conv3
16 x 16
8
20

---

Table 6: The details of the decoder of the frame-wise VAE SEVIR frames. It decodes latent x² 216x16x4 back to a frame in pixel space x E 12 Conv3 is the 2D convolutional layer with 3 x kernel. GroupNormm22 is the Group Normalization (GN) layer [64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function SiLU(z) = x sigmoid(r)
The Attention is the self attention layer [54] that first maps the input to queries Q, keys K and values V by three Linear layers, and then does self attention operation: Attention(z) = Softmax(QK" /VC)).
Block
Layer
Resolution
Channels
Input zd
16 x 16
4
Conv3 x 3
16 x 16
4
CNN
Conv3 3
16 x 16
4
512
GroupNorm32
16 x 16
512
Self Attention Block
Attention
16 x 16
512
Linear
16 x 16
512
GroupNorm32
16 x 16
512
Conv3 x 3
16 x 16
512
ResNet Block x3
GroupNorm32
16 x 16
512
Conv3 3
x
512
SiLU
16 x 16
512
Upsampler
Conv3 3
16 16 >> 32 32
512
GroupNorm32
32 x 32
512
Conv3 x
x
512
ResNet Block x3
GroupNorm32
32 x 32
512
Conv3 3
32 x 32
512
SiLU
32 x 32
512
Upsampler
Conv3 3
32
32 64 64
512
GroupNorm32
64 x 64
512
Conv3 3
64 x 64
512 -> 256, 256, 256
ResNet Block x3
GroupNorm32
64 x 64
256
Conv3 x 3
64 x 64
256
SiLU
64 x
256
Upsampler
Conv3 x
64 x 64 + 128 128
256
GroupNorm32
128 x 128
256
Conv3 x 3
128 x 128
256 -
128, 128, 128
ResNet Block x3
GroupNorm32
128 x
128
Conv3 x 3
128 x 128
128
SiLU
128 x 128
128
GroupNorm32
128 x 128
128
Output Block
SiLU
128 x 128
128
Conv3 x 3
128 x 128
128 -
21

---

Table 7:
The details of the discriminator for the adversarial loss of on N-body MNIST and SEVIR frames. Conv4 4 is the 2D convolutional layer with 4 x kernel, or 1 x stride, and x padding. BatchNorm the Batch Normalization (BN) layer 25 is 0.2.
The negative slope in LeakyReLU
Resolution
Block
Layer
Channels
N N-body MNIST
SEVIR
Input x)
64 64
128 128
2D CNN
Conv4 4
64 64 +
32 x 32
128 128 + 64
+ 64
LeakyReLU
32 32
64 64
64
Downsampler
Conv4 4
32
16
64 64
32
64 + 128
BatchNorm
16 16
32 32
128
LeakyReLU
16 16
32 32
128
Downsampler
Conv4
16 16 + 8
32 32
16
128 256
BatchNorm
8 8
16 x 16
256
LeakyReLU
8 8
16 16
256
Downsampler
Conv4
16 16
15
256
512
BatchNorm
7 7
15 x 15
512
LeakyReLU Conv4 4
7 x 7 7 6 x
15 x 15
512
Output Block
15 15
14 14
AvgPool
6 6 6 +1
15 15 +1
22

---

Table The details of the Earthformer- UNet as the latent diffusion backbone on N-body MNIST and SEVIR datasets. The ConcatMask layer for the Observation Mask block concatenates one more channel the input indicates whether the input the encoded observation Acond or the noisy latent 2,. for Scond and for 2,. Conv3 X 3 is the 2D convolutional layer with 3 x 3 kernel. GroupNorm32 the Group Normalization (GN) layer 64 with 32 groups. If the number of the input data channels is smaller than 32, then the number of groups is set to the number of channels. SiLU is the Sigmoid z sigmoid(r:. The negative slope in LeakyReLU is 0.1. Dropout is the dropout layer [21] with the probability 0.1 to drop element to zeroed. The FFN consists of two Linear layers separated by GeLU activation layer [18] PosEmbed is the positional embedding layer [54] that adds learned positional embeddings to the input. TEmbed is the embedding layer [22] that embeds the denoising step t. PatchMerge splits p patches and merges the spatial
Linear Unit activation layer [8] with function Si)U(x)
2D input tensor with C' channels into N non-overlapping dimensions into channels, gets N 1 x patches with p C² channels and concatenates them back along spatial dimensions. Residual connections [17] are added from blocks in the downsampling phase to corresponding blocks the upsampling phase.
Block
Layer
Spatial Resolution
Channels N-body MNIST
SEVIR
Input Foond
16 16
3
Observation Mask
ConcatMank
16 x
16
3 3-
5
GroupNorm32
16 16
S1LU
16
Conv3
16 16
4
256
256
Projector
GroupNorn32
16 16
256
S1LU
Dropout
16
256
Conv3 x
16 16
256
Positional Embedding
PosEmbed
16 16
256
TEabed
x
256
LayerNorm Cuboid(T, 1,
256
16 16
256
FFN
16 16
256
LayerNorm Cuboid(1,
Cuhoid Attention Block x4
256
FFN
16 16
256
LayerNorm
16 16
256
FFN
16
256
PatchMerge
16 16 8
256
1024
Downsampler
LayerNorm
8
1024
Linear
1024
TEnbed
1024
LayerNorm
x 8
1024
Cuboid(T, 1,
8 8
1024
FFN
1024
LayerNorm Cuboid(1,
1024
Cuboid Attention Block x8
8 x 8
1024
FFN
1024
LayerNorm Cuboid(1, 1.
1024
1024
FEN
8 8
1024
Upsampler
NearestleeighoorIncerp Conv3 x
8 - 16 x 16
1024
16 16
1024 - 256
TEnbed
1 16
256
LayerNorm
16
256
Cuboid(T,
16 16
256
FFNI
16 16
256
LayerNorm Cuboid(1,H,
Cuboid Attention Block x4
16
256
FFN
16 16
256
LayerNorm
16 16
256
(1boid(1,1.W)
256
FFN
16 16
256
Output Block
Linear
16 16
256 -
256 4
23

---

Table 9: The details Earthformer encoders for the parameterization networks Uz Rcond) on N-body MNIST and SEVIR datasets. The ConcatMask layer for the Observation Mask block concatenates one more channel to the imput to indicates whether the input the encoded observation Acond the noisy latent 24. 1 for 2codd and 0 for 2,. Conv3 x 3 is the 2D convolutional layer with 3 x kernel. GroupNorm32 is the Group Normalization (GN) layer [64] with 32 groups. If the number of the input data channels is smaller than 32, then the number of groups is set to the number of channels. SiLU the Sigmoid Linear Unit activation layer [8] with function (i)
knowledge alignment
sigmoid(x). The negative slope in LeakyReLU is 0.1. Dropout is the dropout layer 21 with the probability 0.1 to drop an element to be zeroed. The FFN consists of two Linear layers separated GeLU activation layer [8]. PosEmbed the positional embedding layer (54] that adds learned positional embeddings to the input. TEmbed is the embedding layer [22] that embeds the denoising step PatchMerge splits 2D input tensor with C' channels into N non-overiapping
p patches and merges the spatial dimensions into channels, gets N 1 x patches with p². C' channels and concatenates them back along spatial dimensions. Residual connections [17] are added from blocks in the downsampling phase to corresponding blocks in the upsampling phase. The Attention is the self attention layer [54] with an extra "cls" token for information aggregation. It first flattens the input and concatenates it with the "cls" token. Then it maps the concatenated input to queries Q, keys K and values V by three Linear layers, and then does self attention operation: Attention on (x) = Sof toftmaxx(Ki /VC)V). Finally, the value of the "cls" token after self attention
operation serves the layer's output.
Channels
Block
Layer
Spatial Resolution
N-body MNIST
SEVIR
Input Reond, 2t
16 x 16
4
Observation Mask
ConcatMask
16
4 - 5
GroupNorm32
16
5
SiLU
16 x 16
5
Conv3 x
16 x 16
4 +
64
5 + 64
Projector
GroupNorm32
16 16
64
SiLU
16 16
64
Dropout
16
Conv3 x
16 x 16
64
Positional Embedding
PosEmbed
16 x 16
64
TEmbed
16 16
LayerNorm
16
Cuboid
16
FFN
16 x 16
64
LayerNorm
16 16
64
Cuboid Attention Block
Cuboid(1,H,
16 16
64
FFN
16
LayerNorm
16
64
Cuboid
16 x 16
64
FFN
16 x 16
64
PatchMerge
16
8
64 + 256
Downsampler
LayerNorm
8
256
Linear
8 8 8
256
TEmbed
8 8 8
256
LayerNorm
8 x 8
256
Cuboid 1,
8 8
256
FFN
LayerNorm
8 x 8
256
Cuboid Attention Block
Cuboid(1,H,
8 x 8
256
FFN
8 8 8 8
256
LayerNorm
256
Cuboid(1,1,W)
8 8
256
FFN
8 8
256
GroupNorm32
8 8
256
Output Pooling Block
Attention
8 8
256
Linear
256

---

Table 10: Hyperparameters the Adam optimizer for training frame-wise VAEs and discriminators on N-body MNIST and SEVIR datasets.
Haper-parameter of VAE
Value
Learning rate
4.5 10-6
B1
0,5
B2
0.9
Weight decay
10-2
Batch size
512
Training epochs
200
Hyper-parameter of discriminator
Value
Learning rate
4.5 10 6
B1
0.5
B3
0.9
Weight decay
10
Batch size
512
Training epochs
200
Training start step
50000
Table 11: Hyperparameters of the AdamW optimizer for training LDMs N-body MNIST and
SEVIR datasets.
Hyper-parameter of VAE
Value
Learning rate
1.0 x 10-3
31
0.9
0.999
Weight decay
10-5
Batch size
64
Training epochs
1000
Warm up percentage
10%
Learning rate decay
Cosine
Table 12: Hyperparameters of the AdamW optimizer for training knowledge alignment networks on N N-body MNIST and SEVIR datasets.
of VAE
Value
Learning rate
1.0 x 10
31
0.9
0.999
Weight decay
10-5
Batch size
64
Training epochs
200
Warm up percentage
10%
Learning rate decay
Cosine
25

---

B.2 Baselines
We train baseline algorithms following their officially released configurations and the learning rate, learning rate scheduler, working resolution, etc., to optimize their performance on each dataset. We list the modifications we applied the baselines for each dataset in Table I3
Table 13: Implementation details of baseline algorithms. Modifications based on the officially released implementations are listed according to different datasets. "_" means modification is applied. "reverse enc-dec" means adopting the reversed encoder-decoder architecture proposed in 48 Other terms listed are the hyperparameters in their officially released implementations.
Model
N-body MNIST
SEVIR
UNet [55]
reverse enc-dec 48
enc-dec 48
conv cony_kemels- (3)33)
cony cony_kernels (5.5).
ConvLSTM 47
deconv kemnels = (4444)
accon_kernels = ((6.6))(4.4))(4444
channels=196, 128, 256]
channels=[96, 128, 256]
PredRNN 61
PhyDNet I|
convcell_hiddes [256, 256, 256, 64]
E3D-LSTM 60
factors= downscaling_factors-22... 2, 2, 2| hidden dim=32
downscaling__fatooss444 2,
Rainformer I
heads=(4, 4, 8. 16]
head_dim=tt
Earthformer
context context_steps 7
DGMR 41
forecast steps
vqvae n_nnn____ co_ccoee=== 512
vqvae_n __cod_s
VideoGPT 65
vqvae downsample [1, 4, 4]
vqvae downsample [1, 8, 8]
vae: x > 16 3
vae: 128 128 - 16 x 16
LDM 42
conv 1ii_dii== 3
conv jomy_dim= =
godel_channels 256
model_ ceaenelchnnels 256
26

---

C Derivation of the Approximation to Knowledge Alignment Guidance
We derive the approximation the knowledge alignment guided denoising transition (
following 4 We rewrite 5 to (8) using a normalization constant Z that normalizes z f
dzc 1:
pe, (2 = pe pe(zelze+1 Zoond) Ze Aur -
(8)
In what follows, we abbreviate pe 2++1 1: t, Zcond) as yo, and Eg(2 (2(22-1: Zeond) as Eg for brevity. We use C,, i = fl..... 7) to denote constants.
(=(2e Acand) N(uo Ed);
log (2) 2++1, Zoond) = 7(21 pe) 20)² p (zt - pe) + C1.
(9)
log Ze A= -= ->>FlUu t.g) Fo(y)) C2.
By assuming that log Ze A> (U)()tt:))) has low curvature compared to Mg which is reason- 0), can approximate by a Taylor expansion
able in the limit of infinite diffusion steps [| at 2t 140
log Ze Ax=1CCo(4dd.))-Fuu0
vs ->=1 L, Fo(y)lla. (2: ua)A= v. v. Uol t, - Fo
(10)
(zt - no)g C3.
where g Ax (.UUllU(( Fo(y)lllz=pp- By taking the log of (8) and applying the results
from and (TO), we get
log (2el2e11)1 = loogopppz (2 + log Ze -
(2t ma) T 2 (z+ ue) + (ze - (44) + C4
1
1 E + C3 (11)
2 (at M4 Σag) Eg (z1 μrl Eeg)
2
2 (2t 14 Σeg) TEg (zt
p4 - Eng) + Cs
= loog p(2) + C7, 2 ~ Nips + Eeg, Me)
Therefore, the transition distribution under the guidance of knowledge alignment shown (5) can be approximated by Gaussian similar to the transition without knowledge guidance, but with its mean shifted by Eg.
27

---

D More Quantitative Results on SEVIR
D.1 Quantitative Analysis of BIAS on SEVIR
Similar Critical Success Index (CSI) introduced in Sec. 3.2 BIAS #Bita+ Alame
# is
calculated by counting the #Hits (truth=1, pred=1), #Misses (truth=1, pred=0) and #F #F.Alarms (truth=O, pred=1) of the predictions binarized at thresholds [16, 74, 133, 160, 181, 219]. This mea- surement assesses the model's inclination towards either F. F.Alarms Misses.
The results from Table 14 demonstrate that deterministic spatiotemporal forecasting models, such as UNet [55], ConvLSTM 147,, PredRNN [61], PhyINNtt [11]. E3D-LSTM [60], and Earthformer (8), tend to produce predictions with lower intensity. These models prioritize avoiding high-intensity predictions that have higher chance of being incorrect due their limited ability handle such un- certainty effectively. On the other hand, probabilistic spatiotemporal forecasting baselines, including DGMR [41). VideoGPT [65] and LDM
42 demonstrate more daring approach by predicting pos- sible high-intensity signals, even if it results in lower CSI scores, as depicted in Table 2 Among these baselines, PreDiff achieves the best performance in BIAS. It consistently achieves BIAS scores closest to 1, irrespective of the chosen threshold. These results demonstrate that PreDiff has effectively learned to unbiasedly capture the distribution of intensity.
Table 14:
Quantitative Analysis of BIAS on SEVIR. The BIAS is calculated at differ- precipitation thresholds and denoted BIAS-thresh. BIAS-m reports the mean of BIAS-[16, 74, 133, 160, 181, 219). A BIAS score closer to indicates that the model less biased to either F.Alarms or Misses. The best BIAS score is in boldface while the second best underscored.
Metrics
Model
BIAS-m
BIAS-219
BIAS-181
BIAS-160
BIAS-133
BIAS-74
BIAS-16
Persistence
1.0177
1.0391
1.0323
1.0258
1.0099
1.0016
0.9983
UNet [55]
0.6658
0.2503
0.4013
0.5428
0.7665
0.9551
1.0781
CSnvL.TTM
0.8341
0.5344
0.6811
0.7626
0.9643
0.9957
1.0663
PredRNN 61
0.6605
0.2565
0.4377
0.4909
0.6806
0,9419
1.1554
PhyDNet |1
0.6798
0.3970
0,6593
0.7312
1.0543
1.0553
1.2238
E3D-LSTM 660
0.6925
0,2696
0.4861
0,5686
0.8352
0,9887
1.0070
Earthformer
g
0.7043
0.2423
0,4605
0,5734
0.8623
0.9733
1.1140
DGMR [4
0.7302
0.3704
0.5254
0.6495
0.8312
0.9594
1.0456
VideoGPT 65
0.8594
0.6106
0.7738
0,8629
0.9606
0,9681
0,9805
LDM 42
1.2951
1.4534
1.3525
1.3827
1.3154
1.1817
1.0847
PreDiff
0.9769
0.9647
0.9268
0.9617
0.9978
1.0047
1.0058
28

---

D.2 CSI at Varying Thresholds on SEVIR
We include representative deterministic methods ConvLSTM and Earthformer, and all studied probabilistic methods to compare CSI, CSI, CSI-poo14 and CSI-poo116 at varying thresholds. It is important to note that CSI tends favor conservative predictions, especially in situations with high levels of uncertainty. To ensure fair comparison, we calculated the CSI scores by averaging the samples for each model, while scores in other metrics are averaged over the scores of each sample. The results presented in Table 15
16 17 demonstrate that our PreDiff achieves competitive CSI
scores and outperforms baselines in CSI scores at pooling scale x and 16 16, particularly at higher thresholds.
Table 15: CSI thresholds [16, 74, 133, 160, 181, 219 on SEVIR.
Model
Metrics
CBI-219
09I-181
CSI-160
CSI-133
09I-16
ConvI.STM 0.4185
0.1220
0,2381
0.2905
0.4135
0,6846
0.7510
Earthformer
0.4419
0.1791
0.2848
0.3232
0.4271
0.6860
0.7513
DGMR [
0.2675
0.0151
0.0537
0.0970
0.2184
0.5500
0.6710
VideoGPT
0.3653
0.1029
0.1997
0.2352
0.3432
0,6062
0.7045
LDM 42
0.3580
0.1019
0.1894
0.2340
0.3537
0,5848
0.6841
Table 16: CSI-poo14 at thresholds [16, 74, 133, 160, 181, 219) on SEVIR.
Model
CSI:poo14-219
CSI poo14-181
Metries
CSI-poc14-133>
pocL44747
CoMLSSTM R7
0.4152
0.1850
0.2861
0.3245
0,4502
0.6694
0,7556
Earthformer
0.4567
0.1484
0.2972
0.3341
0,4911
0.7000
0,7892
DGMR [T]
0,3431
00114
0.1191
0,1950
0.3452
0.6302
0.7273
1696
0.442
0.6529
Preliff
0.4624
0,2065
0.3130
0.3613
0,4807
0.6691
0.7438
Table 17: CSI-poo116 166 at thresholds [16, 74, 133, 160, 181, 219] on SEVIR.
Model
CaI pool16
CSI pool16
CBE poo116 181
CSI poo216 180
Menics
CBI pool16 133
CS2 pool16 74
CSI pool16 16
Com-LSTMM
DGMR
0.4543
29

---

E More Qualitative Results on N-body MNIST
Fig. to Fig. 13 show several sets of example predictions on the N-body MNIST In each figure, visualizations from top to bottom are context sequence target sequence æ, predictions by ConvLSTM 47], Earthformer [8, VideoGPT [5], LDM 42,, PreDiff, PreDifI-KA. E.MSE denotes the average error between the total energy (the sum of kinetic energy and potential energy) of the predictions ( and the total energy of the last step context Ely
Contexl:
5
5
Target:
5
E.MSE=0.0094
2_
Convi
0.S0=000090
Barthformer
E.SSE-0.100
VideoGPT
E.MSE 0.0488
IDM
EMSE=0.0115
ProDief
EMSE-000120
PreDiT-KA
Mepp
Arp
wept
Mpp
Figure 6: A set of example predictions on the N-body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "1" in the last frame.
1
50
ConvLSTM
B.SE-8.0198
Earthformer
0.SE=0.0241
VideoGPT
F. 0.MSE-000006
LDM
EMSE-B.037
PreDict
RMSE=000256
65
PreDid-KA F. J.MSE- -0001
m98
mm3
mg4
mo
mBc
mep
mid
inpg
Figure 7: A set of example predictions the N N-body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "0" in the last frame.
30

---

Context:
0
Targel: x
E.MSE-0.0128
ConvLSTM
3
B.MEE0000106
3
E.MSE-0.006
8
EMSE=0.0201
PreDict B.MSE-0.112
3
3
PreDiiTKA E.ME-0.0008
3
3
4
Aep
Figure 8: A set of example predictions on the N-body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "0" in the last frame.
Context: y
%
796
Targel: x E.MSE-B.0155
CoMvLSTM
EMSE-8.016
Barthformer
H.MSE= 0.0113
VideoGPT
7
7
7
E.MSEBB.00..
Figure 9: A set of example predictions on the body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "8" in the last frame.
31

---

Context:
Targel: x
E.MSE-0.0160
ConvLSTM
J
E.SE0000117
9
3
3
3
E.MSE-0.0488
2
3
3
3
3
3
PreDict B.MSE-0.0154
5
3
3
3
3
3
3
PreDitIKA E.MSE-0.002
33
9
Figure 10: set of example predictions on the N-body MNIST set. The red dashed line is to help the reader to judge the position of the digit "4" in the last frame
Context: y
a
Targel: x F.MSE-B.O0
COMVLSTM
H.MEE 0.02000024
Barthformer H.MSE
3
VideoGPT
E.MSE -.000-
#a
13
8
13
3
E.MSE-0.0012
Figure 11: A set of example predictions on the N-body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "1" in the last frame.
32

---

Context: y
7
E.MSE-0.0184
Targel: x
10
8
9
E.MSE00.0151
ConvLSTM
70
B
7
9
6
10
B
3
9
07
E.MSE-0.0797
P
9
6
LO
E.SE-000362
00
L1
B.MSE-0.0189
PreDict
0
4
PreDiTTKA E.MSE-0.000
7
4
9
01
set of example predictions on the N-body MNIST set. The red dashed line is to help the reader to judge the position of the digit "7" in the last frame.
Figure 12:
Context: y
Targel: x
E.MSE-B.00D
w
CoRvLSTM
EMSE-8.001
2
P
E.MS 0.0025
Barthformer
72
3
7
V
#
33
VideoGPT
F.MSE-B.0447
29
Figure 13: A set of example predictions on the N-body MNIST test set. The red dashed line is to help the reader to judge the position of the digit "7" in the last frame.
33

---

More Qualitative Results on SEVIR
Fig. 14 Fig. 19 show several of example predictions on the SEVIR test set. In subfigure (a) of each higure, visualizations from top to bottom are context sequence y, target sequence x, predictions
subfigure (b) of each figure, visualizations from top to bottom are context sequence y, target sequence predictions by PreDiff-Ke -A with anticipated average future intensity H4, + 71 4, 2, 0 2, -4.
by ConvLSTM 477, Earthformer [8), VideoGPT 655 LDM 42, PreDiff, PreDiff-KA.
Tept
Madia
Maddes
Marda
Mali
Figure 14: A set of example predictions on the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDiff-KA under the guidance of anticipated average intensity.
CE-Aa Gigie-C
Di PEEEEE
Incheac
Figure 15: A set of example predictions the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDiff-KA under the guidance of anticipated average intensity.
34

---

D4 piu
Figure 16: A set of example predictions the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDiff-KA KA under the guidance of anticipated average intensity.
Typs
-2gg
Gigic-44
1W
Of DEpIEEET
Ind
Figure 17: A set of example predictions the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDifT-KA under the guidance of anticipated average intensity.
35

---

DI piu
Figure 18: A set of example predictions the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDiff-KA under the guidance of anticipated average intensity.
Of pIEEE IN
(b)
Figure 19: A set of example predictions the SEVIR test set. (a) Comparison of PreDiff with baselines. (b) Predictions by PreDifT-KA under the guidance of anticipated average intensity.
36


Reviews:
[{"review_id": "1qvPM5sy10", "reviewer": "Reviewer_Ntzp", "summary": "This work proposes a latent diffusion model called PreDiff for spatiotemporal forecasting. The main contributions of this paper are two-fold: 1) they use a conditional latent diffusion model to handle the uncertainty of spatiotemporal forecasting. 2) They introduce an explicit knowledge control mechanism to align the model with domain-specific prior knowledge.\n", "questions": "Please see the weakness.", "limitations": "Please see the weakness.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper is well-written and easy to follow.\n2. The design of the knowledge control mechanism is interesting. It is decoupled with the latent diffusion model, thus providing more flexibility to the whole framework.\n", "weaknesses": "1. My main concern is the experiments. This paper does not conduct enough ablation studies, and the effectiveness of the diffusion model cannot be supported by experiments shown in Tables 1 & 2.\n2. Table 2 only reports the mean of CSI at different thresholds, which is insufficient to prove the superiority of the proposed method. CSI with a high threshold correlates with strong weather, thus, it is important in this work to prove that PreDiff can generate better forecasts, which is also one of the main claims of this paper. I suggest authors list CSI at different thresholds in Table 2 to make a comprehensive comparison.\n3. The parameter number of PreDiff is significantly higher than that of EarthFormer (about 15 times). Thus the comparison in Tables 1 & 2 is unfair. So is it possible to reduce the number of parameters to make a fair comparison with other methods?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for providing detailed feedback. We hope that our response below can address your main concerns.\n\n**W1.1: ablation studies**\n\nwe have conducted **several experiments serving the role of ablatiion studies**. We first ablate the two aspects that differentiate PreDiff from LDM [1]:\n+ We replace the UNet backbone in LDM with Earthformer-UNet, derived from Earthformer's encoder [2]. Prior work has demonstrated Earthformer's advantages over UNet for spatiotemporal modeling [2]. We leverage these optimized configurations from [2] and adjust the latent channels to match LDM, without further tuning. As mentioned in Section 3.1.1, **the performance gap between LDM and PreDiff served as an ablation study** that highlights the importance of the latent backbone's spatiotemporal modeling capacity (line 211).\n+ We propose knowledge control to guide PreDiff to generate forecasts better aligned with domain knowledge. The performance gap between PreDiff and PreDiff-KC on E.MSE and E.MAE in Table 1 serves as an ablation study demonstrating the quantitative effectiveness of knowledge control.\n\nFurthermore, we conduct **two additional experiments that ablate the knowledge control mechanism**:\n+ We evaluate trained PreDiff on a cleaner N-body MNIST dataset (N-body MNIST Clean) with less energy noise. As shown in Table R1, PreDiff still suffers from energy noise, fails to complies with the law of energy conservation, and therefore generates forecasts deviating from the ground truth, which does comply with the law of energy conservation. This results in poor MSE, MAE, and SSIM. In contrast, PreDiff-KC better complies with the law of energy conservation, which holds in the new cleaner dataset. By generating forecasts closer to the ground truth, PreDiff-KC achieves better pixel-wise metrics.\n+ We evaluate PreDiff-KC using the ground-truth future average intensity as guidance on SEVIR. As shown in Table R2, PreDiff-KC-GT outperforms PreDiff on CRPS and CSI.\n\nThese ablation studies demonstrate that knowledge control guides the purely data-driven model to generate forecasts better aligned with ground-truth. We will mark out the ablation study clearly in our revision.\n\n**W1.2: effectiveness of the diffusion model**\n\nWe include LDM [1] as a strong diffusion model baseline to study the effectiveness of diffusion models on spatiotemporal forecasting. As shown in Tables 1 and 2, LDM outperforms other probabilistic baselines including DGMR and VideoGPT. This demonstrates the **superiority of diffusion models over GANs (DGMR) and autoregressive models (VideoGPT)** on this task. LDM also achieves much lower FVD than deterministic baselines, highlighting its **ability to generate high-fidelity forecasts**. However, LDM does not surpass deterministic baselines on pixel-wise metrics. This suggests limitations in LDM's modeling of spatiotemporal correlations. \n\nTo address this, we propose PreDiff, which replaces LDM's latent backbone with an Earthformer-UNet. PreDiff achieves state-of-the-art performance compared to both deterministic and probabilistic baselines **on both FVD and pixel-wise metrics**.\n\n**W2: list CSI at different thresholds**\n\nThis is a good suggestion. As advised, **we have added Figure R1,R2,R3** in the attached PDF in our General Response, which presents CSI for our method and baselines at varying thresholds. These figures show that our PreDiff consistently achieves competitive CSI across all studied thresholds. Notably, PreDiff excels at capturing the data distribution within a local area, resulting in higher CSI-pool4 and CSI-pool16 scores across all thresholds. We will include these results in our revision.\n\n**W3: the number of parameters**\n\nOur PreDiff framework replaces the UNet backbone in LDM [1] with an Earthformer-UNet, while retaining the rest configurations. This substitution alone reduces the number of parameters in LDM from 410.3M to 120.7M, **a significant decrease**. As shown in Tables 1 and 2, probabilistic models like DGMR, VideoGPT and LDM typically require more parameters to sufficiently model uncertainty. \n\nOn the other hand, **naively scaling up deterministic models like Earthformer is not straightforward**. Our empirical studies show that increasing either the latent channels (from 128 to 256, 512) or model depth (from [1, 1] to [2, 2], [4, 4]) results in inferior performance. For instance, an **Earthformer** variant with **570M parameters** (channels=512, depth=[4, 4]) **performs worse** than an **Earthformer** with **15M parameters** in Table 2 (0.4332 versus 0.4419). More parameters do not necessarily improve performance. \n\n[1] Rombach, Robin, et al. \"High-resolution image synthesis with latent diffusion models.\" Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.\n\n[2] Gao, Zhihan, et al. \"Earthformer: Exploring space-time transformers for earth system forecasting.\" Advances in Neural Information Processing Systems 35 (2022): 25390-25403."}, {"review_id": "mzQKTh7b8C", "reviewer": "Reviewer_sDg9", "summary": "This paper introduces a new method for doing precipitation nowcasting. The method is based on building a conditional diffusion model in a global latent space, which is found separately in an autoencoding fashion. Additionally, a “knowledge control” mechanism is introduced to help ensure that the generated samples satisfy known physical constraints. Numerical validations are conducted on a synthetic and a real precipitation dataset.", "questions": "* Ln 47 - “AR models” refer to autoregressive models? I didn’t see this acronym defined \n* Equation (1) - what is $p_\\mathcal{E}$ exactly? I presume $y$ gets encoded through $\\mathcal{E}$ and how do you draw $z_\\text{cond}$?\n* Ln 196: “.. by comparing the statistics of feature vectors extracted from real data and generations” - what is this real data? The references cited on ln 201 (in audio and chemical applications) both use embeddings trained on their respective domain-specific datasets. It is meaningless to compare against references not related to the specific tasks since the feature distributions will be completely different.\n*  Ln 209-210: “.. generate predictions with sharp and clear digits” - this is too handwavy to me. Although the statement looks qualitatively true, a quantitative measure would make it much more rigorous. A potential option is to look at Fourier spectrum  -‘blurry’ images tend to show underestimated energy in high frequency modes.\n* Ln 299-300: “.. highlighting possible extreme cases like rainstorms and droughts” - this is not easy to verify by looking at Figure 3 for someone without domain knowledge. Additional labels or captions would help.\n* I do not fully understand the purpose of knowledge control discussed in section 3.2.2. Is it only meant to encourage the generation of extreme scenarios? Do you have ways to evaluate them besides FVD?", "limitations": "The authors mention limitations related to potential misalignment between visual quality and scores, as well as dependence on data quality and quantity.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The problem of precipitation nowcasting is highly significant and has a lot of potential for real-life impact.\n* The paper is generally well written. I find it easy to understand. Tables and figures are informative.\n* Lots of details for reproducibility.", "weaknesses": "* The model itself is not the most novel. Both latent diffusion and knowledge control (similar to classifier guidance) are not new. This is somewhat compensated by the fact that the model is being applied to a specific and impactful domain. \n* Despite the seemingly rich set of baselines being compared, some of the most recent SOTA models developed specifically for the precipitation nowcasting problem are missing (despite being mentioned in the paper): Metnet2 (ref 5), DGMR (ref 28) - these models were applied to real weather data, but one should be able to adapt them to the 2nd benchmark dataset considered here. It is important to know where the model stands compared to these, especially considering that the methodology alone does not qualify for accept due to lack of novelty.\n* The use of the FVD metric is questionable (see questions below).\n* Value of knowledge control on the precipitation dataset is unclear.\n* A few other places that need further clarifications (see questions section).\n\nI feel in its current state these outweigh the strengths so I vote for borderline reject.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for providing detailed feedback. We hope our response below can address your main concerns.\n\n**W1: the model itself is not the most novel**\n\nOur literature review covers related works from three perspectives: deep learning for Earth system forecasting (line 30-42), diffusion models (line 43-58), and controllable generative models (line 59). To the best of our knowledge, we are the first to propose a controllable latent diffusion model for precipitation nowcasting. \n\nAs for our knowledge control mechanism, there are **fundamental differences from classifier guidance**. While classifier guidance requires data-label pairs for training the classifier, we directly formulate the domain knowledge as Eq.(4), without the need of additional categorical labels. The training and inference of our model thus can be **entirely label-free**.\n\nMost importantly as emphasized in our General Response, we propose a **novel two-stage pipeline** for training data-driven Earth system forecasting models. This pipeline incorporates both data-driven learning and constraints from domain knowledge, representing a **promising direction for combining data-driven methods with physical alignment**.\n\n**W2: some SOTA models are missing**\n\nWe apologize for misspelling DGMR as DGRM in Table 2. Actually, we **did evaluate DGMR (ref 28)** for SEVIR precipitation nowcasting tasks and included it for performance comparison in Table 2.\n\nDGMR is outperformed by other probabilistic baselines and our PreDiff by a large margin. The causes can be two-fold:\n+ GAN-based models are well known to be difficult to optimize. The training of DGMR is highly unstable and struggles to converge to an optimal solution.\n+ The basic building blocks of the DGMR generator are ConvGRU, which is an outdated architecture and may limit its capacity.\n\nAs for MetNet-2, we excluded it in our paper because the **evaluation methodology differs**. Directly evaluating MetNet-2 on SEVIR would be unfair. Specifically, according to its official formulation (Figure 2(a) in the MetNet-2 paper), it is designed to generate forecasts for only a center crop of the observed area. \n\nNevertheless, we adapted MetNet-2 for SEVIR by zero padding the observations with other configurations the same as specified in Appendix D of the MetNet-2 paper. As shown in Table R2, MetNet-2 does not perform well on SEVIR in this setting.\n\n**W3&Q3: FVD metric is questionable**\n\n\"Real data\" refers to the target sequences in the test data.\n\nReferences adapt FID for audios and molecules by training inception networks on domain-specific datasets, as generated examples **cannot be visually evaluated**. In contrast, spatiotemporal forecasting datasets like N-Body MNIST and SEVIR can be directly visualized as videos, where the distribution discrepancy in extracted features is **highly correlated with visual quality**. Thus, calculating FVD with a video-pretrained inception network provides meaningful insights into generation quality. \n\n**W4: value of knowledge control on the precipitation**\n\nTo clarify, in Section 3.2.2, we **train a simple probabilistic time series forecasting model $p_\\tau()$** that predicts the future average intensity $I(x)$ with mean $\\mu_\\tau$ and variance $\\sigma_\\tau^2$ (line 290). The simple $p_\\tau()$ is introduced not as an accurate forecasting tool, but rather to showcase how the knowledge control could inform PreDiff of probable future states. By probabilistically predicting the future average intensity $I(x)$ using $p_\\tau()$, we can simulate how weather unfolds in extreme cases. For example, Figure 3(b)'s Row 3 simulates how weather unfolds in an extreme case (with probability around 0.35%) where the future average intensity $I(x)>\\mu_\\tau+4\\sigma_\\tau$. Such simulation can be valuable for goverments to evalaute possible future damage in extreme cases. \n\n**Q1: definition of \"AR\"**\n\nThanks for catching that omission. We'll clearly define \"autoregressive\" (AR) on first use in our revision.\n\n**Q2: Equation (1)**\n\nAs introduced in paragraph **Frame-wise autoencoder (line 106-107)**, the distribution $p_{\\mathcal{E}}$ is parameterized by the trained frame-wise encoder $\\mathcal{E}$. $z\\sim p_{\\mathcal{E}}(z|x)$ denotes encoding a sequence of frames $x$ in pixel space into latent spatiotemporal states $z$ via $\\mathcal{E}$. Similarly, a context sequence of frames $y$ can be encoded into latent context $z_{\\text{cond}}$ by $\\mathcal{E}$ as shown in Eq.(1).\n\n**Q4: quantitative measure on \"sharp and clear digits\"**\n\nSorry for the confusion. We would like to clarify that we are **not making comparative claims about sharpness** in line 208, but rather highlighting that deterministic models produce blurry forecasts (line 209), while PreDiff retains digits' shape as demonstrated in Figure 2. \n\nOur **quantitative** evaluation focuses on general **visual quality**, and therefore uses **standard metrics FVD and SSIM** to measure video-level and frame-level sharpness respectively. As shown in Table 1, PreDiff achieves significantly lower FVD and higher SSIM scores compared to baselines, verifying that our method generates less blurry and more visually appealing forecasts.\n\n**Q5: clarify extreme cases in Figure 3**\n\nThanks for your suggestion. We will clarify that **\"AvgInt\" labels** refer to \"the average forecast intensity\" in the revision.\n\n**Q6: purpose of knowledge control in section 3.2.2**\n\nYes, the primary motivation is to encourage the generation of extreme yet plausible scenarios; these scenarios are **crucial for robust decision-making and precaution** (line 297-299). \n\nTo further validate the effectiveness of knowledge control, we have added an experiment in **Table R2 with PreDiff-KC-GT**, where PreDiff is guided by ground-truth future average intensity. The improved performance of PreDiff-KC-GT compared with PreDiff indicates that our two-stage pipeline further benefits the data-driven models with knowledge control."}, {"review_id": "41lvh38V40", "reviewer": "Reviewer_yf9m", "summary": "This paper applies latent diffusion models to the problem of precipitation nowcasting, developing a modular and flexible method for enforcing external constraints during the generating process. Constraints are general in that they can be specified in terms of any function F from the data to the reals, F:(x,y)->R. In the proposed \"Knowledge Control\" (KC) method, a separate NN is trained to predict F from the noised latent data. Conveniently this can be trained before training the diffusion model; it is then used while training the diffusion model as a regularizer, punishing the diffusion model for generating samples that violate the constraints. \n\nExperiments are performed on a simulation dataset (MovingMNIST) to demonstrate the advantage of the modeling choices and the KC method, and compare diffusion models with other deep learning models. The KC method is used to enforce conservation of energy, and the experiments demonstrate that KC is effective at constraining the generated samples without interfering with the other performance metrics.\n\nThe model is then applied to the SEVIR precipitation dataset and compared against a number of other deep learning models. Experiments demonstrate the use of the KC method in controlling the overall amount of rainfall.", "questions": "- In experiments, deterministic models are compared against probabilistic models. For the probabilistic models, a single sample was generated from the probabilistic distribution in each forecast. The probabilistic models perform worse than the deterministic models on some metrics. Isn't it likely that the performance of the probabilistic models would be improved if multiple samples were drawn from them?", "limitations": "The authors concisely and fairly address the limitations. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- The knowledge control method will be of high interest to those applying diffusion models to scientific applications, where there are often physical constraints.\n- Experiments are well done and will be valuable to researchers in this area. Experiments are done on both simulated and real data, with multiple performance metrics, and comparisons to strong baselines. \n- The paper is well-written. \n", "weaknesses": "None.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for providing detailed feedback. We hope that our response below can address your main concerns.\n\n**Q1: multiple samples from the probabilistic models and the performance comparison with the deterministic models**\n\nThank you for the thoughtful feedback. To clarify, instead of sampling only once from each probabilistic model for evaluation, we sample **eight forecasts** from each model. We will update the manuscript to clearly reflect this methodology.\n\nYou make an excellent observation that the choice of evaluation metrics can significantly impact how deterministic and probabilistic models are assessed comparatively. **Certain metrics may favor deterministic models** over probabilistic models.\n\nAs a simple illustration, consider a dataset where the true data distribution $p_{\\text{data}}(x)$ is a Bernoulli with equal probability for -1 and 1, i.e., $p_{\\text{data}}(-1)=0.5$, $p_{\\text{data}}(1)=0.5$, and $p_{\\text{data}}(x)=0$ when $x\\neq \\pm1$.\n\nA probabilistic model that perfectly captures this distribution, i.e., $p_{\\text{prob}}(\\hat{x}) = p_{\\text{data}}(x)$, would have $\\mathbb{E}\\_{x\\sim p\\_{\\text{data}}, \\hat{x}\\sim p\\_{\\text{prob}}}\\|x-\\hat{x}\\|^2 = 2$. In contrast, a deterministic model $p\\_{\\text{determ}}(\\hat{x})$ that always predicts 0 (the data mean) would have $\\mathbb{E}\\_{x\\sim p\\_{\\text{data}}, \\hat{x}\\sim p\\_{\\text{determ}}}\\|x-\\hat{x}\\|^2 = 1$. $p_{\\text{determ}}$ outperforms $p_{\\text{prob}}$ in MSE even though $p_{\\text{prob}}$ perfectly captures the true distribution while $p_{\\text{determ}}$ fails.\n\nWhile no single metric is perfect, to provide a nuanced comparison, we include multiple metrics and additionally continuous ranked probability score (CRPS) as shown in Table R2. his will give a more complete view of how the probabilistic and deterministic models compare across different evaluation perspectives."}, {"review_id": "46kVFX9BDy", "reviewer": "Reviewer_zn7Z", "summary": "This paper presents the PreDiff for precipitation nowcasting. PreDiff adopts the latent diffusion for prediction. Besides, a special knowledge control design is used in a classifier guidance way to incorporate auxiliary prior knowledge, which can guide the model to satisfy conservation constrain. In both synthetic and real-world benchmarks, PreDiff achieves a good performance.", "questions": "(1) In precipitation nowcasting, how to decide the hyper-parameter n?\n\n(2) It is interesting to see the results in Figure 3 (b). I want to know the CSI, Bias and FVD metrics for the $4\\sigma$ and $-4\\sigma$ models.", "limitations": "They have discussed limitations.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. This paper presents a practical method for precipitation nowcasting. Using diffusion model can benefit the prediction skill.\n\n2. The proposed knowledge control is reasonable and effective.\n\n3. This paper is well-organized.", "weaknesses": "1. The major concern is about novelty.\n\nThis paper introduces the famous diffusion model techniques into precipitation nowcasting task, including the condition diffusion, latent diffusion model and classifier guidance. \n\nThe key design is the knowledge control. But accoding to the paper, the knowledge control requires maual designed for every specific task. And there does not exist evidence that anticipated precipitation intensity is a right prior knowledge for precipitation nowcasting. More explanation are expected.\n\n2. More comparing baselines are insufficient.\n\nThere are some advanced deep models for precipitation nowcasting that the authors should take into consideration, especially the second baseline that is based on the GAN.\n\n[1] MotionRNN: A Flexible Model for Video Prediction with Spacetime-Varying Motions, CVPR 2021\n\n[2] Skilful precipitation nowcasting using deep generative models of radar, Nature 2021\n\n3. Model efficiency comparisons are also required, including the GPU memory and running time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for providing detailed feedback. We hope our response below can address your main concerns.\n\n**W1: more explanation on novelty and prior knowledge for precipitation nowcasting**\n\nOur literature review covers related works from three perspectives: deep learning for Earth system forecasting (line 30-42), diffusion models (line 43-58), and controllable generative models (line 59). To the best of our knowledge, we are the first to propose a controllable latent diffusion model for precipitation nowcasting. \n\nAs for our knowledge control mechanism, there are **fundamental differences from classifier guidance**. While classifier guidance requires data-label pairs for training the classifier, we directly formulate the domain knowledge as Eq.(4), without the need of additional categorical labels. The training and inference of our model thus can be **entirely label-free**.\n\nMost importantly, we propose a **novel two-stage pipeline** for training data-driven Earth system forecasting models. This pipeline incorporates both data-driven learning and constraints from domain knowledge, representing a promising direction for combining physical alignment with data-driven methods. \n\n**Anticipated Precipitation Intensity as Prior Knowledge.** For precipitation nowcasting on SEVIR, **even without directly available auxiliary information or prior knowledge**, By **flexibly** guiding PreDiff with **anticipated** future average intensity predicted by a simple probabilistic time series forecasting model, we can simulate how weather unfolds in extreme cases. E.g., Figure 3(b)'s Row 3 simulates how weather unfolds in an extreme case (with probability around 0.35%) where $I(x)>\\mu_\\tau+4\\sigma_\\tau$. Such simulation can be valuable for goverments to evalaute possible future damage in extreme cases.\n\nOverall, we believe our work is the first to demonstrate how to effectively and flexibly integrate domain knowledge into controllable diffusion models for Earth science tasks. This paves the way for leveraging knowledge control for physical alignment with constraints and domain expertise to improve data-driven Earth system modeling.\n\n**W2: more baselines: MotionRNN and DGMR**\n\nThanks for mentioning these important references. Actually, we **did evaluate DGMR (ref 28)** for SEVIR precipitation nowcasting tasks and included its results in Table 2. We apologize for misspelling DGMR as DGRM in Table 2. \n\nDGMR underperformed other probabilistic baselines and our PreDiff by a large margin. The causes can be two-fold:\n+ GAN-based models are well known to be difficult to optimize. The training of DGMR is highly unstable and struggles to converge to an optimal solution.\n+ The basic building blocks of the DGMR generator are ConvGRU, which is an outdated architecture and may limit its capacity.\n\nFor MotionRNN, we will cite and discuss it in the final version. We did not include MotionRNN for comparison, as it underperforms PredRNN, which is already one of our baselines (shown in Table 5 of PredRNN paper).\n\n**W3: model efficiency comparisons**\n\nGood suggestion! We have added **GPU memory costs** and **giga floating point operations per second (GFLOPS) in Table R2**. The computational expense of Prediff is moderate among other methods. We also report the **inference costs** here for methods that differ from the training costs:\n\n|Model|Inference GFLOPS|\n|-|-:|\n|VideoGPT|611,378|\n|LDM|1,415,522|\n|PreDiff|89,318|\n|PreDiff-KC|92,257|\n\nWe'll revise Table 1 accordingly in our final revision. The findings indicate that probabilistic models will require more parameters, consequently greater computational costs to model uncertainty.\n\n**Q1: in precipitation nowcasting, how to decide the hyper-parameter $n$**\n\nTo clarify, in Section 3.2.2, we **train a simple probabilistic time series forecasting model $p_\\tau()$** that predicts the future average intensity $I(x)$ with mean $\\mu_\\tau$ and variance $\\sigma_\\tau^2$ (line 290). The simple $p_\\tau()$ is introduced not as an accurate forecasting tool, but rather to showcase how the knowledge control could inform PreDiff of probable future states. By probabilistically predicting the future average intensity $I(x)$ using $p_\\tau()$, we can guide the diffusion process to simulate probable extreme weather events accordingly. **The choice of $n$ is arbitrary**, as it **depends on the desired extremity** of the simulated case. By doing this, we can simulate how weather unfolds in extreme cases. E.g., Figure 3(b)'s Row 3 simulates how weather unfolds in an extreme case (with probability around 0.35%) where the future average intensity $I(x)>\\mu_\\tau+4\\sigma_\\tau$. Such simulation can be valuable for goverments to evalaute possible future damage in extreme cases.\n\n**Q2: quantitative results for $\\pm4\\sigma$ in Figure 3(b)**\n\nWe have added the CSI and FVD for PreDiff-KC ($\\in\\{-4\\sigma_\\tau, 4\\sigma_\\tau\\}$) **in Table R2**. As expected, PreDiff-KC with $\\pm 4\\sigma_\\tau$ deviations from the learned distribution mean results in poor per-pixel accuracy, and hence low CSI. This aligns with the intuition that forecasting such extreme cases should reduce the per-pixel accuracy comparing to the target. On the other hand, since the guided forecasts still follow the learned distribution overall, the FVD does not decrease significantly - it measures discrepancy between complete distributions rather than individual forecast-target pairs. Regarding BIAS, the results shown in the following Table R5 align with the expectation that forecasting extreme cases results in exaggerated False Alarms (rainstorms, with $+4\\sigma_\\tau$) or Misses (droughts, with $-4\\sigma_\\tau$), especially for high thresholds.\n\nTable R5: BIAS of PreDiff-KC ($\\in\\{-4\\sigma_\\tau, 4\\sigma_\\tau\\}$) on SEVIR\n|Model|BIAS-m|BIAS-219|BIAS-181|BIAS-160|BIAS-133|BIAS-74|BIAS-16|\n|:-|:-:|:-:|:-:|:-:|:-:|:-:|:-:|\n|$+4\\sigma_\\tau$|7.2340|10.812|8.4762|8.5841|7.6071|4.6057|3.3127|\n|$-4\\sigma_\\tau$|0.7186|0.5290|0.5275|0.5072|0.4691|0.8349|1.4443|"}]

Author rebuttal:
# General Response

Dear Reviewers,

We appreciate reviewers’ effort in evaluating this work. We are glad that they found our paper "well-written" (Ntzp,sDg9,yf9m,zn7Z), our PerDiff with novel knowledge control mechanism "interesting"/"practical"/"effective" (Ntzp/yf9m/zn7Z), our empirical studies "informative"/"well done and valuable" (sDg9/yf9m), and that our PreDiff "achieves good performance" (zn7Z). 

Below please find a summary of the common questions where we have responded in each reviewer's rebuttal correspondingly:
1. novelty (sDg9,zn7Z).
2. CSI at all thresholds, BIAS, FVD and more experiment details on SEVIR (Ntzp,zn7Z).
3. model size and computational costs (Ntzp,zn7Z).
4. more reference baselines (Dg9,zn7Z).
5. the evaluation metric FVD (sDg9).
6. more details of knowledge control on SEVIR (sDg9,zn7Z).

We want to re-iterate our **main contributions and novelty** here, as the reviewers' summaries indicate some uncertainty about these key points. Our contribution lies in the novel **two-stage pipeline** where we first construct a data-driven model and then leverage knowledge control for physical alignment with constraints and domain expertise. This integration of data-driven deep learning with knowledge control represents a ***promising new paradigm*** for Earth system forecasting that combines the strengths of both approaches. With the drastic growth in the volume of Earth observation data, retraining a model with different prior knowledge can become prohibitively expensive. In contrast, our two-stage pipeline keeps the purely data-driven foundation model unchanged once trained. We only need to update the lightweight knowledge control in the second stage when integrating different prior knowledge. This makes our approach uniquely suited to **leverage the explosion of Earth system data while remaining flexible to new constraints and physical insights**.

In addition, we include additional empirical studies in Table R1,R2 and Figure R1,R2,R3 (attached PDF) to address reviewers' concerns on 
1. the effectiveness of knowledge control (sDg9,zn7Z)
2. more details of experiments on SEVIR (Ntzp,zn7Z)
3. model size and computational costs (Ntzp,zn7Z)
4. additional reference baselines to (sDg9,zn7Z).

Table R1: Notice that the PreDiff in Sec. 3.1 is trained on an **energy-noisy** N-body MNIST. We evaluate this trained PreDiff on a **cleaner** N-body MNIST Clean dataset (Ground-truth has much lower E.MSE and E.MAE). As shown in Table R1, PreDiff still suffers from energy noise and generates forecasts deviating from the ground truth. This results in poorer scores in both frame metrics and energy metrics, since in N-body MNIST Clean the target complies with the law of conservation of energy. In contrast, PreDiff-KC better complies with such knowledge and hence achieves better scores by generating forecasts closer to the ground truth.
|Model|||Frame Metrics||Energy Metrics||
|:-|:-:|:-:|:-:|:-:|:-:|:-:|
||MSE$\downarrow$|MAE$\downarrow$|SSIM$\uparrow$|FVD$\downarrow$|E.MSE$\downarrow$|E.MAE$\downarrow$|
|Ground-truth|0|0|1|0|0.0002|0.0105|
|PreDiff|20.69|40.96|0.9414|1.710|0.0090|0.0559|
|PreDiff-KC|**20.44**|**40.77**|**0.9415**|**1.678**|**0.0032**|**0.0377**|

In Table R2 (see **attached PDF for the full table**), we report **GPU memory** and **GFLOPS** (separate by ";" if there are multiple stages) of single training batch forward for evaluating the computational costs per reviewer zn7Z's request.
In addition to FVD and CSI, we also include the **continuous ranked probability score (CRPS)** for assessing probabilistic forecasts on SEVIR dataset. CRPS measures the discrepancy between the predicted distribution and the true distribution. When the forecast distribution collapses into a single value, as in deterministic forecasting models, CRPS reduces to Mean Absolute Error (MAE). A lower CRPS value indicates higher forecast accuracy. We calculate FVD, CRPS and CSI for an ensemble of eight samples from each model. **PreDiff-KC-GT** denotes PreDiff guided by the ground-truth future intensity. As Table R2 shows, compared to the baselines, our PreDiff model achieves the **lowest CRPS**, indicating it best captures the data distribution. 
Notice that PreDiff improves performance compared to LDM while **reducing computational costs**. With recent fast development of acceleration methods like DPM-solver, the inference time of diffusion models can be significantly reduced from thousands of diffusion steps to just dozens without sacrificing performance. This further facilitates the deployment of PreDiff.
With **guidance from the ground-truth** future intensity, the performance of PreDiff **further improves**. This demonstrates that knowledge control is able to guide the purely data-driven model to generate forecasts better aligned with the ground-truth. 
For convenience, we **excerpt the most concerned results from Table R2** here, including the newly added **MetNet-2** baseline. We also copy the results of **DGMR from our Table 2**. **Bold** and ***italics*** highlight the best and the second best scores. Please refer to the **attached PDF for full details**.
|Model|\#Param.(M)|||Metrics|||
|:-|-:|-:|:-:|:-:|:-:|:-:|
|||FVD$\downarrow$|CRPS$\downarrow$|CSI$\uparrow$|CSI-pool4$\uparrow$|CSI-pool16$\uparrow$|
|Persistence|-||0.0526|0.2613|0.3702|0.4690|
|Earthformer|15.1|690.7|0.0304|**0.4419**|0.4567|0.5005|
|MetNet-2|15.6|589.0|0.0386|0.4046|0.4434|0.4997|
|LDM|438.6|133.0|0.0280|0.3580|0.4022|0.5522|
|PreDiff|220.5|**33.05**|***0.0246***|0.4072|***0.4624***|***0.6244***|
|PreDiff-KC($\in\{-4\sigma_\tau,4\sigma_\tau\}$)|229.4|44.30|-|0.3272|0.4080|0.4787|
|PreDiff-KC-GT|229.4|***34.07***|**0.0240**|0.4080|**0.4633**|**0.6261**|

We are sincerely looking forward to further discussion.

Thanks,

Authors of PreDiff

Meta-review: This paper presents a new generative model for precipitation nowcasting by combining a conditional latent diffusion model with a form of guidance, called “knowledge control”, that penalises latent images departing from a known constraint. The method is benchmarked against several other generative models  on both the N-body MNIST synthetic dataset and on the satellite data-based SEVIR dataset, and it is compared to several transformer-based architectures, conv-LSTM-based baselines and a few more recent algorithms like VideoGPT or GANs-based DGRM. Metrics include deterministic CSI and the Fréchet Video Distance.

Reviewers praised the impactful application (sDg9, zn7Z), the paper’s writing (Ntzp, sDg9, yf9m, zn7Z), the extensive descriptions of implementation and the paper’s reproducibility (sDg9), the extensive experimental evaluation (yf9m), the knowledge control guidance (yf9m, Ntzp).

Reviewer Ntzp had concerns about the number of ablations and the authors provided additional comparisons, notably for the knowledge control guidance. Reviewer Ntzp suggested to list CSI at different thresholds, which the authors did. Reviewer sDg9 had concerns about the FVD metric, and the authors included CRPS in the revision. Reviewers sDg9 and zn7Z asked for some clarifications about knowledge control, which were acted upon. Reviewer zn7Z asked for metrics of model efficiency, which were provided by the authors.

There was additional discussion about model novelty (sDg9, zn7Z), but these considerations are not taken into account for evaluating the paper.

Note from the meta-reviewer: the authors should explain whether they evaluated existing pre-trained baseline models (e.g., DGMR) or if they trained them from scratch on the SEVIR data.

In light of the existing scores (4, 5, 6, 8) and of the discussion with the reviewers, the paper crosses the threshold for acceptance and I recommend accepting it.
Decision: Accept (poster)