peal.sparse_dictionaries.ra_sae_decomposition¶
Pretrained RA-SAE (Relaxed Archetypal SAE) dictionaries of DINOv2 token space.
Wraps https://huggingface.co/matybohacek/RA-SAE-DINOv2-32k – a 32000-concept BatchTopK SAE with an archetypal dictionary (Fel et al. 2025, arXiv:2502.12892), trained on the ~261 tokens per image that dinov2_vitb14_reg produces.
The published checkpoint is a 4.4 GB pickled object graph (torch.load(…, weights_only=False) of a DataParallel(BatchTopKSAE)), which would drag the overcomplete package in as an import-time dependency just to unpickle. Instead the loader below reads it once behind stub classes, fuses the archetypal dictionary, and caches ~300 MB of plain tensors; every later run loads only that cache and never touches the original file or overcomplete.
Inference reproduces overcomplete 0.3.0 exactly:
pre_z = LayerNorm(Linear(x)) # MLPEncoder.final_block z = ReLU(pre_z) # MLPEncoder.final_activation codes = z * (z >= running_threshold) # BatchTopKSAE.encode, eval branch x_hat = codes @ D # RelaxedArchetypalDictionary.forward D = exp(multiplier) * (rowstochastic(relu(W)) @ C + Relax)
running_threshold is not stored in the checkpoint (it is None there); the model card’s config.json supplies it, which is also what the HF wrapper sets before calling .eval().
Classes
Sparse dictionary wrapper around a pretrained RA-SAE checkpoint. |
|
Configuration for a pretrained RA-SAE decomposition. |