[논문 리뷰] What can a Single Attention Layer Learn? A Study Through the Random Features Lens

Hengyu Fu, Tianyu Guo외 2명|arXiv (Cornell University)|2023. 07. 21.

Domain Adaptation and Few-Shot LearningComputer Science인용 수 3

한 줄 요약

이 논문은 쿼리 및 키 행렬을 무작위로 초기화하고 동결한 상태에서 값 행렬만 학습 가능한 랜덤 특징 설정에서 단일 멀티헤드 어텐션 레이어의 표현 및 일반화 성질을 연구한다. 이 모델이 쿼리-키 행렬의 평균이 0이 아닌 경우에 특히 유리하게, 키 벡터의 순열 불변 함수를 효율적으로 학습할 수 있으며, 표준 두층 신경망보다 더 낮은 샘플 복잡도를 달성함을 보여준다. 이는 상관관계 기반 함수의 보다 우수한 근사화를 가능하게 한다.

ABSTRACT

Attention layers -- which map a sequence of inputs to a sequence of outputs -- are core building blocks of the Transformer architecture which has achieved significant breakthroughs in modern artificial intelligence. This paper presents a rigorous theoretical study on the learning and generalization of a single multi-head attention layer, with a sequence of key vectors and a separate query vector as input. We consider the random feature setting where the attention layer has a large number of heads, with randomly sampled frozen query and key matrices, and trainable value matrices. We show that such a random-feature attention layer can express a broad class of target functions that are permutation invariant to the key vectors. We further provide quantitative excess risk bounds for learning these target functions from finite samples, using random feature attention with finitely many heads. Our results feature several implications unique to the attention structure compared with existing random features theory for neural networks, such as (1) Advantages in the sample complexity over standard two-layer random-feature networks; (2) Concrete and natural classes of functions that can be learned efficiently by a random-feature attention layer; and (3) The effect of the sampling distribution of the query-key weight matrix (the product of the query and key matrix), where Gaussian random weights with a non-zero mean result in better sample complexities over the zero-mean counterpart for learning certain natural target functions. Experiments on simulated data corroborate our theoretical findings and further illustrate the interplay between the sample size and the complexity of the target function.

연구 동기 및 목표

• 단일 멀티헤드 어텐션 레이어가 유의미한 샘플 복잡도로 학습할 수 있는 함수 클래스를 이해하는 것.
• 유한한 샘플과 유한하지만 큰 수의 헤드를 가진 랜덤 특징 어텐션(RFA) 모델의 일반화 성능을 분석하는 것.
• 쿼리-키 행렬 초기화, 특히 비영 평균 분포가 상관관계 기반 함수 학습 효율성에 미치는 영향을 조사하는 것.
• RFA의 샘플 복잡도를 표준 두층 랜덤 특징 네트워크(RFMLP)와 비교하고 어텐션의 이점이 나타나는 시나리오를 규명하는 것.
• 시뮬레이션된 데이터를 통한 실험을 통해 이론적 결과를 검증하고, 샘플 크기와 목표 함수 복잡도 간의 상호작용을 분석하는 것.

제안 방법

• 단일 쿼리 벡터와 N개의 키 벡터를 가진 단일 어텐션 레이어 메커니즘을 모델링하며, 쿼리 및 키 행렬은 무작위로 초기화되고 동결되며, 값 행렬은 학습 가능하다.
• 어텐션 메커니즘이 쿼리-키 내적을 기반으로 한 어텐션 점수를 통해 값의 가중 평균을 계산하는 랜덤 특징 어텐션(RFA) 모델을 도입한다.
• 목표 함수 $f_\star$의 복잡도를 측정하는 $B(f_\star)$와 샘플 수 $n$을 포함한 초과 위험 경계 $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$를 도출한다.
• 쿼리-키 행렬 곱의 평균이 0이 아닌(예: 항등행렬) 비편향 RFA 변형(BRFA)을 제안하여 실제 어텐션 초기화 패턴을 모델링한다.
• 이론적 분석을 통해 쿼리-키 벡터 간 상관관계에 의존하는 함수에 대해 BRFA가 표준 RFA보다 더 낮은 샘플 복잡도를 달성함을 보여준다.
• 시뮬레이션된 데이터에서 RFA, BRFA, RFMLP를 다양한 시퀀스 길이 $N$, 입력 차원 $d$, 목표 함수 복잡도에서 비교하여 이론적 샘플 복잡도 경향을 검증한다.

실험 결과

연구 질문

• RQ1 무작위로 초기화되고 동결된 쿼리 및 키 행렬을 가진 단일 멀티헤드 어텐션 레이어는 넓은 범위의 순열 불변 목표 함수를 표현할 수 있는가?
• RQ2 RFA 모델을 사용할 때 이러한 목표 함수를 학습하는 데 필요한 샘플 복잡도는 얼마이며, 표준 두층 랜덤 특징 네트워크와 비교해 볼 때 어떻게 다른가?
• RQ3 쿼리-키 행렬 곱의 초기화 분포(특히 비영 평균 대비 영 평균)가 상관관계 기반 함수 학습 효율성에 미치는 영향은 무엇인가?
• RQ4 이론적 샘플 복잡도 경계는 현실적인 모델 크기와 목표 함수에서 유한 샘플 설정에서도 검증 가능할 수 있는가?
• RQ5 RFA 및 BRFA를 통한 목표 함수 학습에서 시퀀스 길이, 입력 차원, 모델 성능 간의 경험적 상호작용은 어떠한가?

주요 결과

• RFA 모델은 다항식 및 상관관계 기반 함수를 포함한 두 토큰의 함수 평균으로 구성된 넓은 범위의 순열 불변 목표 함수를 표현할 수 있다.
• RFA의 초과 위험 경계는 $\widetilde{\mathcal{O}}(\sqrt{B(f_\star)/n})$이며, 이는 입력 차원에만 의존하고 키 토큰 수 $N$에는 영향을 받지 않아 RFMLP보다 더 낮은 샘플 복잡도를 보인다.
• 쿼리-키 벡터 간 상관관계에 의존하는 함수에 대해 비편향 RFA(BRFA)는 평균이 0이 아닌 쿼리-키 행렬을 사용할 경우 표준 RFA보다 더 낮은 샘플 복잡도를 달성하며, $d \gg p$일 경우 $\text{Poly}(d)$ 요소를 절감할 수 있다.
• 실험 결과 RFA는 특히 시퀀스 길이 $N$이 증가할수록 키 순열 불변 함수 학습에서 RFMLP를 능가함을 확인하여 이론적 샘플 복잡도 우수성 검증됨.
• BRFA는 상관관계 기반 함수 학습에서 RFA를 크게 능가하며, 상관관계 강도($\gamma$로 제어)가 증가할수록 성능 격차가 커져 비영 평균 초기화의 이론적 이점이 확인됨.
• 분석 결과 BERT의 사전 학습된 쿼리-키 행렬은 대각선으로 지배되어 있음을 확인하여 실제 어텐션 메커니즘이 비영 평균 초기화의 인덕티브 바이어스를 이미 활용하고 있음을 시사하며, 이는 본 논문의 이론적 결과와 일치한다.