Scaling and evaluating sparse autoencoders

Sparse autoencoders provide a promising unsupervised approach for extracting interpretable features from a language model by reconstructing activations from a sparse bottleneck layer. Since language models learn many concepts, autoencoders need to be very large to recover all relevant features. However, studying the properties of autoencoder scaling is difficult due to the need to balance reconstruction and sparsity objectives and the presence of dead latents. We propose using k-sparse autoencoders [Makhzani and Frey, 2013] to directly control sparsity, simplifying tuning and improving the reconstruction-sparsity frontier. Additionally, we find modifications that result in few dead latents, even at the largest scales we tried. Using these techniques, we find clean scaling laws with respect to autoencoder size and sparsity. We also introduce several new metrics for evaluating feature quality based on the recovery of hypothesized features, the explainability of activation patterns, and the sparsity of downstream effects. These metrics all generally improve with autoencoder size. To demonstrate the scalability of our approach, we train a 16 million latent autoencoder on GPT-4 activations for 40 billion tokens. We release training code and autoencoders for open-source models, as well as a visualizer.

한 줄 요약

GPT-4 활성화를 400억 토큰으로 학습한 1600만 레이턴트 SAE를 통해 체계적인 스케일링 법칙과 새로운 평가 지표를 제시한다.

핵심 기여도

k-sparse autoencoder를 사용해 sparsity를 직접 제어, reconstruction-sparsity trade-off 개선.
1600만 레이턴트 SAE를 GPT-4 활성화에 학습, 400억 토큰 기반.
Downstream loss, Probe loss, Explainability, Ablation sparsity 등 4가지 새로운 평가 지표 제안.
레이턴트 수 증가에 따라 대부분의 지표가 개선됨을 실험적으로 입증.

핵심 아이디어

기존 SAE는 reconstruction loss와 sparsity loss 간 균형 조절이 어려웠으며, 데드 레이턴트(dead latent) 문제로 인해 대규모 학습이 불안정했다. 본 연구는 k-sparse autoencoder를 도입하여 각 샘플당 정확히 k개의 레이턴트만 활성화되도록 강제함으로써 sparsity를 직접 제어한다. 이를 통해 hyperparameter 조정이 용이해지고, reconstruction-스파르시티 trade-off가 개선된다. 또한, 데드 레이턴트를 최소화하는 기법을 제안해 대규모 SAE 학습을 안정적으로 수행할 수 있도록 했다.

기술적 접근법

**모델**: k-sparse autoencoder (Makhzani and Frey, 2013) 사용.
**데이터**: GPT-4의 residual stream 활성화, 400억 토큰 기반.
**하이퍼파라미터**: 레이턴트 수 16,000,000, k=16.
**학습 기법**: 데드 레이턴트 최소화를 위한 수정된 손실 함수와 초기화 전략 적용.
**평가 지표**: Downstream loss, Probe loss, Explainability, Ablation sparsity.

주요 결과

**Downstream loss**: 레이턴트 수 증가에 따라 GPT-4의 loss 감소 (명시되지 않음).
**Probe loss**: 레이턴트 수 증가에 따라 특정 가설된 특징의 복원률 향상 (명시되지 않음).
**Explainability**: 레이턴트 수 증가에 따라 설명 가능성 향상.
**Ablation sparsity**: 레이턴트 수 증가에 따라 개별 레이턴트 제거 시 downstream logits의 영향이 sparser 해짐.
**레이턴트 수 L₀**: L₀ 증가 시 Probe loss와 Ablation sparsity 향상, 그러나 Explainability 저하. L₀가 모델 차원 $d_{\text{model}}$에 가까워질수록 레이턴트가 밀집됨.

의의 및 한계

본 연구는 SAE의 대규모 학습과 평가 방법론을 체계적으로 제시하며, GPT-4 활성화를 기반으로 1600만 레이턴트 SAE를 성공적으로 학습함으로써 기존 연구의 한계를 극복했다. 또한, 기존의 reconstruction-스파르시티 trade-off에만 의존하지 않고, 실제 응용에 유용한 feature 복원 능력을 평가하는 새로운 지표를 제안한 점이 학술적·실용적 가치를 높인다. 한편, L₀가 $d_{\text{model}}$에 가까워질 경우 레이턴트가 밀집되어 SAE의 해석 가능성 저하라는 한계가 존재한다.

실용적 활용

본 연구는 언어 모델의 내부 작동 원리를 해석하는 데 활용 가능한 SAE 학습 기법을 제시하며, 특히 GPT-4와 같은 대형 모델의 활성화를 분석하는 데 유용하다. 개방된 코드와 시각화 도구를 제공함으로써 연구자들이 모델 해석과 feature 추출을 보다 쉽게 수행할 수 있도록 지원한다.