[논문 리뷰] BatchTopK Sparse Autoencoders

Bart Bussmann, Patrick Leask외 1명|arXiv (Cornell University)|2024. 12. 09.

Generative Adversarial Networks and Image Synthesis인용 수 4

한 줄 요약

이 논문은 표준 TopK SAE의 각 샘플별 상위-k 활성화 제약를 완화하기 위해 전체 배치에서 상위-k 활성화를 선택하는 방식으로, 각 샘플당 변동 가능한 잠재 활성화를 허용하는 BatchTopK 희박 오토에인커(BatchTopK Sparse Autoencoders)를 소개한다. 이 방법은 평균 희박성에 손실을 입히지 않은 채 복원 성능을 향상시키며, 직접적인 희박성 수준 지정이 가능해져 TopK와 비교해 성능이 뛰어나고 JumpReLU SAE와 유사한 복원 성능을 보이며, 비용이 많이 드는 하이퍼파라미터 튜닝을 피할 수 있다.

ABSTRACT

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting language model activations by decomposing them into sparse, interpretable features. A popular approach is the TopK SAE, that uses a fixed number of the most active latents per sample to reconstruct the model activations. We introduce BatchTopK SAEs, a training method that improves upon TopK SAEs by relaxing the top-k constraint to the batch-level, allowing for a variable number of latents to be active per sample. As a result, BatchTopK adaptively allocates more or fewer latents depending on the sample, improving reconstruction without sacrificing average sparsity. We show that BatchTopK SAEs consistently outperform TopK SAEs in reconstructing activations from GPT-2 Small and Gemma 2 2B, and achieve comparable performance to state-of-the-art JumpReLU SAEs. However, an advantage of BatchTopK is that the average number of latents can be directly specified, rather than approximately tuned through a costly hyperparameter sweep. We provide code for training and evaluating BatchTopK SAEs at https://github.com/bartbussmann/BatchTopK

연구 동기 및 목표

• 모든 입력이 동일한 수의 활성 잠재변수를 필요로 하는 고정된 각 샘플별 희박성의 제한을 해결하기 위해, 입력의 복잡성에 관계없이 동일한 수의 활성 잠재변수를 요구하는 TopK SAE의 한계를 해결한다.
• 입력의 복잡성에 따라 각 샘플의 활성화 수를 변동시켜 평균 희박성을 증가시키지 않고도 복원 성능을 향상시키는 것을 목표로 한다.
• 다른 SAE 변종에서 흔히 발생하는 비용이 많이 드는 하이퍼파라미터 스윕을 피하기 위해 평균 희박성 수준을 직접 지정할 수 있도록 하는 것.
• 배치 수준의 활성화 선택을 통해 해석 가능성은 유지하면서 복원 성능을 향상시키는 것.
• GPT-2 Small 및 Gemma 2 2B에서 표준 지표를 사용해 TopK 및 JumpReLU SAE와의 성능 비교를 통해 방법을 평가하는 것.

제안 방법

• BatchTopK는 각 샘플별 상위-k 선택을 배치 수준의 상위-k 선택으로 대체하여, 한 배치의 모든 샘플에서 상위-k 활성화를 복원에 사용한다.
• 학습 중 기울기 전파를 가능하게 하기 위해 상위-k 연산의 미분 가능 근사치를 사용하여 잠재 사전의 엔드 투 엔드 최적화를 가능하게 한다.
• 추론 시, 배치 간 최소 양수 활성화의 평균으로 전역 임계치 θ를 추정하고, 배치 기반 선택을 임계치 기반 활성화 함수로 대체한다.
• 임계치 θ는 다수의 배치에서 가장 작은 양수 활성화의 기대값으로 계산되어 추론 시 일관된 동작을 보장한다.
• 기본 SAE 인코더-디코더 아키텍처에 ReLU 활성화와 L2 복원 손실을 사용하며, 죽은 잠재변수 방지를 위한 희박성 페널티와 보조 손실을 추가로 적용한다.
• BatchTopK의 보조 손실은 TopK SAE와 동일하며, 죽은 잠재변수의 소규모 복원을 통해 그 유용성을 유지한다.

Figure 1: On GPT-2 Small activations, BatchTopK largely achieves better NMSE and CE than standard TopK across different dictionary sizes, for a fixed number of active latents of 32 (Left). JumpReLU SAEs are omitted from this comparison as their L0 cannot be fixed to a value. For fixed dictionary siz

실험 결과

연구 질문

• RQ1 각 샘플별 상위-k 제약를 배치 수준의 제약로 완화함으로써 희박 오토에인커의 복원 성능을 향상시킬 수 있는가?
• RQ2 각 샘플당 활성화 수를 변동시킴으로써 평균 희박성을 증가시키지 않고도 더 나은 복원 성능을 달성할 수 있는가?
• RQ3 BatchTopK는 성능이 뛰어난 최신 기술인 JumpReLU SAE와 유사한 성능을 달성하면서도 평균 희박성 수준을 직접 지정할 수 있는가?
• RQ4 BatchTopK와 고정 상위-k 방법 간의 샘플당 활성화 분포는 어떻게 다른가?
• RQ5 BatchTopK의 향상된 유연성이 교차 엔트로피 감소 측면에서 더 나은 최종 성능을 이끌어내는가?

주요 결과

• GPT-2 Small 및 Gemma 2 2B에서 테스트된 모든 사전 크기와 희박성 수준에서 BatchTopK SAE는 TopK SAE보다 유의미하게 낮은 정규화된 평균 제곱오차(NMSE)를 기록한다.
• 고정된 활성 잠재변수 수(L0=32) 조건에서, 두 모델 모두에서 BatchTopK는 NMSE와 교차 엔트로피(CE) 감소 측면에서 TopK를 능가한다.
• GPT-2 Small에서, BatchTopK는 모든 희박성 수준에서 NMSE와 CE 감소 측면에서 JumpReLU SAE를 능가한다. 이는 JumpReLU가 고활성 영역에서 더 뛰어난 최종 성능 안정성을 보이지만도 불구하고이다.
• Gemma 2 2B에서, BatchTopK는 복원(NMSE) 측면에서 JumpReLU와 동등하거나 이를 초월하지만, CE 감소 측면에서는 오직 가장 낮은 희박성 수준(k=16)에서만 승리한다.
• BatchTopK의 샘플당 활성화 분포는 넓은 변동성을 보이며, 일부 샘플은 단 한 개의 잠재변수만 사용하고 다른 샘플은 80개 이상을 사용함으로써, 이 방법이 적응형 활성화 전략을 구현하고 있음을 확인한다.
• BatchTopK는 JumpReLU SAE가 목표 희박성에 도달하기 위해 다수의 파rameter를 튜닝해야 하는 데 반해, 하이퍼파라미터 스윕 없이도 평균 희박성 수준을 직접 지정할 수 있다.

Figure 2: On Gemma 2 2B activations, BatchTopK outperforms TopK SAEs across different values of k. Although BatchTopK has a better reconstruction performance (left), it only outperforms JumpReLU in terms of downstream CE degradation in the setting where k=16 (right).