在科学、工程和金融领域,许多最重要的风险来自低概率、高影响事件。使用强力蒙特卡罗采样 (使用随机抽取的输入重复运行模型以估计罕见结果的概率) 来预测这些事件的概率可能需要进行过多的模型迭代,尤其是当每个样本都来自基于物理性质的昂贵模型时。
引导扩散模型提供了一种引导罕见事件的新方法,但引导带来了一项关键挑战:如果模型对分布的尾部进行过采样,我们还如何估算这些样本的真实概率?
在一篇新的论文《从扩散模型气候仿真器获得准确的极端事件概率》中,我们通过引导基于扩散的气候仿真器转向热带风暴,并根据引导和非引导概率计算概率比,来探索气候科学中的这个问题。在本文中,我们将展示一个使用热带气旋案例的最小工作示例。
还有很多工作要做:值得尝试新方法,需要性能优化,而且这种影响可能会为科学和经济学领域的新一代风险分析释放巨大的成本效益。
请继续阅读并阅读快速入门指南,了解其实际运作方式。
极端气候的引导扩散
极端天气风险通常取决于罕见的局部事件,而使用传统气候模拟进行采样成本高昂。例如,靠近特定海岸线的热带气旋可能需要非常大的集合,然后才能有足够多的案例以有用的置信度估算其可能性。这给登陆风险分析、气候属性实验、基础设施规划、保险建模和下游影响研究带来了瓶颈。
我们的研究表明,基于扩散的气候仿真器不仅能生成真实的大气状态,还能发挥更大的作用。它们还可以为这些状态提供有用的概率估计。使用基于扩散的气候条件天气状态生成器 NVIDIA cBottle,我们可以将扩散模型引导到热带气旋状态,并计算引导分布和非引导分布之间的概率比,从而对罕见事件进行重要性采样。
这一点很重要,因为仅靠指导是不够的。如果模型转向罕见事件,则会根据设计对其进行过采样。要使用这些样本进行概率估计,我们需要了解该指南对其概率的影响程度。几率比可提供该校正,从而对罕见事件进行重要性采样,同时仍可估算其在原始气候分布下的概率。
通过 CBottleTCGuidance 示例,NVIDIA Earth2Studio (用于 AI 气候和天气工作流的开源推理平台) 提供了此工作流的实施。
从引导样本到几率比
cBottle 热带气旋引导工作流根据用户指定的热带气旋活动生成大气样本。在 Earth2Studio 示例中,制导张量会标记佛罗里达州附近的所需位置,并在飓风季节请求样本。
从概念上讲,该模型比较了生成大气状态 \(x\) 的两个概率:
\(p_{\text{guided}}(x)\)
以及
\(p_{\text{unguided}}(x)\)
对数几率比为:
\(\log o(x) = \log p_{\text{unguided}}(x) – \log p_{\text{guided}}(x)\)
此值量化了样本在无引导模型下的可能性比在引导模型下的可能性低多少。
要使用引导样本估计无引导分布下的概率,相应的重要性权重为:
\(o(x) = \frac{p_{\text{unguided}}(x)}{p_{\text{guided}}(x)} = \exp(\log o(x))\)
这正是重要性采样的实现之处:在罕见事件区域生成更多样本,然后重新加权以估计其在原始气候分布下的概率。
\(P_{IS}(\text{TC}) \approx \frac{1}{K} \sum_{i=1}^{K} I_{\text{TC}}(x_i) o(x_i), \quad \text{where } x_i \sim p_{\text{guided}}\)
在 Earth2Studio 中运行引导式热带气旋采样
Earth2Studio 示例使用内置的 CBottleTCGuidance 诊断模型。此设置会加载默认的 cBottle TC 指导包,该指导包会从 NGC 下载检查点。
from earth2studio.models.dx import CBottleTCGuidance
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load the default model package which downloads the checkpoint from NGC
package = CBottleTCGuidance.load_default_package()
快速路径针对标准引导式推理进行了优化。此处,使用纬度 27.0 度和经度 -82.0 度,将一个制导点放置在佛罗里达州附近。然后,该示例会在飓风季节请求一个时间。
lat = torch.tensor([27.0], device=device) # Near Florida
lon = torch.tensor([-82.0], device=device) # Converted internally to [0, 360)
times = [datetime(2005, 10, 11, 12)]
model = CBottleTCGuidance.load_model(package, seed=0).to(device)
# Create guidance tensor
guidance, coords = model.create_guidance_tensor(lat, lon, times)
guidance = guidance.to(device)
# Run guided sampling
guided_sample, guided_coords = model(guidance, coords)
输出 guided_sample 包含生成的大气状态。
这使得提取 10 米分区风、平均海平面压力或其他生成场等变量进行分析和可视化变得简单。该示例在加勒比海域上绘制了 10 米长的带风分量 u10 米 (图 1 如下图所示) 。此可视化有助于验证生成状态在请求位置附近是否包含连贯的热带气旋式风力结构。
图 1:引导样本显示了带状风中典型的 TC 模式计算对数几率比
对引导的热带气旋进行采样很有用,但主要的研究贡献在于能够估算与非引导的碱基分布相比,该样本在引导分布下的可能性。
比值比计算需要估计生成流的离散,而生成流又需要通过模型的二阶导数。因此,该示例使用 allow_second_order_derivatives=True 重新加载模型。
model = CBottleTCGuidance.load_model(
package,
seed=0,
sampler_steps=2,
allow_second_order_derivatives=True,
).to(device)
log_odds_ratio, forward_latents, latent_coords = model.calculate_odds_ratio(
guidance,
coords,
)
print(f"Log odds ratio: {log_odds_ratio:.4f}")
print(f"Forward latents shape: {tuple(forward_latents.shape)}"
该示例减少了 sampler_steps,以加快运行时。对于质量更高的样本和更稳定的比值,应使用默认采样器设置。
返回的 log_odds_ratio 会告诉您,指导对所生成样本的可能性的影响有多大。负值意味着样本在引导分布下的可能性比在无引导基础模型下的可能性大得多。
forward_latents 张量提供用于比值比计算的引导样本。
为什么这对气候科学很重要
该论文表明,扩散模型气候模拟器可以提供概率信息,而对于完整的大气状态,过去无法以实际形式提供这些信息。
这一点很重要,因为许多气候问题从根本上讲都是概率性问题。研究人员想知道热带气旋在特定海岸附近的可能性有多大,在不同的海面温度下,这种可能性会如何变化,以及引导生成模型是否能比强力采样更高效地产生罕见事件。
通过将引导生成与概率比诊断相结合,可以更频繁地对罕见事件进行采样,同时仍可估算其在基础气候分布下的可能性。该论文认为,与简单的蒙特卡洛采样相比,这有助于对热带气旋状态进行重要性采样,并将标准误差降低 25%。
气候以外的机遇
虽然该示例侧重于热带气旋,但其思路更为广泛:引导生成式模型走向分布中罕见但重要的部分,然后计算该指南对样本概率的改变程度。
这种模式在罕见事件主导风险的领域非常有用,包括电网弹性、航空航天压力测试、金融尾部风险评估、材料发现、机器人边缘案例和分子模拟。
当天真的采样很少发现重要事件时,引导式生成式采样加上比值校正可以使罕见事件估计更加实用。
未来研究方向
目前的结果是早期的,但前景光明。要使这些方法更快、更准确和更广泛地适用,还需要做更多的工作。
扩散采样的计算成本很高,尤其是当比值计算需要二阶导数时。更快的采样器、蒸馏或更少的步长扩散方法可以使罕见事件概率估计更具可扩展性。更好的密度估计和更稳定的对数概率计算将提高下游概率估计的置信度。更精确的引导功能还可以针对不同的事件类型、强度和区域,同时保持物理一致性。
热带气旋是极端事件的一个例子。对于大气河流、热浪、强风暴、干旱和复合极端天气,也可以探索类似的方法。当扩散模拟器在不同的气候条件下进行训练时,引导概率估计还可以通过比较不同边界条件或气候状态下的事件概率来支持新的类似于归属的实验。
开始使用
阅读论文:通过扩散模型气候模拟器获得准确的极端事件概率,并探索 Earth2Studio 示例 05_cbottle_tc_guidance.py。
如需了解更多背景信息,请参阅 cBottle 论文《 Climate in a Bottle 》和 NVIDIA Earth2Studio GitHub 存储库。