sae-lens 6.26.1__py3-none-any.whl → 6.28.2__py3-none-any.whl

sae_lens/synthetic/evals.py

@@ -0,0 +1,141 @@
+"""
+Utilities for training SAEs on synthetic data.
+
+This module provides helpers for:
+
+- Generating training data from feature dictionaries
+- Training SAEs on synthetic data
+- Evaluating SAEs against known ground truth features
+- Initializing SAEs to match feature dictionaries
+"""
+
+from dataclasses import dataclass
+
+import torch
+from scipy.optimize import linear_sum_assignment
+
+from sae_lens.synthetic.activation_generator import ActivationGenerator
+from sae_lens.synthetic.feature_dictionary import FeatureDictionary
+
+
+def mean_correlation_coefficient(
+    features_a: torch.Tensor,
+    features_b: torch.Tensor,
+) -> float:
+    """
+    Compute Mean Correlation Coefficient (MCC) between two sets of feature vectors.
+
+    MCC measures how well learned features align with ground truth features by finding
+    an optimal one-to-one matching using the Hungarian algorithm and computing the
+    mean absolute cosine similarity of matched pairs.
+
+    Reference: O'Neill et al. "Compute Optimal Inference and Provable Amortisation
+    Gap in Sparse Autoencoders" (arXiv:2411.13117)
+
+    Args:
+        features_a: Feature vectors of shape [num_features_a, hidden_dim]
+        features_b: Feature vectors of shape [num_features_b, hidden_dim]
+
+    Returns:
+        MCC score in range [0, 1], where 1 indicates perfect alignment
+    """
+    # Normalize to unit vectors
+    a_norm = features_a / features_a.norm(dim=1, keepdim=True).clamp(min=1e-8)
+    b_norm = features_b / features_b.norm(dim=1, keepdim=True).clamp(min=1e-8)
+
+    # Compute absolute cosine similarity matrix
+    cos_sim = torch.abs(a_norm @ b_norm.T)
+
+    # Convert to cost matrix for Hungarian algorithm (which minimizes)
+    cost_matrix = 1 - cos_sim.cpu().numpy()
+
+    # Find optimal matching
+    row_ind, col_ind = linear_sum_assignment(cost_matrix)
+
+    # Compute mean of matched similarities
+    matched_similarities = cos_sim[row_ind, col_ind]
+    return matched_similarities.mean().item()
+
+
+@dataclass
+class SyntheticDataEvalResult:
+    """Results from evaluating an SAE on synthetic data."""
+
+    true_l0: float
+    """Average L0 of the true feature activations"""
+
+    sae_l0: float
+    """Average L0 of the SAE's latent activations"""
+
+    dead_latents: int
+    """Number of SAE latents that never fired"""
+
+    shrinkage: float
+    """Average ratio of SAE output norm to input norm (1.0 = no shrinkage)"""
+
+    mcc: float
+    """Mean Correlation Coefficient between SAE decoder and ground truth features"""
+
+
+@torch.no_grad()
+def eval_sae_on_synthetic_data(
+    sae: torch.nn.Module,
+    feature_dict: FeatureDictionary,
+    activations_generator: ActivationGenerator,
+    num_samples: int = 100_000,
+) -> SyntheticDataEvalResult:
+    """
+    Evaluate an SAE on synthetic data with known ground truth.
+
+    Args:
+        sae: The SAE to evaluate. Must have encode() and decode() methods.
+        feature_dict: The feature dictionary used to generate activations
+        activations_generator: Generator that produces feature activations
+        num_samples: Number of samples to use for evaluation
+
+    Returns:
+        SyntheticDataEvalResult containing evaluation metrics
+    """
+    sae.eval()
+
+    # Generate samples
+    feature_acts = activations_generator.sample(num_samples)
+    true_l0 = (feature_acts > 0).float().sum(dim=-1).mean().item()
+    hidden_acts = feature_dict(feature_acts)
+
+    # Filter out entries where no features fire
+    non_zero_mask = hidden_acts.norm(dim=-1) > 0
+    hidden_acts_filtered = hidden_acts[non_zero_mask]
+
+    # Get SAE reconstructions
+    sae_latents = sae.encode(hidden_acts_filtered)  # type: ignore[attr-defined]
+    sae_output = sae.decode(sae_latents)  # type: ignore[attr-defined]
+
+    sae_l0 = (sae_latents > 0).float().sum(dim=-1).mean().item()
+    dead_latents = int(
+        ((sae_latents == 0).sum(dim=0) == sae_latents.shape[0]).sum().item()
+    )
+    if hidden_acts_filtered.shape[0] == 0:
+        shrinkage = 0.0
+    else:
+        shrinkage = (
+            (
+                sae_output.norm(dim=-1)
+                / hidden_acts_filtered.norm(dim=-1).clamp(min=1e-8)
+            )
+            .mean()
+            .item()
+        )
+
+    # Compute MCC between SAE decoder and ground truth features
+    sae_decoder: torch.Tensor = sae.W_dec  # type: ignore[attr-defined]
+    gt_features = feature_dict.feature_vectors
+    mcc = mean_correlation_coefficient(sae_decoder, gt_features)
+
+    return SyntheticDataEvalResult(
+        true_l0=true_l0,
+        sae_l0=sae_l0,
+        dead_latents=dead_latents,
+        shrinkage=shrinkage,
+        mcc=mcc,
+    )

sae_lens/synthetic/feature_dictionary.py

@@ -0,0 +1,176 @@
+"""
+Feature dictionary for generating synthetic activations.
+
+A FeatureDictionary maps feature activations (sparse coefficients) to dense hidden activations
+by multiplying with a learned or constructed feature embedding matrix.
+"""
+
+from typing import Callable
+
+import torch
+from torch import nn
+from tqdm import tqdm
+
+FeatureDictionaryInitializer = Callable[["FeatureDictionary"], None]
+
+
+def orthogonalize_embeddings(
+    embeddings: torch.Tensor,
+    num_steps: int = 200,
+    lr: float = 0.01,
+    show_progress: bool = False,
+    chunk_size: int = 1024,
+) -> torch.Tensor:
+    """
+    Orthogonalize embeddings using gradient descent with chunked computation.
+
+    Uses chunked computation to avoid O(n²) memory usage when computing pairwise
+    dot products. Memory usage is O(chunk_size × n) instead of O(n²).
+
+    Args:
+        embeddings: Tensor of shape [num_vectors, hidden_dim]
+        num_steps: Number of optimization steps
+        lr: Learning rate for Adam optimizer
+        show_progress: Whether to show progress bar
+        chunk_size: Number of vectors to process at once. Smaller values use less
+            memory but may be slower.
+
+    Returns:
+        Orthogonalized embeddings of the same shape, normalized to unit length.
+    """
+    num_vectors = embeddings.shape[0]
+    # Create a detached copy and normalize, then enable gradients
+    embeddings = embeddings.detach().clone()
+    embeddings = embeddings / embeddings.norm(p=2, dim=1, keepdim=True).clamp(min=1e-8)
+    embeddings = embeddings.requires_grad_(True)
+
+    optimizer = torch.optim.Adam([embeddings], lr=lr)  # type: ignore[list-item]
+
+    pbar = tqdm(
+        range(num_steps), desc="Orthogonalizing vectors", disable=not show_progress
+    )
+    for _ in pbar:
+        optimizer.zero_grad()
+
+        off_diag_loss = torch.tensor(0.0, device=embeddings.device)
+        diag_loss = torch.tensor(0.0, device=embeddings.device)
+
+        for i in range(0, num_vectors, chunk_size):
+            end_i = min(i + chunk_size, num_vectors)
+            chunk = embeddings[i:end_i]
+            chunk_dots = chunk @ embeddings.T  # [chunk_size, num_vectors]
+
+            # Create mask to zero out diagonal elements for this chunk
+            # Diagonal of full matrix: position (i+k, i+k) → in chunk_dots: (k, i+k)
+            chunk_len = end_i - i
+            row_indices = torch.arange(chunk_len, device=embeddings.device)
+            col_indices = i + row_indices  # column indices in full matrix
+
+            # Boolean mask: True for off-diagonal elements we want to include
+            off_diag_mask = torch.ones_like(chunk_dots, dtype=torch.bool)
+            off_diag_mask[row_indices, col_indices] = False
+
+            off_diag_loss = off_diag_loss + chunk_dots[off_diag_mask].pow(2).sum()
+
+            # Diagonal loss: keep self-dot-products at 1
+            diag_vals = chunk_dots[row_indices, col_indices]
+            diag_loss = diag_loss + (diag_vals - 1).pow(2).sum()
+
+        loss = off_diag_loss + num_vectors * diag_loss
+        loss.backward()
+        optimizer.step()
+        pbar.set_description(f"loss: {loss.item():.3f}")
+
+    with torch.no_grad():
+        embeddings = embeddings / embeddings.norm(p=2, dim=1, keepdim=True).clamp(
+            min=1e-8
+        )
+    return embeddings.detach().clone()
+
+
+def orthogonal_initializer(
+    num_steps: int = 200,
+    lr: float = 0.01,
+    show_progress: bool = False,
+    chunk_size: int = 1024,
+) -> FeatureDictionaryInitializer:
+    def initializer(feature_dict: "FeatureDictionary") -> None:
+        feature_dict.feature_vectors.data = orthogonalize_embeddings(
+            feature_dict.feature_vectors,
+            num_steps=num_steps,
+            lr=lr,
+            show_progress=show_progress,
+            chunk_size=chunk_size,
+        )
+
+    return initializer
+
+
+class FeatureDictionary(nn.Module):
+    """
+    A feature dictionary that maps sparse feature activations to dense hidden activations.
+
+    This class creates a set of feature vectors (the "dictionary") and provides methods
+    to generate hidden activations from feature activations via a linear transformation.
+
+    The feature vectors can be configured to have a specific pairwise cosine similarity,
+    which is useful for controlling the difficulty of sparse recovery.
+
+    Attributes:
+        feature_vectors: Parameter of shape [num_features, hidden_dim] containing the
+            feature embedding vectors
+        bias: Parameter of shape [hidden_dim] containing the bias term (zeros if bias=False)
+    """
+
+    feature_vectors: nn.Parameter
+    bias: nn.Parameter
+
+    def __init__(
+        self,
+        num_features: int,
+        hidden_dim: int,
+        bias: bool = False,
+        initializer: FeatureDictionaryInitializer | None = orthogonal_initializer(),
+        device: str | torch.device = "cpu",
+    ):
+        """
+        Create a new FeatureDictionary.
+
+        Args:
+            num_features: Number of features in the dictionary
+            hidden_dim: Dimensionality of the hidden space
+            bias: Whether to include a bias term in the embedding
+            initializer: Initializer function to use. If None, the embeddings are initialized to random unit vectors. By default will orthogonalize embeddings.
+            device: Device to use for the feature dictionary.
+        """
+        super().__init__()
+        self.num_features = num_features
+        self.hidden_dim = hidden_dim
+
+        # Initialize feature vectors as unit vectors
+        embeddings = torch.randn(num_features, hidden_dim, device=device)
+        embeddings = embeddings / embeddings.norm(p=2, dim=1, keepdim=True).clamp(
+            min=1e-8
+        )
+        self.feature_vectors = nn.Parameter(embeddings)
+
+        # Initialize bias (zeros if not using bias, but still a parameter for consistent API)
+        self.bias = nn.Parameter(
+            torch.zeros(hidden_dim, device=device), requires_grad=bias
+        )
+
+        if initializer is not None:
+            initializer(self)
+
+    def forward(self, feature_activations: torch.Tensor) -> torch.Tensor:
+        """
+        Convert feature activations to hidden activations.
+
+        Args:
+            feature_activations: Tensor of shape [batch, num_features] containing
+                sparse feature activation values
+
+        Returns:
+            Tensor of shape [batch, hidden_dim] containing dense hidden activations
+        """
+        return feature_activations @ self.feature_vectors + self.bias

sae_lens/synthetic/firing_probabilities.py

@@ -0,0 +1,104 @@
+"""
+Helper functions for generating firing probability distributions.
+"""
+
+import torch
+
+
+def zipfian_firing_probabilities(
+    num_features: int,
+    exponent: float = 1.0,
+    max_prob: float = 0.3,
+    min_prob: float = 0.01,
+) -> torch.Tensor:
+    """
+    Generate firing probabilities following a Zipfian (power-law) distribution.
+
+    Creates probabilities where a few features fire frequently and most fire rarely,
+    which mirrors the distribution often observed in real neural network features.
+
+    Args:
+        num_features: Number of features to generate probabilities for
+        exponent: Zipf exponent (higher = steeper dropoff). Default 1.0.
+        max_prob: Maximum firing probability (for the most frequent feature)
+        min_prob: Minimum firing probability (for the least frequent feature)
+
+    Returns:
+        Tensor of shape [num_features] with firing probabilities in descending order
+    """
+    if num_features < 1:
+        raise ValueError("num_features must be at least 1")
+    if exponent <= 0:
+        raise ValueError("exponent must be positive")
+    if not 0 < min_prob < max_prob <= 1:
+        raise ValueError("Must have 0 < min_prob < max_prob <= 1")
+
+    ranks = torch.arange(1, num_features + 1, dtype=torch.float32)
+    probs = 1.0 / ranks**exponent
+
+    # Scale to [min_prob, max_prob]
+    if num_features == 1:
+        return torch.tensor([max_prob])
+
+    probs_min, probs_max = probs.min(), probs.max()
+    return min_prob + (max_prob - min_prob) * (probs - probs_min) / (
+        probs_max - probs_min
+    )
+
+
+def linear_firing_probabilities(
+    num_features: int,
+    max_prob: float = 0.3,
+    min_prob: float = 0.01,
+) -> torch.Tensor:
+    """
+    Generate firing probabilities that decay linearly from max to min.
+
+    Args:
+        num_features: Number of features to generate probabilities for
+        max_prob: Firing probability for the first feature
+        min_prob: Firing probability for the last feature
+
+    Returns:
+        Tensor of shape [num_features] with linearly decaying probabilities
+    """
+    if num_features < 1:
+        raise ValueError("num_features must be at least 1")
+    if not 0 < min_prob <= max_prob <= 1:
+        raise ValueError("Must have 0 < min_prob <= max_prob <= 1")
+
+    if num_features == 1:
+        return torch.tensor([max_prob])
+
+    return torch.linspace(max_prob, min_prob, num_features)
+
+
+def random_firing_probabilities(
+    num_features: int,
+    max_prob: float = 0.5,
+    min_prob: float = 0.01,
+    seed: int | None = None,
+) -> torch.Tensor:
+    """
+    Generate random firing probabilities uniformly sampled from a range.
+
+    Args:
+        num_features: Number of features to generate probabilities for
+        max_prob: Maximum firing probability
+        min_prob: Minimum firing probability
+        seed: Optional random seed for reproducibility
+
+    Returns:
+        Tensor of shape [num_features] with random firing probabilities
+    """
+    if num_features < 1:
+        raise ValueError("num_features must be at least 1")
+    if not 0 < min_prob < max_prob <= 1:
+        raise ValueError("Must have 0 < min_prob < max_prob <= 1")
+
+    generator = torch.Generator()
+    if seed is not None:
+        generator.manual_seed(seed)
+
+    probs = torch.rand(num_features, generator=generator, dtype=torch.float32)
+    return min_prob + (max_prob - min_prob) * probs