小波扩散神经算子与复杂系统建模与控制

原创胡佩炎 2025-04-18 18:04 上海

导语

一方面，模拟和控制现实世界中的复杂物理系统，是当前科学计算和人工智能交叉领域的重要任务。另一方面，扩散模型在不同的领域都展现出了强大的能力。然而，扩散模型应用于仿真和控制复杂物理系统时会遇到难以捕获突变细节、无法跨分辨率泛化的挑战。针对这些挑战，近日，西湖大学特聘研究员吴泰霖课题组提出了“小波扩散神经网络算子（Wavelet Diffusion Neural Operator, WDNO）”。该工作基于小波域建模和多尺度扩散生成机制，在多个控制与模拟任务中实现了领先的性能。该文章已成功入选人工智能顶级国际会议 ICLR 2025。

关键词：复杂系统建模，复杂系统控制，人工智能，扩散模型，小波变换，超分辨率

胡佩炎丨作者

论文题目：Wavelet Diffusion Neural Operator
论文链接：https://arxiv.org/abs/2412.04833
代码链接：https://github.com/AI4Science-WestlakeU/wdno.git

一、背景：扩散模型在何处失效？

在机器学习模拟物理系统的浪潮中，神经算子（Neural Operators）正成为连接数据驱动与机制建模的重要桥梁。相比传统的神经网络，神经算子能够直接学习从函数到函数的映射，天然适配偏微分方程（PDE）建模任务。

近年来，扩散模型（Diffusion Models）[1] 在生成建模领域取得了显著突破，其核心思想是通过逐步注入噪声，将原始样本变换为近似高斯分布的噪声，然后再通过学习一个反向去噪过程恢复原始数据。该机制具备理论上的可逆性与灵活性，已在图像、音频等生成任务中展现出极强的表现力。

图1: 基于扩散模型建模物理系统的过程。[2]

正如图中所示，扩散模型在物理系统建模中的潜力也日益受到关注。左侧展示的是一个有序的状态轨迹与控制序列，而右侧表示通过逐步扩散变为噪声形式的轨迹与控制。扩散模型通过学习的反向去噪过程，可以从高斯分布逐步去噪得到目标数据分布。这一过程为处理高维、多阶段的动力学建模问题提供了新的思路。

伴随扩散模型的崛起，基于其的物理建模方法在高维动态系统、偏微分方程（PDE）求解等任务中展现了巨大的潜力。然而，扩散模型用于建模和控制物理系统时仍会面对以下三类现实挑战：

1. 捕获突变（Abrupt Changes）

如流体中的激波、气象中的锋面，这类现象具有剧烈的局部变化，它们天然较难被神经算子捕捉到，而扩散模型的加噪、去噪机制使得建模这些突变更为困难。

2. 跨分辨率泛化（Cross-Resolution Generalization）

由于神经算子学习函数空间之间的映射，因此一个重要的特征为具有不同分辨率的泛化能力。而扩散模型难以直接具备这种能力，也就是零样本的超分辨率（zero-shot super-resolution）。

3. 长时间演化与控制（Long-Term Prediction & Control）

小误差在多步预测和控制中可能被迅速放大，导致发散。传统模型容易陷入“短期精度好，长期失控”的困境，控制任务尤为严重。

二、核心方法：如何应对挑战？

为解决上述挑战，WDNO 的作者们通过引入扩散模型在小波域建模和多分辨率训练来应对以上的问题，几个模块能够互为补充、自然融合。方法核心包括两部分：

图2：WDNO的核心算法。

1. 小波空间的扩散建模

为了应对细节丢失的挑战，WDNO将多尺度小波变换引入扩散模型架构。作者通过选取只在时间和空间有限区间取值的小波基，使得小波变换能够相较傅立叶变换等同时保留空间位置和频率信息，适配于对突变处动力学的建模。[3] 通过在小波域进行扩散去噪，WDNO能弥补加噪、去噪机制导致的难以捕捉突变的缺陷，生成准确的轨迹或控制策略。与此同时，小波变换的局域性和线性性使得它能够天然适配于下文提出的多分辨率训练。

2. 多分辨率训练

WDNO首先基于训练数据集，通过下采样构造出从细到粗的多分辨训练数据，无需利用任何超分辨率的训练数据。在训练阶段，作者构建了超分辨率模型（Super-Resolution Model, SRM）同样在小波域中习得对分辨率的自适应泛化。在推理时，通过同时利用基础分辨率模型（Base-Resolution Model, BRM）和超分辨率模型（Super-Resolution Model, SRM），就能够实现零样本超分辨率的泛化。

通过这两项相互融合、互为补充的机制，WDNO能够实现强建模能力、高跨分辨率泛化性和长期稳定性。

三、实验评估：从建模到控制任务

WDNO 在五个具有代表性的场景中进行了广泛测试：1D Advection方程，1D Burgers’方程，1D Compressible Navier-Stokes方程，2D流体，ERA5 气象数据集。

主要成果包括：

• 控制和模拟任务中，WDNO优于其他主流神经算子与控制方法。

表1：各个算法在模拟任务上的均方误差（MSE）。粗体表示最优的结果，下划线表示次优的结果。

表2：各个算法在控制任务上的结果。粗体表示最优的结果，下划线表示次优的结果。

• 在超分辨率预测任务中，WDNO领先原始的扩散模型（DDPM）、Fourier Neural Operator和Wavelet Neural Operator，验证其跨尺度泛化能力。

图3：零样本超分辨率实验的均方误差（MSE）结果。所有结果均插值到最细网格上以便于比较。图(a)和(c)为在1D Burgers'方程上的结果，图(b)为在2D流体上的结果。

• 在突变存在的区域，WDNO展现了更准确的建模。

图3：WDNO与DDPM相比，有着更强的捕捉突变细节的能力。

• 长时多步预测下，WDNO 表现出更稳定的误差累积特性。

图4：不同算法的误差随物理时间的变化。

四、总结与展望：通往物理智能的路径

本文提出的小波扩散神经算子（Wavelet Diffusion Neural Operator, WDNO）实现了在连续空间上建模和控制复杂物理过程的能力显著提升，在多个具有挑战性的PDE建模任务中取得了优异的表现。

从更广义的角度来看，WDNO的设计展示了一种可能的范式：将先进的生成建模方法（如扩散模型）与神经算子结合，从而迈向更具表达力、更泛化、可用于控制与反演等任务的物理建模工具。这为通向物理智能提供了新的可能性。

未来，我们希望沿以下几个方向进一步推进该方向的研究：探索更高效的跨分辨率扩散机制，适配更复杂边界条件和几何形状的PDE问题，结合实际物理系统数据，进一步推动其在流体、气象等领域的落地应用。

参考文献：

1. Jonatha Ho et al. Denoising diffusion probabilistic models. NeurIPS 2020.

2. Long Wei et al. DiffPhyCon: A Generative Approach to Control Complex Physical Systems. NeurIPS 2024.

3. Beylkin Alpert, Gregory Beylkin, David Gines, and Lev Vozovoi. Adaptive solution of partial differential equations in multiwavelet bases. Journal of Computational Physics, 182(1):149–190, 2002.

复杂系统自动建模读书会第二季

“复杂世界，简单规则”。

集智俱乐部联合复旦大学智能复杂体系实验室青年研究员朱群喜、浙江大学百人计划研究员李樵风、清华大学电子工程系数据科学与智能实验室博士后研究员丁璟韬、美国东北大学物理系Albert-László Barabási指导的博士后高婷婷、北京大学博雅博士后曹文祺、复旦大学数学科学学院应用数学方向博士研究生赵伯林、北京师范大学系统科学学院博士研究生牟牧云，共同发起「复杂系统自动建模」读书会第二季。

读书会将于9月5日起每周四晚上20:00-22:00进行，探讨四个核心模块：数据驱动的复杂系统建模、复杂网络结构推断、具有可解释性的复杂系统推断（动力学+网络结构）、应用-超材料设计和城市系统，通过重点讨论75篇经典、前沿的重要文献，从黑盒（数据驱动）到白盒（可解释性），逐步捕捉系统的“本质”规律，帮助大家更好的认识、理解、预测、控制、设计复杂系统，为相关领域的研究和应用提供洞见。欢迎感兴趣的朋友报名参与！