[논문 리뷰] $ξ$-torch: differentiable scientific computing library

Muhammad Kasim, S. M. Vinko|arXiv (Cornell University)|2020. 10. 05.

Model Reduction and Neural Networks참고 문헌 14인용 수 4

한 줄 요약

이 논문은 과학 계산을 위한 미분 가능 기능들(예: 근 찾기, 상미분방정식 풀이기구, 적분 방법 등)을 제공하는 PyTorch 기반 라이브러리인 $\xi$-torch를 소개한다. 이 라이브러리는 수치적 안정성과 메모리 효율성을 향상시키기 위해 분석적 도함수를 활용한다. 이는 물리 기반 딥러닝의 엔드 투 엔드 미분 가능성을 가능하게 하며, 일阶 및 고계 도함수를 지원하여 거울 설계 및 분자 동역학 최적화에서 빠른 수렴을 보여준다.

ABSTRACT

Physics-informed learning has shown to have a better generalization than learning without physical priors. However, training physics-informed deep neural networks requires some aspect of physical simulations to be written in a differentiable manner. Unfortunately, some operations and functionals commonly used in physical simulations are scattered, hard to integrate, and lack higher order derivatives which are needed in physical simulations. In this work, we present $ξ$-torch, a library of differentiable functionals for scientific simulations. Example functionals are a root finder and an initial value problem solver, among others. The gradient of functionals in $ξ$-torch are written based on their analytical expression to improve numerical stability and reduce memory requirements. $ξ$-torch also provides second and higher order derivatives of the functionals which are rarely available in existing packages. We show two applications of this library in optimizing parameters in physics simulations. The library and all test cases in this work can be found at https://github.com/xitorch/xitorch/ and the documentation at https://xitorch.readthedocs.io.

연구 동기 및 목표

• 물리 기반 딥러닝을 위한 기존 라이브러리에서 통합된, 미분 가능한 과학 기능의 부족을 해결하기 위해.
• 자동 미분을 사용하는 전진 알고리즘 대신 분석적 도함수를 사용하여 과학 시뮬레이션의 수치적 안정성을 향상시키고 메모리 사용량을 줄이기 위해.
• 진동수 계산과 같은 고급 과학 시뮬레이션을 위해 고계 도함수(예: 2차 도함수)를 효율적이고 정확하게 계산할 수 있도록 하기 위해.
• 순수 함수와 선형 연산자 추상화를 통해 딥러닝과 물리 시뮬레이션을 통합할 수 있도록 단일화되고 모듈러한 라이브러리를 제공하기 위해.
• ODE 풀이기구나 근 찾기와 같은 복잡한 시뮬레이션 구성 요소를 통해 기울기 흐름을 가능하게 하여 물리 시스템의 매개변수 최적화를 지원하기 위해.

제안 방법

• 근 찾기, 상미분방정식 풀이기구(solve_ivp), 적분 방법(mcquad, SQuad)와 같은 기능들을 분석적 도함수로 구현하여 수치적 안정성 향상과 메모리 소비 감소를 도모한다.
• 고계 도함수 계산이 정확하게 이루어지도록 하기 위해 입력으로 순수 함수를 사용하며, 출력에 영향을 주는 상태 변수를 추적하여 상태 기반 객체 메서드를 순수 함수로 변환하는 메커니즘을 도입한다.
• 메서드 출력에 영향을 주는 매개변수 이름을 수동으로 지정할 수 있도록 EditableModule 기본 클래스를 도입하여 복잡한 클래스에서 정확한 기울기 계산을 가능하게 한다.
• 고차원 문제에서 비용이 많이 드는 행렬 구축을 피하기 위해 희소 선형 대수 연산을 위한 LinearOperator 추상화를 제공한다.
• PyTorch의 자동 미분 엔진을 활용하면서도 분석적 도함수를 사용하여 gradgradcheck 검증을 통해 2차 및 고계 도함수를 계산한다.
• SciPy를 영감으로 삼았지만, 미분 가능한 과학 계산에 맞게 조정된 모듈러한 라이브러리 아키텍처를 설계하여 최적화, 통합, 보간, 선형 대수 전용 모듈을 전용으로 제공한다.

실험 결과

연구 질문

• RQ1 분석적 도함수를 활용한 미분 가능한 과학 기능을 구현함으로써 과학 시뮬레이션에서 수치적 안정성을 향상시키고 메모리 사용량을 줄일 수 있는가?
• RQ2 ODE 풀이기구나 근 찾기와 같은 복잡한 기능에 대해 고계 도함수(예: 2차 도함수)를 효율적이고 정확하게 계산할 수 있는가?
• RQ3 통합된 미분 가능한 기능 라이브러리가 과학적 응용 분야에서 물리 기반 딥러닝의 가속화에 얼마나 기여하는가?
• RQ4 상태 기반 객체 메서드를 어떻게 순수 함수로 변환하여 미분 가능한 과학 계산에서 정확한 기울기 계산을 가능하게 할 수 있는가?
• RQ5 표현된 시뮬레이션을 통해 $\xi$-torch가 거울 설계 및 분자 동역학와 같은 물리 시스템에서 효과적인 매개변수 최적화를 가능하게 하는가?

주요 결과

• $\xi$-torch 라이브러리는 알고리즘적 구현을 통한 역전파보다 분석적 표현을 사용하여 과학 기능의 일阶 및 고계 도함수를 성공적으로 계산하며, 수치적 안정성 향상과 메모리 사용 감소를 달성한다.
• gradgradcheck 검증을 통해 2차 도함수의 정확한 계산이 이루어졌으며, 복잡한 미분 가능한 기능에서의 강건성을 입증한다.
• 거울 설계 사례에서는 근 찾기 기능을 통해 신경망 기반 거울 표면을 200회 반복 내에 최적화하여 목표 지점에서 빛의 광선을 정확하게 집중시켰다.
• 분자 동역학 사례에서는 solve_ivp 기능을 통해 초기 속도를 최적화하여 t=1일 때 'OX'라는 문자를 형성했으며, Adam 최적화기를 사용해 7500회 반복 내에 수렴을 달성했다.
• 순수 함수 및 편집 가능한 상태 변수를 지원하는 라이브러리의 모듈러한 설계 덕분에 복잡한 물리 시뮬레이션 파이프라인에서 신뢰할 수 있는 기울기 계산이 가능해졌다.
• 몬테카를로 적분 및 고정 샘플 적분과 같은 미분 가능한 기능의 가용성은 표준 ODE 및 근 찾기 작업을 넘어서 미분 가능한 과학 계산의 범위를 확장시켰다.