Temporal SAE integration

[canrager]

No description provided.

[hijohnnylin]

Loading the temporalSAE safetensor using SAE.from_pretrained or SAE.from_pretrained_with_cfg_and_sparsity fails currently due to:

• W_enc doesn't exist in the state dictionary (checked Gemma 2). I printed out the state_dict_raw keys in the pretrained_sae_loaders.py - temporal_sae_huggingface_loader, and I don't see "E":
  • dict_keys(['D', 'attn_layers.0.c_proj.bias', 'attn_layers.0.c_proj.weight', 'attn_layers.0.k_ctx.bias', 'attn_layers.0.k_ctx.weight', 'attn_layers.0.q_target.bias', 'attn_layers.0.q_target.weight', 'attn_layers.0.v_ctx.bias', 'attn_layers.0.v_ctx.weight', 'b'])
  • @canrager please verify and reupload the safetensor files
• b_dec shape is slightly off (we are expecting [2304] but b_dec is [1, 2304])
  • workaround: doing squeeze on the b_dec while importing in sae.py in the converter

[chanind]

Why is this needed? I'd rather avoid adding config options to the global SAE config if it's just for temporal SAEs. If constant_norm_rescale isn't used currently we should just delete it from the types IMO. Can you just fold the scaling factor into your temporal SAE weights when you load them so this isn't needed as a separate global SAE config option?