msae-32k-2L-mlp_in-v1

2-level Matryoshka BatchTopK SAE (dict=32768, k=48, groups=[1024,31744]) trained on layer 6 mlp_in of vanilla-loss-parity-v2. Inner group (1024) mirrors CayleySAE L0 size.

Training details

• trainer: MatryoshkaBatchTopKTrainer
• dict_size: 32768
• k: 48
• group_sizes: [1024, 31744]
• activation_dim: 1024
• layer: 6
• location: mlp_in
• base_model: hf_cache/vanilla-loss-parity-v2
• steps: 97656
• lr: 5e-05

Usage

from dictionary_learning.trainers.matryoshka_batch_top_k import MatryoshkaBatchTopKSAE
sae = MatryoshkaBatchTopKSAE.from_pretrained('ae.pt', device='cuda')
# encode: sae(x) returns (recon, f, loss_dict)
# decode: sae.decode(f)

Part of a research project comparing Matryoshka SAEs and CayleySAEs on hierarchy properties. See evals/msae/ in sparse-nanogpt-private.