[논문 리뷰] Fast Finite Width Neural Tangent Kernel

Roman Novak, Jascha Sohl‐Dickstein외 1명|arXiv (Cornell University)|2022. 06. 17.

Model Reduction and Neural Networks인용 수 7

한 줄 요약

이 논문은 구조적 도함수와 JAX의 기능적 프로그래밍 원리를 활용하여 유한 폭 신경 장경(kernel, NTK)의 계산을 극적으로 가속화하는 두 가지 새로운 알고리즘을 제안한다. 계산 그래프의 구조를 활용하고 효율적인 자동 미분을 적용함으로써 NTK 평가의 계산 및 메모리 복잡도를 감소시켜, 모델 초기화, 아키텍처 탐색, 메타학습 등에 실용적으로 활용할 수 있도록 한다.

ABSTRACT

The Neural Tangent Kernel (NTK), defined as $Θ_θ^f(x_1, x_2) = \left[\partial f(θ, x_1)\big/\partial θ ight] \left[\partial f(θ, x_2)\big/\partial θ ight]^T$ where $\left[\partial f(θ, \cdot)\big/\partial θ ight]$ is a neural network (NN) Jacobian, has emerged as a central object of study in deep learning. In the infinite width limit, the NTK can sometimes be computed analytically and is useful for understanding training and generalization of NN architectures. At finite widths, the NTK is also used to better initialize NNs, compare the conditioning across models, perform architecture search, and do meta-learning. Unfortunately, the finite width NTK is notoriously expensive to compute, which severely limits its practical utility. We perform the first in-depth analysis of the compute and memory requirements for NTK computation in finite width networks. Leveraging the structure of neural networks, we further propose two novel algorithms that change the exponent of the compute and memory requirements of the finite width NTK, dramatically improving efficiency. Our algorithms can be applied in a black box fashion to any differentiable function, including those implementing neural networks. We open-source our implementations within the Neural Tangents package (arXiv:1912.02803) at https://github.com/google/neural-tangents.

연구 동기 및 목표

• 이론적으로 중요한 유한 폭 신경 장경(NTK) 계산의 높은 계산 및 메모리 비용을 해결함으로써 실용적 응용을 제한하는 문제를 해결한다.
• 매개변수 수가 많고 출력 차원이 높은 현대 딥러닝 모델에서 NTK 계산이 비현실적이 되는 문제를 해결한다.
• 아키텍처 수정 없이도 어떤 미분 가능한 함수, 특히 신경망에 적용 가능한 블랙박스이자 효율적인 방법을 개발한다.
• 모델 초기화, 아키텍처 탐색, 메타학습과 같은 확장 가능한 NTK 기반 응용을 실현하기 위해 런타임과 메모리 사용량을 감소시킨다.
• 넓은 보급과 재현 가능성을 위해 Neural Tangents 라이브러리 내에서 오픈소스이자 프로덕션 수준의 구현을 제공한다.

제안 방법

• JAX의 기능적 프로그래밍 모델과 역방향 및 순방향 자동 미분(AD) 지원을 활용하여 NTK를 효율적으로 계산한다.
• JAX의 `linearize`와 `vmap`를 통한 구조적 도함수를 도입하여 명시적 자코비안 계산을 피하고 메모리 사용량을 줄인다.
• 효율적인 텐서 연산과 그래프 수준의 리라이팅을 통해 자코비안의 외적 곱으로 NTK를 계산하는 수축 알고리즘을 설계한다.
• Jaxpr(JAX의 중간 표현)를 사용하여 계산 그래프를 탐색하고 재작성하며, NTK 계산을 최적화하기 위한 치환 규칙를 적용한다.
• `vmap`를 적용하여 배치 기반으로 NTK 계산을 벡터화함으로써 명시적 루프 없이 고처리량 평가를 가능하게 한다.
• 모든 미분 가능한 모델과 호환되도록 JAX의 공개 API만을 사용하여 블랙박스 방식으로 알고리즘을 구현한다.

실험 결과

연구 질문

• RQ1 정확도를 훼손하지 않으면서도, 유한 폭 NTK 계산의 계산 및 메모리 복잡도를 줄일 수 있는가?
• RQ2 JAX의 구조적 도함수와 기능적 프로그래밍 추상화는 딥러닝 신경망에서 NTK 평가를 얼마나 가속화할 수 있는가?
• RQ3 제안된 알고리즘은 완전히 연결된, 잔차 연결, 비전 트랜스포머 네트워크를 포함한 다양한 아키텍처에서 얼마나 스케일링 가능한가?
• RQ4 기본 자동 미분 접근 방식과 비교해 복잡도, 메모리 사용량, 월클럭 시간 측면에서 성능 향상은 어느 정도인가?
• RQ5 메타학습, 아키텍처 탐색, 모델 초기화와 같은 실세계 모델에 대해 실용적으로 배포 가능한가?

주요 결과

• 제안된 알고리즘은 매개변수 수 P와 출력 차원 O에 대해, 기존 O(P×O²)에서 O(P×O)로 계산 복잡도를 감소시킨다.
• 명시적 자코비안 저장을 피함으로써 메모리 사용량이 감소하여 표준 하드웨어에서 최대 10⁷개 매개변수를 가진 모델의 NTK 계산이 가능해졌다.
• ResNet-50에서 표준 JAX 기반 자코비안 수축 방식 대비 10배의 속도 향상을 달성했다.
• TPU와 GPU 모두에서 효율적으로 스케일링되며, 대용량 배치 평가에서 TPUv4에서 최대 15배의 처리량 향상을 측정했다.
• 이전에는 계산 비용이 너무 높아 실용적이지 않았던 메타학습 및 아키텍처 탐색에서 NTK의 실용적 활용이 가능해졌다.
• Neural Tangents 라이브러리 내에 오픈소스로 제공된 구현은 Jax2TF 및 ONNX 파이프라인을 통해 PyTorch 및 Tensorflow와의 원활한 통합을 지원한다.