[논문 리뷰] Neural Tangents: Fast and Easy Infinite Neural Networks in Python

Roman Novak, Lechao Xiao외 5명|arXiv (Cornell University)|2019. 12. 05.

Model Reduction and Neural Networks참고 문헌 28인용 수 57

한 줄 요약

Neural Tangents는 분석적 NNGP/NTK 커널 또는 유한 폭 근사를 통해 무한 너비 신경망을 정의, 학습, 분석하는 JAX 기반 API를 제공합니다. CPU/GPU/TPU 지원 확장 가능.

ABSTRACT

Neural Tangents is a library designed to enable research into infinite-width neural networks. It provides a high-level API for specifying complex and hierarchical neural network architectures. These networks can then be trained and evaluated either at finite-width as usual or in their infinite-width limit. Infinite-width networks can be trained analytically using exact Bayesian inference or using gradient descent via the Neural Tangent Kernel. Additionally, Neural Tangents provides tools to study gradient descent training dynamics of wide but finite networks in either function space or weight space. The entire library runs out-of-the-box on CPU, GPU, or TPU. All computations can be automatically distributed over multiple accelerators with near-linear scaling in the number of devices. Neural Tangents is available at www.github.com/google/neural-tangents. We also provide an accompanying interactive Colab notebook.

연구 동기 및 목표

• Neural Tangents를 소개하는 무한 너비 신경망 연구를 위한 오픈 소스 라이브러리.
• 복잡한 구조를 지정하고 해석적 무한 너비 커널(NNGP/NTK)을 얻기 위한 고수준 API를 제공.
• 무한 너비 한계에서의 정확한 Bayesian 추론과 경사 하강 유사 학습을 가능하게 함.
• 해석적 형태가 없는 커널에 대한 몬테카를로 근사를 지원.
• 다양한 아키텍처와 데이터셋에서 확장성 및 실용성을 입증합니다.

제안 방법

• 표준 텐서 연산(밀집, 합성곱, 비선형성)을 NNGP 및 NTK용 커널 연산으로 매핑하는 변환 정의.
• 무한 폭 네트워크에 대한 해석적 커널 계산 및 해석적 형태가 없는 경우 MC 기반 커널 근사를 제공합니다.
• 정확한 커널을 통해 무한 너비 한계에서 베이지안 추론을 수행하거나 경사하강 유사 학습을 위한 API를 제공합니다.
• Taylor 급수 근사를 통한 연속 시간 학습 역학 및 유한 폭 가중치 공간 탐색을 포함합니다.
• CPU/GPU/TPU 실행을 위한 XLA/JAX를 활용하여 자동 장치 분산 및 배치를 제공합니다.

실험 결과

연구 질문

• RQ1 무한 너비 한계에서 광범위한 신경망 아키텍처에 대해 NNGP 및 NTK 커널을 어떻게 효율적으로 계산하고 활용할 수 있는가?
• RQ2 해석적 또는 몬테카를로 커널 추정을 사용하여 유한 폭 네트워크를 무한 폭 대응으로 정확하게 분석할 수 있는가?
• RQ3 다양한 디바이스와 아키텍처에서 해석적 커널 계산의 성능과 확장성은 어떠한가?
• RQ4 완전 연결, CNN, WideResNet 등 다양한 아키텍처가 무한 너비 추론 및 커널 기반 학습에서 어떻게 비교되는가?
• RQ5 무한 너비 신경망 연구를 가속화하는 실용적인 도구와 워크플로우는 무엇인가?

주요 결과

• Neural Tangents는 많은 아키텍처에 대해 분석적 NNGP/NTK 커널을 통한 무한 너비 네트워크에 대한 정확한 추론을 가능하게 한다.
• 해석적 커널이 이용 불가능한 경우 몬테카를로 커널 추정을 통해 유연한 근사를 제공하며, 폭 또는 샘플 수가 커질수록 해석적 값으로 수렴한다.
• 실험에서 무한 너비 추론이 합성 데이터에서 유한 폭 앙상블 학습과 근접하게 일치할 수 있어 실용적 충실도를 보여준다.
• CIFAR-10에서의 평가에서 아키텍처 간 성능의 수렴 및 계층 구조를 시연하며, WideResNet과 유사한 구조가 테스트된 변형들 중 최상의 성능을 보인다.
• 추가 가속기와의 거의 선형 스케일링, 배치 처리, 연산 동합(op fusion), 2D-합성곱 기반 공분산 계산을 통한 효율성 확보를 구현한다.