기술 요약: 언어 모델에서 해석 가능한 프롬프트별 회로 찾기

문제 제기

대규모 언어 모델 (LLM) 의 내부 계산 메커니즘을 이해하는 것은 기계적 해석 가능성 분야에서 여전히 핵심적인 과제로 남아 있습니다. 이전 연구들은 모델이 잔여 스트림 (residual stream) 에서 정보를 쓰고 읽는 구성 요소들의 인과적 경로인 '회로'를 통해 작동함을 입증했습니다. 그러나 이러한 회로를 추적하는 기존 방법들은 중대한 한계에 직면해 있습니다. 전통적인 접근법은 주로 활성화 패치 (activation patching), 반사실적 입력, 또는 대체 모델을 사용하는데, 이는 계산 비용이 많이 들고 여러 번의 순전파 (forward passes) 를 요구하며 노이즈나 인공물을 도입할 수 있습니다. 더 나아가, 이전의 회로 추적 방법들은 종종 통신에 사용되는 특정 저차원 신호를 흐리게 만드는 크고 노이즈가 많은 그래프를 생성하여, 프롬프트 단위에서 모델 행동을 인간이 이해할 수 있는 설명으로 도출하기 어렵게 만듭니다.

방법론: ACC++

본 논문은 주의 - 인과적 통신 (Attention-Causal Communication, ACC) 원리를 기반으로 한 개선된 회로 추적 방법인 **ACC++(Attention-Causal Communication++)**를 소개합니다. ACC++ 는 특정 어텐션 헤드가 특정 토큰 쌍에 주의를 기울이게 만드는 특정 신호 (저차원 부분 공간의 내용) 를 식별하도록 설계되었습니다.

핵심 메커니즘

ACC++ 는 대체 모델, 패치, 또는 반사실적 입력 없이 단일 순전파로 작동합니다. 이 방법은 목적지 - 소스 토큰 쌍 $(d, s)$에 주의를 기울이는 주어진 어텐션 헤드 $(\ell, a)$에 대해 다음 네 단계로 진행됩니다:

1. 분해: 목적지 잔여 스트림 $x_d$를 상류 구성 요소 (임베딩, 이전 어텐션 헤드, MLP) 의 출력들의 합으로 분해합니다.
2. 후보 신호 구성: 헤드의 쿼리 - 키 (QK) 행렬 $\Omega = W_Q W_K^\top$의 특이값 분해 (SVD) 를 사용하여, 상류 구성 요소의 출력을 특이 벡터 방향으로 투영합니다. 이는 각 신호가 특정 상류 구성 요소가 특정 특이 방향 (랭크 1 후보) 으로 쓰는 것에 대응하는 후보 신호 집합을 생성합니다.
   • 기술적 진전: 모델 이질성 (예: 편향 항과 회전 위치 임베딩/RoPE) 을 처리하기 위해, ACC++ 는 편향과 RoPE 회전을 토큰 표현에 포함시켜 어텐션 계산을 단일 이차 형식 $\tilde{x}_d^\top \Omega \tilde{x}_s$로 통합합니다. 이를 통해 아키텍처에 관계없이 헤드당 단일 SVD 분석이 가능해집니다.
3. 귀속 및 점수 매기기: 각 후보 신호가 소프트맥스 후 어텐션 가중치 $A_{ds}$에 기여하는 정도를 통합 기울기 (Integrated Gradients, IG) 를 사용하여 점수화합니다. 소프트맥스 전 점수에 대한 선형 대리 모델을 사용했던 이전 연구와 달리, ACC++ 는 신호 제거가 전체 소프트맥스 분포에 미치는 비선형 영향을 명시적으로 평가합니다.
4. 탐욕적 반사실적 탐색: 전체 후보 집합에서 시작하여, 어텐션 가중치 $A_{ds}$가 임계값 $\tau$ 미만으로 떨어질 때까지 점수가 가장 높은 신호들을 탐욕적으로 제거합니다. 제거된 집합은 해당 특정 어텐션 결정에 대한 인과적 회로를 구성합니다.

회로 구성

프롬프트에 대한 완전한 회로는 ACC++ 를 재귀적으로 적용하여 구성됩니다: 목표 로그 (logit) 에 직접적인 영향을 미치는 구성 요소에서 시작하여, 더 이상 상류 구성 요소가 남지 않을 때까지 각 어텐션 발화의 인과적 신호를 식별하며 상류 어텐션 발화를 추적합니다.

주요 기여

1. 개념적 및 기술적 진전

• 정확한 신호 식별: ACC++ 는 구성 요소와 채널을 독립적으로 식별하는 대신, 상류 구성 요소와 특이 벡터 방향의 외적 (cross-product) 을 신호로 식별합니다. 이는 더 저차원적이고 정확한 신호를 산출합니다.
• 비선형 어텐션 계상: 선형 대리 모델이 아닌 소프트맥스 후 어텐션 가중치에 직접 반사실적 목적 함수를 정의함으로써, ACC++ 는 가짜 구성 요소를 피하고 더 작고 정확한 회로를 식별합니다.
• 통합된 QK 분석: 이 방법은 토큰을 변환하여 편향과 RoPE 를 처리하므로, 다양한 모델 아키텍처 (예: GPT-2, Pythia, Gemma) 간에 일관된 SVD 기반 분석을 가능하게 합니다.

2. 자동화된 해석 가능성

저자들은 ACC++ 신호의 상당 부분이 해석 가능함을 입증합니다. 그들은 다음을 수행하는 자동화된 파이프라인 ( [18] 에서 적응됨) 을 사용합니다:

• 데이터셋 (The Pile) 에서 각 신호에 대한 고활성도 컨텍스트를 검색합니다.
• LLM 을 사용하여 신호에 대한 짧은 자연어 설명을 생성합니다.
• LLM 을 심사자 (judge) 로 활용하는 프로토콜을 사용하여 설명을 검증하고 신호와 노이즈를 구분합니다.
  결과는 많은 신호가 짧고 일관된 자연어 설명 (예: "병렬 쌍의 두 번째 항목", "고유명사", "구조적 구분자") 을 허용함을 보여줍니다.

3. 프롬프트별 회로 분석

ACC++ 는 프롬프트 단위에서 회로를 분석할 수 있게 하여, 모델 행동이 정적이지 않고 프롬프트 구조와 언어에 따라 크게 변함을 드러냅니다.

결과 및 통찰

1. 회로의 컴팩트함과 품질

세 가지 모델 (GPT-2 Small, Pythia-160M, Gemma-2 2B) 과 세 가지 작업 (간접 목적어 식별, Greater-Than, 성대명사) 에서 ACC++ 회로는 원래 ACC 방법이 생성한 회로보다 2 배에서 10 배 더 작습니다.

• 차원성: ACC++ 신호의 대부분은 단일 특이 방향을 사용하는 랭크 1 이며, 이전 연구의 고차원 신호와 대조됩니다.
• 인과적 타당성: 제거 실험 (ablation studies) 은 ACC++ 신호를 제거하는 것이 목표 작업에서 모델 성능에 상당한 영향을 미치면서도 잔여 스트림을 최소로 교란 (코사인 유사도 $\approx 0.99$) 시킨다는 것을 확인하여, 그 인과적 역할을 입증합니다.

2. 프롬프트 구조에 대한 민감성 (IOI 작업)

간접 목적어 식별 (IOI) 작업에서, 논문은 회로가 프롬프트 템플릿에 매우 민감함을 발견합니다:

• GPT-2: 회로는 고수준 템플릿 구조 (이름의 ABBA 대 BABA 순서) 에 따라 엄격하게 군집화됩니다. 동일한 "이름 이동 (name-mover)" 헤드 (9,9) 는 템플릿에 따라 근본적으로 다른 입력 신호를 사용합니다 (예: "쌍의 두 번째 항목" 대 일반적인 구조적 단서).
• Pythia: 회로는 저수준 단어 변경 (예: "While" 대 "Then") 에 따라 군집화되어, 구조적 순서뿐만 아니라 표면적 문구에 대한 민감성을 나타냅니다.
• Gemma-2: 더 복잡하고 템플릿 의존성이 적은 군집화를 보여주며, 더 세밀한 프롬프트 세부 사항에 의존함을 시사합니다.

3. 다국어 메커니즘 (mIOI 작업)

다국어 IOI 환경 (영어, 스페인어, 프랑스어, 포르투갈어) 에서:

• 구성 요소 재사용: 모델은 언어 간에 동일한 어텐션 헤드와 MLP 를 재사용합니다.
• 언어별 신호: 구성 요소는 공유되지만, 이를 통해 흐르는 신호는 종종 언어별로 다릅니다.
• 언어학적 상관관계: 서로 다른 언어의 회로 간 거리는 해당 언어 간의 유형론적 거리 (구문적 및 계통적 거리) 와 강하게 상관관계가 있습니다. 예를 들어, 스페인어와 포르투갈어 회로는 영어나 프랑스어 회로보다 서로 더 유사합니다.

중요성과 주장

본 논문은 ACC++ 가 값비싼 반사실적 개입에 의존하지 않는 실용적이고 확장 가능한 기계적 해석 가능성 도구를 제공한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

1. 확장성: 단일 순전파만 요구함으로써 프롬프트별 회로 분석을 대규모로 실현 가능하게 합니다.
2. 해석 가능성: 저차원 통신 채널과 인간이 읽을 수 있는 자연어 설명 간의 직접적인 연결을 제공하여, 수학적 메커니즘과 의미적 이해 사이의 간극을 메웁니다.
3. 미묘한 이해: 모델 "회로"가 단일한 것이 아니라 프롬프트 구조와 언어에 적응하는 구성 요소와 신호의 역동적 집합임을 드러냅니다. 이는 작업에 대한 단일 보편적 회로라는 관점을 도전하고, 대신 프롬프트별 메커니즘의 스펙트럼을 보여줍니다.

저자들은 임계값 $\tau$의 휴리스틱적 성격, FFN 으로 인한 신호에 대한 추적 부재, 그리고 자동화된 해석이 결론적이기보다는 시사적이라는 점 등 한계를 지적합니다. 그들은 다국어 연구 결과가 소수의 언어 쌍에 기반하므로 결정적인 증거라기보다는 더 넓은 경향성을 나타내는 것으로 간주되어야 함을 강조합니다.