sae-lens 5.10.3__py3-none-any.whl → 6.0.0rc2__py3-none-any.whl

sae_lens/saes/standard_sae.py

@@ -0,0 +1,185 @@
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+from jaxtyping import Float
+from numpy.typing import NDArray
+from torch import nn
+from typing_extensions import override
+
+from sae_lens.saes.sae import (
+    SAE,
+    SAEConfig,
+    TrainCoefficientConfig,
+    TrainingSAE,
+    TrainingSAEConfig,
+    TrainStepInput,
+)
+from sae_lens.util import filter_valid_dataclass_fields
+
+
+@dataclass
+class StandardSAEConfig(SAEConfig):
+    """
+    Configuration class for a StandardSAE.
+    """
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "standard"
+
+
+class StandardSAE(SAE[StandardSAEConfig]):
+    """
+    StandardSAE is an inference-only implementation of a Sparse Autoencoder (SAE)
+    using a simple linear encoder and decoder.
+
+    It implements the required abstract methods from BaseSAE:
+      - initialize_weights: sets up simple parameter initializations for W_enc, b_enc, W_dec, and b_dec.
+      - encode: computes the feature activations from an input.
+      - decode: reconstructs the input from the feature activations.
+
+    The BaseSAE.forward() method automatically calls encode and decode,
+    including any error-term processing if configured.
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(self, cfg: StandardSAEConfig, use_error_term: bool = False):
+        super().__init__(cfg, use_error_term)
+
+    @override
+    def initialize_weights(self) -> None:
+        # Initialize encoder weights and bias.
+        super().initialize_weights()
+        _init_weights_standard(self)
+
+    def encode(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        """
+        Encode the input tensor into the feature space.
+        For inference, no noise is added.
+        """
+        # Preprocess the SAE input (casting type, applying hooks, normalization)
+        sae_in = self.process_sae_in(x)
+        # Compute the pre-activation values
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+        # Apply the activation function (e.g., ReLU, tanh-relu, depending on config)
+        return self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+
+    def decode(
+        self, feature_acts: Float[torch.Tensor, "... d_sae"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """
+        Decode the feature activations back to the input space.
+        Now, if hook_z reshaping is turned on, we reverse the flattening.
+        """
+        # 1) linear transform
+        sae_out_pre = feature_acts @ self.W_dec + self.b_dec
+        # 2) hook reconstruction
+        sae_out_pre = self.hook_sae_recons(sae_out_pre)
+        # 4) optional out-normalization (e.g. constant_norm_rescale or layer_norm)
+        sae_out_pre = self.run_time_activation_norm_fn_out(sae_out_pre)
+        # 5) if hook_z is enabled, rearrange back to (..., n_heads, d_head).
+        return self.reshape_fn_out(sae_out_pre, self.d_head)
+
+
+@dataclass
+class StandardTrainingSAEConfig(TrainingSAEConfig):
+    """
+    Configuration class for training a StandardTrainingSAE.
+    """
+
+    l1_coefficient: float = 1.0
+    lp_norm: float = 1.0
+    l1_warm_up_steps: int = 0
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "standard"
+
+
+class StandardTrainingSAE(TrainingSAE[StandardTrainingSAEConfig]):
+    """
+    StandardTrainingSAE is a concrete implementation of BaseTrainingSAE using the "standard" SAE architecture.
+    It implements:
+      - initialize_weights: basic weight initialization for encoder/decoder.
+      - encode: inference encoding (invokes encode_with_hidden_pre).
+      - decode: a simple linear decoder.
+      - encode_with_hidden_pre: computes pre-activations, adds noise when training, and then activates.
+      - calculate_aux_loss: computes a sparsity penalty based on the (optionally scaled) p-norm of feature activations.
+    """
+
+    b_enc: nn.Parameter
+
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        _init_weights_standard(self)
+
+    @override
+    def get_coefficients(self) -> dict[str, float | TrainCoefficientConfig]:
+        return {
+            "l1": TrainCoefficientConfig(
+                value=self.cfg.l1_coefficient,
+                warm_up_steps=self.cfg.l1_warm_up_steps,
+            ),
+        }
+
+    def encode_with_hidden_pre(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        # Process the input (including dtype conversion, hook call, and any activation normalization)
+        sae_in = self.process_sae_in(x)
+        # Compute the pre-activation (and allow for a hook if desired)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)  # type: ignore
+        # Add noise during training for robustness (scaled by noise_scale from the configuration)
+        if self.training and self.cfg.noise_scale > 0:
+            hidden_pre_noised = (
+                hidden_pre + torch.randn_like(hidden_pre) * self.cfg.noise_scale
+            )
+        else:
+            hidden_pre_noised = hidden_pre
+        # Apply the activation function (and any post-activation hook)
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre_noised))
+        return feature_acts, hidden_pre_noised
+
+    def calculate_aux_loss(
+        self,
+        step_input: TrainStepInput,
+        feature_acts: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        sae_out: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        # The "standard" auxiliary loss is a sparsity penalty on the feature activations
+        weighted_feature_acts = feature_acts * self.W_dec.norm(dim=1)
+
+        # Compute the p-norm (set by cfg.lp_norm) over the feature dimension
+        sparsity = weighted_feature_acts.norm(p=self.cfg.lp_norm, dim=-1)
+        l1_loss = (step_input.coefficients["l1"] * sparsity).mean()
+
+        return {"l1_loss": l1_loss}
+
+    def log_histograms(self) -> dict[str, NDArray[np.generic]]:
+        """Log histograms of the weights and biases."""
+        b_e_dist = self.b_enc.detach().float().cpu().numpy()
+        return {
+            **super().log_histograms(),
+            "weights/b_e": b_e_dist,
+        }
+
+    def to_inference_config_dict(self) -> dict[str, Any]:
+        return filter_valid_dataclass_fields(
+            self.cfg.to_dict(), StandardSAEConfig, ["architecture"]
+        )
+
+
+def _init_weights_standard(
+    sae: SAE[StandardSAEConfig] | TrainingSAE[StandardTrainingSAEConfig],
+) -> None:
+    sae.b_enc = nn.Parameter(
+        torch.zeros(sae.cfg.d_sae, dtype=sae.dtype, device=sae.device)
+    )

sae_lens/saes/topk_sae.py

@@ -0,0 +1,294 @@
+"""Inference-only TopKSAE variant, similar in spirit to StandardSAE but using a TopK-based activation."""
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import torch
+from jaxtyping import Float
+from torch import nn
+from typing_extensions import override
+
+from sae_lens.saes.sae import (
+    SAE,
+    SAEConfig,
+    TrainCoefficientConfig,
+    TrainingSAE,
+    TrainingSAEConfig,
+    TrainStepInput,
+)
+from sae_lens.util import filter_valid_dataclass_fields
+
+
+class TopK(nn.Module):
+    """
+    A simple TopK activation that zeroes out all but the top K elements along the last dimension,
+    then optionally applies a post-activation function (e.g., ReLU).
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(
+        self,
+        k: int,
+        postact_fn: Callable[[torch.Tensor], torch.Tensor] = nn.ReLU(),
+    ):
+        super().__init__()
+        self.k = k
+        self.postact_fn = postact_fn
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        1) Select top K elements along the last dimension.
+        2) Apply post-activation (often ReLU).
+        3) Zero out all other entries.
+        """
+        topk = torch.topk(x, k=self.k, dim=-1)
+        values = self.postact_fn(topk.values)
+        result = torch.zeros_like(x)
+        result.scatter_(-1, topk.indices, values)
+        return result
+
+
+@dataclass
+class TopKSAEConfig(SAEConfig):
+    """
+    Configuration class for a TopKSAE.
+    """
+
+    k: int = 100
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "topk"
+
+
+class TopKSAE(SAE[TopKSAEConfig]):
+    """
+    An inference-only sparse autoencoder using a "topk" activation function.
+    It uses linear encoder and decoder layers, applying the TopK activation
+    to the hidden pre-activation in its encode step.
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(self, cfg: TopKSAEConfig, use_error_term: bool = False):
+        """
+        Args:
+            cfg: SAEConfig defining model size and behavior.
+            use_error_term: Whether to apply the error-term approach in the forward pass.
+        """
+        super().__init__(cfg, use_error_term)
+
+    @override
+    def initialize_weights(self) -> None:
+        # Initialize encoder weights and bias.
+        super().initialize_weights()
+        _init_weights_topk(self)
+
+    def encode(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        """
+        Converts input x into feature activations.
+        Uses topk activation from the config (cfg.activation_fn == "topk")
+        under the hood.
+        """
+        sae_in = self.process_sae_in(x)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+        # The BaseSAE already sets self.activation_fn to TopK(...) if config requests topk.
+        return self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+
+    def decode(
+        self, feature_acts: Float[torch.Tensor, "... d_sae"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """
+        Reconstructs the input from topk feature activations.
+        Applies optional finetuning scaling, hooking to recons, out normalization,
+        and optional head reshaping.
+        """
+        sae_out_pre = feature_acts @ self.W_dec + self.b_dec
+        sae_out_pre = self.hook_sae_recons(sae_out_pre)
+        sae_out_pre = self.run_time_activation_norm_fn_out(sae_out_pre)
+        return self.reshape_fn_out(sae_out_pre, self.d_head)
+
+    @override
+    def get_activation_fn(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        return TopK(self.cfg.k)
+
+
+@dataclass
+class TopKTrainingSAEConfig(TrainingSAEConfig):
+    """
+    Configuration class for training a TopKTrainingSAE.
+    """
+
+    k: int = 100
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "topk"
+
+
+class TopKTrainingSAE(TrainingSAE[TopKTrainingSAEConfig]):
+    """
+    TopK variant with training functionality. Injects noise during training, optionally
+    calculates a topk-related auxiliary loss, etc.
+    """
+
+    b_enc: nn.Parameter
+
+    def __init__(self, cfg: TopKTrainingSAEConfig, use_error_term: bool = False):
+        super().__init__(cfg, use_error_term)
+
+    @override
+    def initialize_weights(self) -> None:
+        super().initialize_weights()
+        _init_weights_topk(self)
+
+    def encode_with_hidden_pre(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        """
+        Similar to the base training method: cast input, optionally add noise, then apply TopK.
+        """
+        sae_in = self.process_sae_in(x)
+        hidden_pre = self.hook_sae_acts_pre(sae_in @ self.W_enc + self.b_enc)
+
+        # Inject noise if training
+        if self.training and self.cfg.noise_scale > 0:
+            hidden_pre_noised = (
+                hidden_pre + torch.randn_like(hidden_pre) * self.cfg.noise_scale
+            )
+        else:
+            hidden_pre_noised = hidden_pre
+
+        # Apply the TopK activation function (already set in self.activation_fn if config is "topk")
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre_noised))
+        return feature_acts, hidden_pre_noised
+
+    def calculate_aux_loss(
+        self,
+        step_input: TrainStepInput,
+        feature_acts: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        sae_out: torch.Tensor,
+    ) -> dict[str, torch.Tensor]:
+        # Calculate the auxiliary loss for dead neurons
+        topk_loss = self.calculate_topk_aux_loss(
+            sae_in=step_input.sae_in,
+            sae_out=sae_out,
+            hidden_pre=hidden_pre,
+            dead_neuron_mask=step_input.dead_neuron_mask,
+        )
+        return {"auxiliary_reconstruction_loss": topk_loss}
+
+    @override
+    def get_activation_fn(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        return TopK(self.cfg.k)
+
+    @override
+    def get_coefficients(self) -> dict[str, TrainCoefficientConfig | float]:
+        return {}
+
+    def calculate_topk_aux_loss(
+        self,
+        sae_in: torch.Tensor,
+        sae_out: torch.Tensor,
+        hidden_pre: torch.Tensor,
+        dead_neuron_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        """
+        Calculate TopK auxiliary loss.
+
+        This auxiliary loss encourages dead neurons to learn useful features by having
+        them reconstruct the residual error from the live neurons. It's a key part of
+        preventing neuron death in TopK SAEs.
+        """
+        # Mostly taken from https://github.com/EleutherAI/sae/blob/main/sae/sae.py, except without variance normalization
+        # NOTE: checking the number of dead neurons will force a GPU sync, so performance can likely be improved here
+        if dead_neuron_mask is None or (num_dead := int(dead_neuron_mask.sum())) == 0:
+            return sae_out.new_tensor(0.0)
+        residual = (sae_in - sae_out).detach()
+
+        # Heuristic from Appendix B.1 in the paper
+        k_aux = sae_in.shape[-1] // 2
+
+        # Reduce the scale of the loss if there are a small number of dead latents
+        scale = min(num_dead / k_aux, 1.0)
+        k_aux = min(k_aux, num_dead)
+
+        auxk_acts = _calculate_topk_aux_acts(
+            k_aux=k_aux,
+            hidden_pre=hidden_pre,
+            dead_neuron_mask=dead_neuron_mask,
+        )
+
+        # Encourage the top ~50% of dead latents to predict the residual of the
+        # top k living latents
+        recons = self.decode(auxk_acts)
+        auxk_loss = (recons - residual).pow(2).sum(dim=-1).mean()
+        return scale * auxk_loss
+
+    def _calculate_topk_aux_acts(
+        self,
+        k_aux: int,
+        hidden_pre: torch.Tensor,
+        dead_neuron_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Helper method to calculate activations for the auxiliary loss.
+
+        Args:
+            k_aux: Number of top dead neurons to select
+            hidden_pre: Pre-activation values from encoder
+            dead_neuron_mask: Boolean mask indicating which neurons are dead
+
+        Returns:
+            Tensor with activations for only the top-k dead neurons, zeros elsewhere
+        """
+        # Don't include living latents in this loss (set them to -inf so they won't be selected)
+        auxk_latents = torch.where(
+            dead_neuron_mask[None],
+            hidden_pre,
+            torch.tensor(-float("inf"), device=hidden_pre.device),
+        )
+
+        # Find topk values among dead neurons
+        auxk_topk = auxk_latents.topk(k_aux, dim=-1, sorted=False)
+
+        # Create a tensor of zeros, then place the topk values at their proper indices
+        auxk_acts = torch.zeros_like(hidden_pre)
+        auxk_acts.scatter_(-1, auxk_topk.indices, auxk_topk.values)
+
+        return auxk_acts
+
+    def to_inference_config_dict(self) -> dict[str, Any]:
+        return filter_valid_dataclass_fields(
+            self.cfg.to_dict(), TopKSAEConfig, ["architecture"]
+        )
+
+
+def _calculate_topk_aux_acts(
+    k_aux: int,
+    hidden_pre: torch.Tensor,
+    dead_neuron_mask: torch.Tensor,
+) -> torch.Tensor:
+    # Don't include living latents in this loss
+    auxk_latents = torch.where(dead_neuron_mask[None], hidden_pre, -torch.inf)
+    # Top-k dead latents
+    auxk_topk = auxk_latents.topk(k_aux, sorted=False)
+    # Set the activations to zero for all but the top k_aux dead latents
+    auxk_acts = torch.zeros_like(hidden_pre)
+    auxk_acts.scatter_(-1, auxk_topk.indices, auxk_topk.values)
+    # Set activations to zero for all but top k_aux dead latents
+    return auxk_acts
+
+
+def _init_weights_topk(
+    sae: SAE[TopKSAEConfig] | TrainingSAE[TopKTrainingSAEConfig],
+) -> None:
+    sae.b_enc = nn.Parameter(
+        torch.zeros(sae.cfg.d_sae, dtype=sae.dtype, device=sae.device)
+    )

sae_lens/training/activations_store.py

@@ -23,12 +23,12 @@ from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
 from sae_lens import logger
 from sae_lens.config import (
-    DTYPE_MAP,
     CacheActivationsRunnerConfig,
     HfDataset,
     LanguageModelSAERunnerConfig,
 )
-from sae_lens.sae import SAE
+from sae_lens.constants import DTYPE_MAP
+from sae_lens.saes.sae import SAE, T_SAE_CONFIG, T_TRAINING_SAE_CONFIG
 from sae_lens.tokenization_and_batching import concat_and_batch_sequences
 
 
@@ -91,7 +91,8 @@ class ActivationsStore:
     def from_config(
         cls,
         model: HookedRootModule,
-        cfg: LanguageModelSAERunnerConfig | CacheActivationsRunnerConfig,
+        cfg: LanguageModelSAERunnerConfig[T_TRAINING_SAE_CONFIG]
+        | CacheActivationsRunnerConfig,
         override_dataset: HfDataset | None = None,
     ) -> ActivationsStore:
         if isinstance(cfg, CacheActivationsRunnerConfig):
@@ -128,13 +129,15 @@ class ActivationsStore:
             hook_layer=cfg.hook_layer,
             hook_head_index=cfg.hook_head_index,
             context_size=cfg.context_size,
-            d_in=cfg.d_in,
+            d_in=cfg.d_in
+            if isinstance(cfg, CacheActivationsRunnerConfig)
+            else cfg.sae.d_in,
             n_batches_in_buffer=cfg.n_batches_in_buffer,
             total_training_tokens=cfg.training_tokens,
             store_batch_size_prompts=cfg.store_batch_size_prompts,
             train_batch_size_tokens=cfg.train_batch_size_tokens,
             prepend_bos=cfg.prepend_bos,
-            normalize_activations=cfg.normalize_activations,
+            normalize_activations=cfg.sae.normalize_activations,
             device=device,
             dtype=cfg.dtype,
             cached_activations_path=cached_activations_path,
@@ -149,9 +152,10 @@ class ActivationsStore:
     def from_sae(
         cls,
         model: HookedRootModule,
-        sae: SAE,
+        sae: SAE[T_SAE_CONFIG],
+        dataset: HfDataset | str,
+        dataset_trust_remote_code: bool = False,
         context_size: int | None = None,
-        dataset: HfDataset | str | None = None,
         streaming: bool = True,
         store_batch_size_prompts: int = 8,
         n_batches_in_buffer: int = 8,
@@ -159,25 +163,37 @@ class ActivationsStore:
         total_tokens: int = 10**9,
         device: str = "cpu",
     ) -> ActivationsStore:
+        if sae.cfg.metadata.hook_name is None:
+            raise ValueError("hook_name is required")
+        if sae.cfg.metadata.hook_layer is None:
+            raise ValueError("hook_layer is required")
+        if sae.cfg.metadata.hook_head_index is None:
+            raise ValueError("hook_head_index is required")
+        if sae.cfg.metadata.context_size is None:
+            raise ValueError("context_size is required")
+        if sae.cfg.metadata.prepend_bos is None:
+            raise ValueError("prepend_bos is required")
         return cls(
             model=model,
-            dataset=sae.cfg.dataset_path if dataset is None else dataset,
+            dataset=dataset,
             d_in=sae.cfg.d_in,
-            hook_name=sae.cfg.hook_name,
-            hook_layer=sae.cfg.hook_layer,
-            hook_head_index=sae.cfg.hook_head_index,
-            context_size=sae.cfg.context_size if context_size is None else context_size,
-            prepend_bos=sae.cfg.prepend_bos,
+            hook_name=sae.cfg.metadata.hook_name,
+            hook_layer=sae.cfg.metadata.hook_layer,
+            hook_head_index=sae.cfg.metadata.hook_head_index,
+            context_size=sae.cfg.metadata.context_size
+            if context_size is None
+            else context_size,
+            prepend_bos=sae.cfg.metadata.prepend_bos,
             streaming=streaming,
             store_batch_size_prompts=store_batch_size_prompts,
             train_batch_size_tokens=train_batch_size_tokens,
             n_batches_in_buffer=n_batches_in_buffer,
             total_training_tokens=total_tokens,
             normalize_activations=sae.cfg.normalize_activations,
-            dataset_trust_remote_code=sae.cfg.dataset_trust_remote_code,
+            dataset_trust_remote_code=dataset_trust_remote_code,
             dtype=sae.cfg.dtype,
             device=torch.device(device),
-            seqpos_slice=sae.cfg.seqpos_slice,
+            seqpos_slice=sae.cfg.metadata.seqpos_slice or (None,),
         )
 
     def __init__(
@@ -448,7 +464,7 @@ class ActivationsStore:
         ):
             # temporalily set estimated_norm_scaling_factor to 1.0 so the dataloader works
             self.estimated_norm_scaling_factor = 1.0
-            acts = self.next_batch()[:, 0]
+            acts = self.next_batch()[0]
             self.estimated_norm_scaling_factor = None
             norms_per_batch.append(acts.norm(dim=-1).mean().item())
         mean_norm = np.mean(norms_per_batch)