peal.sparse_dictionaries.batch_topk.BatchTopKSAE

class peal.sparse_dictionaries.batch_topk.BatchTopKSAE

Bases: SparseDictionary

Sparse dictionary backed by the vendored BatchTopK SAE.

Activations are centred by mu (the training mean) before encoding and the mean is added back after decoding, so encode/decode work on raw activations of size config.act_size.

Parameters:

config (BatchTopKSAEConfig) – Hyper-parameters; config.cfg may override the whole internal config dict passed to the dependency.

Variables:

• sae (peal.sparse_dictionaries.batch_topk_network.BatchTopKNetwork) – The underlying model with W_enc, W_dec, b_enc, b_dec and (after training) threshold.

• mu (torch.Tensor) – Shape (act_size,): mean of the training activations.

Methods

decode	Reconstruct raw activations from sparse codes (adds mu back).
encode	Sparse codes of raw activations (thresholded, as at inference).
encode_no_threshold	Pre-activation codes without the learned threshold applied.
fit	Train the SAE with train_batch_topk_saes.
fit_from_activations	Everything fit_from_dataloaders does once the encoder has been run.
fit_from_dataloaders	Run feature_extractor over the dataloaders and fit on the result.
fit_with_evaluation	Train the SAE with periodic F1 evaluation against ground truth labels.
get_components	Return the encoder weight W_enc of shape (act_size, dict_size).
load_from_disk	Load a checkpoint written by save_on_disk onto config.device.
save_on_disk	Save W_enc, W_dec, b_enc, b_dec, mu and threshold (if present) as a CPU torch checkpoint.

Attributes

config

__init__(config=pydantic.BaseModel)

fit(X)

Train the SAE with train_batch_topk_saes.

Parameters:

X (torch.Tensor or ActivationStore) – Activations of shape (N, act_size) (wrapped into an ActivationStore) or an object with next_batch. Note that X is used as is; centring by mu is the caller’s job (fit_from_activations does it).

fit_with_evaluation(activation_store, ground_truth_labels, log_dir=None)

Train the SAE with periodic F1 evaluation against ground truth labels. Logs F1 scores for the first 10 components to TensorBoard.

Parameters:

• activation_store – ActivationStore wrapping the centered activations

• ground_truth_labels – Tensor (N, K) of ground truth labels

• log_dir – Directory for TensorBoard logs. If None, uses config base_path.

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

Run feature_extractor over the dataloaders and fit on the result.

Parameters:

• dataloaders (iterable of torch.utils.data.DataLoader) – Batches whose first element is the input and, optionally, whose second element holds ground-truth labels for evaluation.

• feature_extractor (callable) – Maps an input batch on config.device to activations of shape (B, act_size).

fit_from_activations(X, y=None)

Everything fit_from_dataloaders does once the encoder has been run.

Sets mu to the activation mean, wraps the centred activations in an ActivationStore and trains with fit_with_evaluation when labels are given, otherwise with fit.

Parameters:

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

• y (torch.Tensor, optional) – Ground-truth labels of shape (N,) or (N, K) used for the final SAE evaluation.

get_components()

Return the encoder weight W_enc of shape (act_size, dict_size).

save_on_disk(path)

Save W_enc, W_dec, b_enc, b_dec, mu and threshold (if present) as a CPU torch checkpoint.

Parameters:

path (str) – Target file; parent directories are created.

load_from_disk(path)

Load a checkpoint written by save_on_disk onto config.device.

Parameters:

path (str) – Checkpoint file.

encode(x)

Sparse codes of raw activations (thresholded, as at inference).

Parameters:

x (torch.Tensor) – Shape (B, act_size).

Returns:

Shape (B, dict_size).

Return type:

torch.Tensor

encode_no_threshold(x)

Pre-activation codes without the learned threshold applied.

Parameters:

x (torch.Tensor) – Shape (B, act_size).

Returns:

Shape (B, dict_size); first output of sae.compute_activations on the centred input.

Return type:

torch.Tensor

decode(acts)

Reconstruct raw activations from sparse codes (adds mu back).

Parameters:

acts (torch.Tensor) – Shape (B, dict_size).

Returns:

Shape (B, act_size).

Return type:

torch.Tensor