SAEs for use with the SAELens library

This repository contains the following batch topk Matryoshka SAEs for Gemma-2-2b. All SAEs have 32k width and are trained with k=40 on 750M tokens from the Pile using SAELens. The SAEs were trained with Matryoshka layers of width 128, 512, 2048, 8192, and 32768. This release contains both standard Matryoshka SAE and snap loss Matryoshka SAEs.

This repository contains the following SAEs:

Snap Loss Matryoshka SAEs

layer	SAE ID	width	l0	explained variance
0	snap/blocks.0.hook_resid_post	32768	40	0.919964
1	snap/blocks.1.hook_resid_post	32768	40	0.863969
2	snap/blocks.2.hook_resid_post	32768	40	0.858767
3	snap/blocks.3.hook_resid_post	32768	40	0.815844
4	snap/blocks.4.hook_resid_post	32768	40	0.821094
5	snap/blocks.5.hook_resid_post	32768	40	0.797083
6	snap/blocks.6.hook_resid_post	32768	40	0.79815
7	snap/blocks.7.hook_resid_post	32768	40	0.78946
8	snap/blocks.8.hook_resid_post	32768	40	0.779236
9	snap/blocks.9.hook_resid_post	32768	40	0.759022
10	snap/blocks.10.hook_resid_post	32768	40	0.743998
11	snap/blocks.11.hook_resid_post	32768	40	0.731758
12	snap/blocks.12.hook_resid_post	32768	40	0.725974
13	snap/blocks.13.hook_resid_post	32768	40	0.727936
14	snap/blocks.14.hook_resid_post	32768	40	0.727065
15	snap/blocks.15.hook_resid_post	32768	40	0.757408
16	snap/blocks.16.hook_resid_post	32768	40	0.751874
17	snap/blocks.17.hook_resid_post	32768	40	0.763654
18	snap/blocks.18.hook_resid_post	32768	40	0.77644
19	snap/blocks.19.hook_resid_post	32768	40	0.768622
20	snap/blocks.20.hook_resid_post	32768	40	0.761658
21	snap/blocks.21.hook_resid_post	32768	40	0.765593
22	snap/blocks.22.hook_resid_post	32768	40	0.741098
23	snap/blocks.23.hook_resid_post	32768	40	0.729718
24	snap/blocks.24.hook_resid_post	32768	40	0.754838

Standard Matryoshka SAEs

layer	SAE ID	width	l0	explained variance
0	standard/blocks.0.hook_resid_post	32768	40	0.91832
1	standard/blocks.1.hook_resid_post	32768	40	0.863454
2	standard/blocks.2.hook_resid_post	32768	40	0.841324
3	standard/blocks.3.hook_resid_post	32768	40	0.814794
4	standard/blocks.4.hook_resid_post	32768	40	0.820418
5	standard/blocks.5.hook_resid_post	32768	40	0.796252
6	standard/blocks.6.hook_resid_post	32768	40	0.797322
7	standard/blocks.7.hook_resid_post	32768	40	0.787601
8	standard/blocks.8.hook_resid_post	32768	40	0.779433
9	standard/blocks.9.hook_resid_post	32768	40	0.75697
10	standard/blocks.10.hook_resid_post	32768	40	0.745011
11	standard/blocks.11.hook_resid_post	32768	40	0.732177
12	standard/blocks.12.hook_resid_post	32768	40	0.726209
13	standard/blocks.13.hook_resid_post	32768	40	0.719405
14	standard/blocks.14.hook_resid_post	32768	40	0.719056
15	standard/blocks.15.hook_resid_post	32768	40	0.756888
16	standard/blocks.16.hook_resid_post	32768	40	0.742889
17	standard/blocks.17.hook_resid_post	32768	40	0.757294
18	standard/blocks.18.hook_resid_post	32768	40	0.76921
19	standard/blocks.19.hook_resid_post	32768	40	0.766661
20	standard/blocks.20.hook_resid_post	32768	40	0.760939
21	standard/blocks.21.hook_resid_post	32768	40	0.759883
22	standard/blocks.22.hook_resid_post	32768	40	0.740612
23	standard/blocks.23.hook_resid_post	32768	40	0.729678
24	standard/blocks.24.hook_resid_post	32768	40	0.747313

Load these SAEs using SAELens as below:

from sae_lens import SAE

sae, cfg_dict, sparsity = SAE.from_pretrained("chanind/gemma-2-2b-batch-topk-matryoshka-saes-w-32k-l0-40", "<sae_id>")