peal.sparse_dictionaries.msae_decomposition.MSAEDecomposition

class peal.sparse_dictionaries.msae_decomposition.MSAEDecomposition

Bases: SparseDictionary

Sparse dictionary wrapper around MSAE (Matryoshka Sparse Autoencoder).

Supports pretrained MSAE checkpoints as well as fitting from scratch on dataloaders. When config.ckpt_url is set the constructor downloads and loads the pretrained model immediately, falling back to trainable mode (with a printed message) if that fails.

Parameters:

config (MSAEDecompositionConfig, optional) – Checkpoint URLs and training hyper-parameters.

Variables:

• sae (torch.nn.Module or None) – The MSAE SAE wrapper (pretrained) or a trained Autoencoder.

• vocabulary (list of str or None) – One concept name per latent, from the concept interpreter.

• mu (torch.Tensor or None) – Activation mean subtracted before encoding a self-trained model.

Methods

encode	Encodes activations into sparse MSAE activations.
fit	Train MSAE from raw activation tensor X [N, act_size].
fit_from_dataloaders	Encode every batch of the loaders with feature_extractor, then fit.
get_components	Returns decoder weights of shape [n_inputs, n_components].
get_vocabulary	Return the per-latent concept names, or None when not loaded.
load_from_disk	Rebuild a self-trained autoencoder from a save_on_disk() file.
save_on_disk	Write state_dict, mu and the config to path with torch.save.

Attributes

config

__init__(config=pydantic.BaseModel)

Pick a device and, if config.ckpt_url is set, load the pretrained MSAE.

A failure to download or load the checkpoint is caught and reported; the instance then stays in trainable mode with sae unset.

get_components()

Returns decoder weights of shape [n_inputs, n_components].

Return type:

torch.Tensor

get_vocabulary()

Return the per-latent concept names, or None when not loaded.

Return type:

List[str]

encode(x, topk=64)

Encodes activations into sparse MSAE activations.

Parameters:

• x (torch.Tensor) – Activations of shape (n, act_size).

• topk (int, optional) – Number of active latents for the pretrained SAE.encode; ignored by a self-trained autoencoder, which instead subtracts mu first. Default 64.

Returns:

Sparse latents of shape (n, n_components).

Return type:

torch.Tensor

fit(X, y=None)

Train MSAE from raw activation tensor X [N, act_size].

Centres X (storing the mean in mu), builds an Autoencoder or MatryoshkaAutoencoder sized by config.n_components (default 2 * act_size) and trains it with AdamW and an MSE SAELoss for config.epochs epochs, printing the loss every five epochs.

Parameters:

• X (torch.Tensor) – Activations of shape (N, act_size).

• y (torch.Tensor, optional) – Ignored; accepted for interface symmetry.

fit_from_dataloaders(dataloaders, feature_extractor=<function MSAEDecomposition.<lambda>>)

Encode every batch of the loaders with feature_extractor, then fit.

Parameters:

• dataloaders (iterable of torch.utils.data.DataLoader) – Loaders whose first batch element is the input.

• feature_extractor (callable, optional) – Maps a batch to (batch, act_size) activations; identity by default.

save_on_disk(path)

Write state_dict, mu and the config to path with torch.save.

Parameters:

path (str)

load_from_disk(path)

Rebuild a self-trained autoencoder from a save_on_disk() file.

Does nothing when path does not exist. The architecture is re-derived from the current config (act_size, n_components, activation, use_matryoshka), not from the saved config.

Parameters:

path (str)