SlowGuess/NeurIPS_2023_OCR huggingface dataset:
All experiments are conducted on machines with NVIDIA A10G 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 VQVAEs) 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 \times 4$ . A frame  $x^{j} \in \mathbb{R}^{64 \times 64 \times 1}$  is encoded to  $z^{j} \in \mathbb{R}^{16 \times 16 \times 3}$  by parameterizing  $p_{\mathcal{E}}(z^{j}|x^{j}) = \mathcal{N}(\mu_{\mathcal{E}}(x^{j})|\sigma_{\mathcal{E}}(x^{j}))$ . On SEVIR dataset, the spatial downsampling ratio is  $8 \times 8$ . A frame  $x^{j} \in \mathbb{R}^{128 \times 128 \times 1}$  is encoded to  $z^{j} \in \mathbb{R}^{16 \times 16 \times 4}$  similarly.
The detailed configurations of the encoder and decoder of the VAE on  $N$ -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 5 and Table 6. The discriminators for adversarial training on  $N$ -body MNIST and SEVIR datasets share the same configurations, which are shown in Table 7.
Latent diffusion model that instantiates  $p_{\theta}(z_{t-1}|z_t, z_{\mathrm{cond}})$  Stemming from Earthformer [8], we build Earthformer-UNet, which is a hierarchical UNet with self cuboid attention [8] layers as basic building blocks, as shown in Fig. 5.
On  $N$ -body MNIST, it takes the concatenation along the temporal dimension (the sequence length axis) of  $z_{\mathrm{cond}} \in \mathbb{R}^{10 \times 16 \times 16 \times 3}$  and  $z_t \in \mathbb{R}^{10 \times 16 \times 16 \times 3}$  as input, and outputs  $z_{t-1} \in \mathbb{R}^{10 \times 16 \times 16 \times 3}$ . On SEVIR, it takes the concatenation along the temporal dimension (the sequence length axis) of  $z_{\mathrm{cond}} \in \mathbb{R}^{7 \times 16 \times 16 \times 4}$  and  $z_t \in \mathbb{R}^{6 \times 16 \times 16 \times 4}$  as input, and outputs  $z_{t-1} \in \mathbb{R}^{6 \times 16 \times 16 \times 4}$ . Besides, we add the embedding of the denoising step  $t$  to the state in front of each cuboid attention block via an embedding layer TEmbed, following [22]. The detailed configurations of the Earthformer-UNet is described in Table 8.
Knowledge alignment networks A knowledge alignment network parameterizes  $U_{\phi}(z_t, t, y)$  to predict  $\mathcal{F}(\widehat{x}, y)$  using the noisy latent  $z_t$ . In practice, we build an Earthformer encoder [8] with a final pooling block as the knowledge alignment network to parameterize  $U_{\phi}(z_t, t, z_{\mathrm{cond}})$ , which takes  $t$ , and the concatenation of  $z_{\mathrm{cond}}$  and  $z_t$ , instead of  $t, y$  and  $z_t$  as the inputs. We find this implementation accurate enough when  $t$  is small. The detailed configurations of the knowledge alignment network is described in Table 9
Figure 5: Earthformer-UNet architecture. PreDiff employs an Earthformer-UNet as the backbone for parameterizing the latent diffusion model  $p_{\theta}(z_t|z_{t + 1},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) along the temporal dimension (the sequence length axis) as input, and outputs  $z_{t}$ . (Best viewed in color).
**Optimization** We train the frame-wise VAEs using the Adam optimizer [27] following [7]. We train the latent Earthformer-UNet and the knowledge alignment network using the AdamW optimizer [31] following [8]. Detailed configurations are shown in Table 10, Table 11 and Table 12 for the frame-wise VAE, the latent Earthformer-UNet 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.
Table 3: The details of the encoder of the frame-wise VAE on  $N$ -body MNIST frames. It encodes an input frame  $x^{j} \in \mathbb{R}^{64 \times 64 \times 1}$  into a latent  $z^{j} \in \mathbb{R}^{16 \times 16 \times 3}$ . Conv3  $\times$  3 is the 2D convolutional layer with  $3 \times 3$  kernel. GroupNorm32 is the Group Normalization (GN) layer [64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function  $\operatorname{SiLU}(x) = x \cdot \operatorname{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  $(x) = \operatorname{Softmax}(QK^T / \sqrt{C})V$ .
Block	Layer	Resolution	Channels
Input xj	-	64 × 64	1
2D CNN	Conv3 × 3	64 × 64	1 → 128
ResNet Block × 2	GroupNorm32	64 × 64	128
Conv3 × 3	64 × 64	128
GroupNorm32	64 × 64	128
Conv3 × 3	64 × 64	128
SiLU	64 × 64	128
Downsampler	Conv3 × 3	64 × 64 → 32 × 32	128
ResNet Block × 2	GroupNorm32	32 × 32	128
Conv3 × 3	32 × 32	128 → 256, 256
GroupNorm32	32 × 32	256
Conv3 × 3	32 × 32	256
SiLU	32 × 32	256
Downsampler	Conv3 × 3	32 × 32 → 16 × 16	256
ResNet Block × 2	GroupNorm32	16 × 16	256
Conv3 × 3	16 × 16	256 → 512, 512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU	16 × 16	512
Self Attention Block	GroupNorm32	16 × 16	512
Attention	16 × 16	512
Linear	16 × 16	512
ResNet Block × 2	GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU	16 × 16	512
Output Block	GroupNorm32	16 × 16	512
SiLU	16 × 16	512
Conv3 × 3	16 × 16	512 → 6
Conv3 × 3	16 × 16	6
Table 4: The details of the decoder of the frame-wise VAE on  $N$ -body MNIST frames. It decodes a latent  $z^{j} \in \mathbb{R}^{16 \times 16 \times 3}$  back to a frame in pixel space  $x^{j} \in \mathbb{R}^{64 \times 64 \times 1}$ . Conv3 × 3 is the 2D convolutional layer with  $3 \times 3$  kernel. GroupNorm32 is the Group Normalization (GN) layer [64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function  $\operatorname{SiLU}(x) = x \cdot \operatorname{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(x) = \operatorname{Softmax}(QK^{T}/\sqrt{C})V$ .
Block	Layer	Resolution	Channels
Input zj	-	16 × 16	3
2D CNN	Conv3 × 3	16 × 16	3
Conv3 × 3	16 × 16	3 → 512
Self Attention Block	GroupNorm32	16 × 16	512
Attention	16 × 16	512
Linear	16 × 16	512
ResNet Block ×3	GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU	16 × 16	512
Upsampler	Conv3 × 3	16 × 16 → 32 × 32	512
ResNet Block ×3	GroupNorm32	32 × 32	512
Conv3 × 3	32 × 32	512 → 256, 256, 256
GroupNorm32	32 × 32	256
Conv3 × 3	32 × 32	256
SiLU	32 × 32	256
Upsampler	Conv3 × 3	32 × 32 → 64 × 64	256
ResNet Block ×3	GroupNorm32	64 × 64	256
Conv3 × 3	64 × 64	256 → 128, 128, 128
GroupNorm32	64 × 64	128
Conv3 × 3	64 × 64	128
SiLU	64 × 64	128
Output Block	GroupNorm32	64 × 64	128
SiLU	64 × 64	128
Conv3 × 3	64 × 64	128 → 1
Table 5: The details of the encoder of the frame-wise VAE on SEVIR frames. It encodes an input frame  $x^{j} \in \mathbb{R}^{128 \times 128 \times 1}$  into a latent  $z^{j} \in \mathbb{R}^{16 \times 16 \times 4}$ . Conv3 × 3 is the 2D convolutional layer with  $3 \times 3$  kernel. GroupNorm32 is the Group Normalization (GN) layer [64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function  $\mathrm{SiLU}(x) = x \cdot \mathrm{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(x) = \mathrm{Softmax}(QK^T / \sqrt{C})V$ .
Block	Layer	Resolution	Channels
Input xj	-	128 × 128	1
2D CNN	Conv3 × 3	128 × 128	1 → 128
ResNet Block ×2	GroupNorm32	128 × 128	128
Conv3 × 3	128 × 128	128
GroupNorm32	128 × 128	128
Conv3 × 3	128 × 128	128
SiLU	128 × 128	128
Downsampler	Conv3 × 3	128 × 128 → 64 × 64	128
ResNet Block ×2	GroupNorm32	64 × 64	128
Conv3 × 3	64 × 64	128 → 256, 256
GroupNorm32	64 × 64	256
Conv3 × 3	64 × 64	256
SiLU	64 × 64	256
Downsampler	Conv3 × 3	64 × 64 → 32 × 32	256
ResNet Block ×2	GroupNorm32	32 × 32	256
Conv3 × 3	32 × 32	256 → 512, 512
GroupNorm32	32 × 32	512
Conv3 × 3	32 × 32	512
SiLU	32 × 32	512
Downsampler	Conv3 × 3	32 × 32 → 16 × 16	512
ResNet Block ×2	GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU	16 × 16	512
Self Attention Block	GroupNorm32	16 × 16	512
Attention	16 × 16	512
Linear	16 × 16	512
ResNet Block ×2	GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU.	16 × 16	512
Output Block	GroupNorm32	16 × 16	512
SiLU	16 × 16	512
Conv3 × 3	16 × 16	512 → 8
Conv3 × 3	16 × 16	8
Table 6: The details of the decoder of the frame-wise VAE on SEVIR frames. It decodes a latent  $z^{j} \in \mathbb{R}^{16 \times 16 \times 4}$  back to a frame in pixel space  $x^{j} \in \mathbb{R}^{128 \times 128 \times 1}$ . Conv3 × 3 is the 2D convolutional layer with  $3 \times 3$  kernel. GroupNorm32 is the Group Normalization (GN) layer [64] with 32 groups. SiLU is the Sigmoid Linear Unit activation layer [18] with function  $\operatorname{SiLU}(x) = x \cdot \operatorname{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(x) = \text{Softmax}(QK^{T}/\sqrt{C})V$ .
Block	Layer	Resolution	Channels
Input zj	-	16 × 16	4
2D CNN	Conv3 × 3	16 × 16	4
Conv3 × 3	16 × 16	4 → 512
Self Attention Block	GroupNorm32	16 × 16	512
Attention	16 × 16	512
Linear	16 × 16	512
ResNet Block ×3	GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
GroupNorm32	16 × 16	512
Conv3 × 3	16 × 16	512
SiLU	16 × 16	512
Upsampler	Conv3 × 3	16 × 16 → 32 × 32	512
ResNet Block ×3	GroupNorm32	32 × 32	512
Conv3 × 3	32 × 32	512
GroupNorm32	32 × 32	512
Conv3 × 3	32 × 32	512
SiLU	32 × 32	512
Upsampler	Conv3 × 3	32 × 32 → 64 × 64	512
ResNet Block ×3	GroupNorm32	64 × 64	512
Conv3 × 3	64 × 64	512 → 256, 256, 256
GroupNorm32	64 × 64	256
Conv3 × 3	64 × 64	256
SiLU	64 × 64	256
Upsampler	Conv3 × 3	64 × 64 → 128 × 128	256
ResNet Block ×3	GroupNorm32	128 × 128	256
Conv3 × 3	128 × 128	256 → 128, 128, 128
GroupNorm32	128 × 128	128
Conv3 × 3	128 × 128	128
SiLU	128 × 128	128
Output Block	GroupNorm32	128 × 128	128
SiLU	128 × 128	128
Conv3 × 3	128 × 128	128 → 1
Table 7: The details of the discriminator for the adversarial loss of on  $N$ -body MNIST and SEVIR frames. Conv4  $\times$  4 is the 2D convolutional layer with  $4\times 4$  kernel,  $2\times 2$  or  $1\times 1$  stride, and  $1\times 1$  padding. BatchNorm is the Batch Normalization (BN) layer [25]. The negative slope in LeakyReLU is 0.2.
Block	Layer	Resolution	Channels
N-body MNIST	SEVIR
Input xj	-	64 × 64	128 × 128	1
2D CNN	Conv4 × 4	64 × 64 → 32 × 32	128 × 128 → 64 × 64	1 → 64
Downsampler	LeakyReLU	32 × 32	64 × 64	64
Conv4 × 4	32 × 32 → 16 × 16	64 × 64 → 32 × 32	64 → 128
BatchNorm	16 × 16	32 × 32	128
Downsampler	LeakyReLU	16 × 16	32 × 32	128
Conv4 × 4	16 × 16 → 8 × 8	32 × 32 → 16 × 16	128 → 256
BatchNorm	8 × 8	16 × 16	256
Downsampler	LeakyReLU	8 × 8	16 × 16	256
Conv4 × 4	8 × 8 → 7 × 7	16 × 16 → 15 × 15	256 → 512
BatchNorm	7 × 7	15 × 15	512
Output Block	LeakyReLU	7 × 7	15 × 15	512
Conv4 × 4	7 × 7 → 6 × 6	15 × 15 → 14 × 14	1
AvgPool	6 × 6 → 1	15 × 15 → 1	1
Table 8: 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 to the input to indicate whether the input is the encoded observation  $z_{\mathrm{cond}}$  or the noisy latent  $z_{t}$ . 1 for  $z_{\mathrm{cond}}$  and 0 for  $z_{t}$ . Conv3  $\times$  3 is the 2D convolutional layer with  $3 \times 3$  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 is the Sigmoid Linear Unit activation layer [18] with function  $\operatorname{SiLU}(x) = x \cdot \operatorname{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 by a 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 a 2D input tensor with  $C$  channels into  $N$  non-overlapping  $p \times p$  patches and merges the spatial dimensions into channels, gets  $N \times 1$  patches with  $p^2 \cdot 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.
Block	Layer	Spatial Resolution	Channels
N-body MNIST	SEVIR
Input [zcond,zt]	-	16 × 16	3	4
Observation Mask	ConcatMask	16 × 16	3 → 4	4 → 5
Projector	GroupNorm32	16 × 16	4	5
SiLU	16 × 16	4	5
Conv3 × 3	16 × 16	4 → 256	5 → 256
GroupNorm32	16 × 16	256
SiLU	16 × 16	256
Dropout	16 × 16	256
Conv3 × 3	16 × 16	256
Positional Embedding	PosEmbed	16 × 16	256
Cuboid Attention Block ×4	TEmb	16 × 16	256
LayerNorm	16 × 16	256
Cuboid(T,1,1)	16 × 16	256
FFN	16 × 16	256
LayerNorm	16 × 16	256
Cuboid(1,H,1)	16 × 16	256
FFN	16 × 16	256
LayerNorm	16 × 16	256
Cuboid(1,1,W)	16 × 16	256
FFN	16 × 16	256
Downsampler	PatchMerge	16 × 16 → 8 × 8	256 → 1024
LayerNorm	8 × 8	1024
Linear	8 × 8	1024
Cuboid Attention Block ×8	TEmb	8 × 8	1024
LayerNorm	8 × 8	1024
Cuboid(T,1,1)	8 × 8	1024
FFN	8 × 8	1024
LayerNorm	8 × 8	1024
Cuboid(1,H,1)	8 × 8	1024
FFN	8 × 8	1024
LayerNorm	8 × 8	1024
Cuboid(1,1,W)	8 × 8	1024
FFN	8 × 8	1024
Upsampler	NearestNeighborInterp	8 × 8 → 16 × 16	1024
Conv3 × 3	16 × 16	1024 → 256
Cuboid Attention Block ×4	TEmb	16 × 16	256
LayerNorm	16 × 16	256
Cuboid(T,1,1)	16 × 16	256
FFN	16 × 16	256
LayerNorm	16×16	256
Cuboid(1,H,1)	16 × 16	256
FFN	16 × 16	256
LayerNorm	16 × 16	256
Cuboid(1,1,W)	16 × 16	256
FFN	256
Output Block	Linear	16 × 16	256 → 3	256 → 4
Table 9: The details of the Earthformer encoders for the parameterization of the knowledge alignment networks  $U_{\phi}(z_t, t, z_{\mathrm{cond}})$  on  $N$ -body MNIST and SEVIR datasets. The ConcatMask layer for the Observation Mask block concatenates one more channel to the input to indicate whether the input is the encoded observation  $z_{\mathrm{cond}}$  or the noisy latent  $z_t$ . 1 for  $z_{\mathrm{cond}}$  and 0 for  $z_t$ . Conv3 × 3 is the 2D convolutional layer with  $3 \times 3$  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 is the Sigmoid Linear Unit activation layer [18] with function  $\mathrm{SiLU}(x) = x \cdot \mathrm{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 by a 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 a 2D input tensor with  $C$  channels into  $N$  non-overlapping  $p \times p$  patches and merges the spatial dimensions into channels, gets  $N \times 1$  patches with  $p^2 \cdot 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  $(x) = \text{Softmax}(QK^T / \sqrt{C})V$ . Finally, the value of the "cls" token after self attention operation serves as the layer's output.
Block	Layer	Spatial Resolution	Channels
N-body MNIST	SEVIR
Input [zcond,zt]	-	16 × 16	3	4
Observation Mask	ConcatMask	16 × 16	3 → 4	4 → 5
Projectile	GroupNorm32	16 × 16	4	5
SiLU	16 × 16	4	5
Conv3 × 3	16 × 16	4 → 64	5 → 64
GroupNorm32	16 × 16	64
SiLU	16 × 16	64
Dropout	16 × 16	64
Conv3 × 3	16 × 16	64
Positional Embedding	PosEmbed	16 × 16	64
Cuboid Attention Block	TEmbd	16 × 16	64
LayerNorm	16 × 16	64
Cuboid(T,1,1)	16 × 16	64
FFN	16 × 16	64
LayerNorm	16 × 16	64
Cuboid(1,H,1)	16 × 16	64
FFN	16 × 16	64
LayerNorm	16 × 16	64
Cuboid(1,1,W)	16 × 16	64
FFN	16 × 16	64
Downsampler	PatchMerge	16 × 16 → 8 × 8	64 → 256
LayerNorm	8 × 8	256
Linear	8 × 8	256
Cuboid Attention Block	TEmbd	8 × 8	256
LayerNorm	8 × 8	256
Cuboid(T,1,1)	8 × 8	256
FFN	8 × 8	256
LayerNorm	8 × 8	256
Cuboid(1,H,1)	8 × 8	256
FFN	8 × 8	256
LayerNorm	8 × 8	256
Cuboid(1,1,W)	8 × 8	256
FFN	8 × 8	256
Output Pooling Block	GroupNorm32	8 × 8	256
Attention	8 × 8 → 1	256
Linear	1	256 → 1
Table 10: Hyperparameters of the Adam optimizer for training frame-wise VAEs and discriminators on  $N$ -body MNIST and SEVIR datasets.
Hyper-parameter of VAE	Value
Learning rate	4.5 × 10-6
β1	0.5
β2	0.9
Weight decay	10-2
Batch size	512
Training epochs	200
Hyper-parameter of discriminator	Value
Learning rate	4.5 × 10-6
β1	0.5
β2	0.9
Weight decay	10-2
Batch size	512
Training epochs	200
Training start step	50000
Table 11: Hyperparameters of the AdamW optimizer for training LDMs on  $N$ -body MNIST and SEVIR datasets.
Hyper-parameter of VAE	Value
Learning rate	1.0 × 10-3
β1	0.9
β2	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$ -body MNIST and SEVIR datasets.
Hyper-parameter of VAE	Value
Learning rate	1.0 × 10-3
β1	0.9
β2	0.999
Weight decay	10-5
Batch size	64
Training epochs	200
Warm up percentage	10%
Learning rate decay	Cosine
We train baseline algorithms following their officially released configurations and tune the learning rate, learning rate scheduler, working resolution, etc., to optimize their performance on each dataset. We list the modifications we applied to the baselines for each dataset in Table 13.
Table 13: Implementation details of baseline algorithms. Modifications based on the officially released implementations are listed according to different datasets. “-” means no 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]	-	-
ConvLSTM [47]	reverse enc-dec [48]
conv_kernels = [(7,7),(5,5),(3,3)]
deconv_kernels = [(6,6),(4,4),(4,4)]
channels=[96, 128, 256]	reverse enc-dec [48]
conv_kernels = [(7,7),(5,5),(3,3)]
deconv_kernels = [(6,6),(4,4),(4,4)]
channels=[96, 128, 256]
PredRNN [61]	-	-
PhyDNet [11]	-	convcell Hidden = [256, 256, 256, 64]
E3D-LSTM [60]	-	-
Rainformer [1]	downscalingFactors=[2, 2, 2, 2]
hidden_dim=32
heads=[4, 4, 8, 16]
head_dim=8	downscalingFactors=[4, 2, 2, 2]
-
-
-
Earthformer [8]	-	-
DGMR [41]	-	context_steps = 7
forecast_steps = 6
VideoGPT [65]	vqvae_n_codes = 512
vqvae_downsample = [1, 4, 4]	vqvae_n_codes = 512
vqvae_downsample = [1, 8, 8]
LDM [42]	vae: 64 × 64 × 1 → 16 × 16 × 3
conv_dim = 3
model_channels = 256	vae: 128 × 128 × 1 → 16 × 16 × 4
conv_dim = 3
model_channels = 256
C Derivation of the Approximation to Knowledge Alignment Guidance
We derive the approximation to the knowledge alignment guided denoising transition (5) following [4]. We rewrite (5) to (8) using a normalization constant  $Z$  that normalizes  $Z\int e^{-\lambda_{\mathcal{F}}\| U_{\phi}(z_t,t,y) - \mathcal{F}_0(y)\|}dz_t = 1$ :
$$
p _ {\theta , \phi} \left(z _ {t} \mid z _ {t + 1}, y, \mathcal {F} _ {0}\right) = p _ {\theta} \left(z _ {t} \mid z _ {t + 1}, z _ {\text {c o n d}}\right) \cdot Z e ^ {- \lambda_ {\mathcal {F}} \left\| U _ {\phi} \left(z _ {t}, t, y\right) - \mathcal {F} _ {0} (y) \right\|}. \tag {8}
$$
In what follows, we abbreviate  $\mu_{\theta}(z_{t + 1},t,z_{\mathrm{cond}})$  as  $\mu_{\theta}$ , and  $\Sigma_{\theta}(z_{t + 1},t,z_{\mathrm{cond}})$  as  $\Sigma_{\theta}$  for brevity. We use  $C_i,i = \{1,\ldots ,7\}$  to denote constants.
$$
p _ {\theta} \left(z _ {t} \mid z _ {t + 1}, z _ {\text {c o n d}}\right) = \mathcal {N} \left(\mu_ {\theta}, \Sigma_ {\theta}\right),
$$
$$
\log p _ {\theta} \left(z _ {t} \mid z _ {t + 1}, z _ {\text {c o n d}}\right) = - \frac {1}{2} \left(z _ {t} - \mu_ {\theta}\right) ^ {T} \Sigma_ {\theta} ^ {- 1} \left(z _ {t} - \mu_ {\theta}\right) + C _ {1}, \tag {9}
$$
$$
\log Z e ^ {- \lambda_ {\mathcal {F}} \| U _ {\phi} (z _ {t}, t, y) - \mathcal {F} _ {0} (y) \|} = - \lambda_ {\mathcal {F}} \| U _ {\phi} (z _ {t}, t, y) - \mathcal {F} _ {0} (y) \| + C _ {2},
$$
By assuming that  $\log Ze^{-\lambda_{\mathcal{F}}\| U_{\phi}(z_t,t,y) - \mathcal{F}_0(y)\|$  has low curvature compared to  $\Sigma_{\theta}^{-1}$ , which is reasonable in the limit of infinite diffusion steps ( $\| \Sigma_{\theta}\| \to 0$ ), we can approximate it by a Taylor expansion at  $z_{t} = \mu_{\theta}$
$$
\begin{array}{l} \log Z e ^ {- \lambda_ {\mathcal {F}} \| U _ {\phi} (z _ {t}, t, y) - \mathcal {F} _ {0} (y) \|} \approx - \lambda_ {\mathcal {F}} \| U _ {\phi} (z _ {t}, t, y) - \mathcal {F} _ {0} (y) \| | _ {z _ {t} = \mu_ {\theta}} \\ - \left(z _ {t} - \mu_ {\theta}\right) \lambda_ {\mathcal {F}} \nabla_ {z _ {t}} \| U _ {\phi} (z _ {t}, t, y) - \mathcal {F} _ {0} (y) \| | _ {z _ {t} = \mu_ {\theta}} \tag {10} \\ = \left(z _ {t} - \mu_ {\theta}\right) g + C _ {3}, \\ \end{array}
$$
where  $g = -\lambda_{\mathcal{F}}\nabla_{z_t}\| U_\phi (z_t,t,y) - \mathcal{F}_0(y)\| |_{z_t = \mu_\theta}$ . By taking the log of (8) and applying the results from (9) and (10), we get
$$
\begin{array}{l} \log p _ {\theta , \phi} \left(z _ {t} \mid z _ {t + 1}, y, \mathcal {F} _ {0}\right) = \log p _ {\theta} \left(z _ {t} \mid z _ {t + 1}, z _ {\text {c o n d}}\right) + \log Z e ^ {- \lambda_ {\mathcal {F}} \left\| U _ {\phi} \left(z _ {t}, t, y\right) - \mathcal {F} _ {0} (y) \right\|} \\ \approx - \frac {1}{2} (z _ {t} - \mu_ {\theta}) ^ {T} \Sigma_ {\theta} ^ {- 1} (z _ {t} - \mu_ {\theta}) + (z _ {t} - \mu_ {\theta}) g + C _ {4} \\ = - \frac {1}{2} \left(z _ {t} - \mu_ {\theta} - \Sigma_ {\theta} g\right) ^ {T} \Sigma_ {\theta} ^ {- 1} \left(z _ {t} - \mu_ {\theta} - \Sigma_ {\theta} g\right) + \frac {1}{2} g ^ {T} \Sigma_ {\theta} g + C _ {5} \tag {11} \\ = - \frac {1}{2} \left(z _ {t} - \mu_ {\theta} - \Sigma_ {\theta} g\right) ^ {T} \Sigma_ {\theta} ^ {- 1} \left(z _ {t} - \mu_ {\theta} - \Sigma_ {\theta} g\right) + C _ {6} \\ = \log p (z) + C _ {7}, z \sim \mathcal {N} \left(\mu_ {\theta} + \Sigma_ {\theta} g, \Sigma_ {\theta}\right). \\ \end{array}
$$
Therefore, the transition distribution under the guidance of knowledge alignment shown in (5) can be approximated by a Gaussian similar to the transition without knowledge guidance, but with its mean shifted by  $\Sigma_{\theta}$ .
D More Quantitative Results on SEVIR
D.1 Quantitative Analysis of BIAS on SEVIR
Similar to Critical Success Index (CSI) introduced in Sec. 3.2,  $\text{BIAS} = \frac{\# \text{Hits} + \# \text{F.Alarms}}{\# \text{Hits} + \# \text{Misses}}$  is calculated by counting the  $\# \text{Hits}$  (truth=1, pred=1),  $\# \text{Misses}$  (truth=1, pred=0) and  $\# \text{F.Alarms}$  (truth=0, pred=1) of the predictions binarized at thresholds [16, 74, 133, 160, 181, 219]. This measurement assesses the model's inclination towards either F.Alarms or Misses.