SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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 1: Performance comparison on  $N$ -body MNIST. We report conventional frame quality metrics (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  $E_{\mathrm{det}}(\widehat{x})$  and the initial energy  $E(y^{L_{\mathrm{in}}})$ . Lower values on the energy metrics indicate better compliance with conservation of energy.
Model	#Param. (M)	Frame Metrics	Energy Metrics
MSE ↓	MAE ↓	SSIM ↑	FVD ↓	E.MSE ↓	E.MAE ↓
Target	-	0.000	0.000	1.0000	0.000	0.0132	0.0697
Persistence	-	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 [1]	19.2	38.89	96.47	0.8036	163.5	-	-
Earthformer [8]	7.6	14.82	39.93	0.9538	6.798	-	-
VideoGPT [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	120.7	9.492	25.01	0.9716	0.987	0.0226	0.1083
PreDiff-KA	129.4	21.90	43.57	0.9303	4.063	0.0039	0.0443
energy. The total energy of the whole system  $E(x^{j})$  at any future time step  $j$  during evolution should equal the total energy at the last observation time step  $E(y^{L_{\mathrm{in}}})$ .
We impose the law of conservation of energy for the knowledge alignment on  $N$ -body MNIST in the form of (4):
$$
\mathcal {F} (\widehat {x}, y) \equiv \left[ E \left(\widehat {x} ^ {1}\right), \dots , E \left(\widehat {x} ^ {L _ {\text {o u t}}}\right) \right] ^ {T}, \tag {6}
$$
$$
\mathcal {F} _ {0} (y) \equiv \left[ E \left(y ^ {L _ {\mathrm {i n}}}\right), \dots , E \left(y ^ {L _ {\mathrm {i n}}}\right) \right] ^ {T}. \tag {7}
$$
The ground-truth values of the total energy  $E(y^{L_{\mathrm{in}}})$  and  $E(x^j)$  are directly accessible since  $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  $U_{\phi}$  is trained following Alg. 1 to guide the PreDiff to generate forecasts  $\hat{x}$  that conserve the same energy as the initial step  $E(y^{L_{\mathrm{in}}})$ .
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  $E_{\mathrm{det}}(\widehat{x})^4$  that detects the total energy of the forecasts  $\widehat{x}$ . We evaluate the energy error between the forecasts and the initial energy using E.MSE  $(\widehat{x},y)\equiv \mathrm{MSE}(E_{\mathrm{det}}(\widehat{x}),E(y^{L_{\mathrm{in}}}))$  and E.MAE  $(\widehat{x},y)\equiv \mathrm{MAE}(E_{\mathrm{det}}(\widehat{x}),E(y^{L_{\mathrm{in}}}))$ . In this evaluation, 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.
As illustrated in Table 1, PreDiff-KA substantially outperforms all baseline methods and PreDiff without knowledge alignment in E.MSE and E.MAE. 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 1 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 1 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 (E.MSE = 0.0261), PreDiff-KA successfully
4The test MSE of the energy detector is  $5.56\times 10^{-5}$  , which is much smaller than the scores of E.MSE shown in Table 1. This indicates that the energy detector has high precision and reliability for verifying energy conservation in the model forecasts.
Figure 3: A set of example forecasts from baselines and PreDiff on the SEVIR test set. From top to bottom: context sequence  $y$ , target sequence  $x$ , forecasts from ConvLSTM [47], Earthformer [8], VideoGPT[65], 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 a much lower E.MSE = 0.0086. On the contrary, none of the evaluated baselines can 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 of the energy results in changing the velocities and positions of the digits. A minor change in the position of a 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 a spatiotemporal Earth observation dataset which consists of  $384\mathrm{km}\times 384\mathrm{km}$  image sequences spanning over 4 hours. Images in SEVIR are sampled and aligned across five different data types: three channels (C02, C09, C13) from the GOES-16 advanced baseline imager, NEXRAD Vertically Integrated Liquid (VIL) mosaics, and GOES-16 Geostationary Lightning Mapper (GLM) flashes. The SEVIR benchmark supports scientific research on multiple meteorological applications including precipitation nowcasting, synthetic radar generation, front detection, etc. Due to computational resource limitations, we adopt a 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 a spatial resolution of  $128\times 128$ , i.e.  $x\in \mathbb{R}^{6\times 128\times 128\times 1}$ ,  $y\in \mathbb{R}^{7\times 128\times 128\times 1}$ .
Evaluation. Following [55, 8], we adopt the Critical Success Index (CSI) for evaluation, which is commonly used in precipitation nowcasting and is defined as CSI = #Hits/#Hits + #Misses + #F.Alarms. To count the #Hits (truth=1, pred=1), #Misses (truth=1, pred=0) and #F.Alarms (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 \times 4$  and  $16 \times 16$ , which evaluate the performance on neighborhood aggregations at multiple spatial scales. These pooled CSI metrics assess the models' ability to capture local pattern distributions. Additionally, we incorporate FVD [51] and continuous ranked probability score
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  $(\mu_{\tau} + n\sigma_{\tau})$ , where  $n$  takes the values of 4, 2, -2, -4.
(CRPS) [9] 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 a single value, as in deterministic models, CRPS reduces to Mean Absolute Error (MAE). A lower CRPS value indicates higher forecast accuracy. Scores for all involved metrics are calculated using an ensemble of eight samples from each model.
3.2.1 Comparison to the State of the Art
We adjust the configurations of involved baselines accordingly and tune some of the hyperparameters for adaptation on the SEVIR dataset. More implementation details of baselines are provided in Appendix B.2. The experiment results listed in Table 2 show that probabilistic spatiotemporal forecasting methods are not good at 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-poo116. Qualitative results shown in Fig. 3 demonstrate that CSI is not aligned with human perceptual judgement. For such a complex system, deterministic methods give up capturing the real patterns and resort to 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 a local area, resulting in higher CSI-poo116, lower CRPS, and succeed in keeping the correct local patterns, which can be 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 usually 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 a data sequence as  $I(x) \in \mathbb{R}^{+}$ . In order to estimate the conditional quantiles of future intensity, we train a simple probabilistic time series forecasting model with a parametric (Gaussian) distribution  $p_{\tau}(I(x) | [I(y^j)]) = \mathcal{N}(\mu_{\tau}([I(y^j)]), \sigma_{\tau}([I(y^j)]))$  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  $s = 16$  report the CSI at pooling scales of  $4 \times 4$  and  $16 \times 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.
Model	#Param. (M)	Metrics
FVD ↓	CRPS ↓	CSI ↑	CSI-pool14 ↑	CSI-pool16 ↑
Persistence	-	525.2	0.0526	0.2613	0.3702	0.4690
UNet [55]	16.6	753.6	0.0353	0.3593	0.4098	0.4805
ConvLSTM [47]	14.0	659.7	0.0332	0.4185	0.4452	0.5135
PredRNN [61]	46.6	663.5	0.0306	0.4080	0.4497	0.5005
PhyDNet [11]	13.7	723.2	0.0319	0.3940	0.4379	0.4854
E3D-LSTM [60]	35.6	600.1	0.0297	0.4038	0.4492	0.4961
Rainformer [1]	184.0	760.5	0.0357	0.3661	0.4232	0.4738
Earthformer [8]	15.1	690.7	0.0304	0.4419	0.4567	0.5005
DGMR [41]	71.5	485.2	0.0435	0.2675	0.3431	0.4832
VideoGPT [65]	99.6	261.6	0.0381	0.3653	0.4349	0.5798
LDM [42]	438.6	133.0	0.0280	0.3580	0.4022	0.5522
PreDiff	220.5	33.05	0.0246	0.4100	0.4624	0.6244
PreDiff-KA (∈[−2στ, 2στ])	229.4	34.18	-	-	-	-
the distribution of the average future intensity  $I(x)$  given the average intensity of each context frame  $[I(y^{j})]_{j = 1}^{L_{\mathrm{in}}}$  (abbreviated as  $[I(y^{j})]$ ). By incorporating  $\mathcal{F}(\widehat{x},y)\equiv I(\widehat{x})$  and  $\mathcal{F}_0(y)\equiv \mu_\tau +n\sigma_\tau$  for knowledge alignment, PreDiff-KA gains the capability of generating forecasts for potential extreme cases, e.g., where  $I(\widehat{x})$  falls outside the typical range of  $\mu_{\tau}\pm \sigma_{\tau}$ .
Fig. 4 shows a set of generations from PreDiff and PreDiff-KA with anticipated future intensity  $\mu_{\tau} + n\sigma_{\tau}$ ,  $n \in \{-4, -2, 2, 4\}$ . 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.
4 Conclusions and Broader Impacts
In this paper, we propose PreDiff, a 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$ -body MNIST and SEVIR datasets.
Our work has certain limitations: 1) Benchmark datasets and evaluation metrics for precipitation nowcasting and Earth system forecasting are still maturing compared to 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. 3) 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 ability to 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 a 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.
[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. Conditional image generation with pixelCNN 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 NeurIPS, volume 30, 2017.
[55] Mark Veillette, Siddharth Samsi, and Chris Mattioli. SEVIR: A storm event imagery dataset for deep learning applications in radar and satellite meteorology. Advances in Neural Information Processing Systems, 33:22009–22019, 2020.
[56] Vikram Voleti, Alexia Jolicoeur-Martineau, and Christopher Pal. Masked conditional video diffusion for prediction, generation, and interpolation. arXiv preprint arXiv:2205.09853, 2022.
[57] Hao Wang, SHI Xingjian, and Dit-Yan Yeung. Natural-parameter networks: A class of probabilistic 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, 28(12):3395-3408, 2016.
[59] Hao Wang and Dit-Yan Yeung. A survey on bayesian deep learning. CSUR, 53(5):1-37, 2020.
[60] Yunbo Wang, Lu Jiang, Ming-Hsuan 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 NeurIPS, 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 vq-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 IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
[67] Ye Yuan, Jiaming Song, Umar Iqbal, Arash Vahdat, and Jan Kautz. Physdiff: Physics-guided 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-Hua Zhang. A self-attention-based neural network for three-dimensional multivariate modeling and its skillful enso predictions. Science Advances, 9(10):eadb2827, 2023.
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 nowcasting 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 [61] builds upon ConvLSTM by incorporating a spatiotemporal memory flow structure. E3D-LSTM [60] integrates 3D CNN to LSTM to enhance long-term high-level relation modeling. PhyDNet [11] incorporated partial differential equation (PDE) constraints in the latent space. MetNet [49] and its successor, MetNet-2 [5], propose architectures based on ConvLSTM and dilated CNN, enabling skillful precipitation forecasts up to twelve hours ahead. DGMR [41] 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 ERA5 [19] dataset. Several DL-based models have emerged in this area. FourCastNet [37] proposes an architecture with Adaptive Fourier Neural Operators (AFNO) [12] as building blocks for autoregressive weather forecasting. FengWu [3] introduces a multi-model Transformer-based global medium-range weather 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 [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-specific prior knowledge.
Diffusion models Diffusion models (DMs) [22] are a 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 a variant of DMs that are trained on latent vector outputs from a 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 [56] 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 low-dimensional latent space as 2D vectors, and presents a joint training of unconditional and frame conditional video generations. LFDM [36] 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 a pretrained DM by freezing the base model and training a copy
end-to-end 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 [67] trains a physics-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 a powerful technique in various domains, its application in DL for precipitation nowcasting remains an unexplored area.