sae-lens 5.11.0__py3-none-any.whl → 6.0.0__py3-none-any.whl

sae_lens/saes/standard_sae.py

@@ -0,0 +1,178 @@
+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, 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)
+        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
+        # Apply the activation function (and any post-activation hook)
+        feature_acts = self.hook_sae_acts_post(self.activation_fn(hidden_pre))
+        return feature_acts, hidden_pre
+
+    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,300 @@
+"""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 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)
+
+    @override
+    @torch.no_grad()
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError(
+            "Folding W_dec_norm is not safe for TopKSAEs, as this may change the topk activations"
+        )
+
+
+@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)
+
+        # 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))
+        return feature_acts, hidden_pre
+
+    @override
+    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
+    @torch.no_grad()
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError(
+            "Folding W_dec_norm is not safe for TopKSAEs, as this may change the topk activations"
+        )
+
+    @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/activation_scaler.py

@@ -0,0 +1,53 @@
+import json
+from dataclasses import dataclass
+from statistics import mean
+
+import torch
+from tqdm.auto import tqdm
+
+from sae_lens.training.types import DataProvider
+
+
+@dataclass
+class ActivationScaler:
+    scaling_factor: float | None = None
+
+    def scale(self, acts: torch.Tensor) -> torch.Tensor:
+        return acts if self.scaling_factor is None else acts * self.scaling_factor
+
+    def unscale(self, acts: torch.Tensor) -> torch.Tensor:
+        return acts if self.scaling_factor is None else acts / self.scaling_factor
+
+    def __call__(self, acts: torch.Tensor) -> torch.Tensor:
+        return self.scale(acts)
+
+    @torch.no_grad()
+    def _calculate_mean_norm(
+        self, data_provider: DataProvider, n_batches_for_norm_estimate: int = int(1e3)
+    ) -> float:
+        norms_per_batch: list[float] = []
+        for _ in tqdm(
+            range(n_batches_for_norm_estimate), desc="Estimating norm scaling factor"
+        ):
+            acts = next(data_provider)
+            norms_per_batch.append(acts.norm(dim=-1).mean().item())
+        return mean(norms_per_batch)
+
+    def estimate_scaling_factor(
+        self,
+        d_in: int,
+        data_provider: DataProvider,
+        n_batches_for_norm_estimate: int = int(1e3),
+    ):
+        mean_norm = self._calculate_mean_norm(
+            data_provider, n_batches_for_norm_estimate
+        )
+        self.scaling_factor = (d_in**0.5) / mean_norm
+
+    def save(self, file_path: str):
+        """save the state dict to a file in json format"""
+        if not file_path.endswith(".json"):
+            raise ValueError("file_path must end with .json")
+
+        with open(file_path, "w") as f:
+            json.dump({"scaling_factor": self.scaling_factor}, f)