SAELens

The transition to v6 introduced several breaking changes to simplify the API and decouple LLM-specific code from SAE training:

• Renamed Runner: SAETrainingRunner is now LanguageModelSAETrainingRunner (though the old name is temporarily supported).
• Config Renaming:
  • JumpReLU L0 coefficient is now l0_coefficient.
  • TopK k is now an explicit parameter in the config.
• Default Changes:
  • JumpReLU bandwidth default increased from 0.001 to 0.05.
  • JumpReLU starting threshold increased from 0.001 to 0.01.
• Removed Options: Several legacy options were removed, including expansion_factor, hook_layer, ghost grads, b_dec init options, decoder init options, MSE loss normalization, decoder normalization, finetuning_tokens/finetuning_method, noise_scale, and activation_fn/activation_fn_kwargs.
• Metadata Move: SAE.cfg now only contains essential keys. Non-essential keys like prepend_bos have been moved to SAE.cfg.metadata.

When collecting activations sequentially (e.g., processing documents one at a time), consecutive activations are highly correlated, which can negatively impact training. The mixing_buffer utility shuffles activations to ensure each batch contains data from many different contexts.

How it works:

1. Accumulates activations until reaching buffer_size.
2. Randomly shuffles the buffer.
3. Yields half as batches while keeping the other half.
4. Refills with new activations and repeats.

Use mixing_buffer by passing a buffer_size, a batch_size, and an activations_loader (an iterator yielding tensors).

from sae_lens.training.mixing_buffer import mixing_buffer
from collections.abc import Iterator
import torch

def my_sequential_activations() -> Iterator[torch.Tensor]:
    """Yields activations in document order (correlated)."""
    while True:
        # Process a document and yield its activations
        yield torch.randn(1024, 768)  # (tokens_in_doc, d_in)

# Wrap with mixing_buffer
shuffled_provider = mixing_buffer(
    buffer_size=100_000,      # Total activations to buffer
    batch_size=4096,          # Output batch size
    activations_loader=my_sequential_activations(),
)

# Now batches are shuffled!
for batch in shuffled_provider:
    print(batch.shape)  # (4096, 768)
    break

SAELens SAEs are standard PyTorch modules. While they offer deep integration with TransformerLens via HookedSAETransformer, they are compatible with any framework that provides activation tensors of the correct shape.

• Hugging Face Transformers: Extract activations using PyTorch hooks and pass them to sae.encode() or sae.decode().
• nnsight: Use the nnsight tracing API for intervention patterns.
• Generic PyTorch: Pass any tensor with the correct input dimension to the core methods.

import torch
from sae_lens import SAE

sae = SAE.from_pretrained(
    release="gemma-scope-2b-pt-res-canonical",
    sae_id="layer_12/width_16k/canonical",
    device="cuda"
)

# SAEs work with any activation tensor of the right shape
# Gemma 2 2B has d_model=2304
activations = torch.randn(1, 128, 2304, device="cuda")  # From any source
features = sae.encode(activations)
reconstructed = sae.decode(features)

SAELens uses a modular architecture system that separates the training process from the inference process. To create a custom architecture, you must define three components:

1. Configuration classes: Define hyperparameters and settings (inheriting from SAEConfig or TrainingSAEConfig).
2. SAE classes: Implement the neural network logic (inheriting from SAE or TrainingSAE).
3. Registration: Making the architecture available to the training system via register_sae_training_class.

Training vs Inference SAEs

• TrainingSAE: Used during the training phase to implement custom training logic (e.g., specialized loss functions).
• SAE (Inference SAE): The final model saved after training, designed for deployment and efficient inference.

Note that some training architectures might save as a different inference type. For example, BatchTopKTrainingSAE is saved as a JumpReLU SAE for inference.

The MultiSAETrainingRunner allows you to train a sweep of SAEs (e.g., different l1_coefficient values or different layers) simultaneously. The runner performs a single LLM forward pass per batch and multiplexes the resulting activations to all configured SAEs, which is more efficient than training them sequentially.

Key Constraints (V1):

• Does not support CLI/argparse, cached activations, or from_pretrained_path per entry.
• compile_sae=True is not supported, but compile_llm=True is.
• SAEs sharing a hook must agree on d_in and hook_head_index.

Configuration:

• saes: A dictionary mapping names to SAE configurations (e.g., StandardTrainingSAEConfig or TopKTrainingSAEConfig).
• hook_names: A dictionary mapping SAE names to their specific hook points, or a single string if all SAEs share the same hook.

Checkpointing: Checkpoints include per-SAE subdirectories. To resume, use resume_from_checkpoint=<checkpoint_dir> and ensure the keys in cfg.saes match the subdirectory names.

from sae_lens import (
    MultiSAETrainingRunner,
    MultiSAETrainingRunnerConfig,
    StandardTrainingSAEConfig,
    TopKTrainingSAEConfig,
    LoggingConfig,
)

cfg = MultiSAETrainingRunnerConfig(
    saes={
        "h5_l1_low":  StandardTrainingSAEConfig(d_in=768, d_sae=16 * 1024, l1_coefficient=2.0),
        "h5_l1_high": StandardTrainingSAEConfig(d_in=768, d_sae=16 * 1024, l1_coefficient=5.0),
        "h10_topk":   TopKTrainingSAEConfig(d_in=768, d_sae=16 * 1024, k=64),
    },
    hook_names={
        "h5_l1_low":  "blocks.5.hook_resid_pre",
        "h5_l1_high": "blocks.5.hook_resid_pre",
        "h10_topk":   "blocks.10.hook_resid_pre",
    },
    # OR a single string when every SAE shares one hook:
    # hook_names="blocks.5.hook_resid_pre",

    model_name="gpt2",
    dataset_path="apollo-research/Skylion007-openwebtext-tokenizer-gpt2",
    training_tokens=int(1e8),
    train_batch_size_tokens=4096,
    output_path="output/sweep_run_1",
    logger=LoggingConfig(log_to_wandb=True, wandb_project="multi_sae_sweep"),
)

trained_saes = MultiSAETrainingRunner(cfg).run()  # dict[name, TrainingSAE]