[논문 리뷰] Multi-Grid Tensorized Fourier Neural Operator for High-Resolution PDEs

Kossaifi, Jean, Kovachki, Nikola외 2명|arXiv (Cornell University)|2021. 01. 01.

Model Reduction and Neural Networks인용 수 10

한 줄 요약

이 논문은 다중 격자 도메인 분할과 푸리에 도메인 내 텐서화된 파arameterization을 조합함으로써 고해상도 PDE에 적합한 확장성 있고 메모리 효율적인 신경 연산자인 Multi-Grid Tensorized Fourier Neural Operator (MG-TFNO)를 제안한다. 입력 도메인을 계층적 글로벌 컨텍스트를 갖춘 국소 영역으로 분할하고, 저질서 텐서 분해(예: 터커)를 통해 가중치를 표현함으로써 MG-TFNO는 표준 FNO 대비 오차의 절반 이하로 150배 이상의 파라미터 압축과 7배의 입력 크기 감소를 달성한다. 난류 나비에-스토크스 방정식에서 성능을 입증한다.

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의 연산자 학습을 제한하는 높은 메모리 복잡성과 데이터 부족 문제를 해결한다.
• 물리적 및 파rameter 공간 내 局소 및 글로벌 구조를 활용하여 푸리에 신경 연산자(FNOs)의 차원의 저주와 메모리 증가 문제를 극복한다.
• 다중 격자 계층적 컨텍스트를 갖춘 국소 영역으로 도메인을 분할하여 입력 샘플 간 병렬 처리를 가능하게 한다.
• 푸리에 도메인 내 고차원 텐서 분해를 통해 모델 파라미터와 메모리 프로파일을 감소시키되 정확도를 유지한다.
• 실제 과학 계산 문제, 예를 들어 기상 예측 및 유체역학 문제에 대한 일반화 및 확장성을 향상시킨다.

제안 방법

• 입력 도메인을 국소 영역으로 분할하는 다중 격자 도메인 분할 전략을 적용하며, 각 영역은 관심 영역에서의 거리에 따라 점차 다운샘플링된 글로벌 컨텍스트로 풍부화된다.
• 국소 영역은 고해상도로 처리되고 먼 영역은 저해상도로 처리되는 계층적 다중 해상도 표현을 사용하여 입력 크기를 최대 7배까지 감소시킨다.
• 푸리에 도메인에서 FNO의 컨볼루션 가중치를 저질서 텐서 형식(예: 터커, CP, TT)으로 표현하는 텐서화된 푸리에 연산자(TFNO)를 도입하여 파라미터를 100배 이상 압축한다.
• 중간 압축 비율에서 성능과 압축의 균형을 이루는 데 터커 분해를 사용하고, 고압축을 위해 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는 난류 나비에-스토크스 방정식이나 기상 예측과 같은 실제 고해상도 과학 문제에 얼마나 확장 가능한가?
• RQ5 텐서화된 파rameterization은 훈련 데이터가 제한된 상황에서도 일반화 성능을 향상시키는 정규화 효과를 가지는가?

주요 결과

• MG-TFNO는 나비에-스토크스 방정식에서 정확도에 손실 없이 모델 파라미터를 150배 감소시키고 입력 도메인 크기를 7배 줄였다.
• 나비에-스토크스 벤치마크에서 MG-TFNO는 상대적 $L^2$ 테스트 오차를 0.85%로 줄였으며, 표준 FNO의 2.54% 오차보다 절반 이하였다.
• 터커 기반 TFNO 변종은 10배의 모델 압축을 달성했고, 550만 개의 파라미터를 유지하면서도 0.85% 오차율을 유지했다.
• CP 분해를 통해 482배의 압축이 가능했지만, 터커 대비 높은 오차(2.24%)를 보여 압축과 정확도 사이의 상충 관계를 확인했다.
• TT 분해를 통해 574배의 압축이 이루어졌고, 117만 개의 파라미터와 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