Precipitation Downscaling with Spatiotemporal Video Diffusion

Status: completed

Authors: Prakhar Srivastava, Ruihan Yang, Gavin Kerrigan, Gideon Dresdner, Jeremy McGibbon, Christopher Bretherton, Stephan Mandt

Venue / Year: Advances in Neural Information Processing Systems (NeurIPS), 2024

Links: NeurIPS proceedings PDF · arXiv:2312.06071 · OpenReview forum · official code · proceedings DOI

Tags: [[气象降尺度]] [[降水]] [[概率降尺度]] [[Diffusion Model]] [[Video Super-Resolution]] [[频谱一致性]] [[极端事件]]

Screening Assessment

维度	分数	依据
创新性	24/25	首次将多帧视频超分辨率与条件残差扩散结合到全球降水降尺度；deterministic base + stochastic residual 和时空分解注意力都直接服务气象结构。
严谨性	24/25	给出条件分布、v-parameterization、训练/采样算法和注意力复杂度设计；对上下文长度、额外变量和采样步数有消融。
主题相关性	25/25	直接评价降水分布、99.999% 极端分位数、空间自相关和高频谱，而非只用视觉指标。
实验	15/15	与六个强图像/视频 SR 基线比较，并报告 CRPS、EMD、MSE、尾部误差、SAE、区域年平均和谱图。
复现/迁移	8/10	官方代码、环境、数据路径和训练入口公开；未发布预训练权重，且配置中的扩散步数/模型开关与论文实验存在小差异。
总分	96/100	严谨性与主题相关性均高于 15，无否决项，准予完整笔记。

One-Sentence Summary

STVD 先用带时空分解注意力的 UNet 对五帧粗分辨率序列做确定性上采样，再对高分辨率残差进行条件 DDPM/DDIM 采样，从而同时保持天气系统的时间连贯性、降水空间高频和极端尾部概率。

Story and Section Missions

• §1 Introduction： 将降水降尺度定位为 one-to-many 条件分布问题，指出 MSE 模式平均会抹去暴雨，GAN 又可能模式坍塌。
• §2 Problem/Method： 先给出 $p(y_{0:T}\mid x_{0:T})$，再拆成 deterministic mean 和 residual diffusion；说明五帧联合推理如何避免逐帧不连续。
• §2.1.4/A.1： 解释两套 UNet、跨模块特征拼接、时间/空间 factorized attention 以及 v-parameterization。
• §3 Experiments： 用 FV3GFS 11-member ensemble、200 km→25 km 配对资料，与 SwinIR、VRT、RVRT、PSRT、SwinIR-Diff、VDM 做统一比较。
• §5 Conclusion/Limitations： 强调分布和极端事件收益，同时承认必须使用配对 LR-HR，不能直接滚动生成无限长序列。

Problem

设五帧输入 $x_{0:T}\in\mathbb{R}^{C\times H\times W}$，目标降水帧 $y_{0:T}\in\mathbb{R}^{1\times sH\times sW}$，本文 $T+1=5$、$C=12$、$s=8$。由于未解析天气尺度存在多种合理细节，目标不是单一回归函数，而是

$$p_{\theta}(y_{0:T}\mid x_{0:T})\approx p(y_{0:T}\mid x_{0:T}).$$

训练假设每个高分辨率序列都和对应的粗场由面积加权聚合得到，因而可形成 paired $(x_{0:T},y_{0:T})$。这个假设适合 FV3GFS 同一模拟的网格 coarsening，却不适用于只有粗模式、没有高分辨率参考的现实 ESM 资料。

Method

Residual diffusion decomposition

确定性下采样器先输出

$$\bar y_{0:T}=\mu_{\phi}(x_{0:T}),$$

然后定义稀疏、易建模的残差 $r_{0:T}=y_{0:T}-\bar y_{0:T}$。最终样本为

$$\hat y_{0:T}=\bar y_{0:T}+r_{0:T},\qquad r_{0:T}\sim p_{\psi}(r_{0:T}\mid x_{0:T},\bar y_{0:T}).$$

前向扩散取

$$r_n=\alpha_n r_0+\sigma_n\epsilon,\qquad v=\alpha_n\epsilon-\sigma_n r_0,$$

其中 $\epsilon\sim\mathcal N(0,I)$。网络 $M_{\psi}$ 以 $r_n$、扩散步 $n$、粗场与 base 特征为条件，最小化

$$\mathcal L(\psi,\phi)=\mathbb E\sum_{t=0}^{T}\left\|v_t-M_{\psi}(r_{n,t},n,c)\right\|_2^2.$$

论文采用 $N=1400$ 的训练深度和 DDIM 30 步推理；官方配置文件把 diffusion_steps 默认写为 1500，复现实验需以论文设置为准并记录该差异。

Architecture and data flow

模块	数据流、维度与内部结构	作用
Deterministic UNet	输入先逐帧 bicubic 到 $sH\times sW$，经 6 层 encoder/decoder；高分辨率 tile 为 $384\times384$，base channel 64，multipliers $(1,1,2,2,3,4)$。输出五帧 $\bar y$。	将粗尺度天气状态与地形/多变量信息变成可用的均值场。
ResBlock	两个 weight-standardized 3×3 Conv，GroupNorm（8 groups）、SiLU，并以 1×1 Conv 调整通道；encoder/decoder 之间有 skip connection。	稳定高分辨率特征提取，并保留局部梯度。
Factorized attention	bottleneck 使用空间和时间 quadratic attention；大 token 层使用 patch 内 linear spatial attention 和按通道 factorized temporal attention；4 heads，每 head 32 维。	让五个时刻相互注意，同时避免在完整视频立方体上做昂贵的全局注意力。
Diffusion UNet	输入拼接 $x_{0:T}$、$\bar y_{0:T}$ 和 $r_n$；接受 deterministic UNet encoder 特征的跨模块拼接，并通过扩散步 MLP 注入每个 ResBlock；输出 $v$。	对均值难以表达的细粒度残差进行条件随机采样。
时间条件	32 维随机 Fourier/sinusoidal features → Linear → GELU → Linear	把离散噪声步编码为各层条件；不是天气时间戳。

训练时两套 UNet 端到端共同优化；推理先计算一次 base，再从高斯噪声开始用 DDIM 逐步生成完整五帧残差，帧之间通过 temporal attention 并发耦合，而非逐帧自回归。EMA decay 为 0.995，训练 batch size 为 1，梯度裁剪 1.0（代码 Trainer）。

Downscaling Data Audit

项目	论文中的精确情况	审计判断
降尺度前数据	NOAA FV3GFS 全球大气模式的 200 km cubed-sphere 粗场，每 tile $48\times48$；输入 12 个通道，包括粗降水、地形、风、温度、气压和水汽相关变量。	低分辨率与目标来自同一模式模拟，非独立观测。
降尺度后数据	FV3GFS 25 km fine grid，每 tile $384\times384$；目标为 3-hourly average precipitation。	线性倍率 $s=8$，且粗细网格尺寸为 $48\rightarrow384$。
Ground Truth 精确类型	同一 FV3GFS 11-member initial-condition ensemble、13-month simulation 的 fine-grid PRATEsfc；首月丢弃 spin-up，10 年训练、1 年验证。	GT 是动力模式高分辨率模拟，不是雨量计或独立再分析观测。
LR–HR 对齐	通过同一 fine-grid 序列按 factor 8 coarsening 得到 paired coarse sequence；五个连续 3 小时帧联合采样。	paired=是，空间聚合严格同源；时间上只覆盖 5 帧，不支持长 rollout。
训练与应用差异	训练随机抽取六个 cubed-sphere tile；测试时完整 tile/局地 Himalaya、California 区域分析。输入状态仍来自 FV3GFS 分布。	没有跨模式、跨观测域或真实未来气候 OOD。
实际意义	中（RCM/高分辨率模式 emulation）：§3 Table 1、Fig. 4–5 和 Appendix Fig. 11 证明可重建 FV3GFS fine-grid 的概率尾部、时间结构与频谱；但 LR 由同一 HR 模式场粗化，GT 不是独立观测。	对廉价复现同一模式的小尺度统计有明确价值，不能据此声称改善真实洪水/水资源估计。
主要风险	仅使用配对同源 FV3GFS，未强制降水质量守恒；扩散样本可能与粗场总量不闭合；数据分布迁移会低估洪水/干旱风险。	部署到观测或其他模式前需再训练并检查校准和守恒。

通用 CRPS、EMD、空间相关、PSNR 等定义见 papers/downscaling/general_introduction，这里保留 STVD 采用的特定尾部与谱证据。

Claims → Evidence

Claim	类型	证据定位	强度	风险
STVD 学习条件概率而非单一平滑解	Construction	§2.1 Eq. (1)–(2)、Algorithm 1–2；随机 residual diffusion	强	概率分布仍由有限 FV3GFS ensemble 近似。
多帧上下文改善时空一致性	Empirical	Table 1：STVD-3 CRPS 1.96、STVD-1 2.05，高于 STVD 1.85；Fig. 2、Fig. 5	强	只比较长度 1/3/5，未测更长序列。
覆盖极端降水尾部	Empirical	Table 1 PE 1.2 vs Swin-IR 23.4；Fig. 4 分布；§3 99.999% percentile error	中-强	尾部阈值依赖模式资料与 10 个样本实现，未报告重现期校准。
恢复小尺度空间频谱	Empirical	Appendix Fig. 11：基线高频谱衰减更快，STVD 最接近 GT	中-强	频谱是二阶统计，不等于风暴相位或动力守恒。
优于六个强基线	Empirical	Table 1：STVD 在 CRPS/MSE/EMD/PE/SAE 均列最佳	强（FV3GFS 条件下）	基线训练代码、算力和随机样本数不完全对等；未含后续更强扩散模型。
可用于山地水资源/灾害分析	Empirical + application	Fig. 5 Himalaya/California 年平均降水和 topography	中	区域图是模拟内插结果，非独立流域观测验证。

Data and Experiments

FV3GFS 使用 11 个初始条件成员，各进行 13 个月模拟；去掉每个成员首月后得到约 11 年 3-hourly 数据，前 10 年训练、末年验证。fine grid 为 25 km，coarsening 后 coarse 为 200 km，全球 cubed-sphere 六 tile。每个样本包含连续 5 帧，输入除 precipitation 外还提供地形、表面/高空风温和水汽等 11 个辅助状态。降水做对数变换并归一化到 $[-1,1]$。

论文 Table 1（单位分别为 $10^{-5}$、$10^{-8}$、$10^{-6}$、$10^{-3}$、$10^{-6}$）报告 STVD：CRPS 1.85、MSE 0.59、EMD 2.49、PE 1.2、SAE 4.00；PSRT 为 2.15、0.66、4.21、3.8、6.24；VRT 为 3.58、1.74、4.61、4.0、7.39；VDM 为 2.21、0.73、12.70、6.4、8.84。STVD-single（只给 precipitation）和 STVD-1（单帧）分别退化至 CRPS 1.81 和 2.05，说明额外状态与时间上下文的贡献并非只来自模型容量。

训练使用 Adam，学习率从 $10^{-4}$ cosine 降到 $5\times10^{-7}$，N=1400，约 1,000,000 steps、单 RTX A6000 约 7 天。评估 CRPS 使用 10 个随机样本；EMD 比较全球降水分布，PE 直接看 99.999% 尾部，SAE 比较空间自相关，Appendix Fig. 11 另看功率谱。

Baseline fairness and missing experiments

基线覆盖 image SR（Swin-IR、Swin-IR-Diff）、video SR（VRT、RVRT、PSRT）和 video diffusion（VDM），并有去掉多变量、时间上下文的消融，设计相当完整。潜在不公平点是 VRT/RVRT/PSRT 的预处理和训练预算来自不同官方实现，论文没有逐模型参数量与 wall-clock 对齐；STVD 使用 10 个随机样本计算 CRPS，而其余确定性基线的分布指标退化为单点。缺失实验包括：粗细降水质量/面积守恒误差、概率可靠性图或 rank histogram、跨 FV3GFS 物理方案/观测验证、长时段自回归 rollout、不同 ensemble 成员的严格留出和真实极端重现期。

Conservation, Spectrum, Probability and Extremes Check

• 守恒： 模型只学习 $\bar y+r$，没有将每个 coarse cell 的降水总量投影回输入；因此即使 CRPS 和 EMD 很好，样本仍可能改变粗场的面积平均降水。质量守恒是最需要补充的后处理约束。
• 频谱： Appendix Fig. 11 显示 STVD 高频谱最接近 GT，基线衰减过快；这是直接频谱证据，但未给出跨季节/天气型置信区间。
• 概率： 条件 DDPM/DDIM 生成 10 个样本，CRPS 和 EMD 反映分布；但论文未给逐尺度 reliability diagram、coverage 或 rank histogram，因此“分布真实”仍不等于完全校准。
• 极端： 99.999% percentile error、Fig. 4 细网格降水分布和冷锋/山地案例支持尾部改善；没有极值理论重现期、台风中心气压或连续暴雨事件持续时间评估。

Reviewer Feedback

论文是 NeurIPS 2024 正式会议稿，公开 proceedings、arXiv 与 OpenReview 论坛均可定位到论文；本次本地核验以会议 PDF、补充材料和官方代码为主。没有把无法在静态 PDF 中逐条重建的评审对话写成“无评审”；若需复核 OpenReview 动态树，应在有稳定论坛 API 时按 invitation/replyto 读取 Official Review、Rebuttal 和 Decision。

Limitations

STVD 的 paired c48→c384 输入由同一 FV3GFS fine-grid 模拟粗化，属于模式内 perfect downscaling／emulation；其极端、频谱和 CRPS 都以模式世界为 reference。五帧窗口不支持长 rollout，模型没有粗网格降水回代或质量守恒，也缺少跨模式、雷达/站点观测和概率可靠性验证。因而它是强概率架构与评价基线，而不是已经验证的真实观测降尺度系统。

Critical Assessment

Strengths

1. residual diffusion 把确定性结构与随机小尺度分开，既降低扩散建模难度又保留多模态细节；这是比直接 pixel-space diffusion 更适合降水的工程分解。
2. 时空 factorized attention 与五帧并发采样针对 3-hourly 天气连续性，STVD-1/3 消融证明上下文确有作用。
3. 评价不局限于 MSE：CRPS、EMD、99.999% 尾部、SAE 和谱图共同覆盖概率、极端和空间统计，且与六个强基线比较。

Weaknesses & Risks

1. 训练/GT 都来自 FV3GFS 模式，因此“真实极端”是模式极端；模式偏差会被扩散模型忠实地学习，不能替代雨量计或雷达校准。
2. 没有硬质量守恒，随机残差可能与粗场降水总量不一致；这对流域水量和连续预测是实质风险。
3. 配对的 25 km FV3GFS PRATEsfc 模拟昂贵，且每次只处理 5 帧；作者明确把长时段 rollout 留作 future work，不能直接用于多年自回归气候生成。
4. 官方仓库没有预训练权重，环境是 CUDA/Linux 导出；config 中 1500 步与论文 N=1400 不同，复现者需锁定 commit 和参数。

Reviewer Recommendation

Accept. 对概率降水降尺度的技术与评价增量很清楚，实验也覆盖了最关心的频谱和极端指标；部署前建议加入面积加权守恒投影、校准曲线、跨模式/观测留出和长序列稳定性。

Innovation Score: 9.1/10

• Novelty：★★★★★（4.8/5）——将视频上下文、残差扩散和因子化时空注意力组合为专门的降水框架。
• Rigor：★★★★☆（4.6/5）——算法、消融和分布指标全面，但真实性依赖单一模式 ensemble。
• Impact：★★★★☆（4.7/5）——为概率、极端和山地空间结构提供可迁移基线；计算成本与配对资料是现实门槛。

Reusable Ideas

• 对 one-to-many 降尺度，先学习确定性低频/地形结构，再让扩散模型只采样 residual，可显著降低条件生成难度。
• 把时间上下文作为模型维度而非逐帧后处理，并用 factorized attention 在全局时间一致性与显存之间折中。
• 将 MSE 与 CRPS、尾部 percentile、空间自相关和 power spectrum 同时报告，避免“视觉锐利但统计错误”。
• 在 STVD 输出后添加面积加权 mass-conservation projection，检验它是否损害 CRPS/尾部而非默认守恒层一定有益。

Related Work

DSFNO 以确定性 FNO 解决任意倍率和水量约束；STVD 的贡献是条件概率、五帧时空一致性和显式谱/尾部评价。CorrDiff 也采用 deterministic mean + residual diffusion，但当前 STVD 更强调视频上下文和全球降水分布；后续 consistency model 和 R2-D2 则关注更快采样或多变量/动力下采样。视觉 SR 中 SwinIR、VRT、RVRT、PSRT 是本论文直接借鉴和对照的架构。

Personal Notes

这篇论文最值得移植的是评价观：在气象降尺度中，概率分布、极端尾部和频谱比单张图的 PSNR 更接近实际用途。若用于真实观测，我会先把 residual diffusion 改成守恒条件生成，并按天气型和流域做时空 block split；否则模式内随机 tile 的高分只能说明模型会还原 FV3GFS 的统计。

Code Verification

官方仓库 mandt-lab/STVD 已浅克隆（HEAD=3620f82），完成静态核验，未下载大体量 FV3GFS zarr 或运行 GPU 训练。

Claim	论文描述	代码实现核验	状态
paired c48→c384 数据	11-member FV3GFS，coarse 输入与 fine PRATEsfc 目标	data/vsrdata_ensemble.py 读取 c48_precip_plus_more_ave、c48_atmos_ave、c384_precip_ave 和 c384_topo，并形成 LR/HR	匹配
多变量条件	precipitation 加风、温度、气压、水汽和地形	数据加载器分别列出 6 个 surface、5 个 atmosphere 和高分辨率 topography 通道；config.py 默认单变量开关，multi 配置可切换	匹配（需改配置）
residual diffusion	$\hat y=\bar y+r$，高分辨率 residual 走 Gaussian diffusion	GaussianDiffusion 有 upsample=scale_factor=8、model_predictions 与 sample；Trainer 同时取 lres,hres 并以 diffusion loss 训练	匹配
spatiotemporal attention	空间/时间因子化、skip 和跨模块条件	model.py 定义 Attention、LinearAttention、LACross、ACross；UNet encoder feature 作为 diffusion context 拼接	匹配
EMA 与采样	EMA 0.995、训练多步、推理少步	Trainer 默认 ema_decay=0.995；config.py 默认 diffusion_steps=1500，sampling 配置可用 DDIM	部分匹配（论文 N=1400；需锁定实验配置）

复现风险：README 明确写出“Pretrained models will be released later”，因此不能仅克隆仓库复现论文 Table 1；还需准备公开/内部 FV3GFS zarr、统计 pickle 和 CUDA 环境。代码中的 config.py 默认 multi=False，若不显式开启会得到 STVD-single 而非论文主模型。