sae-lens 6.26.1__tar.gz → 6.27.0__tar.gz

{sae_lens-6.26.1 → sae_lens-6.27.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sae-lens
-Version: 6.26.1
+Version: 6.27.0
 Summary: Training and Analyzing Sparse Autoencoders (SAEs)
 License: MIT
 License-File: LICENSE

{sae_lens-6.26.1 → sae_lens-6.27.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sae-lens"
-version = "6.26.1"
+version = "6.27.0"
 description = "Training and Analyzing Sparse Autoencoders (SAEs)"
 authors = ["Joseph Bloom"]
 readme = "README.md"

{sae_lens-6.26.1 → sae_lens-6.27.0}/sae_lens/__init__.py

@@ -1,5 +1,5 @@
 # ruff: noqa: E402
-__version__ = "6.26.1"
+__version__ = "6.27.0"
 
 import logging
 

{sae_lens-6.26.1 → sae_lens-6.27.0}/sae_lens/config.py

@@ -148,6 +148,7 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
         seqpos_slice (tuple[int | None, ...]): Determines slicing of activations when constructing batches during training. The slice should be (start_pos, end_pos, optional[step_size]), e.g. for Othello we sometimes use (5, -5). Note, step_size > 0.
         disable_concat_sequences (bool): Whether to disable concatenating sequences and ignore sequences shorter than the context size. If True, disables concatenating and ignores short sequences.
         sequence_separator_token (int | Literal["bos", "eos", "sep"] | None): If not `None`, this token will be placed between sentences in a batch to act as a separator. By default, this is the `<bos>` token.
+        activations_mixing_fraction (float): Fraction of the activation buffer to keep for mixing with new activations (default 0.5). Higher values mean more temporal shuffling but slower throughput. If 0, activations are served in order without shuffling (no temporal mixing).
         device (str): The device to use. Usually "cuda".
         act_store_device (str): The device to use for the activation store. "cpu" is advised in order to save VRAM. Defaults to "with_model" which uses the same device as the main model.
         seed (int): The seed to use.
@@ -217,6 +218,7 @@ class LanguageModelSAERunnerConfig(Generic[T_TRAINING_SAE_CONFIG]):
     sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = (
         special_token_field(default="bos")
     )
+    activations_mixing_fraction: float = 0.5
 
     # Misc
     device: str = "cpu"

{sae_lens-6.26.1 → sae_lens-6.27.0}/sae_lens/loading/pretrained_sae_loaders.py

@@ -959,7 +959,7 @@ def get_dictionary_learning_config_1_from_hf(
     architecture = "standard"
     if trainer["dict_class"] == "GatedAutoEncoder":
         architecture = "gated"
-    elif trainer["dict_class"] == "MatryoshkaBatchTopKSAE":
+    elif trainer["dict_class"] in ["MatryoshkaBatchTopKSAE", "BatchTopKSAE"]:
         architecture = "jumprelu"
 
     return {

{sae_lens-6.26.1 → sae_lens-6.27.0}/sae_lens/training/activations_store.py

@@ -148,6 +148,7 @@ class ActivationsStore:
             exclude_special_tokens=exclude_special_tokens,
             disable_concat_sequences=cfg.disable_concat_sequences,
             sequence_separator_token=cfg.sequence_separator_token,
+            activations_mixing_fraction=cfg.activations_mixing_fraction,
         )
 
     @classmethod
@@ -222,6 +223,7 @@ class ActivationsStore:
         exclude_special_tokens: torch.Tensor | None = None,
         disable_concat_sequences: bool = False,
         sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = "bos",
+        activations_mixing_fraction: float = 0.5,
     ):
         self.model = model
         if model_kwargs is None:
@@ -269,6 +271,7 @@ class ActivationsStore:
         self.sequence_separator_token: int | Literal["bos", "eos", "sep"] | None = (
             sequence_separator_token
         )
+        self.activations_mixing_fraction = activations_mixing_fraction
 
         self.n_dataset_processed = 0
 
@@ -708,6 +711,7 @@ class ActivationsStore:
             buffer_size=self.n_batches_in_buffer * self.training_context_size,
             batch_size=self.train_batch_size_tokens,
             activations_loader=self._iterate_filtered_activations(),
+            mix_fraction=self.activations_mixing_fraction,
         )
 
     def next_batch(self) -> torch.Tensor:

{sae_lens-6.26.1 → sae_lens-6.27.0}/sae_lens/training/mixing_buffer.py

@@ -8,15 +8,19 @@ def mixing_buffer(
     buffer_size: int,
     batch_size: int,
     activations_loader: Iterator[torch.Tensor],
+    mix_fraction: float = 0.5,
 ) -> Iterator[torch.Tensor]:
     """
     A generator that maintains a mix of old and new activations for better training.
-    It stores half of the activations and mixes them with new ones to create batches.
+    It keeps a portion of activations and mixes them with new ones to create batches.
 
     Args:
-        buffer_size: Total size of the buffer (will store buffer_size/2 activations)
+        buffer_size: Total size of the buffer
         batch_size: Size of batches to return
         activations_loader: Iterator providing new activations
+        mix_fraction: Fraction of buffer to keep for mixing (default 0.5).
+                      Higher values mean more temporal mixing but slower throughput.
+                      If 0, no shuffling occurs (passthrough mode).
 
     Yields:
         Batches of activations of shape (batch_size, *activation_dims)
@@ -24,6 +28,8 @@ def mixing_buffer(
 
     if buffer_size < batch_size:
         raise ValueError("Buffer size must be greater than or equal to batch size")
+    if not 0 <= mix_fraction <= 1:
+        raise ValueError("mix_fraction must be in [0, 1]")
 
     storage_buffer: torch.Tensor | None = None
 
@@ -35,10 +41,12 @@ def mixing_buffer(
         )
 
         if storage_buffer.shape[0] >= buffer_size:
-            # Shuffle
-            storage_buffer = storage_buffer[torch.randperm(storage_buffer.shape[0])]
+            if mix_fraction > 0:
+                storage_buffer = storage_buffer[torch.randperm(storage_buffer.shape[0])]
 
-            num_serving_batches = max(1, storage_buffer.shape[0] // (2 * batch_size))
+            num_serving_batches = max(
+                1, int(storage_buffer.shape[0] * (1 - mix_fraction)) // batch_size
+            )
             serving_cutoff = num_serving_batches * batch_size
             serving_buffer = storage_buffer[:serving_cutoff]
             storage_buffer = storage_buffer[serving_cutoff:]