Loading... # PreDiff:Precipitation Nowcasting with Latent Diffusion Models ## introduction 这部分首先介绍了降水预报的重要性，并说明随着数据量的增长，一些深度学习方法的出现已经对齐传统数值方法的性能；接着指出现有的深度学习方法的缺陷：由于地球气候系统的复杂性，现有模型无法捕捉多种可能的未来状态，并倾向于生成忽略关键细节的模糊预报，为此，能够表征系统固定不确定性的概率模型应运而生。接着作者介绍了扩散模型相较于对抗生成网络和自回归模型的优势以及其在地球预报系统中的起步阶段的现实，指出了现在这种纯数据驱动的研究缺乏来自动态系统先验知识约束的问题，并由此提出了作者的基于Earthformer的pre-diff模型。 ## Method 这节作者介绍了他的思路：将降水预报问题建模为时空预测问题。 - 输入：$y = [y^j]_{j = 1} ^{L_{in}}\in \mathbb{R}^{L_{in}\times H\times W\times C}$ - 输出：$x = [x^j]_{j = 1} ^{L_{out}}\in \mathbb{R}^{L_{out}\times H\times W\times C}$ - 模型：在给定输入$y$下建立其条件概率分布$p(x|y)$ 同时，作者还给出了模型的推理过程示意图： <img src="https://y0k1n0-1323330522.cos.ap-beijing.myqcloud.com/image-20251128204012197.png" alt="image-20251128204012197" style="zoom:67%;" style=""> 模型将输入序列通过编码得到对应的潜在向量$z_{cond}$，然后将潜在向量$z_{cond}$和纯噪声一起送入模型，通过逐步去噪获得预测图像$\hat{x}$。其中先验知识对齐到相应章节再聊。 ### Preliminary: Diffusion Models 这部分简要介绍了一下扩散模型： - 前向加噪：$q(x_t|x_{t-1}) = \mathcal{N}\left(x_t; \sqrt{\alpha_t}x_{t-1}, (1-\alpha_t)\boldsymbol{I}\right)(1 \leq t \leq T)$，其中$x_0\sim p(x)$是真实数据，$x_T \sim \mathcal{N}(0, \boldsymbol{I})$ 是随机噪声 - 反向去噪：$p_{\theta}(x_{0:T}) = p(x_T)\prod_{t=1}^{T} p_{\theta}(x_{t-1}\mid x_t)$，这也是一个马尔科夫链过程，其中反向去噪的每一步都是一个高斯分布$p_{\theta}(x_{t-1}\mid x_t) = \mathcal{N}\left(\mu_{\theta}(x_t, t),\, \Sigma_{\theta}(x_t, t)\right)$，表示模型从$x_t$步预测得到$x_{t-1}$步的过程。 为了方便理解，本文在重新推导一下扩散模型的前向过程和反向过程。 #### 重参数化技巧 首先需要理解VAE所引入的重参数化技巧： - 在VAE中，编码器的输出是隐变量的分布参数，然后从分布中采样隐变量$z\sim q(z|x)=\mathcal{N}(\mu(x), \sigma^2(x)I)$ - 解码器基于$z$重构输入$x$，损失函数包含重构损失和KL散度损失 - 如果直接采样 $z \sim q(z|x)$，则梯度无法从$z$传递到编码器的参数$\mu$和$\sigma^2$，因为采样操作会 “阻断” 梯度流。 - 为此可以为隐变量的分布引入可训练参数，假设$\epsilon \sim \mathcal{N}(0, I)$， 编码器输出的隐变量的分布为$z=\mu(x)+\sigma(x)*\epsilon$，这样即可保证输出的分布与原分布一致的前提下，不丢失梯度信息 #### 高斯分布的条件概率公式 #### 前向过程 最直接的前向加噪过程的公式如下所示： $$ x_t = \sqrt{\alpha_t}\, x_{t-1} + \sqrt{1-\alpha_t}\, \epsilon_t,\quad \epsilon_t \sim \mathcal{N}(0,\boldsymbol{I}) $$ 其中，$\alpha_t$为权重参数 将$x_{t-1}$用$x_{t-1} = \sqrt{\alpha_t}x_{t-2} + \sqrt{1-\alpha_t}$替代，以此迭代可得： $$ x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\,\epsilon_t,\quad \epsilon_t \sim \mathcal{N}(0,\boldsymbol{I}), \bar{\alpha}_t = \prod_{s=1}^{t} \alpha_s $$ 考虑重参数化，实际上就是在均值上偏移$\sqrt{\bar{\alpha}_t}\, x_0$, 在方差上偏移$(\sqrt{1-\bar{\alpha}_t})^2$，即： $$ x_t \sim p(x_t|x_0) = \mathcal{N}\left(\sqrt{\bar{\alpha}_t}\,x_0,\ (1-\bar{\alpha}_t)\boldsymbol{I}\right) $$ #### 反向过程 反向过程实际上就是在求$p(x_{t-1}|x_t)$，根据贝叶斯定理有： $$ p(x_{t-1}\mid x_t) = \frac{p(x_t\mid x_{t-1})\,p(x_{t-1})} {p(x_t)} $$ 其中$p(x_t\mid x_{t-1})$的分布是已知的，但$p(x_{t-1})$和$p(x_t)$的分布是未知的，但考虑到这个反向过程实际上是在已知$x_0$的情况下的，因此原式可以写为： $$ p(x_{t-1}\mid x_t,x_0) = \frac{p(x_t\mid x_{t-1},x_0)\,p(x_{t-1}\mid x_0)} {p(x_t\mid x_0)} $$ 由此我们已知三个概率的概率分布： - $p(x_{t-1} \mid x_0) \sim \mathcal{N}\Big(x_{t-1}; \, \sqrt{\bar{\alpha}_{t-1}} x_0, \, (1 - \bar{\alpha}_{t-1}) I \Big)$ - $p(x_t \mid x_0) \sim \mathcal{N}\Big(x_t; \, \sqrt{\bar{\alpha}_t} x_0, \, (1 - \bar{\alpha}_t) I \Big)$ - $p(x_t \mid x_{t-1}, x_0) = p(x_t \mid x_{t-1}) \sim \mathcal{N}\Big(x_t; \, \alpha_t x_{t-1}, \, (1-\alpha_t) I \Big)$ 这里将他们展开成数学表达式的形式，然后进行运算可得： $$ p(x_{t-1}\mid x_t, x_0)=\mathcal{N}\left( x_{t-1}; \frac{\sqrt{\alpha_t(1-\bar{\alpha}_{t-1})}}{1-\bar{\alpha}_{t}}x_t+\frac{\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t)}{1-\bar{\alpha}_t}x_0, \left( \frac{\sqrt{1-\alpha_t}\sqrt{1-\bar{\alpha}_{t-1}}}{\sqrt{1-\bar{\alpha}_t}} \right)^2 \right) $$ 我们希望$x_{t-1}$只和$x_t$有关，因此这里考虑使用$x_0$和$x_t$的关系消去$x_0$：$x_0=\frac{x_t-\sqrt{1-\bar{\alpha}_t}\epsilon}{\sqrt{\bar{\alpha}_t}}$ 最终可得： $$ \begin{eqnarray} p(x_{t-1}\mid x_t, x_0) &=&\mathcal{N}\left( x_{t-1}; \frac{\sqrt{\alpha_t(1-\bar{\alpha}_{t-1})}}{1-\bar{\alpha}_{t}}x_t+\frac{\sqrt{\bar{\alpha}_{t-1}}(1-\alpha_t)}{1-\bar{\alpha}_t} \times \frac{x_t-\sqrt{1-\bar{\alpha}_t}\epsilon}{\sqrt{\bar{\alpha}_t}}, \left( \frac{\sqrt{1-\alpha_t}\sqrt{1-\bar{\alpha}_{t-1}}}{\sqrt{1-\bar{\alpha}_t}} \right)^2 \right)\\ \end{eqnarray} $$ ### Conditional Diffusion in Latent Space 这里作者讲了条件扩散模型的训练方式：采用两段式训练： - 先训练一个逐帧变分自编码器 - 然后再训练一个条件扩散模型 #### Frame-wise autoencoder 这里讲了逐帧编码器的输入输出和损失函数： - 输入：单帧像素数据$x^j\in\mathbb{R}^{H\times W\times C}$ - 潜在变量：$z=[z^j]$，其中$z^j = \mathcal{E}(x^j)\in \mathcal{R}^{H_z\times W_z\times C_z}$ - 输出：原始像素帧$\hat{x}^j = \mathcal{D}(z^j)$ - 损失函数：像素损失+对抗损失 根据附录中的训练细节可知，这里引入了判别器从而增强生成帧的感知真实性 #### Latent diffusion 有了上面的基础知识，我们对于扩散模型的正向过程和反向过程有了更深入的了解。 我们知道，对于给定的输入上下文$y$，由编码器编码可得潜在变量$z_{cond}\sim p_{\mathcal{E}}(z_{cond\mid y})$，而反向去噪过程就是基于马尔科夫链：$p_{\theta}(z_{0:T\mid z_{cond}}) = p(z_T)\prod_{t=1}^{T} p_{\theta}(z_{t-1}\mid z_t, z_{cond})$。 扩散模型并没有直接预测$z_{t-1}$，而是学习第$t$步的转移噪声$\epsilon_\theta(z_t, t)$。而模型在前向过程中，第$t$步添加的噪声为$\epsilon = \frac{z_t - \sqrt{\bar{\alpha}_t}z_0}{\sqrt{1-\bar{\alpha}_t}}$。损失函数计算的就是预测噪声与真实噪声的差： $$ \mathcal{L}_{\mathrm{CLDM}} = \mathbb{E}_{x, y, t, \epsilon \sim \mathcal{N}(0, I)} \left[ \left\| \epsilon - \epsilon_{\theta}(z_t, t, z_{\mathrm{cond}}) \right\|_2^2 \right] $$ 作者采用的*Earthformer-UNet*模型结果如下所示： <img src="https://y0k1n0-1323330522.cos.ap-beijing.myqcloud.com/image-20251129193837068.png" alt="image-20251129193837068" style="zoom:50%;" style=""> #### Instantiating $p_{\theta}(z_{t-1}\mid z_t, z_{cond})$ 这部分作者讲了选择*Earthformer-UNet*作为Unet的原因：对于更高维度的时空观测数据，该编码器能够更好的捕捉其中广泛的时空依赖关系。 ### Incorporating Knowledge Alignment 本节首先指出了扩散模型可能存在的两个问题： - 训练数据未必完全符合领域知识，从分布中采样得到的结果未必能保证物理可实现性 - 混沌系统本身的随机性和去噪步骤中的近似误差 为此作者提出知识对齐策略，将辅助先验知识融入扩散生成过程，利用领域专业知识约束生成结果： $$ F(\hat{x}, y)=F_0(y)\in\mathbb{R}^d $$ 其中，损失函数的偏差项$\left\| F(x_b, y) - F_0(y) \right\|$用于量化预测结果和先验知识之间的偏离程度，从而抑制生成高偏差结果的概率。 本文模型采用知识对齐网络$U_\phi(z_t, t, y)$，基于加噪步骤$t$时刻的中间潜态$z_t$去估计预测结果$\hat{x}$对应的约束值$F(\hat{x}, y)$，采用$\left\| F(x_b, y) - F_0(y) \right\|$作为损失函数引导知识对齐网络的训练。 <img src="https://y0k1n0-1323330522.cos.ap-beijing.myqcloud.com/image-20251129184917032.png" alt="image-20251129184917032" style="zoom:50%;" style=""> 如上图所示，知识对齐网络通过调整潜在空间的均值从而影响采样结果，从而减少可能的高偏差结果，其数学表达式如下所示： $$ p_{\theta,\phi}(z_t \mid z_{t+1}, y, F_0) \propto p_{\theta}(z_t \mid z_{t+1}, z_{\mathrm{cond}}) \cdot e^{-\lambda_F \left\| U_{\phi}(z_t, t, y) - F_0(y) \right\|} $$ 本文模型中采用基于*Earthformer-UNet*的encoder去拟合$U_\phi(z_t, t, y)$ ## Experiments 这部分讲了模型在N-body MNIST和SEVIR上的模型性能评价，由于我主要关注临近预报任务，因此我这里主要看了在降水临近预报上的实验 ### SEVIR Precipitation Nowcasting #### Dataset 这部分首先介绍了SEVIR数据集，SEVIR数据集由三部分组成： - 卫星图像数据 - VIL垂直液态降水积分图像 - 雷电数据 作者这里选择的是下采样得到的降水数据，以70分钟的VIL作为输入，去预测60分钟的输出。 #### Evaluation 作者这里选择了三个评价指标： - CSI：关键成功指数。 - 作者通过将预测结果和真实值重缩至0-255的区间后采用分段阈值处理数据 - 此外，还引入了在4x4、16x16尺度下的CSI，以捕捉局部模式分布 - FVD：弗雷歇视频距离 - 用于评估研究方法的视觉质量 - CRPS：连续排序概率得分 - 用于评估不确定性建模能力 #### Comparison to the State of the Art 本节是不同工作之间的横向对比。作者在对比的过程中指出CSI判别方法在临近预报任务中的不可取之处：确定性模型为了提高指标会抛弃真实模式，而是采用平均化的方法给出一个大致预测；而概率模型倾向于捕获数据分布。 #### Knowledge Alignment: Anticipated Average Intensity 这节讨论了知识对齐网络在模型中的设计和作用： - 知识对齐网络设计： - 将序列的平均强度记作$I(x)\in \mathbb{R}^+$ - 将上下文平均强度序列$[I(y^j)]_{j=1}$记作$[I(y^j)]$ - 训练一个简单的概率型时间序列预测模型，基于目前的序列平均强度预测未来的强度$p_{\tau}(I(x) | [I(y^j)]) = \mathcal{N}\left(\mu_{\tau}([I(y^j)]), \sigma_{\tau}([I(y^j)])\right)$ - 应用到*Earthformer-UNet*： - 给定初始条件$y$下的序列平均强度为$\mathcal{F}_0(y)\equiv\mu_\tau + n\sigma_\tau$ - 在给定初始条件$y$下的预测强度$\mathcal{F}(\hat{x}, y)\equiv I(\hat{x})$ - 根据损失函数$\mathcal{L}_u$可知，$\mathcal{F}(\hat{x}, y)$应该趋近于$\mathcal{F}_0(y)$ - 当$n\in\{-1, 0, 1\}$时，模型即可预测常规情况 - 当$n\notin\{-1, 0, 1\}$时，模型即可预测极端天气情况 ## 总结和疑问 这篇文章好像作者再用diffusion去解决天气预测的问题，但又好像不只是想解决天气预测问题。作者似乎是想要提出一个通用的范式，通过可嵌入式的知识对齐网络＋适用于复杂非线性的时空特征输入的网络去解决高限制的现实问题。这样做的好处是对通用模型进行了一定的探究，但是缺点是模型太复杂了.... 但是作者关于diffusion进行临近预测有几点我十分认同： - CSI指标令判别模型生成更加平均和模糊，生成模型具备更大的可能性 - 知识对齐应用的优势：增强了生成结果的合理性和可解释性，提高了生成质量 最后修改：2025 年 11 月 30 日 © 允许规范转载 赞 如果觉得我的文章对你有用，请随意赞赏