기술 요약: Head Pursuit: 멀티모달 트랜스포머의 주의 집중 전문화 탐색

문제 정의

대규모 생성 모델(LLM과 VLM을 포함한 단일 모드 언어 모델 및 멀티모달 시각-언어 모델 모두)은 다양한 작업에서 놀라운 성능을 달아가고 있습니다. 그러나 이들의 내부 메커니즘은 여전히 부분적으로만 이해되고 있습니다. Logit Lens와 같은 기존의 해석 가능성 도구들은 은닉 상태(hidden states)에서 의미론적 정보를 드러내 주었지만, 이는 종종 개별 예시에 대해 휴리스틱하게 적용되므로, 발견된 내용을 일반화하거나 특정 구성 요소(예: 어텐션 헤드)가 모델 출력 형성에 미치는 중요성을 정량화하기 어렵습니다. 따라서 특정 어텐션 헤드가 특정 의미적 또는 시각적 속성에 특화되어 있는지 식별하고, 추가 학습 없이 이러한 헤드를 조작함으로써 모델 행동을 목표에 맞게 제어할 수 있는지 결정할 수 있는 원칙적인 방법이 필요합니다.

방법론

저자들은 Head Pursuit라고 불리는 프레임워크를 제안하며, 이는 중간 활성화(intermediate activations)를 조사하는 관행을 희소 신호 처리(sparse signal processing)의 관점에서 재해석한 것입니다. 핵심 방법론은 다음 단계들을 포함합니다:

1. SOMP를 통한 희소 분해: 저자들은 고차원 신호를 기저 요소(basis elements)의 희소 선형 조합으로 근사하기 위한 탐욕 알고리즘인 **동시 직교 매칭 추적(Simultaneous Orthogonal Matching Pursuit, SOMP)**을 사용합니다.

   • 입력: 레이어 $l$에 있는 주어진 어텐션 헤드 $h$에 대해, $n$개의 샘플을 나타내는 활성화 행렬 $H_{h,l}$을 분석합니다.
   • 사전(Dictionary): 활성화로부터 사전을 학습하는 대신, 이 방법은 모델의 언임베딩 행렬(unembedding matrix) $D$를 고정된 사전으로 사용합니다. $D$의 각 행은 어휘 사전의 토큰에 대응하며, 이는 인간이 해석 가능한 의미론적 개념과 직접 연결됩니다.
   • 과정: SOMP는 데이터셋 전체의 모든 샘플에 걸쳐 헤드 잔차(head residuals)와 최대 상관관계를 갖는 사전 원자(dictionary atoms, 즉 토큰)를 반복적으로 선택합니다. 이를 통해 $H \approx W^*D$를 만족하는 희소 계수 행렬 $W^*$를 산출합니다.
   • 일반화: 단일 잔차 벡터를 투영하는 Logit Lens와 달리, SOMP는 여러 샘플을 동시에 처리하며 여러 개의 사전 방향을 선택함으로써, 헤드의 기능적 역할을 데이터셋 수준에서 견고하게 특징짓습니다.

2. 헤드 선택 및 순위 지정:

   • 특정 타겟 개념(예: "색상", "독성", "숫자")과 관련된 헤드를 식별하기 위해, 언임베딩 사전은 해당 개념과 연관된 토큰들로 제한됩니다.
   • SOMP는 이 개념 특화 사전에 적용됩니다.
   • 헤드들은 SOMP 재구성에 의해 설명되는 활성화의 **분산 비율(fraction of variance)**을 기준으로 순위가 매겨집니다. 높은 분산 설명력은 해당 개념에 대한 강한 전문화를 나타냅니다.

3. 개입(Intervention):

   • 상위 $k$개의 특화된 헤드가 식별되면, 저자들은 잔류 스트림(residual stream)에 대한 이들의 기여도를 재조정(rescaling)하여 개입합니다.
   • 억제 개입(Inhibitory Intervention): 타겟 속성을 억제하기 위해 헤드 표현의 부호를 반전시킵니다 ($\alpha = -1$).
   • 강화 개입(Enhancing Intervention): 타겟 속성의 생성을 촉진하기 위해 헤드 표현을 증폭시킵니다 ($\alpha > 1$).
   • 이 개입은 모델 가중치를 업데이트하지 않고 순전파(forward pass) 중에 수행됩니다.

주요 기여

1. 이론적 프레이밍: 본 논문은 매칭 추적(Matching Pursuit, MP)과 그 다중 샘플 변형인 SOMP를 생성 모델 해석 가능성의 맥맥락 내에서 프레이밍하여, 여러 샘 샘플과 여러 의미적 방향을 동시에 처리하도록 일반화함으로써 Logit Lens와 같은 표준 도구들과의 공식적인 연결을 확립했습니다.
2. 전문화의 발견: 이 방법은 단일 모드(Mistral-7B) 및 멀티모달(LLaVA-NeXT, Gemma3, Qwen2.5-VL) 트랜스포머 전반에 걸친 일관된 헤드 전문화 패턴을 밝혀냈습니다. 특정 헤드들이 의미론적으로 일관된 토큰 그룹(예: 국가 이름, 월, 숫자, 감정)을 안정적으로 생성한다는 것이 발견되었습니다.
3. 제어 가능성: 저자들은 특화된 헤드의 매우 적은 부분(전체의 1%만큼도 안 되는 수준)에 개입하는 것만으로도 타겟 개념을 안정적으로 억제하거나 강화할 수 있음을 입증했습니다. 이는 어텐션 레이어가 해석 가능하고 조작 가능한 선형 구조를 포함하고 있음을 시사합니다.

결과

저자들은 다음과 같은 여러 작업에서 접근 방식을 검증했습니다:

• 질의응답 (TriviaQA): "국가 이름"에 특화된 단 8개의 헤드(모델의 0.8%)의 부호를 반전시키자, 다른 질문에는 최소한의 영향을 미치면서 국가 관련 질문에 대한 성능은 크게 저하되었습니다. 무작위 헤드 선택이나 Logit Lens 기반 선택은 이러한 표적 효과를 보여주지 못했습니다.
• 독성 완화 (Toxicity Mitigation): 독성 키워드 목록을 사용하여 헤드를 선택함으로써, 저자들은 RealToxicityPrompts (RTP) 및 Thoroughly Engineered Toxicity (TET) 데이터셋에서 독성 생성 빈도를 성공적으로 줄였습니다. 32개의 특화된 헤드를 반전시켰을 때, 무작위 개입이 오히려 독성을 높이는 경우가 있었던 것과 대조적으로, 베이스라인 대비 독성 생성을 약 30-50% 감소시켰습니다.
• 이미지 분류 (LLaVA): 특정 데이터셋(예: MNIST, SVHN)을 위해 SOMP로 선택된 상위 32개 헤드를 반전시키면 분류 정확도가 크게 저하되었습니다. 연구 결과, 관련된 의미를 가진 데이터셋(예: 숫자 인식 작업)은 겹치는 헤드 집합을 공유하며, 이는 도메인 간의 기능적 중첩을 확인시켜 줍니다.
• 이미지 캡셔닝 (Flickr30k):
  • 억제: "색상", "감정", 또는 "수량"과 같은 속성에 대해 16개의 특화된 헤드를 반전시키면 캡션 품질(CIDEr 점수)을 높게 유지하면서도 해당 키워드를 거의 완전히 제거할 수 있었습니다.
  • 강화: 이 헤드들을 증폭($\alpha=5$)하면 캡션 품질의 미미한 저하와 함께 타겟 개념의 출현이 60% 이상 증가했습니다.

의의 및 주장

본 논문은 대규모 생성 모델의 어텐션 레이어가 **구조화되고 개념에 정렬된 선형 부분 공간(linear subspace)**을 보유하고 있다고 주장합니다. 이 발견은 고차원 개념이 근사적으로 선형 잔류 부분 공간에 표현된다는 가설을 뒷받침합니다.

이 연구의 의의는 대규모 모델을 이해하고 편집하기 위한 단순하고 학습이 필요 없는 도구를 제공한다는 점에 있습니다. 소수의 유닛이 광범위한 의미적 제어를 행사할 수 있음을 밝힘으로써, 저자들은 휴리스틱한 프롬프트 기반 개입이나 미세 조정(fine-tuning)에 대한 원칙적인 대안을 제시합니다. 이 접근 방식은 다음을 가능하게 합니다:

• 해석 가능성: 모델의 행동을 해석 가능한 작은 요소 세트로 분해합니다.
• 제어: 일반적인 모델 능력에 미치는 영향을 최소화하면서 특정 속성(예: 독성, 특정 시각적 개념)을 표적하여 억제하거나 강화합니다.

저자들은 이 방법이 선형성을 가정하고(이는 모든 비선형 구조를 포착하지 못할 수 있음), 의미론적 사전에 의존하며, 미세한 스티어링보다는 단순한 전역적 재조정을 사용한다는 점을 언급하며 한계점에 대해 겸허한 태도를 유지했습니다. 이들은 향후 연구에서 위치별 개입을 탐구하거나 이미지 생성 작업에 이러한 기술을 적용할 수 있을 것이라고 제안했습니다.