[论文解读] U-NO: U-shaped Neural Operators

Md Ashiqur Rahman, Zachary E. Ross+1 more|arXiv (Cornell University)|Apr 23, 2022

Neural Networks and Applications被引用 32

一句话总结

U-NO 引入一种内存高效的 U 形架构用于神经算子，使模型更深，在 DarcyFlow 和 Navier–Stokes 基准测试上相较 FNO 显示出显著的准确性提升，包括 3D 时空任务。

ABSTRACT

Neural operators generalize classical neural networks to maps between infinite-dimensional spaces, e.g., function spaces. Prior works on neural operators proposed a series of novel methods to learn such maps and demonstrated unprecedented success in learning solution operators of partial differential equations. Due to their close proximity to fully connected architectures, these models mainly suffer from high memory usage and are generally limited to shallow deep learning models. In this paper, we propose U-shaped Neural Operator (U-NO), a U-shaped memory enhanced architecture that allows for deeper neural operators. U-NOs exploit the problem structures in function predictions and demonstrate fast training, data efficiency, and robustness with respect to hyperparameters choices. We study the performance of U-NO on PDE benchmarks, namely, Darcy's flow law and the Navier-Stokes equations. We show that U-NO results in an average of 26% and 44% prediction improvement on Darcy's flow and turbulent Navier-Stokes equations, respectively, over the state of the art. On Navier-Stokes 3D spatiotemporal operator learning task, we show U-NO provides 37% improvement over the state of art methods.

研究动机与目标

• 以在无限维函数空间之间学习算子为动机，而不是在有限维向量之间。
• 通过收缩并扩展函数定义域（U 形编码/解码）来设计内存高效的深层算子架构。
• 在如 Darcy flow 与 Navier–Stokes 等 PDE 基准上，展示对最先进的 Fourier Neural Operator (FNO) 的经验性提升。
• 展示内存效率和数据效率带来的好处，使得对参数过多的模型也能训练。

提出的方法

• 采用提升算子 P 将输入映射到高维特征空间。
• 应用一系列非线性积分算子 G_i（在傅里叶域通过 R_i 进行卷积），并在编码器与解码器之间设置跳跃连接。
• 通过逐步收缩输入函数域同时扩展共域维度进行编码。
• 通过将域展开回目标并在减小共域的同时进行解码，并具有来自相应编码阶段的跳跃连接。
• 使用投影算子 Q 将最终编码表示映射到输出函数空间。
• 对于 3D 时空任务，实现自回归（2D+时间递归）和直接的 3D 时空（空间-时间）变体。

实验结果

研究问题

• RQ1 U 形神经算子是否能够在不造成过高内存消耗的情况下实现更深的架构？
• RQ2 内存高效的 U-NO 变体是否在经典 PDE 基准上提升了相对于最先进神经算子的预测精度？
• RQ3 使用 U-NO 设计的 3D 时空神经算子在可行性和优势方面是否优于 2D+时间方法？
• RQ4 更为激进的域/共域收缩/扩张如何影响性能和鲁棒性？
• RQ5 U-NO 是否对超参数选择具有鲁棒性，且具备零-shot 超分辨能力？

主要发现

内存 (MB)	模型	需求	# 参数量	s=421	s=211	s=141	s=85
166	U-NO†	166	8.2	0.57^{±1.4e-2}	0.58^{±0.4e-2}	0.60^{±0.5e-2}	0.73^{±0.1e-2}
214	FNO	214	1.1	0.78^{±4.0e-2}	0.84^{±5.5e-2}	0.84^{±1.1e-2}	0.87^{±3.0e-2}

• U-NO 及其变体 U-NO† 在 Darcy Flow 和 Navier–Stokes 基准测试中相对于 FNO 取得显著的准确性提升（例如在 Darcy Flow 的平均提升为 26%，在 Navier–Stokes 为 44%）。
• U-NO 使得算子更深且参数化程度更高，同时训练内存更低（例如 Darcy Flow 的内存比 FNO 少 23%，但参数量多 7.5 倍）。
• 在 3D 时空设置中，U-NO 相比 FNO 提升约 37%，且 3D 算子所需内存显著更低。
• U-NO 支持零-shot 超分辨率，并且由于其内存高效的编码/解码设计，能够在单个 GPU 上对高分辨率数据进行训练。
• 更激进的 U-NO 变体（U-NO†）具备更强的收缩/扩张，在提高内存效率和性能的同时，仍对超参数保持鲁棒性。
• 在 2D 和 3D 设置中，U-NO 能在显著增加参数量和更深架构的前提下，达到有竞争力甚至更高的准确性，相较于 FNO，使得在可控的内存增长下实现深度成为可能。