[논문 리뷰] Gated recurrent neural networks discover attention

Nicolas Zucchet, Seijin Kobayashi외 5명|arXiv (Cornell University)|2023. 09. 04.

Neural Networks and ApplicationsComputer Science인용 수 3

한 줄 요약

이 논문은 곱셈 게이팅 메커니즘을 갖춘 게이팅된 순환 신경망(RNN)이 학습된 가중치 설정을 통해 선형 자기주의를 정확히 구현할 수 있음을 보여주며, 실제로 경사하강법이 주의 기반 계산을 발견함을 드러낸다. 문맥 내 학습 작업에서 훈련된 RNN은 선형 자기주의와 동일한 주의 기반 알고리즘을 암묵적으로 채택한다. 이는 RNN이 내부적으로 주의 메커니즘을 수행하고 있음을 시사한다.

ABSTRACT

Recent architectural developments have enabled recurrent neural networks (RNNs) to reach and even surpass the performance of Transformers on certain sequence modeling tasks. These modern RNNs feature a prominent design pattern: linear recurrent layers interconnected by feedforward paths with multiplicative gating. Here, we show how RNNs equipped with these two design elements can exactly implement (linear) self-attention, the main building block of Transformers. By reverse-engineering a set of trained RNNs, we find that gradient descent in practice discovers our construction. In particular, we examine RNNs trained to solve simple in-context learning tasks on which Transformers are known to excel and find that gradient descent instills in our RNNs the same attention-based in-context learning algorithm used by Transformers. Our findings highlight the importance of multiplicative interactions in neural networks and suggest that certain RNNs might be unexpectedly implementing attention under the hood.

연구 동기 및 목표

• 게이팅된 순환 신경망(RNN)이 곱셈 게이팅을 갖추고 있을 때 선형 자기주의 메커니즘을 구현할 수 있는지 조사하기.
• 실제로 RNN에서 경사하강법이 주의 유사 계산으로 이어지는 방식을 이해하기.
• 특히 깊이 있는 선형 RNN, LSTMs, GRUs와 같은 다양한 RNN 아키텍처의 인덕티브 바이어스가 주의 구현에 어떻게 영향을 미치는지 비교하기.
• 훈련된 RNN을 역설계하여, 문맥 내 학습 작업에서 선형 자기주의와 동일한 알고리즘을 학습하는지 확인하기.
• 곱셈 상호작용이 RNN이 주의 유사 계산을 수행하는 데 어떻게 기여하거나 제약을 가하는지 평가하기.

제안 방법

• 저자는 요소별 곱셈을 사용하는 입력 및 출력 게이팅을 갖춘 게이팅된 대각선 선형 RNN의 클래스를 정의한다: $ h_{t+1} = \lambda \odot h_t + g^{\text{in}}(x_t) $, $ y_t = D g^{\text{out}}(h_t) $, 여기서 $ g^{\text{in/out}}(x_t) = (W_m^{\text{in/out}} x_t) \odot (W_x^{\text{in/out}} x_t) $.
• 이러한 RNN이 특정 매개변수 설정 하에 선형 자기주의를 정확히 구현할 수 있음을 도출하며, 핵심 통찰은 순환 상태가 $ \sum_t v_t k_t^\top $를 누적한다는 점이다. 이는 주의 가중치 행렬과 유사하다.
• 선형 회귀 및 문맥 내 학습 작업에서 훈련된 RNN을 역설계하여 학습된 가중치를 이론적 주의 구성과 비교한다.
• 입력-출력 게이팅과 사이드 게이팅을 갖춘 RNN을 분석하여, 특정 가중치 설정 하에 오직 사이드 게이팅 버전만 선형 자기주의를 성공적으로 모방함을 보였다.
• 저자는 훈련에 JAX와 Flax를 사용하였으며, 분석 및 가시화에 Numpy, Scikit-learn, Matplotlib를 사용하였다.
• RNN에서 학습된 $ M(x,y) $ 행렬을 선형 자기주의의 것과 비교하여, 양의 피크와 대칭적인 음의 피크를 갖는 랭크-1 성분이 있음을 발견하였으며, 이는 원-핫 인코딩 덕분에 분류 성능에 영향을 주지 않는다.

실험 결과

연구 질문

• RQ1 특정 매개변수 설정 하에 곱셈 게이팅을 갖춘 게이팅된 RNN이 선형 자기주의를 정확히 구현할 수 있는가?
• RQ2 특히 주의 구조를 명시적으로 설계하지 않은 상황에서, 훈련된 RNN에서 경사하강법이 주의 기반 계산을 발견하는가?
• RQ3 왜 깊이 있는 선형 RNN은 주의를 성공적으로 실현하지만, LSTMs와 GRUs는 동일한 설정 하에 실패하는가?
• RQ4 문맥 내 학습 작업에서 훈련된 RNN이 선형 자기주의와 동일한 알고리즘을 어느 정도 채택하는가?
• RQ5 RNN 내 곱셈 상호작용이 주의 유사 행동의 발생을 어떻게 촉진하거나 제약하는가?

주요 결과

• 특정 매개변수 설정 하에 곱셈 게이팅을 갖춘 게이팅된 RNN은 선형 자기주의를 정확히 구현할 수 있으며, 순환 상태가 $ \sum_t v_t k_t^\top $를 누적한다.
• 문맥 내 학습 작업에서 훈련된 RNN은 선형 자기주의와 동일한 주의 기반 알고리즘을 학습한다. 이는 경사하강법이 주의를 암묵적으로 발견함을 시사한다.
• 사이드 게이팅 RNN 아키텍처는 선형 자기주의를 성공적으로 모방하며, 학습된 가중치는 이론적 구성과 일치한다: $ W_x^{\text{in}} = W^{\text{side}} $, $ W_m^{\text{in}} = [0|B] $, $ D = B^\top $.
• RNN의 출력 행렬 $ M(x,y) $ 는 랭크 1이며, $ (i,j) $ 위치에 양의 피크와 대칭적인 음의 피크를 포함한다. 이는 원-핫 인코딩 덕분에 분류 성능에 영향을 주지 않는다.
• LSTM은 동일한 설정 하에 주의 구조를 성공적으로 실현하지 못하며, 이는 깊이 있는 선형 RNN에 비해 주의 구현에 대한 인덕티브 바이어스가 더 약하다는 것을 시사한다.
• 본 연구는 곱셈 상호작용이 RNN이 주의를 구현하는 데 필수적임을 드러내며, 이러한 메커니즘이 현대 RNN의 시퀀스 모델링 성공에 기여할 수 있음을 시사한다.