Review

• Link: https://arxiv.org/abs/2504.03022

• Cited: 14

• Github: https://dualroute.baulab.info/

• Related paper: In-context Learning and Induction Heads (Enthropic’22) | Cited: 899
  • TL;DR ICL 능력의 핵심 메커니즘은 induction head 때문이고 2 layers로 작동한다!

TL;DR

Preliminary

• In-Context Copying

  • 정의 prompt 안에 있는 패턴을 보고 그대로(or 의미 기반으로) 따라 생성하는 태스크
    • 새로 학습한 게 아닌, prompt 안에 있는 패턴을 보고 즉석에서 복사함!

      

  • e.g.,

    Input:
    A B C → D E F
    G H I → ???

    • In-Context Copying output: J K L

  • 활용 분야
    • 의미적 복사(concept copying)
      • 번역 패턴 복사

        Input:
        cat → 고양이
        dog → ???
        Output:개

• Induction Head

  • 정의 앞에서 본 패턴을 찾아 그 다음 토큰을 이어서 생성하게 만드는 attention head
    • 일반 attention head 와의 차이점
      • attention head: 관련 있는 토큰을 참고
      • induction head: 과거에 등장한 동일한 토큰을 찾아, 그 뒤에 이어졌던 토큰을 현재 위치에서 예측하도록 만드는 attention 패턴(특화된 Attention head)

  • 작동원리

    [A] [B] ... (다른 텍스트들) ... [A]

    1. 모델이 현재 [A]를 읽고 있을 때, Induction Head 가 과거의 문맥에서 첫 번째 [A]를 찾아냄
    2. 과거의 [A] 바로 다음에 왔던 토큰이 [B] 였다는 사실을 Attention 함
    3. 결과: 다음 출력으로 [B]를 예측

    ⇒ 현재 토큰 → 과거 동일 토큰 찾기 → 그 다음 토큰 가져오기

  • ⭐ 2-layer transformer 구조임!!
    • 단일 Attention Head로만 작동하는 것이 아닌 두 개의 Attention Head Layer로 작동함
      • next-token 정보를 hidden state에 저장해두고 그걸 꺼내는 구조
      1. 1st Layer
         • 이전 토큰 정보 저장
         • 현재 토큰이 바로 이전 토큰이 무엇인지 정보를 가져와 기억
      2. 뒷 층 헤드(실제 Induction Head):
         • 현재 시점 토큰 [A]가 등장했을 때, 과거의 문맥을 뒤져서 똑같은 토큰 [A]를 찾고, 과거 [A]가 가지고 있던 바로 다음 토큰 [B]에 강한 어텐션(가중치) 부여

• Activation Patching

  • 모델 내부의 특정 활성화값(activation)을 조작해서 어떤 부분이 특정 출력에 얼마나 영향을 주는지 분석하는 기법
    • 특정 뉴런/레이어가 특정 역할을 했는지 테스트하는 방법
    • 단순 상관관계가 아니라, 해당 head가 실제로 출력에 인과적 영향(causal effect)을 주는지 확인하는 방법
    • 특정 attention head가 어떤 역할을 하는지 분석할 때 사용됨

  • 작동원리
    • 모델에게 줄 2개의 입력이 있다고 가정
      1. 정상 입력 (clean input) → 원하는 정답을 잘 맞춤
      2. 오염된 입력 (corrupted input) → 일부 정보가 망가져서 정답을 못 맞춤
    • 이때 특정 레이어나 뉴런의 activation을 오염된 입력의 실행 중간에 정상 입력에서 나온 activation으로 덮어씌우고(patch) 출력이 다시 좋아지는지 확인
      • 성능이 회복되면 → 해당 activation의 attention head가 핵심 정보를 담고 있구나!!
      • 성능 회복 안되면 → 다른 head를 알아보자..

      

Introduction

Background & Motivation

• 뇌 과학자들은 사람의 뇌가 무언가를 읽을 때 2가지 병렬 경로로 작동한다고 함
  1. 글자를 하나씩 해독하는 경로 sublexical(token)
  2. 단어의 의미를 전체 단위로 접근하는 경로 lexical

• Token 기반 In-Context Copying의 문제점
  • 사람에게 아래 두 문장을 외우라고 해보자!
    1. oane dnn t ephzawfeew eausr lthii → 뜻이 없기에 알파벳 별로 외우게 됨 token-based
    2. the false azure in the windowpane → 문장의 의미를 외우게 됨 concept-based

• 기존 induction head의 문제점
  • 앞에서 본 패턴을 보고 다음 토큰 복사함 ⇒ 토큰 단위 복사임
  • but, 번역 같은 느슨한(fuzzy) 복사 작업은 토큰 단위로 외우지 않고 concept 단위로 외움
    • 번역을 하면 토큰 길이가 완전 달라지기도 하기 때문에 토큰 단위로는 설명이 부족함!
      • e.g., potato (영어) → pommes de terre (프랑스어)

• RQ 모델도 사람의 뇌처럼 토큰 단위 복사와 의미 단위 복사를 같이 할 수는 없을까?

Contribution

• token induction head 이외에 concept induction head도 존재함
  • Token Induction: 토큰 단위 복사
  • Concept Induction: 의미 단위 복사

    복사 = token head + concept head

    

Concept Induction Head를 찾아보자

• 목표 어떤 head가 token/concept 단위의 copying을 하는지 구분해보자
  • concept head를 찾기 위해, multi-token 단어를 복사하는 attention head를 causal intervention을 통해 찾음
  • 특정 head가 multi-token 개념의 미래 토큰 확률을 증가시킨다면,
  • ⇒ 해당 head가 전체 개념을 복사한다고 가정(패턴이 없는데도 복원이 되면 Concept-head 때문)
  • 실험 과정(중요)

    1. 기본 induction prompt 만들기(기존 Induction Head 실험 셋업과 같음)

       x1 x2 ... xn c1...cm | x1' x2' ... xn' c1'

       • 앞부분과 뒷부분이 거의 같은 반복 구조( |는 newline)
       • 끝에 multi-token 개념 c1...cm을 붙임
       • 마지막 토큰은 항상 c1' (c1과 동일한 토큰)
       • e.g.,

         random_tokens... "Paul Cham bers" | random_tokens... "Paul"
                                                                  ↑
         		                                              마지막 토큰 c1'

       • 모델 입장에서는 Paul 다음에 Cham이 와야 한다는 걸 앞 문맥에서 유추해야 함

    2. 두 개의 프롬프트 준비
       • $p_{clean}$ (Step1과 동일한 프롬프트):

       x1...xn c1...cm | x1'...xn' c1'   ← 앞뒤가 일치

       • $p_{corrupt}$ (오염된 프롬프트):

       y1...yn+m      | x1'...xn' c1'   ← 앞부분을 전혀 다른 랜덤 토큰으로 교체

       • $p_{corrupt}$에서는 앞부분이 망가졌기 때문에, 모델이 c2를 예측할 단서가 없음

    3. Activation Patching: $p_{clean}$에서 특정 head의 activation을 뽑아서, $p_{corrupt}$의 같은 위치에 이식(patch)하고 실제로 예측이 되는지 실험(Preliminary 참고..)
       • 만약 특정 attention head $a^{(l,h)}$ $(layer\ l, head\ h)$의 activation을 이식했더니 c2 확률이 올라간다면 → 이 head가 개념 복사를 담당한다는 증거가 됨..!

         

    4. 2가지 Score로 평가
       • (1) Concept Copying Score (개념 단위 복사):

         $$\text{CONCEPTCOPYING}(l,h) = P(c_2 \mid \text{patch}) - P(c_2 \mid p_{\text{corrupt}})$$

         • c2는 개념의 두 번째 토큰 (e.g., " Cham")
         • 이미 c1'("Paul")은 입력에 있으니, c2를 맞추려면 개념 전체를 기억해야 함
         • Concept Copying Score가 높으면 ⇒ 개념 수준 복사 head
       • (2) Token Copying Score (토큰 단위 복사):

         $$\text{TOKENCOPYING}(l,h) = P(r_1 \mid \text{patch}) - P(r_1 \mid p_{\text{corrupt}})$$

         • 개념 대신 랜덤 토큰 사용
         • r1 (바로 다음 토큰) 예측에 집중
         • Token Copying Score가 높으면 ⇒ 토큰 수준 복사 head (전통적 induction head)
  • 실험 결과

    • 각 모델의 모든 attention head에 대해 causal score를 계산

      

      • Head들의 위치 분포

        

        • Concept Copier Heads (🔶)
          • 초중반 레이어에 집중되어 있음
          • 의미 정보는 이후 레이어에서 활용되어야 하므로, 초반 레이어에서 미리 형성되어야 함
        • Token Copier Heads (🔵)
          • 산발적으로 분포, 특히 후반 레이어에 많음
          • 토큰 하나를 복사하는 건 단순한 작업이라 늦게 처리해도 됨
        • 두 헤드의 위치 분포가 거의 겹치지 않음
          • ⇒ 아 각각 다른 역할이구나(둘 다 필요하구나!)

Next-Token & Last-Token Attention Scores

• 목표 Concept copier head는 도대체 어디에 attend하는가?
  • 기존 token induction head는 항상 바로 다음 토큰을 attend 함
  • RQ concept copier head는 개념 전체를 한번에 전달한다면 개념 정보가 저장된 마지막 토큰에 attend하지 않을까?
  • 두 가지 Score 설계
    1. Last-Token Matching Score (concept head용)

       # 프롬프트 구조
       x1...xn c1...cm | x1'...xn'
                                   ↑
                            여기서 attention 측정

       $$\text{LASTTOKENMATCHING}(l,h) = \frac{1}{|C|} \sum_{c \in C} A^{(l,h)}[x'_n, c_m]$$

       • 단어의 마지막 토큰을 바라보는 정도
       • x'n 위치에서 개념의 마지막 토큰 cm에 얼마나 attend하는지 측정
       • 개념은 COUNTERFACT 데이터셋에서 샘플링 (길이 2~5토큰)
    2. Next-Token Matching Score (token head용)

       $$\text{NEXTTOKENMATCHING}(l,h) = \frac{1}{|R|} \sum_{r \in R} A^{(l,h)}[x'_n, r_1]$$

       • 다음 토큰을 바라보는 정도
       • 개념 대신 랜덤 토큰 span 사용
       • 첫 번째 토큰 $r₁$에 얼마나 attend하는지 측정

• 실험 결과
  • 목표 Llama-2-7b의 각 head의 next-token matching score와 last-token matching score

  

  • x축: 각 attention head의 위치(layer 번호.head 번호)
  • 왼쪽 그래프(token copy head): NEXTTOKENMATCHING이 매우 높음
  • 오른쪽 그래프(concept copy head): LASTTOKENMATCHING이 매우 높음

Concept/Token Copier Head 제거 실험

• RQ Concept induction head와 Token induction head를 강제로 끄면 어떤 task에 영향을 주는가?
  • Concept head → fuzzy copying (의미 기반 복사) 에 중요
    • 번역, 동의어, 반의어 등 의미적으로 관련된 단어 복사
  • Token head → verbatim copying (그대로 복사) 에 중요
    • 완전히 동일한 단어를 그대로 옮기는 작업

• 실험 설계
  • Fuzzy copying Tasks (의미 기반 tasks)
    • Translations (5개 언어의 단어 → 영어 번역)
    • uppercasing (대문자 변환)
    • synonyms
    • antonyms
  • Baseline
    • Verbatim task (의미(semantic)가 없는 토큰 복사)
      • English copying: 영어 단어 목록을 그대로 반복
      • Nonsense copying: 랜덤 토큰을 그대로 반복
        • e.g., apple → apple, xyz → ?
  • Metric
    • First-token accuracy: 첫 토큰 맞추면 개념 이해했다고 판단
  • Ablation 방식
    • 특정 Head의 activation을 mean activation (무의미한 상태)으로 바꿈
    • 어떤 Head를 지우냐?
      1. concept head 제거: concept score가 높은 head들
      2. token head 제거: token score가 높은 head들

• 실험 결과

  

  • (b) ablating Token Induction Head
    • Nonsense copying 실패
    • 나머지 task는 concept head가 대신 수행 가능
    • → Token head는 의미 없는 패턴의 그대로 복사에 필요
  • (c) ablating Concept Induction Head
    • 번역, 동의어, 반의어 정확도 하락
    • 표면적 복사 task에는 영향 없음
    • → Concept head는 의미 기반 복사에만 필요
  • 두 ablation 모두 영향 없는 task
    • English copying: 의미로도, 패턴으로도 풀 수 있어서 어느 쪽을 꺼도 나머지 경로가 대체
    • Uppercasing: 마찬가지로 두 경로 모두로 해결 가능

  → Concept head를 끄면 의미 기반 복사가 망가지고, Token head를 끄면 verbatim 복사가 망가짐

  ⇒ 두 경로가 실제로 독립적으로 존재함

Concept Induction is Language-Agnostic

• Concept induction head의 의미적 복사는 언어와 무관한 추상적 의미 표현의 복사다!
  • e.g.,
    • 영어 waxwing
    • 러시아어 свиристель (여새)

    

  • 두 단어는 표면적으로 완전히 다르지만, 같은 개념을 가리킴

    ⇒ concept head가 의미를 다루면 두 경우 모두 동일한 activation을 가질 것!

• 실험 설계
  1. 두 개의 서로 다른 번역 프롬프트를 만듦
     	Source prompt (s)	Base prompt (b)
     번역 방향	스페인어 → 이탈리아어	일본어 → 중국어
     번역 단어	$w_s$	$w_b$

  2. patching 실험 진행
     • Source prompt에서 top-k concept head의 activation 추출
     • Base prompt의 같은 위치에 이식(patch)
     • 모델이 $w_b$ 대신 $w_s$의 의미를 출력하는가?
       • 원래 이탈리어로 번역할 스페인어 $w_s$를 중국어의 의미로 번역할건지?

• 실험 결과

  

  • (a) patching을 하니 source prompt의 $w_s$가 번역됨(figure (a) 참고)
  • (b)
    • 빨간선($w_s$ 출력): patch한 source 단어의 의미가 출력되는 정확도
    • 회색선($w_b$ 출력): 원래 base 단어의 의미가 출력되는 정확도
    • k가 증가할수록:
      • 회색선 하락 → 원래 답이 밀려남
      • 빨간선 상승 → source 의미가 대신 나타남
      • k=80 근처에서 효과가 가장 강력
        • ⇒ 80개 정도의 concept heads가 가장 효과적이구나