peal.sparse_dictionaries.batch_topk.BatchTopKSAEConfig

class peal.sparse_dictionaries.batch_topk.BatchTopKSAEConfig

Bases: SparseDictionaryConfig

Config of BatchTopKSAE.

Parameters:

• n_components (int) – Dictionary size (number of SAE latents).

• top_k (int) – Number of active latents kept per batch on average (BatchTopK).

• lr (float) – Adam learning rate and the L1 / orthogonality penalty weights.

• l1_coeff (float) – Adam learning rate and the L1 / orthogonality penalty weights.

• orth_coeff (float) – Adam learning rate and the L1 / orthogonality penalty weights.

• base_path (str) – Directory for TensorBoard logs and evaluation outputs.

• weights_path (str) – Where fit_with_evaluation saves the trained weights.

• device (str) – Device the SAE trains and lives on.

• batch_size (int) – Activation batch size and total number of activations processed; num_tokens // batch_size is the number of training steps.

• num_tokens (int) – Activation batch size and total number of activations processed; num_tokens // batch_size is the number of training steps.

• n_batches_to_dead (int, int, float) – Dead-latent detection window and the auxiliary loss that revives dead latents.

• top_k_aux (int, int, float) – Dead-latent detection window and the auxiliary loss that revives dead latents.

• aux_penalty (int, int, float) – Dead-latent detection window and the auxiliary loss that revives dead latents.

• input_unit_norm (bool) – Whether the dependency normalises inputs to unit norm.

• eval_epoch_interval (int) – Step interval of the (currently disabled) periodic F1 evaluation.

• cfg (dict) – If non-empty, used verbatim as the dependency’s config instead of being assembled from the fields above.

Methods

Attributes

aux_penalty
base_path
batch_size
cfg
comp_max
comp_min
device
dict_size
dtype
eval_epoch_interval
input_unit_norm
l1_coeff
lr
model_batch_size
n_batches_to_dead
n_components
num_tokens
orth_coeff
seed
sparse_dictionaries_type
top_k
top_k_aux
weights_path

n_components: int = 10016

dict_size: int | None = None

sparse_dictionaries_type: str = 'BatchTopKSAE'

top_k: int = 12

lr: float = 0.0003

l1_coeff: float = 0.0

orth_coeff: float = 0.0

base_path: str = None

weights_path: str = None

comp_min: int = -1

comp_max: int = -1

device: str = 'cuda'

model_batch_size: int = 32

batch_size: int = 4096

num_tokens: int = 1000000

n_batches_to_dead: int = 100

top_k_aux: int = 4

aux_penalty: float = 0.03125

dtype: str = 'torch.float32'

seed: int = 42

input_unit_norm: bool = False

eval_epoch_interval: int = 50

cfg: dict = {}