[논문 리뷰] Infinite attention: NNGP and NTK for deep attention networks

Jiri Hron, Yasaman Bahri외 2명|arXiv (Cornell University)|2020. 06. 18.

Gaussian Processes and Bayesian Inference인용 수 29

한 줄 요약

이 논문은 다중 헤드 어텐션 네트워크가 무한한 너비 근처에서 가우시안 프로세스(NNGP)로 수렴하고, 신경 미분 커널(NTK) 행동을 보임을 입증하며, 단일 헤드 설정에서 어텐션의 비가우시안 출력으로 인해 발생하는 이론적 과제를 해결한다. 유한한 넓이와 무한한 넓이 모델 모두에서 성능을 향상시키는 수정된 어텐션 메커니즘과 위치 인코딩 전략을 제안하며, 데이터 증강 없이 CIFAR-10에서 최신 기준 성능을 달성한다.

ABSTRACT

There is a growing amount of literature on the relationship between wide neural networks (NNs) and Gaussian processes (GPs), identifying an equivalence between the two for a variety of NN architectures. This equivalence enables, for instance, accurate approximation of the behaviour of wide Bayesian NNs without MCMC or variational approximations, or characterisation of the distribution of randomly initialised wide NNs optimised by gradient descent without ever running an optimiser. We provide a rigorous extension of these results to NNs involving attention layers, showing that unlike single-head attention, which induces non-Gaussian behaviour, multi-head attention architectures behave as GPs as the number of heads tends to infinity. We further discuss the effects of positional encodings and layer normalisation, and propose modifications of the attention mechanism which lead to improved results for both finite and infinitely wide NNs. We evaluate attention kernels empirically, leading to a moderate improvement upon the previous state-of-the-art on CIFAR-10 for GPs without trainable kernels and advanced data preprocessing. Finally, we introduce new features to the Neural Tangents library (Novak et al., 2020) allowing applications of NNGP/NTK models, with and without attention, to variable-length sequences, with an example on the IMDb reviews dataset.

연구 동기 및 목표

• 다중 헤드 어텐션 메커니즘을 갖춘 깊은 신경망의 무한 너비 근처를 엄밀하게 특성화하는 것.
• 표준 교환 가능성 또는 조건부 가정을 만족하지 못하는 어텐션 레이어로 확장된 가우시안 프로세스 및 신경 미분 커널 프레임워크를 이론적으로 해결하는 것.
• 유한한 넓이와 무한한 넓이 네트워크 모두에서 성능을 향상시키는 수정된 어텐션 메커니즘과 위치 인코딩 전략을 제안하는 것.
• 제안된 어텐션 커널을 실증적으로 평가하고, GP 추론을 사용하여 CIFAR-10과 같은 벤치마크 데이터셋에서 향상된 성능을 보여주는 것.
• Neural Tangents 라이브러리에 가변 길이 시퀀스를 지원하는 기능을 추가하여, 텍스트와 같은 순차적 데이터에 NNGP/NTK 응용을 가능하게 하는 것.

제안 방법

• d^{-1/2} 및 d^{-1} 스케일링 하에서 다중 헤드 어텐션 네트워크의 NNGP 및 NTK 커널을 유도하며, 무한 너비 근처에서 가우시안 프로세스로 수렴함을 보임.
• d^{-1} 스케일링과 보간 기반 공분산 업데이트를 갖춘 수정된 어텐션 메커니즘을 도입하여 표현력을 유지하면서도 이론적 분석이 가능하도록 함.
• 무한 너비 영역에서 일반화 및 성능 향상을 향상시키는 구조적 위치 인코딩을 제안하며, NNGP/NTK 추론과 조합되어도 효과를 발휘함.
• 잔차 연결 및 레이어 정규화를 다룰 수 있도록 이론 프레임워크를 적응시켜, 커널 수렴 및 모델 행동에 미치는 영향을 분석함.
• 학습 가능한 위치 임베딩의 기울기 기여를 통합하는 새로운 커널 업데이트 규칙을 도입하며, 하이브리드 보간 메커니즘을 통해 NTK를 수정함.
• Neural Tangents 라이브러리에 가변 길이 시퀀스를 지원하는 코드 확장을 배포하여, IMDb와 같은 데이터셋에서 종단 간 NNGP/NTK 추론을 가능하게 함.

실험 결과

연구 질문

• RQ1 넓은 신경망에서 다중 헤드 어텐션은 무한 너비 근처에서 가우시안 프로세스로 수렴하는가? 만약 그렇다면 어떤 조건에서 그러한 수렴이 이루어지는가?
• RQ2 왜 표준 NNGP 및 NTK 유도 기법은 단일 헤드 어텐션에선 실패하며, 이를 다중 헤드 아키텍처에 어떻게 적응시킬 수 있는가?
• RQ3 위치 인코딩과 레이어 정규화는 무한 너비 어텐션 네트워크에서 커널 행동과 일반화 성능에 어떤 영향을 미치는가?
• RQ4 d^{-1} 스케일링을 갖는 수정된 어텐션 메커니즘은 표현력을 유지하면서도 엄밀한 NNGP/NTK 분석을 가능하게 할 수 있는가?
• RQ5 실제 벤치마크, 예를 들어 CIFAR-10에 대해 유도된 NNGP 및 NTK 커널을 적용했을 때, 데이터 증강 없이 어떤 성능 향상이 달성될 수 있는가?

주요 결과

• 다중 헤드 어텐션 네트워크는 무한 너비 근처에서 가우시안 프로세스로 수렴하지만, 단일 헤드 어텐션은 공유된 어텐션 가중치로 인한 통계적 의존성으로 인해 그렇지 않다.
• d^{-1/2} 및 d^{-1} 스케일링 하에서 다중 헤드 어텐션에 대한 NNGP 및 NTK 커널의 엄밀한 유도를 제공하며, 어텐션 메커니즘으로의 이론적 프레임워크 확장을 가능하게 한다.
• d^{-1} 스케일링과 보간 기반 공분산 업데이트를 갖는 수정된 어텐션 메커니즘은 표현력을 유지하면서 안정적인 커널 수렴을 가능하게 하며, 이전 방법에서 관찰된 풀링 유사 행동을 피한다.
• 위치 인코딩은 무한 너비 근처에서 실증적 성능 향상에 크게 기여하며, 모델이 훈련되지 않은 상태에서도 일반화 능력을 향상시킴을 보여준다.
• 제안된 방법은 데이터 증강 없이도 GP 추론을 사용하여 이전 최고 성능에 비해 중간 정도의 성능 향상을 달성한다.
• 저자들은 Neural Tangents 라이브러리에 가변 길이 시퀀스를 지원하는 코드 확장을 배포하여, 텍스트와 같은 순차적 데이터에서 NNGP/NTK 추론을 가능하게 하며, IMDb 데이터셋에서의 적용 사례를 제시한다.