SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhihan Gao*
Hong Kong University of Science and Technology zhihan.gao@connect.ust.hk
Xingjian Shi†
Boson AI
xshiab@connect.ust.hk
Boran Han
boranhan@amazon.com
AWS AI Labs
howngz@amazon.com
Xiaoyong Jin
jxiaoyon@amazon.com
Danielle Maddix
dmmaddix@amazon.com
Yi Zhu†
yi@boson.ai
Mu Li†
mu@boson.ai
Yuyang Wang
yuyawang@amazon.com
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 or neglect domain-specific prior 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 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.
Earth's intricate climate system significantly influences daily life. Precipitation nowcasting, tasked with delivering accurate rainfall forecasts for the near future (e.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 in this field, rivaling conventional numerical methods [8, 5] with their advantages of being more skillful [5], efficient [37], and scalable [3]. However, accurately predicting the future rainfall remains challenging for data-driven algorithms.
*Work conducted during an internship at Amazon. †Work conducted while at Amazon.
The state-of-the-art Earth system forecasting algorithms [47, 61, 41, 37, 8, 69, 2, 29, 3] typically generates blurry predictions. This is caused by the high variability and complexity inherent to Earth's climatic system. Even minor differences in initial conditions can lead to vastly divergent outcomes that are difficult to predict. Most methods adopt a 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.
The emergence of diffusion models (DMs) [22] has enabled powerful probabilistic frameworks for generative modeling. DMs have shown remarkable capabilities in generating high-quality images [40, 45, 43] and videos [15, 23]. As a likelihood-based model, DMs do not exhibit mode collapse or training instabilities like GANs [10]. Compared to autoregressive (AR) models [53, 46, 63, 39, 65] that generate images pixel-by-pixel, DMs can produce higher resolution images faster and with higher quality. They are also better at handling uncertainty [62, 34, 57-59] without drawbacks like exposure bias [13] in AR models. Latent diffusion models (LDMs) [42, 52] further improve on DMs by separating the model into two phases, only applying the costly diffusion in a 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, 56], its application to precipitation nowcasting and Earth system forecasting is in early stages [16]. 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 [11, 1, 37]. This enables them to be aware of prior knowledge and generate physically plausible forecasts. However, these approaches still face challenges, such as requiring to 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.
Inspired by recent success in controllable generative models [68, 24, 4, 33, 6], we propose a 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 [8]. 2) In the second stage, we inject prior knowledge of the Earth system by training a 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 our main contributions as follows:
- 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-art performance on the  $N$ -body MNIST [8] dataset and attains state-of-the-art perceptual quality on the SEVIR [55] dataset.
We follow [47, 48, 55, 1, 8] to formulate precipitation nowcasting as a spatiotemporal forecasting problem. The  $L_{\mathrm{in}}$ -step observation is represented as a spatiotemporal sequence  $y = [y^{j}]_{j=1}^{L_{\mathrm{in}}} \in \mathbb{R}^{L_{\mathrm{in}} \times H \times W \times C}$ , where  $H$  and  $W$  denote the spatial resolution, and  $C$  denotes the number of measurements at each space-time coordinate. Probabilistic forecasting aims to model the conditional probabilistic distribution  $p(x|y)$  of the  $L_{\mathrm{out}}$ -step-ahead future  $x = [x^{j}]_{j=1}^{L_{\mathrm{out}}} \in \mathbb{R}^{L_{\mathrm{out}} \times H \times W \times C}$ , given the observation  $y$ . In what follows, we will present the parameterization of  $p(x|y)$  by a controllable LDM.
Figure 1: Overview of PreDiff inference with knowledge alignment. An observation sequence  $y$  is encoded into a latent context  $z_{\mathrm{cond}}$  by the frame-wise encoder  $\mathcal{E}$ . The latent diffusion model  $p_{\theta}(z_t|z_{t + 1},z_{\mathrm{cond}})$ , which is parameterized by an Earthformer-UNet, then generates the latent future  $z_0$  by autoregressively denoising Gaussian noise  $z_{T}$  conditioned on  $z_{\mathrm{cond}}$ . It takes the concatenation of the latent context  $z_{\mathrm{cond}}$  (in the blue border) and the previous-step noisy latent future  $z_{t + 1}$  (in the cyan border) as input, and outputs  $z_{t}$ . The transition distribution of each step from  $z_{t + 1}$  to  $z_{t}$  can be further refined as  $p_{\theta ,\phi}(z_t|z_{t + 1},y,\mathcal{F}_0)$  via knowledge alignment, according to auxiliary prior knowledge. This denoising process iterates from  $t = T$  to  $t = 0$ , resulting in a denoised latent future  $z_0$ . Finally,  $z_0$  is decoded back to pixel space by the frame-wise decoder  $\mathcal{D}$  to produce the final prediction  $\hat{x}$ . (Best viewed in color).
2.1 Preliminary: Diffusion Models
Diffusion models (DMs) learn the data distribution  $p(x)$  by training a model to reverse a predefined noising process that progressively corrupts the data. Specifically, the noising process is defined as  $q(x_{t}|x_{t - 1}) = \mathcal{N}(x_{t};\sqrt{\alpha_{t}} x_{t - 1},(1 - \alpha_{t})I),1\leq t\leq T$ , where  $x_0\sim p(x)$  is the true data, and  $x_{T}\sim \mathcal{N}(0,I)$  is random noise. The coefficients  $\alpha_{t}$  follow a fixed schedule over the timesteps  $t$ . DMs factorize and parameterize the joint distribution over the data  $x_0$  and noisy latents  $x_{i}$  as  $p_{\theta}(x_{0:T}) = p(x_T)\prod_{t = 1}^{T}p_{\theta}(x_{t - 1}|x_t)$ , where each step of the reverse denoising process is a Gaussian distribution  $p_{\theta}(x_{t - 1}|x_t) = \mathcal{N}(\mu_\theta (x_t,t),\Sigma_\theta (x_t,t))$ , which is trained to recover  $x_{t - 1}$  from  $x_{t}$ .
To apply DMs for spatiotemporal forecasting,  $p(x|y)$  is factorized and parameterized as  $p_{\theta}(x|y) = \int p_{\theta}(x_{0:T}|y)dx_{1:T} = \int p(x_T)\prod_{t=1}^T p_{\theta}(x_{t-1}|x_t,y)dx_{1:T}$ , where  $p_{\theta}(x_{t-1}|x_t,y)$  represents the conditional denoising transition with the condition  $y$ .
2.2 Conditional Diffusion in 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 of lower-dimensional latent representations. The two sequential phases of the PreDiff training are: 1) Training a frame-wise variational autoencoder (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 a combination of the pixel-wise loss (e.g. L2 loss) and an adversarial loss. Different from [7], we exclude the perceptual loss since there are no standard pretrained models for perception on Earth observation data. Specifically, the encoder  $\mathcal{E}$  is trained to encode a data frame  $x^{j} \in \mathbb{R}^{H \times W \times C}$  to a latent representation  $z^{j} = \mathcal{E}(x^{j}) \in \mathbb{R}^{H_{z} \times W_{z} \times C_{z}}$ . The decoder  $\mathcal{D}$  learns to reconstruct the data frame  $\widehat{x}^{j} = \mathcal{D}(z^{j})$  from
the encoded latent. We denote  $z \sim p_{\mathcal{E}}(z|x) \in \mathbb{R}^{L \times H_z \times W_z \times C_z}$  as equivalent to  $z = [z^j] = [\mathcal{E}(x^j)]$ , representing encoding a sequence of frames in pixel space into a latent spatiotemporal sequence. And  $x \sim p_{\mathcal{D}}(x|z)$  denotes decoding a latent spatiotemporal sequence.
Latent diffusion. With the context  $y$  being encoded by the frame-wise encoder  $\mathcal{E}$  into the learned latent space as  $z_{\mathrm{cond}} \in \mathbb{R}^{L_{\mathrm{in}} \times H_z \times W_z \times C_z}$  as (1). The conditional distribution  $p_{\theta}(z_{0:T}|z_{\mathrm{cond}})$  of the latent future  $z_i \in \mathbb{R}^{L_{\mathrm{out}} \times H_z \times W_z \times C_z}$  given  $z_{\mathrm{cond}}$  is factorized and parameterized as (2):
$$
z _ {\text {c o n d}} \sim p _ {\mathcal {E}} \left(z _ {\text {c o n d}} | y\right), \tag {1}
$$
$$
p _ {\theta} \left(z _ {0: T} \mid z _ {\text {c o n d}}\right) = p \left(z _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(z _ {t - 1} \mid z _ {t}, z _ {\text {c o n d}}\right). \tag {2}
$$
where  $z_{T} \sim p(z_{T}) = \mathcal{N}(0, I)$ . As proposed by [22, 45], an equivalent parameterization is to have the DMs learn to match the transition noise  $\epsilon_{\theta}(z_{t}, t)$  of step  $t$  instead of directly predicting  $z_{t-1}$ . The training objective of PreDiff is simplified as shown in (3):
$$
L _ {\mathrm {C L D M}} = \mathbb {E} _ {(x, y), t, \epsilon \sim \mathcal {N} (0, I)} \| \epsilon - \epsilon_ {\theta} \left(z _ {t}, t, z _ {\text {c o n d}}\right) \| _ {2} ^ {2}. \tag {3}
$$
where  $(x,y)$  is a sampled context sequence and target sequence data pair, and given that,  $z_{t} \sim q(z_{t}|z_{0})p_{\mathcal{E}}(z_{0}|x)$  and  $z_{\mathrm{cond}} \sim p_{\mathcal{E}}(z_{\mathrm{cond}}|y)$ .
Instantiating  $p_{\theta}(z_{t-1}|z_t, z_{\mathrm{cond}})$ . Compared to images, modeling spatiotemporal observation data in precipitation nowcasting poses greater challenges due to their higher dimensionality. We propose replacing the UNet backbone in LDM [42] with Earthformer-UNet, derived from Earthformer's encoder [8], which is known for its ability to model intricate and extensive spatiotemporal dependencies in the Earth system.
Earthformer-UNet adopts a hierarchical UNet architecture with self cuboid attention [8] as the building blocks, excluding the bridging cross-attention in the encoder-decoder architecture of Earthformer. More details of the architecture design of Earthformer-UNet are provide in Appendix B.1. We find Earthformer-UNet to be more stable and effective at modeling the transition distribution  $p_{\theta}(z_{t - 1}|z_t,z_{\mathrm{cond}})$ . It takes the concatenation of the encoded latent context  $z_{\mathrm{cond}}$  and the noisy latent future  $z_{t}$  along the temporal dimension as input, and predicts the one-step-ahead noisy latent future  $z_{t - 1}$  (in practice, the transition noise  $\epsilon$  from  $z_{t}$  to  $z_{t - 1}$  is predicted as shown in (3)).
2.3 Incorporating Knowledge Alignment
Though DMs hold great promise for diverse and realistic generation, the generated predictions may violate physical constraints, or disregard domain-specific prior knowledge, thereby fail to give plausible and non-trivial results [14, 44]. One possible reason for this is that DMs are not necessarily 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
Algorithm 1 One training step of the knowledge alignment network  $U_{\phi}$
1:  $(x,y)$  sampled from data
2:  $t\sim \mathrm{Uniform}(0,T)$
3:  $z_{t} \sim q(z_{t}|z_{0})p_{\mathcal{E}}(z_{0}|x)$
4:  $L_{U}\gets \| U_{\phi}(z_{t},t,y) - \mathcal{F}(x,y)\|$
main physically realizable. The causes may also stem from the stochastic nature of chaotic systems, the approximation error in denoising steps, etc.
To address this issue, we propose knowledge alignment to incorporate auxiliary prior knowledge:
$$
\mathcal {F} (\widehat {x}, y) = \mathcal {F} _ {0} (y) \in \mathbb {R} ^ {d}, \tag {4}
$$
into the diffusion generation process. The knowledge alignment imposes a constraint  $\mathcal{F}$  on the forecast  $\widehat{x}$ , optionally with the observation  $y$ , based on domain expertise. E.g., for an isolated physical system, the knowledge  $E(\widehat{x},\cdot) = E_0(y^{L_{\mathrm{in}}})\in \mathbb{R}$  imposes the conservation of energy by enforcing the generation  $\widehat{x}$  to keep the total energy  $E(\widehat{x},\cdot)$  the same as the last observation  $E_0(y^{L_{\mathrm{in}}})$ . The violation  $\| \mathcal{F}(\widehat{x},y) - \mathcal{F}_0(y)\|$  quantifies the deviation of a prediction  $\widehat{x}$  from prior knowledge. The larger violation indicates  $\widehat{x}$  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  $x$  from training data may violate the knowledge, i.e.  $\mathcal{F}(x,y)\neq \mathcal{F}_0(y)$ , due to noise in data collection or simulation.
Inspired by classifier guidance [4], we achieve knowledge alignment by training a knowledge alignment network  $U_{\phi}(z_t,t,y)$  to estimate  $\mathcal{F}(\widehat{x},y)$  from the intermediate latent  $z_{t}$  at noising step  $t$ . The key idea is to adjust the transition probability distribution  $p_{\theta}(z_{t - 1}|z_t,z_{\mathrm{cond}})$  in (2) during each latent denoising step to reduce the likelihood of sampling  $z_{t}$  values expected to violate the constraints:
$$
p _ {\theta , \phi} \left(z _ {t} \mid z _ {t + 1}, y, \mathcal {F} _ {0}\right) \propto p _ {\theta} \left(z _ {t} \mid z _ {t + 1}, z _ {\text {c o n d}}\right) \cdot e ^ {- \lambda_ {\mathcal {F}} \left\| U _ {\phi} \left(z _ {t}, t, y\right) - \mathcal {F} _ {0} (y) \right\|}, \tag {5}
$$
where  $\lambda_{\mathcal{F}}$  is a guidance scale factor. The knowledge alignment network is trained by optimizing the objective  $L_{U}$  in Alg. 1. According to [4], (5) can be approximated by shifting the predicted mean of the denoising transition  $\mu_{\theta}(z_{t + 1},t,z_{\mathrm{cond}})$  by  $-\lambda_{\mathcal{F}}\Sigma_{\theta}\nabla_{z_t}\| U_\phi (z_t,t,y) - \mathcal{F}_0(y)\|$ , where  $\Sigma_{\theta}$  is the variance of the original transition distribution  $p_{\theta}(z_t|z_{t + 1},z_{\mathrm{cond}}) = \mathcal{N}(\mu_{\theta}(z_{t + 1},t,z_{\mathrm{cond}}),\Sigma_{\theta}(z_{t + 1},t,z_{\mathrm{cond}}))$ . Detailed derivation is provided in Appendix C.
The training procedure of knowledge alignment is outlined in Alg. 1. The noisy latent  $z_{t}$  for training the knowledge alignment network  $U_{\phi}$  is sampled by encoding the target  $x$  using the frame-wise encoder  $\mathcal{E}$  and the forward noisig process  $q(z_{t}|z_{0})$ , eliminating the need for an inference sampling process. This makes the training of the knowledge alignment network  $U_{\phi}$  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  $U_{\phi}$  to flexibly explore various constraints and domain knowledge, without the need for retraining the entire model. This stands as a key advantage over incorporating constraints into model architectures or training losses.
We conduct empirical studies and compare PreDiff with other state-of-the-art spatiotemporal forecasting models on a synthetic dataset  $N$ -body MNIST [8] and a real-world precipitation nowcasting benchmark SEVIR[2] [55] to verify the effectiveness of PreDiff in handling the dynamics and uncertainty in complex spatiotemporal systems and generating high quality, accurate forecasts. We impose data-specific knowledge alignment: energy conservation on  $N$ -body MNIST and anticipated precipitation intensity on SEVIR. Experiments demonstrate that PreDiff under the guidance of knowledge alignment (PreDiff-KA) is able to generate predictions that comply with domain expertise much better, without severely sacrificing fidelity.
3.1  $N$ -body 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$ -body MNIST<sup>3</sup>, which is an extension of MovingMNIST [50]. The dataset contains sequences of digits moving subject to the gravitational force from other digits. The governing equation for the motion is  $\frac{d^2\boldsymbol{x}_i}{dt^2} = -\sum_{j\neq i}\frac{Gm_j(\boldsymbol{x}_i - \boldsymbol{x}_j)}{(|\boldsymbol{x}_i - \boldsymbol{x}_j| + d_{\mathrm{soft}})^r}$ , where  $\boldsymbol{x}_i$  is the spatial coordinates of the  $i$ -th digit,  $G$  is the gravitational constant,  $m_j$  is the mass of the  $j$ -th digit,  $r$  is a constant representing the power scale in the gravitational law, and  $d_{\mathrm{soft}}$  is a small softening distance that ensures numerical stability. The motion occurs within a  $64\times 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  $x\in \mathbb{R}^{10\times 64\times 64\times 1}$  given the length-10 context  $y\in \mathbb{R}^{10\times 64\times 64\times 1}$ . 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.
Dataset is available at https://sevir.mit.edu/
$^3$ Code available at https://github.com/amazon-science/earth-forecasting-transformer/tree/main/src/earthformer/datasets/nbody
Figure 2: A set of example predictions on the  $N$ -body MNIST test set. From top to bottom: context sequence  $y$ , target sequence  $x$ , 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(\hat{x}^j)$  and the total energy of the last context frame  $E(y^{L_{\mathrm{in}}})$ . The red dashed line is to help the reader to judge the position of the digit “2” in the last frame.
Evaluation. In addition to standard metrics MSE, MAE and SSIM, we also report the scores of Fréchet Video Distance (FVD) [51], a 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$ -body MNIST and SEVIR datasets cannot be directly compared with those on natural video datasets. Nevertheless, the relative ranking of the FVD scores remains a meaningful indicator of model ability to achieve high visual quality, as FVD has shown consistency with expert evaluations across various domains beyond natural images [38, 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], ConvLSTM [47], PredRNN [61], PhyDNet [11], E3D-LSTM [60], Rainformer [1] and Earthformer [8], as well as two probabilistic spatiotemporal forecasting models: VideoGPT [65] and LDM [42]. All baselines are trained following the default configurations in their officially released code. More implementation details of baselines are provided in Appendix B.2. Results in Table 1 show that PreDiff outperforms these baselines by a large margin in both conventional video prediction metrics (i.e., MSE, MAE, SSIM), and a perceptual quality metric, FVD. The example predictions in Fig. 2 demonstrate that PreDiff generate predictions with sharp and clear digits in accurate positions. In 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 an ablation study that highlights the importance of the latent backbone's spatiotemporal modeling capacity. Specifically, the EarthformerUNet utilized in PreDiff demonstrates superior performance compared to the UNet in LDM [42].