[논문 리뷰] Guaranteed Approximation Bounds for Mixed-Precision Neural Operators

Renbo Tu, Colin White외 6명|arXiv (Cornell University)|2023. 07. 27.

Model Reduction and Neural Networks인용 수 4

한 줄 요약

이 논문은 메모리 사용량을 최대 50% 감소시키고 학습 처리량을 최대 58% 향상시키며 정확도 손실을 최소화하는 푸리에 신경 연산자(Fourier Neural Operators, FNO)를 위한 혼합 정밀도 학습 루틴을 제안한다. tanh 전활성화와 주파수 모드 절단을 적용함으로써, 복소수 푸리에 도메인에서 반정밀도 학습을 안정화시켜 고해상도 PDE 해를 손실 없이 효율적으로 학습할 수 있도록 한다.

ABSTRACT

Neural operators, such as Fourier Neural Operators (FNO), form a principled approach for learning solution operators for PDEs and other mappings between function spaces. However, many real-world problems require high-resolution training data, and the training time and limited GPU memory pose big barriers. One solution is to train neural operators in mixed precision to reduce the memory requirement and increase training speed. However, existing mixed-precision training techniques are designed for standard neural networks, and we find that their direct application to FNO leads to numerical overflow and poor memory efficiency. Further, at first glance, it may appear that mixed precision in FNO will lead to drastic accuracy degradation since reducing the precision of the Fourier transform yields poor results in classical numerical solvers. We show that this is not the case; in fact, we prove that reducing the precision in FNO still guarantees a good approximation bound, when done in a targeted manner. Specifically, we build on the intuition that neural operator learning inherently induces an approximation error, arising from discretizing the infinite-dimensional ground-truth input function, implying that training in full precision is not needed. We formalize this intuition by rigorously characterizing the approximation and precision errors of FNO and bounding these errors for general input functions. We prove that the precision error is asymptotically comparable to the approximation error. Based on this, we design a simple method to optimize the memory-intensive half-precision tensor contractions by greedily finding the optimal contraction order. Through extensive experiments on different state-of-the-art neural operators, datasets, and GPUs, we demonstrate that our approach reduces GPU memory usage by up to 50% and improves throughput by 58% with little or no reduction in accuracy.

연구 동기 및 목표

• 고해상도 PDE 데이터에서 FNO 학습의 높은 메모리 및 계산 요구량을 해결하기 위해.
• 표준 혼합 정밀도 기법이 적용되지 않는 복소수 푸리에 도메인에서 혼합 정밀도 학습을 가능하게 하기 위해.
• 반정밀도 FNO 학습에서의 수치 안정성과 정밀도 간의 상충 관계를 연구하기 위해.
• 정확도를 유지하면서 자원 사용량을 줄이는 강력하고 고처리량의 학습 파이프라인을 개발하기 위해.
• 실제 응용 분야의 PDE, 예를 들어 나비에-스토크스 및 다르시 흐름 문제에서 본 방법의 효과를 입증하기 위해, 특히 제로샷 슈퍼레졸루션을 포함하여.

제안 방법

• 성능 저하 요인과 잠재적 성능 향상 요소를 파악하기 위해 전정밀도 및 혼합 정밀도 학습의 프로파일링을 수행한다.
• 푸리에 변환 이전에 tanh 전활성화를 적용하여 반정밀도에서 수치적 값을 안정화시킨다.
• 고주파 모드를 절단하여 반정밀도 표현 오차의 영향을 줄인다.
• 학습 중에 반정밀도에서 혼합 정밀도로, 다시 전정밀도로 전환하는 정밀도 스케줄링을 구현한다.
• 반정밀도 복소수 연산을 벡터 캐스팅을 통해 효율적으로 처리하는 cuFFT를 사용한다.
• 다양한 해상도와 정밀도 설정에서 나비에-스토크스 및 다르시 흐름 데이터셋을 대상으로 방법을 검증한다.

Figure 1: Overview of the half-precision FNO block. We cast the input tensor to half precision and then use a tanh pre-activation for numerical stability. We use cuFFT to run the FFT, and we compute the tensor contraction in half precision by casting each complex number to a vector.

실험 결과

연구 질문

• RQ1 복소수 값이고 함수 공간 연산을 수행하는 FNO에 대해, 반정밀도 학습을 안전하게 적용할 수 있는가?
• RQ2 반정밀도 FNO 학습에서 발생하는 수치적 불안정성은 무엇이며, 어떻게 완화할 수 있는가?
• RQ3 주파수 모드 절단은 반정밀도 FNO의 정확도와 안정성에 어떤 영향을 미치는가?
• RQ4 정확도를 손상시키지 않고 메모리 사용량과 학습 처리량을 얼마나 줄일 수 있는가?
• RQ5 제안된 방법은 제로샷 슈퍼레졸루션을 위한 FNO의 해상도 불변성을 유지하는가?

주요 결과

• 제안된 혼합 정밀도 FNO 학습 루틴은 나비에-스토크스 및 다르시 흐름 문제에서 GPU 메모리 사용량을 최대 50% 감소시킨다.
• 최적화된 혼합 정밀도 파이프라인을 사용할 경우 학습 처리량이 최대 58% 향상된다.
• tanh 전활성화와 주파수 모드 절단의 조합이 반정밀도 학습에서 수치적 오버플로우 및 언더플로우를 효과적으로 방지한다.
• 반정밀도 FNO는 전정밀도 대비 정확도 손실가 매우 작으며, 예를 들어 1024×1024 나비에-스토크스 문제에서 L2 오차가 0.00213에서 0.00236로 증가할 뿐이다.
• 정밀도 스케줄링 방법은 순수 반정밀도보다 훨씬 높은 정확도를 달성하며, 1024×1024 나비에-스토크스 문제에서 L2 오차를 0.00236에서 0.00170으로 감소시킨다.
• 이 방법은 해상도 불변성을 유지하여, 최대 1024×1024 해상도에서 정확도가 낮게 떨어지지 않는 정확한 제로샷 슈퍼레졸루션을 가능하게 한다.

Figure 2: Training throughput and runtime as a function of the method, on different hardware. For mixed-precision FNO + AMP, we consistently observe an improvement of throughput up to 1.58X over the baseline, and a speedup of up to 36.8% . Our method also improves upon using only AMP in throughput b