[论文解读] Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs

Kossaifi, Jean, Kovachki, Nikola+2 more|arXiv (Cornell University)|Jan 1, 2021

Model Reduction and Neural Networks被引用 10

一句话总结

本文提出多网格张量化傅里叶神经算子（MG-TFNO），一种可扩展、内存高效的神经算子，用于高分辨率PDE问题。该方法结合多网格区域分解与傅里叶域内的张量化参数化。通过将输入区域分解为具有分层全局上下文的局部区域，并利用低秩张量分解（如Tucker）表示权重，MG-TFNO在湍流N-S方程上实现超过150倍的参数压缩和7倍的输入尺寸缩减，且误差小于标准FNO的一半。

ABSTRACT

Memory complexity and data scarcity have so far prohibited learning solution operators of partial differential equations (PDEs) at high resolutions. We address these limitations by introducing a new data efficient and highly parallelizable operator learning approach with reduced memory requirement and better generalization, called multi-grid tensorized neural operator (MG-TFNO). MG-TFNO scales to large resolutions by leveraging local and global structures of full-scale, real-world phenomena, through a decomposition of both the input domain and the operator's parameter space. Our contributions are threefold: i) we enable parallelization over input samples with a novel multi-grid-based domain decomposition, ii) we represent the parameters of the model in a high-order latent subspace of the Fourier domain, through a global tensor factorization, resulting in an extreme reduction in the number of parameters and improved generalization, and iii) we propose architectural improvements to the backbone FNO. Our approach can be used in any operator learning setting. We demonstrate superior performance on the turbulent Navier-Stokes equations where we achieve less than half the error with over 150x compression. The tensorization combined with the domain decomposition, yields over 150x reduction in the number of parameters and 7x reduction in the domain size without losses in accuracy, while slightly enabling parallelism.

研究动机与目标

• 解决高分辨率PDE问题中神经算子学习面临的高内存复杂度与数据稀缺性问题。
• 通过利用物理空间与参数空间中的局部与全局结构，克服傅里叶神经算子（FNOs）在维度与内存增长方面的“维度灾难”。
• 通过将区域分解为具有多网格分层上下文的局部区域，实现在输入样本上的并行化处理。
• 通过在傅里叶域内使用高阶张量分解，减少模型参数与内存占用，同时不损失精度。
• 提升真实科学计算问题（如天气预报与流体动力学）中的泛化性与可扩展性。

提出的方法

• 采用多网格区域分解策略，将输入区域划分为局部区域，每个区域根据与目标区域的距离逐步引入下采样后的全局上下文。
• 使用分层多分辨率表示，其中局部区域以高分辨率处理，远处区域以粗略分辨率处理，从而将输入尺寸最多缩减7倍。
• 引入张量化傅里叶算子（TFNO），在傅里叶域中以低秩张量格式（如Tucker、CP、TT）表示FNO的卷积权重，实现超过100倍的参数压缩。
• 利用Tucker分解在中等压缩比下实现性能与压缩的平衡，CP分解实现高阶压缩，TT分解实现极端压缩下的稳定性。
• 将多网格分解与张量化权重整合进统一架构（MG-TFNO），实现输入区域间的完全并行化。
• 使用促进输出平滑性的损失函数，以确保局部预测拼接后的高质量重建。

Figure 1: Comparison of the performance on the relative $L^{2}$ and $H^{1}$ test errors (lower is better) on a log-scale of our approach, compared with both our improved backbone ( FNO ) and the original FNO, on Navier-Stokes. Our approach enables large compression for both input and parameter, whil

实验结果

研究问题

• RQ1 多网格区域分解能否在保持或提升高分辨率PDE问题精度的同时，减少神经算子的输入尺寸与内存占用？
• RQ2 在傅里叶域中对FNO权重进行高阶张量分解，能否实现极端参数压缩而不引起性能下降？
• RQ3 区域分解与张量化相结合，对泛化性、推理速度与内存效率在算子学习中的影响如何？
• RQ4 MG-TFNO在真实世界高分辨率科学问题（如湍流N-S方程或天气预报）中可扩展到何种程度？
• RQ5 张量化参数化是否起到正则化作用，即使在训练数据有限的情况下也能提升泛化能力？

主要发现

• MG-TFNO在N-S方程上实现150倍的模型参数减少与7倍的输入区域尺寸缩减，且精度无损失。
• 在N-S基准测试中，MG-TFNO将相对$L^2$测试误差降低至0.85%，不足标准FNO 2.54%误差的一半。
• 基于Tucker的TFNO变体实现10倍模型压缩，参数量为550万，同时保持0.85%的误差率。
• CP分解实现482倍压缩，但误差较高（2.24%），相较于Tucker显示出压缩与精度之间的权衡。
• TT分解实现574倍压缩，参数量为11.7万，误差为1.18%，表明在极端压缩比下仍具稳定性。
• 消融实验表明，结合多网格区域分解与张量化权重的性能与内存节省效果，优于单独使用任一技术。

Figure 2: Overview of our approach . First (left), a multi-grid approach is used to create coarse to fine inputs that capture high-resolution details in a local region while still encoding global context. The resulting regions are fed to a tensorized Fourier operator (middle), the parameters of whic