sae-lens 6.16.3__py3-none-any.whl → 6.21.0__py3-none-any.whl

sae_lens/saes/temporal_sae.py

@@ -0,0 +1,372 @@
+"""TemporalSAE: A Sparse Autoencoder with temporal attention mechanism.
+
+TemporalSAE decomposes activations into:
+1. Predicted codes (from attention over context)
+2. Novel codes (sparse features of the residual)
+
+See: https://arxiv.org/abs/2410.04185
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Literal
+
+import torch
+import torch.nn.functional as F
+from jaxtyping import Float
+from torch import nn
+from typing_extensions import override
+
+from sae_lens import logger
+from sae_lens.saes.sae import SAE, SAEConfig
+
+
+def get_attention(query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
+    """Compute causal attention weights."""
+    L, S = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1))
+    attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device)
+    temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril(diagonal=0)
+    attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+    attn_bias.to(query.dtype)
+
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    return torch.softmax(attn_weight, dim=-1)
+
+
+class ManualAttention(nn.Module):
+    """Manual attention implementation for TemporalSAE."""
+
+    def __init__(
+        self,
+        dimin: int,
+        n_heads: int = 4,
+        bottleneck_factor: int = 64,
+        bias_k: bool = True,
+        bias_q: bool = True,
+        bias_v: bool = True,
+        bias_o: bool = True,
+    ):
+        super().__init__()
+        assert dimin % (bottleneck_factor * n_heads) == 0
+
+        self.n_heads = n_heads
+        self.n_embds = dimin // bottleneck_factor
+        self.dimin = dimin
+
+        # Key, query, value projections
+        self.k_ctx = nn.Linear(dimin, self.n_embds, bias=bias_k)
+        self.q_target = nn.Linear(dimin, self.n_embds, bias=bias_q)
+        self.v_ctx = nn.Linear(dimin, dimin, bias=bias_v)
+        self.c_proj = nn.Linear(dimin, dimin, bias=bias_o)
+
+        # Normalize to match scale with representations
+        with torch.no_grad():
+            scaling = 1 / math.sqrt(self.n_embds // self.n_heads)
+            self.k_ctx.weight.copy_(
+                scaling
+                * self.k_ctx.weight
+                / (1e-6 + torch.linalg.norm(self.k_ctx.weight, dim=1, keepdim=True))
+            )
+            self.q_target.weight.copy_(
+                scaling
+                * self.q_target.weight
+                / (1e-6 + torch.linalg.norm(self.q_target.weight, dim=1, keepdim=True))
+            )
+
+            scaling = 1 / math.sqrt(self.dimin // self.n_heads)
+            self.v_ctx.weight.copy_(
+                scaling
+                * self.v_ctx.weight
+                / (1e-6 + torch.linalg.norm(self.v_ctx.weight, dim=1, keepdim=True))
+            )
+
+            scaling = 1 / math.sqrt(self.dimin)
+            self.c_proj.weight.copy_(
+                scaling
+                * self.c_proj.weight
+                / (1e-6 + torch.linalg.norm(self.c_proj.weight, dim=1, keepdim=True))
+            )
+
+    def forward(
+        self, x_ctx: torch.Tensor, x_target: torch.Tensor, get_attn_map: bool = False
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Compute projective attention output."""
+        k = self.k_ctx(x_ctx)
+        v = self.v_ctx(x_ctx)
+        q = self.q_target(x_target)
+
+        # Split into heads
+        B, T, _ = x_ctx.size()
+        k = k.view(B, T, self.n_heads, self.n_embds // self.n_heads).transpose(1, 2)
+        q = q.view(B, T, self.n_heads, self.n_embds // self.n_heads).transpose(1, 2)
+        v = v.view(B, T, self.n_heads, self.dimin // self.n_heads).transpose(1, 2)
+
+        # Attention map (optional)
+        attn_map = None
+        if get_attn_map:
+            attn_map = get_attention(query=q, key=k)
+
+        # Scaled dot-product attention
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=None, dropout_p=0, is_causal=True
+        )
+
+        # Reshape and project
+        d_target = self.c_proj(
+            attn_output.transpose(1, 2).contiguous().view(B, T, self.dimin)
+        )
+
+        return d_target, attn_map
+
+
+@dataclass
+class TemporalSAEConfig(SAEConfig):
+    """Configuration for TemporalSAE inference.
+
+    Args:
+        d_in: Input dimension (dimensionality of the activations being encoded)
+        d_sae: SAE latent dimension (number of features)
+        n_heads: Number of attention heads in temporal attention
+        n_attn_layers: Number of attention layers
+        bottleneck_factor: Bottleneck factor for attention dimension
+        sae_diff_type: Type of SAE for novel codes ('relu' or 'topk')
+        kval_topk: K value for top-k sparsity (if sae_diff_type='topk')
+        tied_weights: Whether to tie encoder and decoder weights
+        activation_normalization_factor: Scalar factor for rescaling activations (used with normalize_activations='constant_scalar_rescale')
+    """
+
+    n_heads: int = 8
+    n_attn_layers: int = 1
+    bottleneck_factor: int = 64
+    sae_diff_type: Literal["relu", "topk"] = "topk"
+    kval_topk: int | None = None
+    tied_weights: bool = True
+    activation_normalization_factor: float = 1.0
+
+    def __post_init__(self):
+        # Call parent's __post_init__ first, but allow constant_scalar_rescale
+        if self.normalize_activations not in [
+            "none",
+            "expected_average_only_in",
+            "constant_norm_rescale",
+            "constant_scalar_rescale",  # Temporal SAEs support this
+            "layer_norm",
+        ]:
+            raise ValueError(
+                f"normalize_activations must be none, expected_average_only_in, layer_norm, constant_norm_rescale, or constant_scalar_rescale. Got {self.normalize_activations}"
+            )
+
+    @override
+    @classmethod
+    def architecture(cls) -> str:
+        return "temporal"
+
+
+class TemporalSAE(SAE[TemporalSAEConfig]):
+    """TemporalSAE: Sparse Autoencoder with temporal attention.
+
+    This SAE decomposes each activation x_t into:
+    - x_pred: Information aggregated from context {x_0, ..., x_{t-1}}
+    - x_novel: Novel information at position t (encoded sparsely)
+
+    The forward pass:
+    1. Uses attention layers to predict x_t from context
+    2. Encodes the residual (novel part) with a sparse SAE
+    3. Combines both for reconstruction
+    """
+
+    # Custom parameters (in addition to W_enc, W_dec, b_dec from base)
+    attn_layers: nn.ModuleList  # Attention layers
+    eps: float
+    lam: float
+
+    def __init__(self, cfg: TemporalSAEConfig, use_error_term: bool = False):
+        # Call parent init first
+        super().__init__(cfg, use_error_term)
+
+        # Initialize attention layers after parent init and move to correct device
+        self.attn_layers = nn.ModuleList(
+            [
+                ManualAttention(
+                    dimin=cfg.d_sae,
+                    n_heads=cfg.n_heads,
+                    bottleneck_factor=cfg.bottleneck_factor,
+                    bias_k=True,
+                    bias_q=True,
+                    bias_v=True,
+                    bias_o=True,
+                ).to(device=self.device, dtype=self.dtype)
+                for _ in range(cfg.n_attn_layers)
+            ]
+        )
+
+        self.eps = 1e-6
+        self.lam = 1 / (4 * self.cfg.d_in)
+
+    @override
+    def _setup_activation_normalization(self):
+        """Set up activation normalization functions for TemporalSAE.
+
+        Overrides the base implementation to handle constant_scalar_rescale
+        using the temporal-specific activation_normalization_factor.
+        """
+        if self.cfg.normalize_activations == "constant_scalar_rescale":
+            # Handle constant scalar rescaling for temporal SAEs
+            def run_time_activation_norm_fn_in(x: torch.Tensor) -> torch.Tensor:
+                return x * self.cfg.activation_normalization_factor
+
+            def run_time_activation_norm_fn_out(x: torch.Tensor) -> torch.Tensor:
+                return x / self.cfg.activation_normalization_factor
+
+            self.run_time_activation_norm_fn_in = run_time_activation_norm_fn_in
+            self.run_time_activation_norm_fn_out = run_time_activation_norm_fn_out
+        else:
+            # Delegate to parent for all other normalization types
+            super()._setup_activation_normalization()
+
+    @override
+    def initialize_weights(self) -> None:
+        """Initialize TemporalSAE weights."""
+        # Initialize D (decoder) and b (bias)
+        self.W_dec = nn.Parameter(
+            torch.randn(
+                (self.cfg.d_sae, self.cfg.d_in), dtype=self.dtype, device=self.device
+            )
+        )
+        self.b_dec = nn.Parameter(
+            torch.zeros((self.cfg.d_in), dtype=self.dtype, device=self.device)
+        )
+
+        # Initialize E (encoder) if not tied
+        if not self.cfg.tied_weights:
+            self.W_enc = nn.Parameter(
+                torch.randn(
+                    (self.cfg.d_in, self.cfg.d_sae),
+                    dtype=self.dtype,
+                    device=self.device,
+                )
+            )
+
+    def encode_with_predictions(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> tuple[Float[torch.Tensor, "... d_sae"], Float[torch.Tensor, "... d_sae"]]:
+        """Encode input to novel codes only.
+
+        Returns only the sparse novel codes (not predicted codes).
+        This is the main feature representation for TemporalSAE.
+        """
+        # Process input through SAELens preprocessing
+        x = self.process_sae_in(x)
+
+        B, L, _ = x.shape
+
+        if self.cfg.tied_weights:  # noqa: SIM108
+            W_enc = self.W_dec.T
+        else:
+            W_enc = self.W_enc
+
+        # Compute predicted codes using attention
+        x_residual = x
+        z_pred = torch.zeros((B, L, self.cfg.d_sae), device=x.device, dtype=x.dtype)
+
+        for attn_layer in self.attn_layers:
+            # Encode input to latent space
+            z_input = F.relu(torch.matmul(x_residual * self.lam, W_enc))
+
+            # Shift context (causal masking)
+            z_ctx = torch.cat(
+                (torch.zeros_like(z_input[:, :1, :]), z_input[:, :-1, :].clone()), dim=1
+            )
+
+            # Apply attention to get predicted codes
+            z_pred_, _ = attn_layer(z_ctx, z_input, get_attn_map=False)
+            z_pred_ = F.relu(z_pred_)
+
+            # Project predicted codes back to input space
+            Dz_pred_ = torch.matmul(z_pred_, self.W_dec)
+            Dz_norm_ = Dz_pred_.norm(dim=-1, keepdim=True) + self.eps
+
+            # Compute projection scale
+            proj_scale = (Dz_pred_ * x_residual).sum(
+                dim=-1, keepdim=True
+            ) / Dz_norm_.pow(2)
+
+            # Accumulate predicted codes
+            z_pred = z_pred + (z_pred_ * proj_scale)
+
+            # Remove prediction from residual
+            x_residual = x_residual - proj_scale * Dz_pred_
+
+        # Encode residual (novel part) with sparse SAE
+        z_novel = F.relu(torch.matmul(x_residual * self.lam, W_enc))
+        if self.cfg.sae_diff_type == "topk":
+            kval = self.cfg.kval_topk
+            if kval is not None:
+                _, topk_indices = torch.topk(z_novel, kval, dim=-1)
+                mask = torch.zeros_like(z_novel)
+                mask.scatter_(-1, topk_indices, 1)
+                z_novel = z_novel * mask
+
+        # Return only novel codes (these are the interpretable features)
+        return z_novel, z_pred
+
+    def encode(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        return self.encode_with_predictions(x)[0]
+
+    def decode(
+        self, feature_acts: Float[torch.Tensor, "... d_sae"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """Decode novel codes to reconstruction.
+
+        Note: This only decodes the novel codes. For full reconstruction,
+        use forward() which includes predicted codes.
+        """
+        # Decode novel codes
+        sae_out = torch.matmul(feature_acts, self.W_dec)
+        sae_out = sae_out + self.b_dec
+
+        # Apply hook
+        sae_out = self.hook_sae_recons(sae_out)
+
+        # Apply output activation normalization (reverses input normalization)
+        sae_out = self.run_time_activation_norm_fn_out(sae_out)
+
+        # Add bias (already removed in process_sae_in)
+        logger.warning(
+            "NOTE this only decodes x_novel. The x_pred is missing, so we're not reconstructing the full x."
+        )
+        return sae_out
+
+    @override
+    def forward(
+        self, x: Float[torch.Tensor, "... d_in"]
+    ) -> Float[torch.Tensor, "... d_in"]:
+        """Full forward pass through TemporalSAE.
+
+        Returns complete reconstruction (predicted + novel).
+        """
+        # Encode
+        z_novel, z_pred = self.encode_with_predictions(x)
+
+        # Decode the sum of predicted and novel codes.
+        x_recons = torch.matmul(z_novel + z_pred, self.W_dec) + self.b_dec
+
+        # Apply output activation normalization (reverses input normalization)
+        x_recons = self.run_time_activation_norm_fn_out(x_recons)
+
+        return self.hook_sae_output(x_recons)
+
+    @override
+    def fold_W_dec_norm(self) -> None:
+        raise NotImplementedError("Folding W_dec_norm is not supported for TemporalSAE")
+
+    @override
+    @torch.no_grad()
+    def fold_activation_norm_scaling_factor(self, scaling_factor: float) -> None:
+        raise NotImplementedError(
+            "Folding activation norm scaling factor is not supported for TemporalSAE"
+        )

sae_lens/training/activation_scaler.py

@@ -1,5 +1,6 @@
 import json
 from dataclasses import dataclass
+from pathlib import Path
 from statistics import mean
 
 import torch
@@ -51,3 +52,9 @@ class ActivationScaler:
 
         with open(file_path, "w") as f:
             json.dump({"scaling_factor": self.scaling_factor}, f)
+
+    def load(self, file_path: str | Path):
+        """load the state dict from a file in json format"""
+        with open(file_path) as f:
+            data = json.load(f)
+            self.scaling_factor = data["scaling_factor"]

sae_lens/training/activations_store.py

@@ -4,6 +4,7 @@ import json
 import os
 import warnings
 from collections.abc import Generator, Iterator, Sequence
+from pathlib import Path
 from typing import Any, Literal, cast
 
 import datasets
@@ -13,8 +14,8 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.utils import HfHubHTTPError
 from jaxtyping import Float, Int
 from requests import HTTPError
-from safetensors.torch import save_file
-from tqdm import tqdm
+from safetensors.torch import load_file, save_file
+from tqdm.auto import tqdm
 from transformer_lens.hook_points import HookedRootModule
 from transformers import AutoTokenizer, PreTrainedTokenizerBase
 
@@ -24,7 +25,7 @@ from sae_lens.config import (
     HfDataset,
     LanguageModelSAERunnerConfig,
 )
-from sae_lens.constants import DTYPE_MAP
+from sae_lens.constants import ACTIVATIONS_STORE_STATE_FILENAME, DTYPE_MAP
 from sae_lens.pretokenize_runner import get_special_token_from_cfg
 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
@@ -318,7 +319,7 @@ class ActivationsStore:
                 )
         else:
             warnings.warn(
-                "Dataset is not tokenized. Pre-tokenizing will improve performance and allows for more control over special tokens. See https://jbloomaus.github.io/SAELens/training_saes/#pretokenizing-datasets for more info."
+                "Dataset is not tokenized. Pre-tokenizing will improve performance and allows for more control over special tokens. See https://decoderesearch.github.io/SAELens/training_saes/#pretokenizing-datasets for more info."
             )
 
         self.iterable_sequences = self._iterate_tokenized_sequences()
@@ -729,6 +730,48 @@ class ActivationsStore:
         """save the state dict to a file in safetensors format"""
         save_file(self.state_dict(), file_path)
 
+    def save_to_checkpoint(self, checkpoint_path: str | Path):
+        """Save the state dict to a checkpoint path"""
+        self.save(str(Path(checkpoint_path) / ACTIVATIONS_STORE_STATE_FILENAME))
+
+    def load_from_checkpoint(self, checkpoint_path: str | Path):
+        """Load the state dict from a checkpoint path"""
+        self.load(str(Path(checkpoint_path) / ACTIVATIONS_STORE_STATE_FILENAME))
+
+    def load(self, file_path: str):
+        """Load the state dict from a file in safetensors format"""
+
+        state_dict = load_file(file_path)
+
+        if "n_dataset_processed" in state_dict:
+            target_n_dataset_processed = state_dict["n_dataset_processed"].item()
+
+            # Only fast-forward if needed
+
+            if target_n_dataset_processed > self.n_dataset_processed:
+                logger.info(
+                    "Fast-forwarding through dataset samples to match checkpoint position"
+                )
+                samples_to_skip = target_n_dataset_processed - self.n_dataset_processed
+
+                pbar = tqdm(
+                    total=samples_to_skip,
+                    desc="Fast-forwarding through dataset",
+                    leave=False,
+                )
+                while target_n_dataset_processed > self.n_dataset_processed:
+                    start = self.n_dataset_processed
+                    try:
+                        # Just consume and ignore the values to fast-forward
+                        next(self.iterable_sequences)
+                    except StopIteration:
+                        logger.warning(
+                            "Dataset exhausted during fast-forward. Resetting dataset."
+                        )
+                        self.iterable_sequences = self._iterate_tokenized_sequences()
+                    pbar.update(self.n_dataset_processed - start)
+                pbar.close()
+
 
 def validate_pretokenized_dataset_tokenizer(
     dataset_path: str, model_tokenizer: PreTrainedTokenizerBase

sae_lens/training/optim.py

@@ -2,6 +2,8 @@
 Took the LR scheduler from my previous work: https://github.com/jbloomAus/DecisionTransformerInterpretability/blob/ee55df35cdb92e81d689c72fb9dd5a7252893363/src/decision_transformer/utils.py#L425
 """
 
+from typing import Any
+
 import torch.optim as optim
 import torch.optim.lr_scheduler as lr_scheduler
 
@@ -150,3 +152,12 @@ class CoefficientScheduler:
     def value(self) -> float:
         """Returns the current scalar value."""
         return self.current_value
+
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "current_step": self.current_step,
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]):
+        for k in state_dict:
+            setattr(self, k, state_dict[k])

sae_lens/training/sae_trainer.py

@@ -1,4 +1,5 @@
 import contextlib
+import math
 from pathlib import Path
 from typing import Any, Callable, Generic, Protocol
 
@@ -10,7 +11,11 @@ from tqdm.auto import tqdm
 
 from sae_lens import __version__
 from sae_lens.config import SAETrainerConfig
-from sae_lens.constants import ACTIVATION_SCALER_CFG_FILENAME, SPARSITY_FILENAME
+from sae_lens.constants import (
+    ACTIVATION_SCALER_CFG_FILENAME,
+    SPARSITY_FILENAME,
+    TRAINER_STATE_FILENAME,
+)
 from sae_lens.saes.sae import (
     T_TRAINING_SAE,
     T_TRAINING_SAE_CONFIG,
@@ -56,6 +61,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
     data_provider: DataProvider
     activation_scaler: ActivationScaler
     evaluator: Evaluator[T_TRAINING_SAE] | None
+    coefficient_schedulers: dict[str, CoefficientScheduler]
 
     def __init__(
         self,
@@ -84,7 +90,9 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
                 range(
                     0,
                     cfg.total_training_samples,
-                    cfg.total_training_samples // self.cfg.n_checkpoints,
+                    math.ceil(
+                        cfg.total_training_samples / (self.cfg.n_checkpoints + 1)
+                    ),
                 )
             )[1:]
 
@@ -93,11 +101,6 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
             sae.cfg.d_sae, device=cfg.device
         )
         self.n_frac_active_samples = 0
-        # we don't train the scaling factor (initially)
-        # set requires grad to false for the scaling factor
-        for name, param in self.sae.named_parameters():
-            if "scaling_factor" in name:
-                param.requires_grad = False
 
         self.optimizer = Adam(
             sae.parameters(),
@@ -210,10 +213,7 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
                 sparsity_path = checkpoint_path / SPARSITY_FILENAME
                 save_file({"sparsity": self.log_feature_sparsity}, sparsity_path)
 
-                activation_scaler_path = (
-                    checkpoint_path / ACTIVATION_SCALER_CFG_FILENAME
-                )
-                self.activation_scaler.save(str(activation_scaler_path))
+                self.save_trainer_state(checkpoint_path)
 
                 if self.cfg.logger.log_to_wandb:
                     self.cfg.logger.log(
@@ -227,6 +227,44 @@ class SAETrainer(Generic[T_TRAINING_SAE, T_TRAINING_SAE_CONFIG]):
         if self.save_checkpoint_fn is not None:
             self.save_checkpoint_fn(checkpoint_path=checkpoint_path)
 
+    def save_trainer_state(self, checkpoint_path: Path) -> None:
+        checkpoint_path.mkdir(exist_ok=True, parents=True)
+        scheduler_state_dicts = {
+            name: scheduler.state_dict()
+            for name, scheduler in self.coefficient_schedulers.items()
+        }
+        torch.save(
+            {
+                "optimizer": self.optimizer.state_dict(),
+                "lr_scheduler": self.lr_scheduler.state_dict(),
+                "n_training_samples": self.n_training_samples,
+                "n_training_steps": self.n_training_steps,
+                "act_freq_scores": self.act_freq_scores,
+                "n_forward_passes_since_fired": self.n_forward_passes_since_fired,
+                "n_frac_active_samples": self.n_frac_active_samples,
+                "started_fine_tuning": self.started_fine_tuning,
+                "coefficient_schedulers": scheduler_state_dicts,
+            },
+            str(checkpoint_path / TRAINER_STATE_FILENAME),
+        )
+        activation_scaler_path = checkpoint_path / ACTIVATION_SCALER_CFG_FILENAME
+        self.activation_scaler.save(str(activation_scaler_path))
+
+    def load_trainer_state(self, checkpoint_path: Path | str) -> None:
+        checkpoint_path = Path(checkpoint_path)
+        self.activation_scaler.load(checkpoint_path / ACTIVATION_SCALER_CFG_FILENAME)
+        state_dict = torch.load(checkpoint_path / TRAINER_STATE_FILENAME)
+        self.optimizer.load_state_dict(state_dict["optimizer"])
+        self.lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
+        self.n_training_samples = state_dict["n_training_samples"]
+        self.n_training_steps = state_dict["n_training_steps"]
+        self.act_freq_scores = state_dict["act_freq_scores"]
+        self.n_forward_passes_since_fired = state_dict["n_forward_passes_since_fired"]
+        self.n_frac_active_samples = state_dict["n_frac_active_samples"]
+        self.started_fine_tuning = state_dict["started_fine_tuning"]
+        for name, scheduler_state_dict in state_dict["coefficient_schedulers"].items():
+            self.coefficient_schedulers[name].load_state_dict(scheduler_state_dict)
+
     def _train_step(
         self,
         sae: T_TRAINING_SAE,