기술 요약: 희소 오토인코더(Sparse Autoencoders)를 통한 트랜스포머 일반화의 한계 추적

문제 정의

사전 학습된 트랜스포머는 놀라운 일반화 능력을 보여주지만, 미세한 오타, 적대적 탈옥(jailbreak), 또는 언어 스타일의 변화와 같은 분포 외(out-of-distribution, OOD) 데이터를 마주할 때 신뢰성이 크게 저하됩니다. 기존 머신러닝 이론은 훈련 데이터와 테스트 데이터가 동일한 분포를 따른다는 가정을 전제로 하지만, 이는 실제 배포 환경에서 거의 충족되지 않는 가정입니다. 대규모 언어 모델(LLM)은 인상적인 제로샷 및 퓨샷 학습 능력을 보여줌에도 불구하고, 분포 변화에 취약하여 불안정한 실패 모드, 잘못된 추론, 그리고 적대적 공격에 대한 취약성을 드러냅니다. 분포 변화를 탐지하기 위한 기존 방법들(예: 최대 소프트맥스 확률, 마할라노비스 거리)은 표면적인 통계나 모델을 고려하지 않은 접근 방식에 의존하는 경우가 많으며, 트랜스포머의 잠재 공간 내에서 발생하는 내부적인 메커니즘 변화를 포착하지 못합니다. 따라서 트랜스포머의 강건성(robustness)의 경계를 체계적으로 획정하고, OOD 입력이 내부 표현을 어떻게 변화시키는지에 대한 메커니즘적 이해를 제공할 수 있는 프레임워크가 절실히 필요합니다.

방법론

저자들은 LLM의 내부 계산 과정을 매핑하기 위한 진단 도구로서 **희소 오토인코더(Sparse Autoencoders, SAE)**를 활용하는 프레임워크를 제안합니다. 인간이 이해할 수 있는 추상 개념이 활성화 공간 내의 선형 방향으로 표현된다는 **선형 표현 가설(linear representation hypothesis)**에 근거하여, SAE는 이러한 개념들을 밀집된 트랜스포머 활성화 값으로부터 분리해내는 데 사용됩니다.

핵심 메커니즘

1. 현미경으로서의 SAE: 저자들은 SAE를 대상 LLM의 잔차 스트림(residual stream) 활성화 값에 대해 학습된 대리 모델로 취급합니다. SAE는 밀집된 활성화 값을 희소하게 활성화되는 과완전 기저(over-complete basis)인 "개념(concepts)"(잠재 특징)으로 투영합니다.
2. 에너지 점수(Energy Score)를 통한 OOD 탐지: 저자들은 분포 변화의 정도를 정량화하기 위해 에너지 점수라는 복합 지표를 도입했습니다. 이 점수는 두 가지 상호 보완적인 신호를 결합합니다:
   • 개념 개수 ($L_0$): SAE 잠재 공간에서 활성화되는 0이 아닌 의미론적 개념의 수입니다. OOD 입력은 이를 표현하기 위해 더 많은 가짜(spurious) 개념을 필요로 할 것이라고 가설을 세웠습니다.
   • 재구성 오차(Reconstruction Error): SAE의 재구성 값과 실제 잔차 스트림 활성화 값 사이의 평균 제곱 오차입니다. 학습된 매니폴드에서 벗어난 입력은 희소 선형 근사기에 의해 제대로 재구성되지 않습니다.
   • 에너지 점수($F(x)$)는 이러한 요소들을 정규화하여 입력이 얼마나 "매니폴드에서 벗어나 있는지(off-manifold)"를 나타내는 단일 스칼라 값을 제공합니다.
3. 실험 설정:
   • 토이 모델: 오타가 없는 TinyStories 코퍼스로 사전 학습된 캐릭터 단위 토큰화 GPT-2 (25M 파라미터). 오타를 다양한 비율로 유도하여 통제된 OOD 입력을 생성합니다.
   • 실제 모델: Llama 3.1 8B (Goodfire의 사전 학습된 SAE 사용) 및 API를 통해 접근 가능한 독점 모델(GPT-4o mini, GPT-5-thinking-nano).
   • 벤치마크: 성능 저하를 측정하기 위한 MMLU(및 오염되지 않은 MMLU-CF)와 안전 정렬 테스트를 위한 WildJailbreak 데이터셋.

주요 기여

1. 가짜 개념의 메커니즘적 식별: 본 연구는 LLM이 OOD 입력(예: 오타 또는 탈옥)을 마주할 때, 분포 내(in-distribution) 입력에 비해 훨씬 더 많은 수의 가짜(spurious), 오프-매니폴드(off-manifold) 개념을 추론한다는 것을 입증합니다. 이는 SAE 활성화의 증가된 $L_0$ 노름을 통해 추적됩니다.
2. 분포 변화의 정량화: 저자들은 SAE에 의해 탐지되는 미세한 분포 변화가 벤치마크 성능 저하와 직접적으로 상관관계가 있음을 보여줍니다. 오타의 비율이 증가함에 따라 활성화되는 개념의 수가 거의 선형적으로 증가하며, MMLU에 대한 모델 정확도는 감소합니다.
3. SAE 기반 미세 조정 전략: 높은 에너지 점수를 가진 샘플에 대해 선택적으로 미세 조정을 수행함으로써 LLM의 강건성을 개선하는 방법을 제시합니다. 오프-매니폴드 활성화(높은 재구성 오차 및 높은 개념 개수)를 유도하는 샘플을 우선시함으로써, 모델은 무작위 또는 낮은 에너지 샘플로 학습할 때보다 OOD 조건에 더 효율적으로 일반화하는 법을 배웁니다.
4. 특징 정렬을 통한 탈옥 완화: 저자들은 성공적인 탈옥 프롬프트가 별도의 "탈옥 전용" 개념들을 활성화한다는 것을 보여줍니다. 성공적인 탈옥의 특징 활성화를 실패한(거부된) 탈옥의 중심점으로 정렬하도록 경량 LoRA 어댑터를 사용함으로써, 일반적인 추론 능력을 희생하지 않고 탈옥 공격에 대한 모델의 취약성을 억제할 수 있습니다.

결과

• 성능 저하: GPT-2 및 Llama 3.1 8B 실험에서, 프롬프트에 오타를 도입하면 MMLU 정확도가 단조 감소했습니다. 예를 들어, Llama 3.1 8B의 정확도는 단 5%의 단어에 오타가 포함되었을 때 66.7%에서 51.01%로 떨어졌습니다. GPT-4o mini 및 GPT-5-thinking-nano와 같은 프런티어 모델들도 상당한 민감도를 보였습니다.
• 개념 확장: 입력이 더욱 OOD가 됨에 따라 활성화되는 고유한 SAE 특징의 수가 증가했습니다. GPT-2 토이 모델에서 모든 단어에 오타가 있는 입력은 깨끗한 입력보다 거의 30% 더 많은 특징을 동원했습니다.
• 효율적인 미세 조정: 에너지 점수의 상위 10% 샘플로 GPT-2를 미세 조정했을 때, 하위 10%를 사용하거나 표준 엔트로피/마할라노비스 기반 선택법을 사용했을 때보다 최종 검증 손실이 낮고 수렴 속도가 빨랐습니다. 고에너지 샘플을 통해 모델은 3분의 2 단계의 훈련 단계만으로도 유사한 손실 수준에 도달할 수 있었습니다.
• 탈옥 방어: Llama 3.1 8B에 SAE 기반 정렬 전략을 적용했을 때, 홀드아웃 테스트 세트에서의 탈옥 성공률이 46%에서 7%로 감소했습니다. 결정적으로, 이러한 개입은 MMLU 벤치마크 성능에 미미한 저하(0.09%)만을 일으켰으며, 정상적인 거부율(over-refusal)을 높이지 않았습니다.
• 스타일 변화에 대한 일반화: 본 방법은 글쓰기 스타일 변화(셰익스피어, 현대 시)를 OOD로 성공적으로 식별하였으며, 표적 미세 조정을 통해 이러한 스타일에 대한 일반화 능력을 향상시켰습니다.

의의 및 주장

본 논문은 트랜스포머 일반화의 한계를 이해하고 완화하기 위한 메커니즘적으로 근거가 있는 프레임워크를 제공한다고 주장합니다. OOD의 개념을 단순한 입력 데이터를 넘어 모델의 사적인 계산 과정으로 확장함으로써, 저자들은 추론 시 감사를 위한 새로운 진단 도구를 제공합니다.

그 의의는 OOD 탐지의 패러다임을 데이터 의존적 통계(엔트로피 등)에서 모델 내재적 매니폴드 분석으로 전환하는 데 있습니다. 저자들은 SAE를 통해 특정 취약점(예: 탈옥)을 "외과적으로" 교정하고, 내부 특징 활성을 재정렬함으로써 일반적인 추론 능력을 해치지 않고 AI 시스템을 강화할 수 있다고 주장합니다. 이 접근 방식은 특히 예상치 못한 조건에서의 신뢰성이 매우 중요한 과학, 비즈니스, 정부 분야의 배치와 같은 안전 필수 환경에서 AI 시스템을 안전하게 만드는 데 있어 중요한 단계로 제시됩니다.

저자들은 탐구한 OOD 시나리오(오타, 탈옥, 스타일)가 광범위하지만, 여전히 탐구되지 않은 많은 맥락이 남아 있음을 인정하며 연구 범위에 대해 겸허한 태도를 유지합니다. 또한, 그들의 프레임워크가 SAE의 전역적 매니폴드 근사 능력에 의존하며, 서로 다른 SAE 초기화에 따라 달라질 수 있는 특정 국소적 특징의 해석 가능성에 의존하지 않는다는 점을 명시합니다.