Learning Multi-Level Features with Matryoshka Sparse Autoencoders

Sparse autoencoders (SAEs) have emerged as a powerful tool for interpreting neural networks by extracting the concepts represented in their activations. However, choosing the size of the SAE dictionary (i.e. number of learned concepts) creates a tension: as dictionary size increases to capture more relevant concepts, sparsity incentivizes features to be split or absorbed into more specific features, leaving high-level features missing or warped. We introduce Matryoshka SAEs, a novel variant that addresses these issues by simultaneously training multiple nested dictionaries of increasing size, forcing the smaller dictionaries to independently reconstruct the inputs without using the larger dictionaries. This organizes features hierarchically - the smaller dictionaries learn general concepts, while the larger dictionaries learn more specific concepts, without incentive to absorb the high-level features. We train Matryoshka SAEs on Gemma-2-2B and TinyStories and find superior performance on sparse probing and targeted concept erasure tasks, more disentangled concept representations, and reduced feature absorption. While there is a minor tradeoff with reconstruction performance, we believe Matryoshka SAEs are a superior alternative for practical tasks, as they enable training arbitrarily large SAEs while retaining interpretable features at different levels of abstraction.

한 줄 요약

Matryoshka SAE는 계층적 특징 학습을 통해 고수준 개념을 유지하면서도 SAE의 확장 문제를 해결하는 새로운 접근법이다.

핵심 기여도

Matryoshka SAE는 중첩된 사전을 동시에 학습하여 고수준 특징의 흡수를 감소시킨다.
Gemma-2-2B와 TinyStories 데이터셋에서 희소 탐색 및 개념 삭제 작업에서 기존 SAE 대비 우수한 성능을 보인다.
최대 코사인 유사도를 기준으로 더 분리된 잠재 표현을 학습한다.
SAE의 확장 시 발생하는 특징 분할 및 병합 문제를 훈련 절차의 간단한 수정으로 완화한다.

핵심 아이디어

기존 SAE는 사전 크기가 커질수록 특징이 분할되거나 흡수되어 고수준 개념이 손실되는 문제가 있었다. Matryoshka SAE는 이 문제를 해결하기 위해 여러 중첩된 사전을 동시에 학습하는 방식을 도입했다. 작은 사전부터 입력을 독립적으로 재구성하도록 강제함으로써, 고수준 개념이 하위 개념에 흡수되지 않고 유지된다. 이는 계층적 특징 학습을 유도하며, ReLU 활성화 함수를 사용하는 기존 SAE와 동일한 아키텍처를 유지하면서도 훈련 오브젝티브를 재구성한다. 핵심 아이디어는 "중첩된 사전이 독립적으로 재구성"이라는 제약을 통해 특징의 계층 구조를 강제하는 것이다.

기술적 접근법

**Matryoshka SAE**: 여러 중첩된 SAE를 동시에 학습. 작은 사전부터 입력을 재구성하도록 훈련.
**Input/Output 차원**: 20, **Latent 수**: 20, **활성화 함수**: ReLU
**중첩 구조**: 각 작은 사전은 전체 잠재 공간의 일부만 사용.
**훈련 절차**: 각 중첩된 SAE는 독립적으로 입력을 재구성해야 하며, 이는 고수준 특징의 흡수를 방지.
**Sparsity 제약**: 기존 SAE와 동일한 희소성 제약을 유지하지만, 계층적 구조로 인해 특징 분할이 감소.

주요 결과

**Gemma-2-2B와 TinyStories 데이터셋**에서 Matryoshka SAE는 기존 SAE 대비 **더 분리된 개념 표현**을 보임.
**최대 코사인 유사도** 기준으로 잠재 표현의 분리도가 향상됨.
**희소 탐색 및 개념 삭제 작업**에서 **성능 개선**을 관찰.
**재구성 오류**는 약간 증가했으나, **고수준 개념의 유지**로 인해 실용적 활용성이 높음.

의의 및 한계

Matryoshka SAE는 SAE의 확장 시 발생하는 고수준 개념 손실 문제를 해결하며, 해석 가능성과 성능을 동시에 유지할 수 있는 새로운 방향을 제시한다. 특히, 기존 SAE가 훈련 손실은 줄이지만 실제 작업 성능이 저하되는 문제를 완화한다. 그러나 재구성 오류가 약간 증가하는 점은 한계로 작용할 수 있다. 또한, 중첩된 사전 구조는 훈련 시간이 증가할 수 있으나, 본 연구에서는 "훈련 시간 증가가 적절한 수준"이라고 언급하며 실용성을 강조한다.

실용적 활용

Matryoshka SAE는 모델 해석, 특징 회로 분석, 개념 기반 모델 조정 등 다양한 해석 가능성 연구에 적용 가능하다. 특히, Gemma-2-2B와 같은 대규모 모델에서 고수준 개념을 유지하면서도 세부 특징을 학습할 수 있어, NLP 및 시각 모델의 해석성 향상에 유용할 것으로 기대된다.