기술 요약: 특징 발견 및 롱 컨텍스트 귀속을 위한 턴 평균 SAE (Turn-Averaged SAEs)

1. 문제 정의

희소 오토인코더(Sparse Autoencoders, SAEs)는 언어 모델로부터 해석 가능한 특징(feature)을 추출하여 회로 분석, 행동 제어, 안전 모니터링을 용이하게 하는 주요 도구로 부상했습니다. 그러나 표준 SAE 아키텍처는 개별 토큰 활성화값에 작동합니다. 이러한 토큰 단위의 입도는 대규모 해석 작업에서 두 가지 중요한 병목 현상을 초래합니다:

1. 확장성 제한: 활성화되는 특징의 수는 컨텍스트 길이에 따라 선형적으로 증가합니다. 롱 컨텍스트 시나리오나 다회차 대화(multi-turn conversations)의 경우, 발생하는 데이터 양이 방대해져 다운스트림 분석(예: 귀속 그래프)을 수행하기 번거롭고 검사하기 어렵게 만듭니다.
2. 고수준 컨텍스트 손실: 많은 해석 목표는 모델이 왜 특정 응답을 생성했는지 이해하는 것을 포함하며, 이는 많은 토큰에 걸쳐 분산된 고수준 속성(예: 주제, 페르소나, 안전 특성)의 식별을 요구합니다. 표준 토큰 단위 SAE는 종종 특정 단어나 구문 패턴에 대한 저수준 탐지기를 학습합니다. 이 모델들은 각 위치를 독립적으로 재구성하는 학습 목적 함수를 가지기 때문에, 이러한 고수준의 스팬(span) 전반에 걸친 특성을 토큰 수준의 노이즈와 분리하도록 강제하는 압력이 없습니다.

2. 방법론

2.1 턴 평균 SAE (Turn-Averaged SAEs)

핵심 제안은 SAE를 개별 토큰 활성화값이 아니라, 토큰의 스팬(특히 단일 인간 또는 어시스턴트 턴)에 대한 **평균 은닉 상태(mean hidden state)**에 대해 학습시키는 것입니다.

• 입력 변환: $T$라는 턴의 토큰들을 $x_t$라고 할 때, 입력은 평균 은닉 상태 $\bar{x} = \frac{1}{|T|} \sum_{t \in T} x_t$입니다.
• 아키텍처: SAE 아키텍처는 희소성을 위해 BatchTopK를 사용하는 기존 구현과 동일하게 유지됩니다. $\mathbb{R}^{d_{model}}$ 내 벡터들의 평균은 개별 벡터와 동일한 차원을 유지하므로, 구조적 변경은 필요하지 않습니다.
• 직관: 평균화는 저역 통과 필터(low-pass filter) 역할을 합니다. 이는 특정 단어 선택이나 구문적 표식과 같은 고주파의 토큰별 세부 사항을 매끄럽게 만드는 동시에, 주제, 스타일, 행동 특성과 같은 광범위한 개념에 해당하는 저주파 성분을 보존합니다.

2.2 하이브리드 아키텍처

평균화가 토큰별 세부 정보를 소실한다는 한계를 해결하기 위해, 저자들은 고수준 구조와 미세한 잔차(residual)를 모두 포착하기 위한 두 가지 하이브리드 접근 방식을 제안합니다.

1. 디커플드 SAE (Decoupled SAE): 두 개의 독립적인 SAE를 학습시킵니다. 하나는 턴 평균 $\bar{x}$를 재구성하고, 다른 하나는 잔차($x_t - \bar{x}$)를 재구성합니다. 이는 신호를 분리하지만, 두 모델이 인과적으로 단절되어 있어 귀속 분석을 복잡하게 만듭니다.
2. 네스티드 SAE (Nested SAE, 마트료시카 스타일): 공유된 인코더와 디코더를 사용하되 특화된 손실 함수를 가진 단일 SAE를 학습시킵니다. 특징 공간은 내부 파티션 $[0:h]$와 외부 파티션 $[h:d_{sae}]$로 나뉩니다.
   • 손실 함수: $L = \sum_t \| \text{dec}(\text{enc}(x_t)) - x_t \|^2 + \alpha \cdot \| \text{dec}_h(\text{enc}(\bar{x})[0:h]) - \bar{x} \|^2$.
   • 내부 파티션은 턴 평균을 재구성하도록 강제되는 반면, 전체 특징 집합은 토큰별 활성화를 재구성합니다. 이를 통해 단일 모델이 고수준의 턴 특성과 토큰 수준의 세부 사항을 동시에 학습할 수 있습니다.

2.3 귀속 그래프 적응 (Attribution Graph Adaptation)

저자들은 귀속 그래프를 토큰 위치의 특징이 아닌, 특정 턴과 레이어에서 활성화되는 특징을 나타내는 턴 단위로 적응시킵니다.

• 에지 가중치(Edge Weights): 소스 턴에 대한 그래디언트 귀속값을 계산하고 턴 길이($1/|T_A|$ 및 $1/|T_B|$)로 정규화하여, 가중치가 턴 길이에 무관하도록 계산합니다.
• 개입 검증: 소스 특징을 턴 내의 모든 토큰에 대해 디코더 방향을 따라 섭동(perturbation)시키고, 타겟 특징의 사전 활성화값 변화를 측정하여 에지 가중치를 검증합니다.

3. 주요 기여

1. 턴 평균 SAE: 턴 평균 은닉 상태에 대해 학습하는 것이 토큰 단위 특징보다 고수준의 턴 특성(주제, 스타일, 안전 특성)을 더 완전하게 포착함을 입증했습니다.
2. 네스티드 SAE 아키텍처: 마트료시카 스타일의 손실 함수를 사용하여 턴 평균 및 토크 단위 특징을 공동으로 학습하는 단일 모델 아키텍처를 도입하였으며, 두 신호 성분에 대해 높은 재구성 품질을 달aram했습니다.
3. 확장 가능한 귀속: 귀속 그래프를 턴 단위로 적응시켜 다회차 대화에 대한 압축된 분석을 가능하게 했습니다. 저자들은 실험적 개입과 충분성 지표(완전성 및 교체)를 통해 이 그래프들을 검증했습니다.

4. 결과

4.1 특징 품질 및 재구성

• 재구성: 네스티드 SAE는 순수 토큰 SAE(0.162)와 유사한 0.167의 토큰 단위 정규화 평균 제곱 오차(nMSE)를 달성하는 동시에, 턴 평균 nMSE 0.116을 달성했습니다. 반면, 표준 토큰 단위 SAE는 턴 평균 입력에 대해 성능이 저조했습니다(nMSE 0.355). 이는 평균값이 토큰 학습 모델에게 분포 외(out-of-distribution) 데이터임을 나타냅니다.
• 정성적 발견: 수학적 오류, 열정적인 기술 스타일, 규제 준수와 관련된 예시에서, 턴 평균 특징은 토큰 단위 특징이 놓쳤거나 부분적으로만 포착했던 고수준의 행동적 특성(예: "오답", "열정적인 기술 스타일")을 성공적으로 식별했습니다.
• 롱 컨텍스트 일반화: 채팅 턴(평균 길이 ~150 토큰)으로 학습된 턴 평균 SAE는 LongBench v2의 훨씬 긴 문서(24k–32k 토크)로도 성공적으로 일반화되어, 토큰 단위 특징이 감지하지 못한 "거시 경제학" 및 "컴퓨터 그래픽스"와 같은 고수준 개념을 식별했습니다.

4.2 정량적 평가

LLM-as-a-judge 프레임워크를 사용하여 특징 설명과 턴의 구조화된 요약을 비교한 결과:

• 변별력 (10-way matching): 토큰 단위 특징(max-pooling 적용)이 턴 평균 특징보다 우세했습니다(95.0% vs 73.9%). 이는 토큰 단위 특징이 텍스트를 구분하기 위해 특정 토큰 수준의 세부 사항에 의존하기 때문입니다.
• 커버리지 (Pairwise 및 5-way ranking): 턴 평균 특징이 토큰 단위 특징을 크게 앞질렀습니다(쌍별 비교에서 87.9% 승리). 이들은 턴의 고수준 특성에 대해 더 완전한 설명을 제공했습니다.
• 임베딩 지표: 턴 평균 특징은 구조화된 요약 필드와 가장 높은 코사인 유사도(0.663)를 기록하여, 의미적 내용에 대한 우수한 커버리지를 확인했습니다.

4.3 귀속 분석

• 그래프 압축성: 턴 평균 그래프는 10회의 턴이 있는 대화에서 노드 수를 약 128K(토큰 단위)에서 ~5K로 줄여 검사가 가능하게 만들었습니다.
• 사례 연구:
  • 시스템 프롬프트 귀속: 턴 평균 그래프는 어떤 시스템 프롬프트 단락이 특정 사용자 쿼리에 영향을 미치는지 정확히 식별했으며, 토큰 단위 그래프와 유사한 귀속 변화를 보여주었습니다.
  • 퇴행 추적 (Degeneration Tracing): 모델이 반복적인 텍스트로 퇴행하는 14회차 대화에서, 턴 평균 그래프는 일관된 인과 사슬을 추적했습니다: "AI 안전 경고문" $\to$ "핵무기" $\to$ "적대적 프롬프트 탐지" $\to$ "퇴행된 출력". 토큰 단위 그래프는 적대적 탐지 특징을 찾지 못하고 관련 없는 "화학 화합물" 같은 개념으로 귀속시켰습니다.
  • 개입 및 충분성: 에지 가중치는 개입 효과와 유의미한 상관관계를 보였습니다. 네스티드 SAE는 순수 토큰 단위 및 사후 집계된 토큰 단위 특징보다 높은 "완전성" 및 "교체" 충분성 점수를 달성했습니다.

5. 의의 및 주장

본 논문은 턴 평균 SAE가 해석 가능성 기술을 롱 컨텍스트 길이에 실용적으로 적용 가능하게 만든다고 주장합니다. 분석 단위를 토큰에서 스팬으로 전환함으로써 이 방법은 다음과 같은 효과를 냅니다:

• 토큰 단위 특징보다 고수준의 특성(주제, 페르소나, 안전성)을 더 완전하게 설명하는 특징을 생성합니다.
• 분석하기에는 너무 클 수 있는 다회차 대화에 대해 압축되고 해석 가능한 귀속 그래프 구축을 가능하게 합니다.
• 학습 신호를 수정(평균화 활용)함으로써, 토큰 수와 무관하게 원하는 입도(고수준 개념)의 특징을 학습하도록 SAE를 유도할 수 있음을 입증합니다.
• 턴 평균 특징이 학습 시 본 것보다 훨씬 긴 스팬에도 일반화됨을 보여줌으로써, 프런티어 모델 분석을 위한 견고함을 시사합니다.

저자들은 현재 실험 규모(7B 모델, LMSYS-Chat-1M 코퍼스)와 평가 지표로서의 LLM 의존도를 한계점으로 언급하며, 프런티어 모델에 대한 적용 및 그라운드 트루스 벤치마크 구축이 향후 과제임을 밝히고 있습니다.