peal.sparse_dictionaries.batch_topk_network¶

PEAL’s own batch top-K sparse autoencoder, its trainer and an activation source.

This replaces the vendored peal.dependencies.matryoshka_sae backend, which was a copy of https://github.com/bartbussmann/matryoshka_sae. That repository publishes no licence file, so all rights stay with its authors by default and PEAL could neither redistribute it nor ship as an openly licensed library while depending on it (see THIRD_PARTY_NOTICES.md).

The method is Bussmann, Leask and Nanda, BatchTopK Sparse Autoencoders (2024); methods are not copyrightable, only a particular expression of them, and this is PEAL’s expression of it, LGPL-3.0 like the rest of PEAL. The two MIT reimplementations (SAELens batchtopk_sae.py and dictionary_learning batch_top_k.py) were not usable as dependencies: both declare requires-python >= 3.10 while PEAL runs on 3.9, and both pull a language model stack (transformer-lens / nnsight, datasets) for what is a 512-dimensional image-representation dictionary.

Compatibility is deliberate and verified: parameter names, shapes, the initialisation order and the forward pass match the previous backend, so checkpoints written before the swap load unchanged and the published Sparse Numbers dictionary keeps its meaning.

What batch top-K does: a plain top-K autoencoder keeps the K largest codes per sample, which forces every sample to use exactly K atoms. Batch top-K instead keeps the K * batch_size largest codes across the whole batch, so the budget is shared and a sample may use more or fewer than K atoms while the average stays K. At inference there is no batch to pool over, so a running threshold learned during training replaces the top-K selection.

Functions

`get_default_cfg`	Return the default trainer/model config dict.
`train_batch_topk_saes`	Train one or more autoencoders on batches drawn from `activation_store`.

Classes

`ActivationStore`	Minimal `next_batch()` source over an in-memory activation tensor.
`BatchTopKNetwork`	Batch top-K sparse autoencoder.
`SparseAutoencoderBase`	Shared parameters and housekeeping of a tied-init sparse autoencoder.