[논문 리뷰] Scaling and evaluating sparse autoencoders

Leo Gao, Tom Dupré la Tour외 7명|arXiv (Cornell University)|2024. 06. 06.

Generative Adversarial Networks and Image Synthesis인용 수 10

한 줄 요약

이 OpenAI 연구는 재구성-희소성 트레이드오프를 분석하기 위한 확장 가능한 TopK-sparse 자동인코더 프레임워크를 개발하고, 새로운 품질 지표를 도입하며, 16M-latent 자동인코더를 GPT-4 활성화에 대해 학습시켜 스케일링 법칙을 연구한다.

ABSTRACT

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-2와 GPT-4를 가로지르는 희소성, 자동인코더 크기, 피실험 모델 크기에 대한 스케일링 법칙을 특성화한다.
• 잠재 특징의 질을 측정하는 새로운 지표를 제안하고 검증한다(특징 회복, 해석 가능성, 다운스트림 영향 포함).

제안 방법

• L1 패널티 없이 희소성을 직접 제어하기 위해 TopK (k-sparse) 자동인코더를 채택하고 순수 L2 손실로 재구성 성능을 평가한다.
• 죽은 잠재를 방지하기 위해 초기화 및 보조 손실을 사용하여 대규모에서도 낮은 죽은 잠재 비율을 달성한다(예: 16M-latent 모델의 경우 7%).
• GPT-2 작은 모델과 GPT-4 계열 활성화에서 자동인코더의 크기(n 잠재)와 희소성(k)을 체계적으로 확장하여 MSE와 희소성에 대한 스케일링 법칙을 도출한다.
• TopK를 ReLU 및 기타 활성화와 비교하여 희소성-재구성 프런티어를 개선하고 활성화 수축을 줄이는 것을 보인다.
• MSE를 넘어선 평가 지표(다운스트림 손실, 프로브 손실, 해석 가능성, 제거-희소성) 정의 및 계산하여 특징의 품질을 평가한다.
• GPT-4 잔류 스트림에서 40B 토큰으로 학습된 16M-latent 자동인코더를 통해 확장성을 입증한다.

Figure 1 : Scaling laws for TopK autoencoders trained on GPT-4 activations. (Left) Optimal loss for a fixed compute budget. (Right) Joint scaling law of loss at convergence with fixed number of total latents $n$ and fixed sparsity (number of active latents) $k$ . Details in Section 3 .

실험 결과

연구 질문

• RQ1 희소성과 자동인코더 크기가 대형 언어 모델 전반의 재구성 품질과 잠재 활용도에 어떤 상호 작용을 보이는가?
• RQ2 더 크고 희소한 자동인코더가 다운스트림 예측 성능과 해석 가능한 특징 회복을 개선하는가?
• RQ3 언어 모델에서 희소 자동인코더가 학습한 특징의 질과 해석 가능성을 가장 잘 포착하는 지표는 무엇인가?
• RQ4 TopK-희소 자동인코더가 대규모에서 죽은 잠재를 완화할 수 있는가, 그리고 이 선택은 대안 활성화와 어떻게 비교되는가?
• RQ5 주제 모델 크기(GPT-2 대 GPT-4)와 고정된 희소성에서 스케일링 법칙은 어떻게 달라지는가?

주요 결과

• TopK 자동인코더는 희소성-재구성 프런티어에서 ReLU 및 ProLU 베이스라인보다 우수하며 활성 예시의 단의-의미성(monosemanticity)이 더 좋다.
• Encoder-Decoder 초기화 및 보조 손실을 사용하여 죽은 잠재를 크게 줄이며, 가장 큰 16M-latent 모델에서 남은 죽은 잠재는 7%에 불과하다.
• MSE에 대해 자동인코더 크기 n과 희소성 k에 대한 깨끗한 스케일링 법칙이 나타나며, GPT-4에 대해 공동 L(n,k) 관계가 포함된다.
• 더 큰 자동인코더는 일반적으로 GPT-2 소형 및 GPT-4 활성화 전반에서 다운스트림 손실, 프로브 손실 및 해석 가능성 지표를 개선한다.
• 40B 토큰에서 학습된 16M-latent GPT-4 자동인코더가 특정 비교에서 GPT-4 사전훈련 컴퓨트의 10%에 비해 다운스트룸 성능에 근접한 결과를 달성한다.
• 새로운 평가 지표(다운스트림 손실, 프로브 손실, 해석 가능성, 제거-희소성)가 더 크고 더 희소한 모델에서 개선을 보인다.

(a) At a fixed number of latents ( $n=32768$ ), TopK has a better reconstruction-sparsity trade off than ReLU and ProLU, and is comparable to Gated.