[논문 리뷰] Feature Learning in Infinite-Width Neural Networks

Greg Yang, J. Edward Hu|arXiv (Cornell University)|2020. 11. 30.

Topic Modeling참고 문헌 34인용 수 20

한 줄 요약

이 논문은 표준 NTK 파arametrization과 달리 초기화 시점에서 특징이 고정되지 않는, 특징 학습을 허용하는 수정된 신경망 파arametrization인 최대 업데이트 파arametrization(µP)을 제안한다. 텐서 프로그램 기법을 사용하여 저자들은 µP의 정확한 무한한 너비 근사를 유도하였으며, 이는 Word2Vec 및 소수 샘플 학습에서 NTK 기반 모델과 유한한 너비 네트워크를 모두 능가함을 보여준다. 유한한 너비 모델은 너비가 증가함에 따라 µP 성능에 점점 가까워진다.

ABSTRACT

As its width tends to infinity, a deep neural network's behavior under gradient descent can become simplified and predictable (e.g. given by the Neural Tangent Kernel (NTK)), if it is parametrized appropriately (e.g. the NTK parametrization). However, we show that the standard and NTK parametrizations of a neural network do not admit infinite-width limits that can learn features, which is crucial for pretraining and transfer learning such as with BERT. We propose simple modifications to the standard parametrization to allow for feature learning in the limit. Using the *Tensor Programs* technique, we derive explicit formulas for such limits. On Word2Vec and few-shot learning on Omniglot via MAML, two canonical tasks that rely crucially on feature learning, we compute these limits exactly. We find that they outperform both NTK baselines and finite-width networks, with the latter approaching the infinite-width feature learning performance as width increases. More generally, we classify a natural space of neural network parametrizations that generalizes standard, NTK, and Mean Field parametrizations. We show 1) any parametrization in this space either admits feature learning or has an infinite-width training dynamics given by kernel gradient descent, but not both; 2) any such infinite-width limit can be computed using the Tensor Programs technique. Code for our experiments can be found at github.com/edwardjhu/TP4.

연구 동기 및 목표

• 과도하게 파aram터화된 네트워크에서 특징 학습을 포착하지 못하는 신경 장경계핵(Neural Tangent Kernel, NTK) 이론의 핵심적 한계를 해결하기 위해.
• 표준 파arametrization 및 NTK 파arametrization이 무한한 너비 근사에서 작동하지 않는 조건에서 특징 학습을 가능하게 하는 파arametrization을 규명하기 위해.
• 텐서 프로그램 기법을 사용하여 임의의 아키텍처와 학습 동역학에 대해 무한한 너비 근사의 정확한 계산 방법을 체계적으로 개발하기 위해.
• 제안된 µP 파arametrization이 MAML 기반의 특징 학습 작업, 예를 들어 Word2Vec 및 소수 샘플 학습에서 NTK 및 유한한 너머의 성능을 뛰어나게 하는지 확인하기 위해.

제안 방법

• 조정 가능한 스케일링 파aram터(a_l, b_l, c)를 통해 표준, NTK, 평균장 파arametrization을 일반화하는 abc-파arametrization 프레임워크를 제안한다.
• 1차 레이어에서 a1 = 0, l ≥ 2일 때 al = 1/2, bl = 0, c = 1/2로 설정된 특별한 abc-파arametrization인 최대 업데이트 파arametrization(µP)을 도입하여, 무한한 너비 근사에서 최대한의 특징 학습을 가능하게 한다.
• 무한한 너비 근사에서 µP 파arametrization의 정확한 근사를 엄밀하게 계산하기 위해 텐서 프로그램 기법을 적용하여, 네트워크 동역학의 정확한 해석적 유도를 가능하게 한다.
• 네트워크의 은닉 표현과 로짓의 무한한 너비 근사를 위한 명시적 공식을 도출하여, 학습 동역학의 정확한 계산을 가능하게 한다.
• Word2Vec 및 MAML 기반 소수 샘플 학습에서 이론적 프레임워크를 실험적으로 검증하며, µP를 NTK, GP 및 유한한 너비 모델과 비교한다.
• PCA 시각화와 단어 유추 정확도를 사용하여 특징 공간의 품질과 일반화 성능을 평가한다.

실험 결과

연구 질문

• RQ1 무한한 너비 신경망은 특징을 학습할 수 있으며, 만약 가능하다면 어떤 파arametrization 조건에서 가능할까?
• RQ2 왜 표준 및 NTK 파arametrization은 무한한 너비 근사에서 특징 학습을 지원하지 못할까?
• RQ3 안정적인 훈련과 비트리비얼한 특징 진화를 동시에 가능하게 하는 파arametrization의 수학적 구조는 무엇일까?
• RQ4 임의의 아키텍처와 학습 동역학에 대해 신경망의 무한한 너비 근사를 정확하고 체계적으로 계산하는 방법은 무엇일까?
• RQ5 제안된 µP 파arametrization은 NTK 및 유한한 너비 모델보다 특징 학습 작업에서 더 뛰어난 성능을 보일까?

주요 결과

• 최대 업데이트 파arametrization(µP)의 무한한 너비 근사는 NTK 근사와 달리 비트리비얼한 특징 학습을 가능하게 하며, 초기화 시점에서 특징이 고정되지 않는다.
• text8 데이터셋에서 µP 무한한 너비 모델은 43.31%의 단어 유추 정확도를 기록하였으며, NTK 기반 모델(0.0%) 및 유한한 너비 네트워크보다 유의미하게 뛰어나다.
• 유한한 너비의 µP 네트워크는 너비가 증가함에 따라 무한한 너비 성능에 수렴하며, 너비 2^6일 때 정확도 33.35%에서 2^10일 때 42.56%로 상승한다.
• 더 큰 fil9 데이터셋에서 µP 무한한 너비 모델은 56.45%의 정확도를 기록하였으며, 다시 한번 NTK 기반 모델(0.0%) 및 유한한 너비 모델을 능가한다.
• 역동적 이분법 정리(Dynamical Dichotomy theorem)는 안정적이고 비트리비얼한 abc-파arametrization이 특징 학습 근사 또는 커널 근사 중 하나를 가져오지만, 둘 다를 동시에 가져오진 않음을 증명한다.
• PCA 시각화 결과, µP 임베딩은 도시와 주를 의미적으로 분리하는 반면, NTK 임베딩은 무한한 너비 근사에서 거의 무작위 상태를 유지한다.