[논문 리뷰] UPS: Efficiently Building Foundation Models for PDE Solving via Cross-Modal Adaptation

Jun-Hong Shen, Tanya Marwah외 1명|arXiv (Cornell University)|2024. 03. 11.

BIM and Construction Integration인용 수 5

한 줄 요약

UPS는 다양한 시공간 PDE 데이터를 공통 표현으로 통합하고 프리트레인된 LLM을 PDE 해법에 적응시키는 트랜스포머 기반 모델을 훈련시켜, PDEBench 작업 전반에서 강한 데이터 효율성과 함께 최첨단 성능을 달성한다.

ABSTRACT

We present Unified PDE Solvers (UPS), a data- and compute-efficient approach to developing unified neural operators for diverse families of spatiotemporal PDEs from various domains, dimensions, and resolutions. UPS embeds different PDEs into a shared representation space and processes them using a FNO-transformer architecture. Rather than training the network from scratch, which is data-demanding and computationally expensive, we warm-start the transformer from pretrained LLMs and perform explicit alignment to reduce the modality gap while improving data and compute efficiency. The cross-modal UPS achieves state-of-the-art results on a wide range of 1D and 2D PDE families from PDEBench, outperforming existing unified models using 4 times less data and 26 times less compute. Meanwhile, it is capable of few-shot transfer to unseen PDE families and coefficients.

연구 동기 및 목표

• 도메인, 차원, 해상도에 걸쳐 다양한 시간 의존 PDE를 해결하기 위한 데이터 효율적이고 기초 모델 기반의 접근법을 제시한다.
• FNO 기반 PDE 임베딩, LLM, 예측 헤드를 결합한 통합 데이터 표현과 3단 구성 아키텍처를 제안한다.
• 여러 PDE 계열에서 교차 모달 적합화의 두 단계로 PDE 임베딩과 텍스트 임베딩을 정렬하고 미세조정한다.
• PDEBench에서 향상된 성능과 샘플 효율성을 입증하고, 강력한 few-shot 및 zero-shot 전이 가능성을 보인다.

제안 방법

• 다양한 PDE 트로젝터리들을 다중 차원, 해상도, 양에 걸친 공유 표현 공간으로 통합한다.
• PDE 특화 임베딩 네트워크와 Fourier Neural Operator(FNO) 계층을 사용하여 PDE 데이터로부터 해상도에 구애받지 않는 특징을 추출한다.
• 텍스트 형태의 PDE 설명으로 메타데이터를 임베딩하고 이를 PDE 특징과 함께 LLM 본체로 처리하기 전에 연결한다.
• 3방향 아키텍처를 채택한다: PDE 임베딩 네트워크, 모델 바디로서의 사전학습된 LLM 계층, 다음 단계 PDE 상태를 예측하는 선형 예측 헤드.
• 두 단계로 학습한다: (i) 모달리티 정렬 손실(MMD 기반)과 작업 손실(nRMSE)을 포함한 임베딩 프리트레이닝; (ii) 여러 PDE 계통에 걸친 다중 작업 파인튜닝.
• 다양한 LLM 백본(RoBERTa, T5, CLIP)과의 호환성을 보여주고 신경 연산자의 격자 및 해상도 불변성을 시연한다.

Figure 1 : UPS outperforms leading competitors on a wide range of tasks from PDEBench, including the single-task neural operators U-Net (Ronneberger et al., 2015 ) , FNO (Li et al., 2020a ) , and ORCA (Shen et al., 2023 ) , and the multi-task MPP-Base (McCabe et al., 2023 ) (MPP only applies to 2D t

실험 결과

연구 질문

• RQ1 하나의 통합 표현과 아키텍처가 서로 다른 PDE 계통, 차원, 해상도 간의 크로스 도메인 전이를 가능하게 할 수 있는가?
• RQ2 사전학습된 LLM의 크로스모달 적응이 다양한 PDE 해결 과제에서 데이터 효율성과 예측 정확도를 향상시키는가?
• RQ3 부분 메타데이터(텍스트 형태의 PDE 설명)가 PDE용 신경 연산자 프레임워크의 학습 및 일반화에 어떤 영향을 미치는가?
• RQ4 미지의 PDE 계통, 계수, 해상도에 대한 few-shot 및 zero-shot 전이에서 UPS의 성능은 어떠한가?

주요 결과

# 매개변수	Advection nRMSE	Burgers nRMSE	Diffusion-Sorption nRMSE	Navier-Stokes nRMSE	Shallow-Water nRMSE	Navier-Stokes Incomp nRMSE	Incomp Navier-Stokes nRMSE
149M	0.0033	0.0399	0.0009	0.0056	0.0019	0.0174	0.104
7.7M	0.67	0.34	0.15	0.72	0.083	5.1	0.1903
466K	0.011	0.042	0.0017	0.068	0.0044	0.36	0.0942
125M	0.0098	0.12	0.0016	0.062	0.006	0.3549	0.1529
116M	-	-	-	-	0.0024	0.0281	-
409M	-	-	-	-	0.0022	0.0208	-
8.5K	1.1	0.96	0.22	-	0.017	-	-

• UPS는 PDEBench 10개 과제 중 8개에서 최첨단 성능을 달성하며, 종종 베이스라인 대비 큰 차이로 우수하다.
• 인디스트리뷰션 과제에서 UPS(RoBERTa-Base)는 7개 데이터세트 중 6개에서 가장 낮은 nRMSE를 달성하고 1D 및 2D 성능이 강하다.
• UPS는 보지 않은 PDE 계통, 계수, 더 높은 해상도에 대한 강력한 few-shot 및 zero-shot 전이를 보여주며, 제한된 데이터에서도 여러 베이스라인을 능가한다.
• 모달리티 정렬 및 작업 학습을 포함한 1단계 임베딩 프리트레이닝은 성능에 필수적이며, 메타데이터의 도입이 결과를 향상시킨다.
• 이 방법은 데이터 및 계산 효율적이며, PDE 계마다 5K 미만의 학습 궤적을 사용하고 단일 GPU 학습을 가능하게 한다.

Figure 2 : To adapt pretrained LLMs for PDE solving, UPS first expresses PDE trajectories of different families, dimensions, channels, and resolutions using a unified data representation (left panel). It then processes the data with a unified transformer architecture that integrates domain-specific