peal.sparse_dictionaries.ra_sae_decomposition.RASAEDecomposition¶

class peal.sparse_dictionaries.ra_sae_decomposition.RASAEDecomposition[source]¶

Bases: SparseDictionary

Sparse dictionary wrapper around a pretrained RA-SAE checkpoint.

Pretrained-only: fit/fit_from_dataloaders are deliberately not implemented, since the point of this dictionary is that it is not fitted on our activations.

Methods

`decode`	Reconstruct token activations from sparse codes.
`encode`	Sparse codes for [N, 768] tokens, as overcomplete computes them in eval.
`fit`	Not supported: the dictionary is pretrained (raises `NotImplementedError`).
`fit_from_dataloaders`	Not supported: the dictionary is pretrained (raises `NotImplementedError`).
`get_components`	Decoder atoms as [n_inputs, n_components], the layout PEAL expects.
`get_vocabulary`	RA-SAE ships no concept names -- unlike MSAE there is no text-space interpreter, so concepts are identified by ground-truth matching or by their maximally-activating images, never by a vocabulary lookup.
`load_from_disk`	Restore tensors written by `save_on_disk()`; a missing path is a no-op.
`save_on_disk`	Save the fused tensors and the threshold to `path` with `torch.save`.

Attributes

config

__init__(config=pydantic.BaseModel)[source]¶

Pick the device, allocate the tensor slots and load the fused weights.

Parameters:: config (RASAEDecompositionConfig, optional) – Download/cache locations and the activation threshold. Its n_components and act_size are overwritten by the values read from the checkpoint.

get_components()[source]¶

Decoder atoms as [n_inputs, n_components], the layout PEAL expects.

Return type:: torch.Tensor

get_vocabulary()[source]¶

RA-SAE ships no concept names – unlike MSAE there is no text-space interpreter, so concepts are identified by ground-truth matching or by their maximally-activating images, never by a vocabulary lookup.

Return type:: List[str] | None

encode(x)[source]¶

Sparse codes for [N, 768] tokens, as overcomplete computes them in eval.

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

decode(codes)[source]¶

Reconstruct token activations from sparse codes.

Parameters:: codes (torch.Tensor)
Return type:: torch.Tensor

fit(X, y=None)[source]¶

Not supported: the dictionary is pretrained (raises NotImplementedError).

Parameters:

X (torch.Tensor)
y (torch.Tensor)

fit_from_dataloaders(dataloaders, feature_extractor=<function RASAEDecomposition.<lambda>>)[source]¶: Not supported: the dictionary is pretrained (raises NotImplementedError).

save_on_disk(path)[source]¶

Save the fused tensors and the threshold to path with torch.save.

Parameters:: path (str) – Destination file; parent directories are created.

load_from_disk(path)[source]¶

Restore tensors written by save_on_disk(); a missing path is a no-op.

Parameters:: path (str) – File written by save_on_disk(). Updates config.n_components, config.act_size and, when stored, the threshold.