LLM 주의 헤드 안정성 측정: 회로 보편성의 새로운 통찰

초록

본 논문은 동일한 GPT‑계열 구조를 서로 다른 초기 시드로 여러 차례 학습시킨 뒤, 각 층의 어텐션 헤드가 생성하는 어텐션 스코어 행렬을 코사인 유사도로 비교해 “안정성”을 정량화한다. 실험 결과, 중간 층 헤드가 가장 불안정하고 표현적으로 다양하며, 모델 깊이가 깊어질수록 이러한 중간‑층 발산이 심해진다. 불안정한 헤드는 깊은 층일수록 기능적 중요도가 높아지고, AdamW와 같은 가중치 감쇠 최적화가 안정성을 크게 향상시킨다. 잔차 스트림은 전반적으로 매우 안정적이다. 연구는 회로(서브컴퓨테이션)의 보편성을 검증하기 위한 필수 전제조건으로 “시드 안정성”을 강조한다.

상세 분석

이 연구는 메커니즘 해석 가능성(mechanistic interpretability) 분야에서 흔히 가정되는 “회로 보편성”을 실증적으로 검증하려는 시도이다. 저자들은 GPT‑2‑small(12 층, 124 M 파라미터)부터 2 층, 4 층, 8 층까지 다양한 깊이와 헤드 수(8 또는 16)를 갖는 26가지 아키텍처를 정의하고, 각 아키텍처마다 50개의 시드(12 층은 5개)로 완전 재학습한다. 데이터는 모델 규모에 따라 OpenWebText(9 B 토큰)와 C4(2 B 토큰)로 나뉘며, 학습 하이퍼파라미터는 대부분 고정한다.

안정성 측정은 다음과 같다. 동일한 프롬프트 집합 P (≈100개)를 입력해 각 헤드 hᵢ 와 다른 시드의 후보 헤드 hⱼ 의 포스트‑소프트맥스 어텐션 스코어 행렬을 벡터화한 뒤, 코사인 유사도 s(m,n)(hᵢ,hⱼ) 를 계산한다. 프롬프트별 유사도를 평균해 \bar{s} 를 얻고, 같은 층 내에서 가장 높은 \bar{s} 값을 선택해 헤드 hᵢ 의 단일‑시드 안정성 s(m,n){hᵢ} 을 정의한다. 마지막으로 모든 쌍 시드에 대해 평균을 취해 최종 안정성 S(m){hᵢ} 을 산출한다. 층별 평균 S(m)_l 을 구하면 층‑단위 안정성 프로파일을 얻으며, 층 제약을 없애고 전 층에서 최적 매치를 찾는 “크로스‑레이어 베스트‑매치”도 별도로 계산한다.

주요 결과는 다섯 가지이다. 첫째, 초기와 최종 층의 헤드가 높은 S 값을 보이는 반면, 중간 층(특히 4‑6층)에서는 급격히 낮아져 평균 ≈0.70 까지 떨어진다. 이는 동일 아키텍처라도 중간 층의 내부 표현이 시드마다 크게 달라진다는 증거다. 둘째, 모델 깊이가 증가할수록 중간‑층 불안정성이 더욱 심화된다; 12 층 모델에서 중간‑층 평균 S 가 0.55 수준까지 감소한다. 셋째, 불안정한 헤드가 기능적으로 더 중요한지를 확인하기 위해, 해당 헤드를 무작위 마스킹하거나 교체했을 때의 퍼플렉시티·로그‑우도 감소량을 측정했으며, 깊은 층일수록 불안정 헤드의 영향력이 크게 나타났다. 넷째, AdamW(분리형 가중치 감쇠)를 사용한 경우, 동일 아키텍처에서 시드 간 S 값이 평균 +0.12 정도 상승해 안정성이 현저히 개선됨을 보였다. 이는 가중치 감쇠가 학습 과정에서 표현 공간을 더 일관되게 수렴하도록 만든다는 가설을 뒷받침한다. 다섯째, 잔차 스트림(Residual Stream)의 토큰‑레벨 표현은 모든 시드에서 코사인 유사도가 0.95 이상으로 매우 안정적이며, 이는 “잔차 기반 정보 전달”이 구조적 변동에 강인함을 시사한다.

이러한 발견은 회로 기반 해석이 실제로 재현 가능한지 여부를 판단하는 데 핵심적인 기준을 제공한다. 특히 안전‑중요 분야에서 모델을 검증·감시하려면, 중간 층의 불안정성을 최소화하거나, 불안정한 헤드가 담당하는 기능을 별도로 검증해야 한다는 실용적 교훈을 얻는다. 또한, 가중치 감쇠와 같은 단순한 최적화 선택이 회로 보편성에 미치는 영향을 정량화함으로써, 향후 대규모 LLM 개발 시 “안정성‑우선” 훈련 전략을 설계할 근거를 제공한다.