[논문 리뷰] Scan and Snap: Understanding Training Dynamics and Token Composition in 1-layer Transformer

Yuandong Tian, Yiping Wang외 2명|arXiv (Cornell University)|2023. 05. 25.

Neural Networks and Applications인용 수 6

한 줄 요약

이 논문은 다음 토큰 예측을 위한 1층 트랜스포머에서 SGD 학습 동역학을 수학적으로 엄밀하게 분석하며, 자기어텐션(self-attention)이 스팀스럽게 다른 토큰들에 초점을 맞추면서도 공통 토큰에는 주의를 줄 줄임으로써 분류적 스캐닝 메커니즘으로 작용함을 밝혀낸다. 핵심 발견은 학습 동역학에서 발생하는 단계 전이(phase transition)로, 이는 승자독점(winner-takes-all)로의 붕괴를 방지하여 안정적이고 희소한 어텐션 패턴을 만들어내며, 학습률 비율에 의해 제어 가능하다.

ABSTRACT

Transformer architecture has shown impressive performance in multiple research domains and has become the backbone of many neural network models. However, there is limited understanding on how it works. In particular, with a simple predictive loss, how the representation emerges from the gradient \emph{training dynamics} remains a mystery. In this paper, for 1-layer transformer with one self-attention layer plus one decoder layer, we analyze its SGD training dynamics for the task of next token prediction in a mathematically rigorous manner. We open the black box of the dynamic process of how the self-attention layer combines input tokens, and reveal the nature of underlying inductive bias. More specifically, with the assumption (a) no positional encoding, (b) long input sequence, and (c) the decoder layer learns faster than the self-attention layer, we prove that self-attention acts as a \emph{discriminative scanning algorithm}: starting from uniform attention, it gradually attends more to distinct key tokens for a specific next token to be predicted, and pays less attention to common key tokens that occur across different next tokens. Among distinct tokens, it progressively drops attention weights, following the order of low to high co-occurrence between the key and the query token in the training set. Interestingly, this procedure does not lead to winner-takes-all, but decelerates due to a \emph{phase transition} that is controllable by the learning rates of the two layers, leaving (almost) fixed token combination. We verify this extbf{\emph{scan and snap}} dynamics on synthetic and real-world data (WikiText).

연구 동기 및 목표

• SGD로 학습하는 동안 트랜스포머의 자기어텐션에서 암묵적인 유도 편향(inductive bias)을 이해하기 위해.
• 특히 위치 인코딩이 없는 상황에서 SGD 하에 어텐션 가중치의 동적 변화를 수식적으로 기술하기 위해.
• 자기어텐션은 증가하는 희소성에도 불구하고 승자독점으로 붕괴되지 않는 이유를 설명하기 위해.
• 디코더 레이어와 자기어텐션 레이어 간의 학습률 비율이 단계 전이의 발생 시점과 어텐션 패턴의 조밀성에 어떻게 영향을 주는지 기술하기 위해.
• 합성 데이터와 실제 데이터(WikiText)를 통해 이론적 예측을 검증하기 위해.

제안 방법

• SGD 최적화 하에 다음 토큰 예측을 위한 1층 트랜스포머를 분석하며, 자기어텐션 레이어 다음에 디코더 레이어를 배치함.
• 매개변수의 진화를 모델링하기 위해 비선형 미분방정식(ODE)을 사용한 연속시간 근사 모델을 도입함.
• 자기어텐션과 디코더의 동역학을 분리하기 위해 어텐션 가중치의 재매개변수화(reparameterization)를 통해 행렬 $Y$와 $Z$를 도입함.
• 장기간 시퀀스 및 느린 어텐션 조건 하에서 어텐션 로짓과 소프트맥스 출력의 해석적 해를 유도함.
• 장기간 시퀀스에서 토큰 공존 빈도의 경험적 추정치를 유한하게 제한하기 위해 허프딩 부등식(Hoeffding's inequality)을 활용함.
• 합성 데이터와 WikiText에서의 실험을 통해 어텐션 엔트로피와 Recall-Threshold 곡선을 측정함으로써 결과를 검증함.

Figure 1: Overall of our setting. (a) A sequence with contextual tokens $\{x_{1},\ldots,x_{T-1}\}$ and query token $x_{T}$ is fed into 1-layer transformer (self-attention, normalization and decoding) to predict the next token $x_{T+1}$ . (b) The definition of sequence classes (Sec. 3.2 ). A sequence

실험 결과

연구 질문

• RQ1 1층 트랜스포머에서 위치 인코딩이 없는 조건에서 자기어텐션은 학습 도중 어떻게 동적으로 진화하는가?
• RQ2 자기어텐션은 어떤 유도 편향을 학습하며, 이는 토큰의 공존 빈도와 고유성과 어떻게 관련되어 있는가?
• RQ3 증가하는 희소성에도 불구하고 자기어텐션은 왜 승자독점 패턴으로 붕괴되지 않는가?
• RQ4 디코더와 자기어텐션 레이어 간의 학습률 비율은 어텐션 패턴의 희소성과 안정성에 어떻게 영향을 주는가?
• RQ5 실제 시퀀스에서 '스캔하고 고정하기(scanf and snap)'의 이론적 동역학을 경험적으로 검증할 수 있는가?

주요 결과

• 자기어텐션은 분류적 스캐닝 행동을 보이며, 쿼리 토큰과 자주 함께 나타나는 고유한 토큰들에 점차로 초점을 맞춤.
• 어텐션 메커니즘은 빈도 편향을 보이며, 훈련 데이터에서 쿼리 토큰과 자주 공존하는 토큰을 선호함.
• 증가하는 희소성에도 불구하고, 기울기 크기의 로그 감소를 유도하는 학습 동역학의 단계 전이로 인해 어텐션은 일항형(one-hot)으로 붕괴되지 않음.
• 디코더 레이어의 학습률이 클수록 단계 전이가 가속화되어 더 조밀한 어텐션 패턴이 되며, 자기어텐션 레이어의 학습률이 클수록 더 희소한 패턴이 됨.
• WikiText에서의 경험적 결과는 높은 학습률이 어텐션 맵 엔트로피를 감소시키고 Recall-Threshold 곡선의 하강 속도를 늦추며, 이는 더 희소한 어텐션을 의미함.
• 학습률 비율에 따라 변하는 어텐션 희소성과 공존 빈도 의존성에 대한 이론적 예측이 Recall-Threshold 곡선 분석을 통해 실증적으로 검증됨.

Figure 2: Overview of the training dynamics of self-attention map. Here $\tilde{c}_{l|m,n}:=\mathbb{P}(l|m,n)\exp(z_{ml})$ is the un-normalized attention score (Eqn. 5 ). (a) Initialization stage. $z_{ml}(0)=0$ and $\tilde{c}_{l|m,n}=\mathbb{P}(l|m,n)$ . Distinct tokens (Sec. 3.2 ) shown in blue, co