Review

TL; DR

Summary

• Interpreting the Repeated Token Phenomenon in Large Language Models, ICLM’25 | Link

• Author

• Citation: 3

Introduction

Preliminaries

• Beginning of Sequence (BoS) 토큰

  • 정의: 문장에서 항상 첫 번째 위치에 있는 토큰
    • e.g., "Once upon a time” → BoS: Once

• Attention Sinks

  Efficient Streaming Language Models with Attention Sinks, ICLR’24

  • 정의: 모델이 한 토큰에 많은 attention을 주는 현상
    • attention 값들은 항상 합이 1이 되어야 하는데, 마땅히 줄 곳이 없으면 특정 토큰에 몰아줌(남는 attention을 버리는 통 ‘sink’)

      → 그 대상이 항상 접근 가능한 초기 토큰인것

    • semantic 때문이 아닌 structural 역할 때문
      • "\n"으로 바꿔도 동일한 패턴 발생

  • Fig.

    • Llama-2-7B에서 256개 문장(각 16토큰)에 대해, 평균 attention logit을 시각화한 heatmap
      • x축: key token 위치
      • y축: query token 위치
        • y = 10, x = 0 이 붉다면 → 10번째 토큰이 0번째 토큰을 매우 강하게 참고한다는 뜻

    

    • Layer2를 넘어갈수록 모든 query 토큰이 초기 토큰(token 0) 에 매우 큰 attention을 줌

  • BoS sink:
    • 정의: BoS (Beginning of Sequence) token이 attention sink 역할을 하는 현상(attention을 버리는 통(sink) 으로 BoS 토큰을 사용한다~)
    • 효과
      • 문장 구조의 기준점 제공 및 문맥 해석의 anchor 역할 수행
      • BoS 토큰은 대표적인 정상 attention sink임
      • attention이 완전히 흩어지지 않도록 중심을 잡아줌(문장의 기준점을 만들어 줌)
      • 모델의 fluency를 높여줌

• Attention Sinks & BoS 간의 관계

  • 모델은 학습 중에 ‘첫 토큰은 문맥의 기준점’으로 배우기 때문에, BoS 토큰은 항상 attention sink 역할을 하도록 학습됨

  • 하지만, 반복 토큰이나 특정 패턴 토큰이 BoS처럼 취급될 때 문제가 발생함. (attnetion sink 가 ‘여러개’인 경우 문제 발생)
    • BoS sink가 여러 개 생기고
    • 문장이 여러 번 새로 시작되는 것처럼 인식되어 모델이 문맥을 잃고 발산함

Background

Repeated Token Divergence Phenomenon

• LLM은 다양한 자연어 태스크에서 뛰어난 성능을 보이고 있음

• 하지만, ‘하나의 단어를 반복하라’라는 단순한 지시를 제대로 수행하지 못하는 경우가 있음

  ⇒ 이를 “repeated token divergence” phenomenon 라고 함

  • e.g.,

    Scalable Extraction of Training Data from Aligned, Production Language Models, ICLR'25

    

    • GPT-3.5-turbo 로 실험 했을 때, 단일 token을 반복하다가 다른 무관한 text를 출력하게 됨.

    • 이렇게 출력되는 text는 모델이 학습 과정에서 언젠간 한번 보았던 데이터라고 함

• ‘반복 토큰 현상’을 분석하는 기존 연구들:
  • 반복 토큰이 첫 토큰(BoS)표현으로 수렴함을 직관적으로 분석 및 관찰
  • 반복 토큰이 이론적으로 문제가 될 수 있다는 분석

  ⇒ 왜 이런 현상이 일어나는지에 대한 구체적인 분석은 충분하지 않았음

In this Paper…

• ‘Repeated token divergence’ 현상을 ‘attention sink’ 현상과 연관이 있다고 보고, 해당 매커니즘과 연결해서 설명하고자 함
  • attention sink: 문장의 첫 토큰(BoS)이 비정상적으로 높은 attention을 받는 현상

  → “반복 토큰 현상은 attention sink를 만드는 동일한 신경 회로가 오작동하면서 발생한다.”

• attention sink를 만드는 신경 회로를 분석하고, 그 회로가 반복 토큰에서 어떻게 오작동하는지 분석하며 그 결과 모델이 왜 붕괴되는지를 설명하고자 함

  ⇒ LLM의 유창성을 가능하게 하는 attention sink 메커니즘이, 동시에 repeat token 취약점을 만들어내는 구조적 원인이다!

Contribution

• Repeated Token Divergence Phenomenon 을 매커니즘 수준에서 처음으로 설명
  • 기존 연구들이 ‘현상’ 만 관찰했다면, 해당 논문은 white-box 세팅에서 분석

• Attention sink의 원인이 되는 sink neuron을 식별하고 인과적으로 검증
  • sink layer에서 norm을 만드는 소수 뉴런을 선별하여 해당 뉴런을 ablation 했을 때 현상이 사라짐을 보임

• 모델이 첫 토큰을 인식하는 first-token detector neuron 구조를 규명
  • 첫 attention layer 이후, first token /non-first token이 선형 분리됨을 증명함

    → LLM이 위치 정보를 명시적 뉴런 하나에 인코딩을 한다!

• Repeated token Attack에 대한 방법 및 mitigation 방안 제안

Mechanistic Analysis of Repeated Token Divergence

• Repeated token divergence현상을 신경 회로 메커니즘 관점으로 설명하고자 함

• Flow
  1. 반복 token의 attention이 BoS의 attention (attention sink 현상)과 유사함을 관찰
  2. Attention sink를 만드는 mechanism 을 파악
  3. 이 mechanism이 repeated token에서 어떻게 재현되어 divergence을 유발하는지 파악

• LLaMa-2 를 활용하여 실험 진행

Large Attention Scores of Repeated Tokens

Q: repeated token이 나타내는 attention score과 attention sink attention score 양상이 유사한가?

• Setting

  • Top panel: 일반문장

  • Bottom panel: ‘the …’를 반복했을 때

  • x축: key token 위치

  • y축: query token 위치

• Bottom panel (일반 문장): 첫 토큰이 가장 높은 attention을 받음 → attention sink

• Top panel (the .. repeat): 여러 위치의 "the" 토큰들이 첫 토큰과 거의 동일한 attention score를 받음

  → 반복되는 토큰들이 문장의 시작 토큰(BoS) 처럼 취급되고 있다는 의미

⇒ 반복된 "the" 토큰들의 attention 점수 분포는 일반 문장에서 첫 토큰이 받는 attention분포과 유사함!

The Attention-Sink Mechanism

Q: 왜 repeated token이 높은 attention score을 받게 되는가?

A: Token이 repeat 될 수록 hidden state norm이 커지는데, norm이 클수록 자동으로 높은 attention을 받게 되기 때문임.

• Transformer는 학습 과정에서 문장 시작 위치(BoS)에 강한 구조적 의미를 부여하도록 학습되는데, 이 과정에서 hidden norm이 항상 크게 형성됨
  • Hidden state norm이 크다 → 해당 토큰 표현이 layer에서 강하게 활성화 되어있다

• 이 Hidden state norm은 attention score에도 영향을 줌

  

  → hidden state norm이 크면, 그 토큰은 자동으로 높은 attention을 받게 되는 구조임

• 따라서, 특정 token을 반복했을 때 해당 token이 BoS token과 같이 hidden state norm양상이 비슷한지 알아보고자 함
  • 논문에서는 attention sink 현상이 처음 강하게 형성되기 시작하는 레이어를 sink layer으로 정의하고, 해당 layer(=1)에서 실험을 진행

• Setting

  • Sink Layer: 1
    • sink layer: attention sink 현상이 처음 강하게 형성되기 시작하는 레이어

• 모든 토큰에서 공통적으로 반복 횟수가 증가할수록 hidden state norm이 지속적으로 증가하며 결국 BoS token norm 수준에 수렴
  • 세 토큰(the, one, es) 모두 속도 차이는 있지만, 방향은 동일함

⇒ Repeated token이 모델 내부 표현 공간에서 BoS token과 거의 동일한 상태로 변해감

Q: 어떤 neuron들이 attention sink 현상을 유발하는가?

• 이러한 현상을 야기하는 특정 neuron이 있다고 가정하고, 이를 ‘sink neurons’으로 정의

• sink neuron 정의 과정

  1. BoS 토큰을 넣었을 때, MLP 출력이 residual stream norm에 가장 크게 기여하는 뉴런들을 top-K로 우선 선별

     

  2. 후보 뉴런들을 하나씩 ablation 진행
     • 이 뉴런을 제거했을 때 반복 토큰 norm이 의미 있게 줄어드는가? 를 확인

  → 1, 2를 모두 만족시키는 neuron을 sink neuron으로 확정

⇒ 반복 토큰의 norm 폭증은 sink neuron이 없으면 발생하지 않음!

• 모델별로 sink neuron IDs

  

Q: 앞서 반복 토큰은 결국 BoS(first token)처럼 취급된다고 하였는데, 그렇다면 모델은 첫 토큰을 언제 인식하는가?
A: 첫 번째 Attention Layer에서 시퀀스에서 첫 토큰과 그 이후 토큰을 구분함

• 첫 attention layer를 지난뒤의 representation을 1차원 축으로 확인하여 분리 정도를 확인하고자 함

• 첫 토큰과 이후 토큰의 분포가 뚜렷하게 분리 되어 있음 → 첫 번째 토큰과 그 외의 토큰 간의 ‘선형 분리’가 일어남

  → 첫 attention layer가 첫 토큰을 식별하고 내부적으로 ‘표식(marking)’ 하는 데 핵심 역할

• 추가적으로 단일 뉴런 하나만으로도 이 분리가 완벽히 된다는 것을 발견
  • LLaMA2: MLP0 gate neuron 912 → Layer 0 MLP의 gate projection에서 912번째 차원 뉴런 하나

Attack and Mitigation

Attack

• 반복 토큰은 ‘취약점’으로 악용될 수 있음
  • 기존 연구에서, 반복 토큰 입력이 모델을 혼란시키고 그 결과 훈련 데이터 유출(training data leakage) 공격에 악용될 수 있다고 함
  • 또한 지시 따르기에서 이탈하게 만들고, 외워둔(memorized) 훈련 데이터를 노출시키는 사례가 있음
    • e.g., Pythia-12B에 ‘as’를 50번 반복하면 모델이 3D 프린팅 설명 같은 텍스트를 출력하는데, 그 출력은 실제 웹사이트에 있는 문장(훈련 데이터에 포함된 것으로 추정되는 텍스트)을 재진술 한 것으로 확인됨
  • Attack 예(Before Patching)

    

• 반복 토큰이 길게 나오면 감지해서 차단하는 방식(Surface-level mitigation)이 있지만, 이런 방식은 근본 원인을 해결하지 못함

  → repeted token이 아니더라도 모델 공격이 가능함

  • Attack detail

    • attention head의 투영 공간을 2개(첫 번째 토큰/ 나머지)가 아니라 여러 개의 자연스러운 그룹으로 분리하여 cluster함

    • 같은 cluster 토큰들을 섞어 넣으면 BoS처럼 취급되는 표현이 생기고 모델이 발산(diverge) 됨

      → 동일한 token 을 repeat하지 않아도 유사한 token들을 모아서 모델을 우회하여 공격이 가능함

Mitigation

• <Fig3. 참고> sink를 유발하는 뉴런의 출력(activation)을 강제로 ‘no-sink’ 상태로 고정하면 repeat token 공격을 막을 수 있음
  • Fig 3

    

  • Mitigation (After Patching)

    

• LLaMA2에서 repeat prompt를 줬을 때 해당 공격이 더 이상 모델을 발산시키지 못함을 보여줌

• <Table 2> 추가적으로, 이 패치가 모델의 기본 능력을 망가뜨리지 않는지 확인 → 패치를 해도, 성능에 미치는 영향은 없었음