[논문 리뷰] Scalable Transformer for PDE Surrogate Modeling

Zijie Li, Dule Shu외 1명|arXiv (Cornell University)|2023. 05. 27.

Model Reduction and Neural Networks인용 수 20

한 줄 요약

논문은 FactFormer를 소개합니다. 이는 PDE 대리 모델링을 위한 확장 가능한 축 방향 분해 트랜스포머로, 고차원 격자(2D/3D 유체 문제)에서 계산을 줄이며 경쟁력 있는 정확도를 달성합니다.

ABSTRACT

Transformer has shown state-of-the-art performance on various applications and has recently emerged as a promising tool for surrogate modeling of partial differential equations (PDEs). Despite the introduction of linear-complexity attention, applying Transformer to problems with a large number of grid points can be numerically unstable and computationally expensive. In this work, we propose Factorized Transformer (FactFormer), which is based on an axial factorized kernel integral. Concretely, we introduce a learnable projection operator that decomposes the input function into multiple sub-functions with one-dimensional domain. These sub-functions are then evaluated and used to compute the instance-based kernel with an axial factorized scheme. We showcase that the proposed model is able to simulate 2D Kolmogorov flow on a $256 imes 256$ grid and 3D smoke buoyancy on a $64 imes64 imes64$ grid with good accuracy and efficiency. The proposed factorized scheme can serve as a computationally efficient low-rank surrogate for the full attention scheme when dealing with multi-dimensional problems.

연구 동기 및 목표

• 고해상도 격자에서 트랜스포머를 이용한 PDE의 확장 가능한 대리 모델링의 동기를 제시합니다.
• 다중 차원 도메인을 다루기 위한 분해된 소프트맥스-프리(attention) 메커니즘을 제안합니다.
• 정확도를 유지하면서 계산 비용을 줄이기 위한 학습 가능한 투영자(projection)와 축 방향 커널 적분을 개발합니다.
• 도전적인 PDE 문제들(2D Kolmogorov 흐름, 3D 연기 부력, 3D 등방성 난류, 2D Darcy 흐름)에서 평가합니다.

제안 방법

• RoPE 기반 상대 위치 지정으로 학습 가능한 커널 적분으로 주의(attention)를 재정의합니다.
• 고차원 입력을 다수의 일차원 도메인 서브 함수로 매핑하는 학습 가능한 투영 연산자를 도입합니다.
• 축별 커널과 차원 간 텐서-행렬 곱을 적용하는 분해된 커널 적분을 정의합니다.
• 자귀 추론 temporal 예측의 안정화를 위해 잠재적 행 march 및 pushforward 기법으로 학습합니다.
• 최종 업데이트를 위한 인스턴스 정규화 및 후(attention) MLP를 사용합니다.

Figure 1 : Model’s prediction (pred.) and reference ground truth (ref.). Left : 2D Kolmogorov flow on $256\times 256$ grid; Right : 3D smoke buoyancy on $64\times 64\times 64$ grid ( $zOy$ cross-section is shown).

실험 결과

연구 질문

• RQ1 축 방향의 분해된 주의(attention) 방식이 고해상도 격자에서 안정적인 PDE 대리 모델링을 제공할 수 있는가?
• RQ2 FactFormer가 다차원 문제에서 효율성을 향상시키면서 FNO 변형, Dil-ResNet 등 최신 신경 연산자와의 정확도에 버금갈 수 있는가?
• RQ3 잠재적 march와 pushforward 같은 학습 전략이 안정성 및 장기 예측 정확도에 어떤 영향을 미치는가?
• RQ4 2D 및 3D 유체 역학 벤치마크( Kolmogorov 흐름, 등방성 난류, 연기 부력, Darcy 흐름)에서 방법의 성능은 어떠한가?

주요 결과

• FactFormer는 2D Kolmogorov 흐름과 3D 연기 부력에서 우수한 효율로 경쟁력 있는 정확도를 달성합니다.
• 3D 등방성 난류에서 Dil-ResNet가 단기 프레임별 정확도에서 더 잘나올 수 있지만, FactFormer는 더 긴 예측 구간에서 여전히 효율적이고 경쟁력이 있습니다.
• 해상도가 달라져도 성능이 상대적으로 안정적으로 유지되어, 해상도 증가에 따라 오차가 증가하는 일부 CNN 기반 모델과 달리 안정성을 보입니다.
• 잠재적 march와 pushforward 학습은 자귀적(autoregressive) 학습에 비해 안정성을 개선하고 롤아웃 오차를 감소시킵니다.
• 전반적으로 FactFormer는 FNO 변형 및 Dil-ResNet에 비해 정확도, 추론 시간, 매개변수 수 간의 우수한 트레이드오프를 제공합니다.

Figure 2 : Schematic of the factorized kernel attention. Upper path : the input is transformed into the Value via a linear transformation. Lower path : the input is first projected into multiple sub-functions with a one-dimensional domain. These sub-functions are then used to derive the Query and Ke