기술적 요약: 스펙트럼 프로브-회로

문제 제기

사전 훈련된 트랜스포머의 기계적 해석 가능성은 일반적으로 사후 분석에 의존합니다. 즉, 모델이 완전히 훈련된 후 헤드들을 제거 (ablating) 하고 고정된 사전 지정된 작업에서의 성능 저하를 관찰하여 어텐션 헤드 회로를 식별하는 방식입니다. 이 워크플로우는 계산적으로 비용이 많이 들며 (모든 조건에 대해 완전히 훈련된 네트워크를 통한 순전파가 필요함), 또한 목표 역량이 탐색 시작 전에 정의되어야 하므로 본질적으로 반응적입니다. furthermore, 기존 방법들은 특히 "어텐션 싱크 (첫 번째 토큰에 주의를 기울이는 헤드)"가 지배적인 더 큰 모델에서, 일반적인 콘텐츠 의존적 계산과 작업별 인과 회로를 구별하는 데 종종 어려움을 겪습니다.

방법론: 3 단계 레시피

이 논문은 헤드별로 작동하며, 작업 레이블이나 귀속 기울기 (attribution gradients) 를 필요로 하지 않고, 훈련 중이나 최종 체크포인트에 적용할 수 있는 어텐션 헤드 회로를 식별하기 위한 3 단계 레시피를 제안합니다.

1 단계: 헤드별 스펙트럼 신호 (일반 지표)

첫 번째 단계는 무작위 행렬 이론에서 유래한 스펙트럼 지표를 사용하여 지속적이고 콘텐츠에 의존적인 계산을 수행하는 헤드를 식별합니다.

• 지표: 각 훈련 체크포인트의 각 헤드에 대해, 헤드별 어텐션 출력 행렬의 **참여 비율 (Participation Ratio, PR)**이 계산됩니다. PR 은 출력 분포의 유효 랭크 (차원성) 를 측정합니다.
• 궤적 특징: 이 방법론은 이 신호의 시간적 진화를 통합합니다: $I(L, H) = \sum_t \max(PR(L, H, t) - 1, 0) \cdot \Delta \log(\text{tokens}_t)$.
• 기능: 이 "PR-적분"은 시간이 지남에 따라 높은 유효 랭크 (많은 방향을 포괄) 를 유지하는 헤드를 순위 매깁니다. 이는 저랭크 (콘텐츠 무관) 상태로 붕괴되거나 노이즈가 많은 헤드를 필터링하는 일반적인 전문화 지표 역할을 합니다. 중요하게도, 이 지표는 헤드가 무엇을 하고 있는지는 구분하지 않고, 전문적인 계산을 수행하고 있다는 사실만을 구분합니다.

2 단계: 작업 패턴 스크리닝 (작업별 필터)

일반적인 스펙트럼 신호는 표준적인 목표에 대한 어텐션 패턴을 측정함으로써 목표 작업에 맞게 전문화됩니다.

• 메커니즘: "전체 헤드" 스크린은 특정 어텐션 패턴 (예: 유도, 이전 토큰, 중복 토큰, 첫 번째 토큰/BOS, 자기 참조, 로컬) 에 대한 모든 헤드의 선택성을 측정합니다.
• 선택성 정의: 선택성은 목표 위치에 대한 어텐션과 균일한 기준선 간의 비율입니다.
• 임계값: 헤드는 선택성이 임계값을 초과할 경우 분류됩니다 (예: 할당 작업의 경우 $\ge 30\times$, 유도 회로 소속의 경우 $\ge 50\times$).
• 기능: 이 단계는 일반적인 "전문화" 신호를 특정 역량을 위한 후보 회로로 변환합니다. 이는 어텐션 싱크가 지배적인 모델에서 1 단계가 작업별 헤드가 아닌 일반적인 콘텐츠 의존적 헤드를 더 높은 순위로 매길 수 있다는 한계를 해결합니다.

3 단계: 인과적 검증 (반증 가능한 주장)

후보 회로는 레이어 구성의 인공물이 아닌 인과적 효과임을 보장하기 위해 엄격한 통제를 갖춘 그룹 제거 (group ablation) 를 통해 검증됩니다.

• 절차: 어텐션 출력 투영 (projection) 에서 잔여 스트림 (residual stream) 의 헤드별 슬라이스가 0 으로 설정됩니다.
• 통제:
  1. 매칭된 무작위 (Matched-Random): 스크린에서 선택되지 않은 동일한 레이어에서 동일한 수의 헤드를 제거합니다. 이는 "이 레이어의 어떤 헤드를 제거해도 성능 저하를 초래한다"는 가설을 배제합니다.
  2. 상한선 (Upper Bound): 선택된 레이어의 모든 헤드를 제거합니다. 이는 해당 레이어의 가능한 최대 효과를 확립합니다.
• 기준: 스크린으로 선택된 헤드가 매칭된 무작위 통제보다 유의미하게 큰 성능 저하를 일으키고, 상한선이 0 에 가깝게 포화될 때 회로는 인과적인 것으로 간주됩니다.

주요 결과

1. 아키텍처 및 규모 간 복제

이 레시피는 8 배의 파라미터 범위 (51M 에서 1B-활성/7B-총계) 와 세 가지 아키텍처 계열 (밀집 GPT-NeoX, 밀집 Llama 스타일, Llama 스타일 MoE) 에서 검증되었으며, 서로 다른 데이터 파이프라인 (TinyStories, FineWeb, Pile, DCLM) 으로 훈련되었습니다.

• 유도 회로: 테스트된 모든 모델에서 소규모 유도 회로 (3~6 개 헤드) 가 합성 유도 작업에 인과적으로 필수적인 것으로 발견되었습니다.
• 아선형 스케일링: 전체 헤드 수는 모델 크기에 따라 선형적으로 증가하는 반면, 유도 회로의 크기는 아선형 (3~11 개 헤드) 으로 유지되어, 특정 역량이 모델 용량에 선형적으로 스케일링되지 않음을 나타냅니다.

2. 보존된 전문화 헤드의 비율

124M 에서 410M 에 이르는 8 배의 규모 범위에서, 식별 가능한 전문 계산을 수행하는 헤드의 비율은 **약 17~19%**로 보존됩니다.

• 이 비율은 두 가지 방법으로 독립적으로 회복되었습니다: (1) 상위 K 개 스펙트럼 선택을 역량 클래스로 분류, (2) 영분포 (null-distribution) 노이즈 바닥 이상의 유도 선택성을 가진 헤드 수 세기.
• 더 큰 1B 급 모델에서는 이 비율이 더 크게 변동 (4.7% 에서 27.3%) 하며, 모델 규모보다는 BOS-인자 (어텐션 싱크) 의 우세와 상관관계가 있어, 보존은 특정 아키텍처/데이터 계열 내의 속성임을 시사합니다.

3. 예측력 및 시드 특이성

• 시드 특이성: 6 개의 서로 다른 무작위 시드로 훈련된 51M 파라미터 모델에서, 이 레시피는 동일한 작업에 대해 서로 다른 시드별 회로를 식별했습니다. 시드 A 에 대해 식별된 회로를 제거하는 것은 시드 B 모델에 영향을 미치지 않았으며, 이는 이 방법이 보편적 패턴이 아니라 모델의 실제 구현을 찾음을 확인시켜 줍니다.
• 예측 타이밍: 스펙트럼 신호 (1 단계) 는 역량 선택성이 기능적 임계값을 넘기 전에 상승합니다. 이 레시피는 역량이 완전히 등장하기 전에 총 훈련 토큰의 작은 부분 (0.3~2%) 만 사용하여 회로를 식별할 수 있습니다.

4. 데이터 변화에 대한 강건성

FineWeb 에서 훈련된 후 OpenWebText 로 계속 훈련된 모델에서, FineWeb 엔드포인트에서 식별된 유도 회로는 데이터 변화를 통해 불변 (Jaccard = 1.00) 으로 유지되었습니다. 헤드는 선택성과 PR 궤적을 유지했으며, 이는 이 방법이 코퍼스별 통계가 아닌 구조적 전문화를 식별함을 시사합니다.

5. 구성적 의존성

이 방법은 "이전 토큰" 회로 (8~9 개 헤드) 를 성공적으로 식별하고 유도에서의 인과적 역할을 입증했습니다. 이전 토큰 회로를 제거하면 유도 성능이 파괴되었으며, 이는 이전 토큰 헤드가 유도 헤드가 사용하는 키 벡터를 구축하는 표준 구성 구조를 확인시켜 줍니다.

중요성 및 주장

이 논문은 3 편으로 구성된 프로그램의 방법론적 앵커로 자신을 위치시킵니다. 주요 주장은 레시피 자체의 유효성에 초점을 맞춘 겸손한 것입니다:

1. 레시피는 레이지이지, 보편적 스크린이 아님: 저자들은 명시적으로 1 단계 (스펙트럼 신호) 가 작업별 감지기가 아닌 일반적인 전문화 지표라고 명시합니다. 2 단계 (스크린) 가 작업별 특이성을 제공하고, 3 단계 (제거) 가 인과적 반증 가능성을 제공합니다. 이 조합이 방법론을 강건하게 만듭니다.
2. 인과적 주장의 정직한 프레임: 이 논문은 "매칭된 무작위" 차이가 인과적 주장을 정직하게 만든다고 강조합니다. 이 통제 없이는 특정 회로와 일반적인 레이어 효과 사이를 구별할 수 없습니다.
3. 발달적 선행: 스펙트럼 신호는 역량 선택성의 등장보다 앞서며, 사전 훈련 중 회로 형성 모니터링을 위한 잠재적 창구를 제공합니다 (동반 논문에서 상세히 다룸).
4. 한계: 저자들은 간접 객체 식별과 같은 구성된 작업의 경우, 주요 어텐션 패턴이 서로 다른 모델 계열 간에 작업-인과 구조와 분리될 수 있으며, 더 넓은 범위의 후보 스크린이 필요하다고 지적합니다. 또한, 이 방법은 고정된 평가 배치에 의존하며, 특정 회로 헤드는 재훈련 시마다 다를 수 있습니다 (비록 전문화 헤드의 비율은 안정적으로 보이지만).

요약하자면, 이 논문은 작고 인과적으로 필수적인 어텐션 헤드 회로를 분리하기 위한 재현 가능하고 레이블이 없는 프레임워크를 제공하며, 모델 용량의 보존된 부분이 모델 크기에 아선형으로 스케일링되는 전문 계산에 할당됨을 입증합니다.