Review

TL; DR

Summary

Motivation

• Transformers를 수학 혹은 알고리즘 태스크에 훈련할 때 보이는 abrupt learning (갑작스러운 학습) 현상
  • : 모델의 성능이 오랫동안 정체되었다가 갑자기 급격하게 향상되는 현상

• 본 논문은 훈련 시 이러한 현상의 보편적인 특성과 기본적 메커니즘을 밝히고자 함

Contribution

• 소형 Transformers를 간단한 알고리즘 태스크로 훈련하여 abrupt learning과 관련된 여러 현상 탐구하고자 함

  • 사용 태스크: moving-window-sum (MWS)

    

    • $x_1, ..., x_n$ 시퀀스가 주어지면, $SEP$ 이후 $y_1, ..., y_n$ 을 출력해야 함
      • $x_i$는 0, 1, 2, …, 17 중 하나의 숫자
    • $y_1$은 $x_1$ 그대로, $y_2$는 $(x_1 + x_2)$를 $p$(=17)로 나눈 나머지, $y_3$은 $(x_2+x_3)\mod p$, …
    • ⇒ ground truth가 잘 알려져 있는 태스크로서 모델의 훈련 진행 정도를 정확히 측정할 수 있음

  • 모델 아키텍쳐: 1-layer, 1-head Transformer
    • 이 구조로도 주어진 태스크 완벽히 수행 가능

      

      • $(s_1, ..., s_L)$ : 입력 토큰 시퀀스
      • $MLP$ : 2-layer NN
      • $Id$ : residual connection
      • $LM$ : linear layer, mapping hidden state to logits
    • greedy decoding 사용
    • ⇒ 작은 모델을 사용하여 모델 내부 메커니즘을 쉽게 분석하고 해석할 수 있음

  • 훈련
    • $(x_1, ..., x_n, SEP, y_1, ..., y_n)$의 전체 시퀀스에 대해 next-token-prediction cross-entropy loss 손실 최소화하도록 훈련
      • 256개 훈련 샘플로 1 에폭 훈련
    • 정확도 측정: output 부분인 $y_1, ..., y_n$의 $n$개 토큰에 대한 예측 정확도 평균

  • abrupt learning
    • 훈련 중 training loss가 sub-optimal 값에서 상당히 많은 step동안 유지되다가 급격한 정확도 증가와 손실 하락
    • ⇒ optimal solution을 갑자기 학습하는 abrupt learning 현상 확인

    

  • attention map
    • 실험 태스크에 대해서 최적의 attention pattern은 각 output token $y_i$가 자신을 계산하는데 관련있는 input token만을 attend하는 것
      • $y_1$은 $x_1$에 attend, $y_2$는 $x_1, x_2$에 attend, …
    • attention progress measure (APM) 사용해 attention pattern 변화 관찰

      

      • $A_{ij}$ : $i$-th output token 계산할 때 $j$-th token에 할당되는 attention score
      • $\Omega$ : 최적 attention map의 position pair set
    • → 훈련 중 APM이 0부터 0.8 정도까지 단조 증가하며, 손실/정확도보다 완만한 변화 곡선 보임
      • step 150 쯤에서 급격한 손실 하락 있는데, 그 이전부터 APM은 이미 상당히 증가함
      • ⇒ 손실 하락은 급격하지만 그 이전부터 attention pattern learning은 점진적으로 진행됨

      

• Transformers 훈련의 초기 정체기 (early loss plateau period) 동안 모델은 종종 partial solution을 학습함

  • e.g., moving-window-sum 태스크에서 첫번째 입력 토큰 $x_1$를 그대로 출력하면 되는 $y_1$ 예측은 빠르게 학습하지만, 전체 손실은 여전히 높으며 이후 토큰에 대한 정확도 떨어짐

    

    • 첫번째 출력 토큰 예측 정확도인 partial solution accuracy가 빠르게 증가
    • 반면 전체 loss는 많은 훈련 스텝 이후 하락
    • ⇒ 초기 정체기 동안 전체 loss는 크게 줄지 않지만 모델의 partial solution 학습 진행됨

• 정체기동안 모델이 반복적 토큰을 출력하는 경향인 repetition bias가 강하게 나타남

  • repetition frequency: repetition bias 정량화 지표

    

  • 결과

    

    • repetition frequency가 훈련 시작 시 작았다가 처음 50 스텝동안 0.8까지 상승
    • ⇒ 초기 정체기의 강한 repetition bias 확인

• output repetition bias는 다른 토큰에 대한 hidden representation이 거의 동등하게 되는 representation collapse를 동반함

  • 출력 위치 $i, j$에서 hidden representation 간 pairwise cosine similarity:

    

    • $h_i$ : $i$th 토큰의 hidden representation (logit 변환 직전)

  • cosine similarity가 훈련 초기 단계에서 급격히 증가

    

    • ⇒ partial solution에서 정확히 예측되는 첫번째 출력 위치를 제외하고는, 훈련 초기 단계에서 여러 출력 위치의 hidden representation이 거의 동등해짐

• attention map learning이 repetition, representation collapse과 loss plateau 형성에도 중요한 역할 함을 보임

  • 최적의 어텐션 맵을 향해(또는 반대로) 편향시켜 반복, 표현 붕괴, 손실 정체가 감소(또는 증폭)되는지 확인

  • $(i,j) \in \Omega$ 에 대해 attention mask $M_{i,j}$를 $c$로 설정: 나머지 경우 $M_{i,j}=1$
    • 기존 attention을 이 attention mask와 hadamard 곱 함으로써 변형하여 훈련과 추론에 사용
    • $c > 1$ 이면 optimal attention map 방향으로 편향시키고, $0<c<1$ 이면 반대 방향으로 편향시키는 것

  • $c > 1$ : hidden state 간 평균 코사인 유사도와 repetition의 감소, 보다 빨리 수렴함

    

  • $0 < c < 1$ : representation collapse 상태에 모델이 보다 오래 머물고 보다 이후에 수렴, repetition frequency 또한 정체기동안 크게 유지됨

    

  • ⇒ attention map learning이 repetition, representation collapse, loss plateau 형성에 중요한 역할 함

• 소형 Transformer만이 아니라 실제 LLM의 사전훈련 초기 단계에서도 이러한 repetition bias, representation collapse가 나타나는지 확인

  • LLM: Pythia, OLMo-2 (open-source)

  • task: ARC-Easy 테스트 데이터에서 랜덤 선정한 100개 질문 (초등학교 수준 과학 객관식 질문)
    • 각 질문에 대해 8개 토큰을 생성하게 하고 hidden representation의 pair-wise cosine similarity 계산

  • 14M, 1B, 1.4B, 2.8B Pythia 모델의 초기 훈련 단계에서 출력 시퀀스의 repetition bias 발견
    • 초기화 시 평균 코사인 유사도가 비교적 낮으나 (0.4~0.65), 모든 사이즈의 모델에 대해 몇 스텝만 학습해도 0.9 이상으로 급격히 증가

    

  • OLMo-2 7B 모델에서도 Pythia와 유사한 representation collapse 현상 관찰
    • 150스텝의 초기 훈련 단계에서 representation 평균 코사인 유사도는 0.93
    • 600스텝에서는 0.43으로 감소

  • ⇒ repetition bias, representation collapse가 LLM 초기 사전훈련 단계에서 실제로 발생하는 현상임

Conclusion

• Transformer 훈련의 초기단계에서 repetition bias와 representation collapse가 발생하며, 이는 loss plateau와 밀접한 관련이 있음
  • loss plateau가 최적의 attention map 탐색하는 과정일 가능성 있음

• 본 논문의 발견에 이어 representation collapse와 같은 현상이나, attention map의 느린 learning의 원인에 대한 향후 연구가 유망