[논문 리뷰] Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-Informed Neural Networks

Zheyuan Hu, Zekun Shi외 2명|arXiv (Cornell University)|2023. 12. 22.

Neural Networks and Applications참고 문헌 57인용 수 7

한 줄 요약

이 논문은 고차원 및 고차 미분방정식(PDE)을 위한 물리기반신경망(PINNs)의 계산 속도를 향상시키기 위해 허친슨 추정법(Hutchinson Trace Estimation, HTE)을 도입한다. 전체 헤시안 행렬 계산을 헤시안-벡터 곱(Hessian-vector products, HVPs)으로 대체함으로써 메모리 사용량을 줄이고, 효율적인 타일러 모드 자동미분을 가능하게 하여, 특히 고분산 상황에서 Stochastic Dimension Gradient Descent(SDGD)보다 분산이 낮고 편향이 없는 수렴 가능한 학습을 달성한다.

ABSTRACT

Physics-Informed Neural Networks (PINNs) have proven effective in solving partial differential equations (PDEs), especially when some data are available by seamlessly blending data and physics. However, extending PINNs to high-dimensional and even high-order PDEs encounters significant challenges due to the computational cost associated with automatic differentiation in the residual loss. Herein, we address the limitations of PINNs in handling high-dimensional and high-order PDEs by introducing Hutchinson Trace Estimation (HTE). Starting with the second-order high-dimensional PDEs ubiquitous in scientific computing, HTE transforms the calculation of the entire Hessian matrix into a Hessian vector product (HVP). This approach alleviates the computational bottleneck via Taylor-mode automatic differentiation and significantly reduces memory consumption from the Hessian matrix to HVP. We further showcase HTE's convergence to the original PINN loss and its unbiased behavior under specific conditions. Comparisons with Stochastic Dimension Gradient Descent (SDGD) highlight the distinct advantages of HTE, particularly in scenarios with significant variance among dimensions. We further extend HTE to higher-order and higher-dimensional PDEs, specifically addressing the biharmonic equation. By employing tensor-vector products (TVP), HTE efficiently computes the colossal tensor associated with the fourth-order high-dimensional biharmonic equation, saving memory and enabling rapid computation. The effectiveness of HTE is illustrated through experimental setups, demonstrating comparable convergence rates with SDGD under memory and speed constraints. Additionally, HTE proves valuable in accelerating the Gradient-Enhanced PINN (gPINN) version as well as the Biharmonic equation. Overall, HTE opens up a new capability in scientific machine learning for tackling high-order and high-dimensional PDEs.

연구 동기 및 목표

• 고차원 및 고차 PDE에서 PINN 학습의 높은 계산 비용과 메모리 부담을 해결한다.
• 차원 간 변동성이 높을 경우 고분산을 보이는 기존 방법(예: SDGD)의 한계를 극복한다.
• 비하모닉 방정식과 같은 4차 PDE에 대해 텐서-벡터 곱(Tensor-vector products, TVPs)을 도입하여 기하급수적으로 커지는 텐서를 관리함으로써 PINN 학습을 4차 미분방정식까지 확장한다.
• 특정 조건 하에서 HTE 기반 손실이 원래 PINN 손실로 이론적으로 수렴하고 편향이 없는지를 보장한다.
• 효율적인 추정을 통해 Gradient-Enhanced PINNs(gPINN) 및 고차 PDE 해법기에서 실제 가속을 가능하게 한다.

제안 방법

• PINN 손실 함수에서 헤시안 행렬의 추정을 위해 허친슨 추정법(HTE)을 적용하여 전체 헤시안 행렬 계산을 대체한다.
• 헤시안 행렬 계산을 헤시안-벡터 곱(HVP)으로 변환함으로써 메모리 소비를 O(d²)에서 O(d)로 감소시키고, JAX에서 효율적인 타일러 모드 자동미분을 가능하게 한다.
• Rademacher 분포를 가진 랜덤 벡터를 사용하여 편향 없는 추정량을 계산한다: $ \frac{1}{V}\sum_{i=1}^{V} \boldsymbol{v}_i^T A \boldsymbol{v}_i \approx \mathrm{Tr}(A) $.
• 헤시안-벡터 곱(HVP)을 일반화하여 텐서-벡터 곱(TVP)으로 확장함으로써 고차 PDE, 예를 들어 비하모닉 방정식에서의 4차 항을 효율적으로 계산할 수 있도록 한다.
• HTE를 PINN 프레임워크에 통합하여, 적절한 조건 하에서 손실 함수가 원래 PINN 손실로 수렴하고 편향이 없는 것을 보장한다.
• 이론적 분산 분 析와 고차원 및 고차 PDE에서의 실증 평가를 통해 HTE와 SDGD를 비교한다.

실험 결과

연구 질문

• RQ1 HTE는 고차원 PINN에서 전체 헤시안 행렬 계산의 스케일러블한 대안이 될 수 있는가? 이는 계산 비용과 메모리 소비를 줄이는가?
• RQ2 고차원 PDE에서 차원 간 변동성이 높을 경우 HTE의 분산은 SDGD보다 어떻게 다른가?
• RQ3 비하모닉 방정식과 같은 4차 PDE에 대해 텐서-벡터 곱(TVPs)을 사용하여 HTE를 확장할 수 있는가?
• RQ4 표준 가정 하에서 HTE 기반 PINN 손실은 원래 PINN 손실로 편향 없이 수렴하는가?
• RQ5 HTE는 Gradient-Enhanced PINNs(gPINN) 및 고차 PDE에서 더 빠르고 메모리 효율적인 학습을 가능하게 하는가?

주요 결과

• HTE는 매트릭스 연산을 스칼라 HVP로 대체함으로써 전체 헤시안에 비해 메모리 소비를 O(d²)에서 O(d)로 줄인다.
• 헤시안의 잔차 손실에 대해 추정기법을 적용할 경우 HTE 기반 손실은 편향이 없고 원래 PINN 손실로 수렴한다.
• 고차원 PDE에서 차원 간 분산이 높은 상황에서 HTE는 SDGD보다 분산이 낮아 복잡한 PDE에 대해 더 안정적인 성능을 보인다.
• 비하모닉 방정식의 경우, HTE와 함께 TVPs를 사용함으로써 O(d⁴) 텐서를 순차적인 텐서-벡터 연산으로 줄여 효율적인 계산이 가능하다.
• 실증 결과에 따르면 HTE는 SDGD와 유사한 수렴 속도를 보이며, 메모리 소비를 크게 줄이고 고차원 환경에서 더 빠른 학습을 가능하게 한다.
• HTE는 gPINN에서 기울기 강화 항을 효율적으로 추정함으로써 모델 표현력에 손상을 주지 않고 학습 속도를 가속화한다.