Learning, Solving and Optimizing PDEs with TensorGalerkin: an efficient high-performance Galerkin assembly algorithm

Motivation

• Known: FEM은 복잡한 도메인 형상과 수학적 보장 측면에서 강력한 PDE 수치해법이지만, inverse design이나 PDE-constrained optimization 같은 many-query 문제에서는 계산 비용이 크다. 한편 PINNs 등 physics-informed ML은 strong-form residual과 automatic differentiation(AD)에 의존하여 ill-conditioning, 깊은 계산 그래프, 큰 메모리 사용량 문제를 겪는다.
• Gap: Python 기반 AD 프레임워크(PyTorch 등)에 표준 FEM assembly를 통합하면, 요소 반복문에 의한 interpreter overhead, 국소 basis index 루프에 의한 계산 그래프 파편화, 그리고 scatter-add 같은 비결정적 atomic 연산으로 인한 재현성 저하 문제가 발생한다. 기존 접근(scikit-fem, JAX-FEM, ad-hoc PyTorch 구현)은 I/O 병목, 동적 메쉬 재컴파일 오버헤드, 비결정성 문제를 각각 부분적으로만 해결한다.
• Why: Galerkin assembly는 PDE 수치해법뿐 아니라 physics-informed learning과 PDE-constrained optimization 전반의 핵심 병목이므로, 이를 O(1) 계산 그래프로 근본적으로 해결하면 세 가지 하위 응용(수치 solver, operator learning, 최적화) 모두에서 효율성과 정확도를 동시에 개선할 수 있다.
• Approach: 변분 구조를 갖는 PDE의 Galerkin discretization을 기반으로, element-wise 연산을 Python 레벨 Map 단계에서 dense tensor contraction으로 융합하고, mesh-induced sparsity graph 상의 message passing을 수행하는 sparse matrix multiplication(SpMM)으로 Reduce 단계를 수행하는 TensorGalerkin을 제안한다.

Achievement

Figure 2. Numerical PDE solver benchmarks on 3D meshes. Panels (a) and (b) report solve-time scaling with the number of

1. O(1) 계산 그래프 달성: Map과 Reduce 단계를 PyTorch autograd 내에서 공동 설계하여, 기존 O(Ek²) (E: 요소 수, k: 국소 DoF 수) 노드 수를 요소 수와 무관한 O(1) 노드로 압축함으로써 backward pass의 파편화 오버헤드를 제거하였다.
2. 세 가지 하위 응용으로의 검증: TensorMesh(고성능 수치 PDE solver), TensorPILS(AD 없이 해석적 shape gradient를 활용하는 physics-informed operator learning), TensorOpt(end-to-end differentiable PDE-constrained optimization) 세 프레임워크를 통해 O(1)-graph 특성을 실증적으로 검증하였다.
3. 다양한 벤치마크에서의 성능 우위: 2D/3D 타원형(elliptic), 포물형(parabolic), 쌍곡형(hyperbolic) PDE의 비정형 메쉬(unstructured mesh) 벤치마크에서 기존 CPU FEM solver 및 baseline 대비 계산 효율성과 정확도 측면에서 유의미한 개선을 달성하였다.

How

Figure 1. Overview of TENSORGALERKIN. Stage I (Batch-Map) computes element-wise operators via a fully tensorized einsum

• 변분 문제 aρ(u,v)=ℓρ(v)를 basis function {ϕi}로 이산화하여 KU=F 형태의 Matrix Equation으로 환원 (Neural Galerkin Discretization)
• Stage I (Batch-Map): 국소 stiffness matrix/load vector 계산을 위한 element-wise 연산을 하나의 dense tensor contraction으로 융합, 단일 GPU 커널에서 실행
• Stage II (Reduce): scatter-add 루프를 precomputed routing 기반의 결정적 sparse matrix multiplication(SpMM)으로 대체하여 mesh sparsity graph 상의 message passing 구현
• 위 두 단계를 PyTorch autograd 그래프 상에서 공동 설계(co-design)하여 backward graph 노드 수를 O(1)로 축소
• 이를 numerical solver(TensorMesh), physics-informed operator learning(TensorPILS, 해석적 shape gradient 사용으로 AD 불필요), PDE-constrained optimization(TensorOpt, end-to-end differentiable)의 세 가지 downstream task에 적용하여 검증

Originality

• 기존 FEM assembly의 element-loop 기반 구조를 Map-Reduce 패러다임으로 재정식화하고 이를 PyTorch autograd와 공동 설계하여 O(1) 그래프 크기를 이론적·실험적으로 규명한 점이 독창적이다.
• 단일 assembly 알고리즘을 numerical solving, physics-informed operator learning, PDE-constrained optimization이라는 세 가지 이질적인 downstream 응용에 통일적으로 적용하는 프레임워크를 제시하였다.
• PINN류 방법이 강한 형태(strong-form) residual과 pointwise AD에 의존하는 것과 달리, 해석적 shape gradient를 활용해 AD 의존성 자체를 줄이는 TensorPILS를 제안하였다.

Limitation & Further Study

• 발췌된 본문은 선형(linear) 정상상태 PDE를 중심으로 서술되어 있으며, 비선형(semi-linear) 시간의존 PDE로의 확장은 보충자료(SM)로 미뤄져 있어 본문만으로는 일반성 검증이 제한적이다.
• 비정형 메쉬(unstructured mesh)에서 sparsity pattern이 매우 불균일하거나 매우 고차 요소(high-order element)일 때 precomputed routing의 메모리 오버헤드나 확장성에 대한 논의가 부족해 보인다.
• 동적 메쉬(adaptive mesh refinement 등) 상황에서 routing 재계산 비용과 JAX-FEM 대비 재컴파일 이점이 실제로 어느 정도인지에 대한 정량적 비교가 발췌문에서는 충분히 제시되지 않았다.
• 후속 연구로는 비선형 PDE, 고차 시간적분 스킴, 대규모 3D 실세계 응용으로의 확장성 검증이 필요할 것으로 보인다.

Evaluation

총평: Galerkin assembly의 근본적인 시스템 레벨 병목을 Map-Reduce 재정식화와 autograd 공동설계로 해결하여 O(1) 계산 그래프를 달성한 기술적으로 탄탄하고 실용적 가치가 큰 연구로, 수치해석과 물리 정보 기반 학습 커뮤니티 모두에 의미 있는 기여를 한다.