기술 요약: 공유 서브서킷을 통한 LLM의 카운트다운 수행 방식

문제 정의

대규모 언어 모델(LLM)은 특정 길이의 문장 작성, ASCII 표 포맷팅, 또는 DNA 서열을 정확한 코돈에서 종료하는 것과 같이 목표 길이에 맞춰 토큰 수를 추적해야 하는 작업을 빈번하게 수행합니다. LLM이 이러한 작업에 능숙함을 보여줌에도 불구하고, 목표 길이에 도달하기까지 남은 "남은 시간"을 추적하는 내부 메커니의는 완전히 이해되지 않았습니다. 특히, 모델이 명시적 대상과 암시적 대상이라는 서로 다른 맥락 전반에 걸쳐 작동하는 일반화된 카운트다운 메커니즘을 보유하고 있는지, 그리고 이 메커니즘이 서로 다른 모델 간에 공유되는지 여부가 불분명합니다.

방법론

저자들은 Llama-3.1-70B-Instruct 모델에 초점을 맞추어 기계론적 해석 가능성(mechanistic interpretability)에 대한 상향식(bottom-up) 접근 방식을 채택했습니다. 방법론은 세 단계로 진행되었습니다:

1. 인과적 국소화 (Causal Localization via Patching):

   • 작업: 모델에게 특정 단어로 끝나는 지정된 길이($N$)의 문장을 작성하도록 하는 "문장 계획(sentence planning)" 프롬프트를 부여했습니다.
   • 기법: 저자들은 잔차 스트림 패칭(residual stream patching)을 수행했습니다. "소스(source)" 프롬프트(목표 길이 $n_{src}$)와 "데스티네이션(destination)" 프롬프트(목표 길이 $n_{dst}$, 여기서 $n_{dst} > n_{src}$)를 구성했습니다. 소스 프롬프트의 목표 길이 토큰과 길이 단위 토큰의 활성화 값을 다양한 레이어에서 데스티네이션 프롬프트로 패칭했습니다.
   • 목표: 특정 레이어와 토큰이 출력 길이를 결정하는 데 인과적으로 충분한 정보를 제공하는지 식별하는 것입니다.

2. 어텐션 헤드 식별:

   • 임계 레이어를 식별한 후, 저자들은 해당 레이어의 어텐션 헤드 출력 노름(output norms)을 분석하여, 길이 단위 토큰으로부터 다운스트림 토큰으로 흐르는 정보의 흐름을 조사했습니다.
   • 이 헤드들이 사용하는 표현의 기하학적 구조(키와 쿼리)를 조사하여, 현재 위치를 목표 길이와 어떻게 비교하는지 분석했습니다.

3. 비지도 프로빙 및 일반화:

   • 암시적 길이: 목표 길이가 명시적으로 언급되지 않고 맥락에서 추론되는 상황(예: Base64 또는 DNA 서열에서의 동일한 길이의 반복 문자열)을 조사했습니다. 패칭을 통해 이러한 문자열의 종료 구분자(ending delimiters)가 길이 정보를 인과적으로 저장하는지 확인했습니다.
   • 기하학적 분석: PCA(주성분 분석)와 코사인 유사도를 사용하여 Llama-3.1-70B에서 식별된 서브서킷과 이전에 식별된 Claude 3.5 Haiku의 "줄 바꿈(line-wrapping)" 서브서킷의 기하학적 매니폴드를 비교했습니다.
   • 비지도 발견: 식별된 "카운트다운 헤드"의 출력 활성화를 사용하여 선형 프로브(linear probes)를 구축했습니다. 이 프로브를 대규모 미라벨링 코퍼스(The Pile 및 LMSYS-1M 채팅 데이터셋)에 실행하여, 모델이 암시적으로 카운트다운을 수행하는 다양하고 높은 활성 사례들을 찾아냈습니다.

주요 기여 및 결과

1. 카운트다운 서브서킷의 식별

본 연구는 남은 토큰을 추적하는 데 책임이 있는 Llama-3.1-70B-Instruct 내의 특정 서브서킷을 성공적으로 격리했습니다:

• 국소화: 길이 정보는 레이어 4 부근부터 "길이 단위 토큰"으로 이동하며, 레이어 18까지 인과적으로 사용됩니다.
• 특정 헤드: 레이어 18의 세 가지 어텐션 헤드(L18H24, L18H28, L18H30)가 "카운트다운 헤드"로 식별되었습니다.
• 메커니즘: 이 헤드들은 중첩된 수용 영역(receptive fields) 패턴을 보입니다. L18H28은 초기에, L18H24는 중간에, L18H30은 종단부에 활성화됩니다. 이들의 결합된 출력 노름은 목표 길이 직전에 정점에 도달하며, 이는 갑작스러운 스파이크가 아닌 연속적인 카운트다운 메커니즘임을 나타냅니다.

2. 암시적 길이 저장 및 보편성

저자들은 목표 길이가 프롬프트에 명시적으로 제공되지 않을 때도 이 서브서킷이 기능함을 입증했습니다:

• 암시적 저장: "반복 문자열" 설정(예: 여러 개의 동일한 길이의 Base64 문자열)에서, 모델은 맥락으로부터 목표 길이를 추론합니다. 패칭 실험을 통해 이러한 문자열의 종료 구분자가 암시적인 목표 길이 정보를 인과적으로 저장함을 확인했습니다.
• 표현 유사성: 코사인 유사도 분석 결과, 암시적 길이(구분자 내)의 키 활성화와 명시적 길이(프롬프트 토큰 내)의 키 활성화가 유사한 매니폴드 위에 놓여 있음을 발견했습니다. 암시적 길이와 명시적 길이 사이의 관계는 매우 선형적($0.93 \le r^2 \le 0.97$)이며, 이는 길이 정보의 통합된 표현을 시사합니다.
• 교차 모델 모티프: Llama-3.1의 카운트다운 서브서킷의 기하학적 구조(특정 정렬 특성을 가진 키와 쿼리의 곡선형 매니폴드)는 이전에 Claude 3.5 Haiku에서 식별된 "줄 바꿈" 서브서킷과 동일한 것으로 나타났습니다. 이는 이 모티프가 서로 다른 모델과 작업 간에 공유되는 것임을 시사합니다.

3. 다양한 응용 분야의 비지도 발견

카운트다운 헤드의 출력을 이용한 비지도 프로빙을 통해, 저자들은 이 서브서킷이 활성화되는 다음과 같은 다양한 작업들을 식별했습니다:

• 임베디드 트윗: 뉴스 텍스트에 포함된 280자 트윗의 종료 예측.
• DNA 서열: 이전 서열을 바탕으로 올바른 코돈에서 서열 종료.
• SHA-256 해시: 64자 헥사(hex) 문자열의 끝 탐지.
• 하이쿠 생성: 음절 수 및 행 바꿈 추적.
• ASCII 표: 특정 너비에 맞춰 셀 정렬.
• 동전 던지기 및 반올림: 생성된 시퀀스의 문자 또는 숫자 개수 추적.

의의 및 주장

본 논문은 단일하고 좁은 작업(문장 계획)에서 서브서킷을 역공학하는 것이 연구자로 하여금 특정 행동이 어떻게 광범위한 작업 및 서로 다른 모델 전반으로 일반화되는지 이해할 수 있게 해준다고 주장합니다.

• 일반화: 본 연구는 단일 서브서킷이 명시적 카운팅(프롬프트 지정)과 암시적 카운팅(맥락 추론) 모두를 위해 재사용됨을 보여줍니다.
• 공유된 메커니즘: Llama-3.1과 Claude 3.5 Haiku에서 동일한 기하학적 모티프를 식별한 것은 시간/길이를 추적하기 위한 특정 계산 프리미티브(computational primitives)가 프런티어 모델들 사이에서 근본적이고 공유되는 것임을 시사합니다.
• 방법론적 접근: 저자들은 "하향식" 기계론적 해석 가능성 워크플로우를 제안합니다: 제어된 작업에서 서브서킷을 격리하고, 그 기하학적 구조를 역공학한 다음, 비지도 프로빙을 사용하여 모델의 전체 운영 공간에서 그 사용처를 매핑하는 방식입니다. 저자들은 이 접근 방식이 피처 딕셔너리(feature dictionary) 분석과 유사하게 자동화되어, LLM 내의 다른 일반화된 행동들을 체계적으로 밝혀낼 수 있을 것이라고 제안합니다.

저자들은 자신들이 공유된 모티프를 식별했음에도 불구하고, 본 연구가 길이 처리의 모든 측면을 해결했다고 주장하기보다는 인과적 개입과 기하학적 구조에 집중하고 있다는 점을 명시하며 신중한 태도를 유지했습니다.