sae-lens 6.26.0__py3-none-any.whl → 6.28.1__py3-none-any.whl

sae_lens/synthetic/activation_generator.py

@@ -0,0 +1,215 @@
+"""
+Functions for generating synthetic feature activations.
+"""
+
+from collections.abc import Callable, Sequence
+
+import torch
+from scipy.stats import norm
+from torch import nn
+from torch.distributions import MultivariateNormal
+
+from sae_lens.util import str_to_dtype
+
+ActivationsModifier = Callable[[torch.Tensor], torch.Tensor]
+ActivationsModifierInput = ActivationsModifier | Sequence[ActivationsModifier] | None
+
+
+class ActivationGenerator(nn.Module):
+    """
+    Generator for synthetic feature activations.
+
+    This module provides a generator for synthetic feature activations with controlled properties.
+    """
+
+    num_features: int
+    firing_probabilities: torch.Tensor
+    std_firing_magnitudes: torch.Tensor
+    mean_firing_magnitudes: torch.Tensor
+    modify_activations: ActivationsModifier | None
+    correlation_matrix: torch.Tensor | None
+    correlation_thresholds: torch.Tensor | None
+
+    def __init__(
+        self,
+        num_features: int,
+        firing_probabilities: torch.Tensor | float,
+        std_firing_magnitudes: torch.Tensor | float = 0.0,
+        mean_firing_magnitudes: torch.Tensor | float = 1.0,
+        modify_activations: ActivationsModifierInput = None,
+        correlation_matrix: torch.Tensor | None = None,
+        device: torch.device | str = "cpu",
+        dtype: torch.dtype | str = "float32",
+    ):
+        super().__init__()
+        self.num_features = num_features
+        self.firing_probabilities = _to_tensor(
+            firing_probabilities, num_features, device, dtype
+        )
+        self.std_firing_magnitudes = _to_tensor(
+            std_firing_magnitudes, num_features, device, dtype
+        )
+        self.mean_firing_magnitudes = _to_tensor(
+            mean_firing_magnitudes, num_features, device, dtype
+        )
+        self.modify_activations = _normalize_modifiers(modify_activations)
+        self.correlation_thresholds = None
+        if correlation_matrix is not None:
+            _validate_correlation_matrix(correlation_matrix, num_features)
+            self.correlation_thresholds = torch.tensor(
+                [norm.ppf(1 - p.item()) for p in self.firing_probabilities],
+                device=device,
+                dtype=self.firing_probabilities.dtype,
+            )
+        self.correlation_matrix = correlation_matrix
+
+    def sample(self, batch_size: int) -> torch.Tensor:
+        """
+        Generate a batch of feature activations with controlled properties.
+
+        This is the main function for generating synthetic training data for SAEs.
+        Features fire independently according to their firing probabilities unless
+        a correlation matrix is provided.
+
+        Args:
+            batch_size: Number of samples to generate
+
+        Returns:
+            Tensor of shape [batch_size, num_features] with non-negative activations
+        """
+        # All tensors (firing_probabilities, std_firing_magnitudes, mean_firing_magnitudes)
+        # are on the same device from __init__ via _to_tensor()
+        device = self.firing_probabilities.device
+
+        if self.correlation_matrix is not None:
+            assert self.correlation_thresholds is not None
+            firing_features = _generate_correlated_features(
+                batch_size,
+                self.correlation_matrix,
+                self.correlation_thresholds,
+                device,
+            )
+        else:
+            firing_features = torch.bernoulli(
+                self.firing_probabilities.unsqueeze(0).expand(batch_size, -1)
+            )
+
+        firing_magnitude_delta = torch.normal(
+            torch.zeros_like(self.firing_probabilities)
+            .unsqueeze(0)
+            .expand(batch_size, -1),
+            self.std_firing_magnitudes.unsqueeze(0).expand(batch_size, -1),
+        )
+        firing_magnitude_delta[firing_features == 0] = 0
+        feature_activations = (
+            firing_features * self.mean_firing_magnitudes + firing_magnitude_delta
+        ).relu()
+
+        if self.modify_activations is not None:
+            feature_activations = self.modify_activations(feature_activations).relu()
+        return feature_activations
+
+    def forward(self, batch_size: int) -> torch.Tensor:
+        return self.sample(batch_size)
+
+
+def _generate_correlated_features(
+    batch_size: int,
+    correlation_matrix: torch.Tensor,
+    thresholds: torch.Tensor,
+    device: torch.device,
+) -> torch.Tensor:
+    """
+    Generate correlated binary features using multivariate Gaussian sampling.
+
+    Uses the Gaussian copula approach: sample from a multivariate normal
+    distribution, then threshold to get binary features.
+
+    Args:
+        batch_size: Number of samples to generate
+        correlation_matrix: Correlation matrix between features
+        thresholds: Pre-computed thresholds for each feature (from inverse normal CDF)
+        device: Device to generate samples on
+
+    Returns:
+        Binary feature matrix of shape [batch_size, num_features]
+    """
+    num_features = correlation_matrix.shape[0]
+
+    mvn = MultivariateNormal(
+        loc=torch.zeros(num_features, device=device, dtype=thresholds.dtype),
+        covariance_matrix=correlation_matrix.to(device=device, dtype=thresholds.dtype),
+    )
+
+    gaussian_samples = mvn.sample((batch_size,))
+    return (gaussian_samples > thresholds.unsqueeze(0)).float()
+
+
+def _to_tensor(
+    value: torch.Tensor | float,
+    num_features: int,
+    device: torch.device | str,
+    dtype: torch.dtype | str,
+) -> torch.Tensor:
+    dtype = str_to_dtype(dtype)
+    device = torch.device(device)
+    if not isinstance(value, torch.Tensor):
+        value = value * torch.ones(num_features, device=device, dtype=dtype)
+    if value.shape != (num_features,):
+        raise ValueError(
+            f"Value must be a tensor of shape ({num_features},) or a float. Got {value.shape}"
+        )
+    return value.to(device, dtype)
+
+
+def _normalize_modifiers(
+    modify_activations: ActivationsModifierInput,
+) -> ActivationsModifier | None:
+    """Convert modifier input to a single modifier or None."""
+    if modify_activations is None:
+        return None
+    if callable(modify_activations):
+        return modify_activations
+    # It's a sequence of modifiers - chain them
+    modifiers = list(modify_activations)
+    if len(modifiers) == 0:
+        return None
+    if len(modifiers) == 1:
+        return modifiers[0]
+
+    def chained(activations: torch.Tensor) -> torch.Tensor:
+        result = activations
+        for modifier in modifiers:
+            result = modifier(result)
+        return result
+
+    return chained
+
+
+def _validate_correlation_matrix(
+    correlation_matrix: torch.Tensor, num_features: int
+) -> None:
+    """Validate that a correlation matrix has correct properties.
+
+    Args:
+        correlation_matrix: The matrix to validate
+        num_features: Expected number of features (matrix should be [num_features, num_features])
+
+    Raises:
+        ValueError: If the matrix has incorrect shape, non-unit diagonal, or is not positive definite
+    """
+    expected_shape = (num_features, num_features)
+    if correlation_matrix.shape != expected_shape:
+        raise ValueError(
+            f"Correlation matrix must have shape {expected_shape}, "
+            f"got {tuple(correlation_matrix.shape)}"
+        )
+
+    diagonal = torch.diag(correlation_matrix)
+    if not torch.allclose(diagonal, torch.ones_like(diagonal)):
+        raise ValueError("Correlation matrix diagonal must be all 1s")
+
+    try:
+        torch.linalg.cholesky(correlation_matrix)
+    except RuntimeError as e:
+        raise ValueError("Correlation matrix must be positive definite") from e

sae_lens/synthetic/correlation.py

@@ -0,0 +1,170 @@
+import random
+
+import torch
+
+
+def create_correlation_matrix_from_correlations(
+    num_features: int,
+    correlations: dict[tuple[int, int], float] | None = None,
+    default_correlation: float = 0.0,
+) -> torch.Tensor:
+    """
+    Create a correlation matrix with specified pairwise correlations.
+
+    Args:
+        num_features: Number of features
+        correlations: Dict mapping (i, j) pairs to correlation values.
+            Pairs should have i < j.
+        default_correlation: Default correlation for unspecified pairs
+
+    Returns:
+        Correlation matrix of shape [num_features, num_features]
+    """
+    matrix = torch.eye(num_features) + default_correlation * (
+        1 - torch.eye(num_features)
+    )
+
+    if correlations is not None:
+        for (i, j), corr in correlations.items():
+            matrix[i, j] = corr
+            matrix[j, i] = corr
+
+    # Ensure matrix is symmetric (numerical precision)
+    matrix = (matrix + matrix.T) / 2
+
+    # Check positive definiteness and fix if necessary
+    # Use eigvalsh for symmetric matrices (returns real eigenvalues)
+    eigenvals = torch.linalg.eigvalsh(matrix)
+    if torch.any(eigenvals < -1e-6):
+        matrix = _fix_correlation_matrix(matrix)
+
+    return matrix
+
+
+def _fix_correlation_matrix(
+    matrix: torch.Tensor, min_eigenval: float = 1e-6
+) -> torch.Tensor:
+    """Fix a correlation matrix to be positive semi-definite."""
+    eigenvals, eigenvecs = torch.linalg.eigh(matrix)
+    eigenvals = torch.clamp(eigenvals, min=min_eigenval)
+    fixed_matrix = eigenvecs @ torch.diag(eigenvals) @ eigenvecs.T
+
+    diag_vals = torch.diag(fixed_matrix)
+    fixed_matrix = fixed_matrix / torch.sqrt(
+        diag_vals.unsqueeze(0) * diag_vals.unsqueeze(1)
+    )
+    fixed_matrix.fill_diagonal_(1.0)
+
+    return fixed_matrix
+
+
+def generate_random_correlations(
+    num_features: int,
+    positive_ratio: float = 0.5,
+    uncorrelated_ratio: float = 0.3,
+    min_correlation_strength: float = 0.1,
+    max_correlation_strength: float = 0.8,
+    seed: int | None = None,
+) -> dict[tuple[int, int], float]:
+    """
+    Generate random correlations between features with specified constraints.
+
+    Args:
+        num_features: Number of features
+        positive_ratio: Fraction of correlations that should be positive (0.0 to 1.0)
+        uncorrelated_ratio: Fraction of feature pairs that should remain uncorrelated (0.0 to 1.0)
+        min_correlation_strength: Minimum absolute correlation strength
+        max_correlation_strength: Maximum absolute correlation strength
+        seed: Random seed for reproducibility
+
+    Returns:
+        Dictionary mapping (i, j) pairs to correlation values
+    """
+    # Use local random number generator to avoid side effects on global state
+    rng = random.Random(seed)
+
+    # Validate inputs
+    if not 0.0 <= positive_ratio <= 1.0:
+        raise ValueError("positive_ratio must be between 0.0 and 1.0")
+    if not 0.0 <= uncorrelated_ratio <= 1.0:
+        raise ValueError("uncorrelated_ratio must be between 0.0 and 1.0")
+    if min_correlation_strength < 0:
+        raise ValueError("min_correlation_strength must be non-negative")
+    if max_correlation_strength > 1.0:
+        raise ValueError("max_correlation_strength must be <= 1.0")
+    if min_correlation_strength > max_correlation_strength:
+        raise ValueError("min_correlation_strength must be <= max_correlation_strength")
+
+    # Generate all possible feature pairs (i, j) where i < j
+    all_pairs = [
+        (i, j) for i in range(num_features) for j in range(i + 1, num_features)
+    ]
+    total_pairs = len(all_pairs)
+
+    if total_pairs == 0:
+        return {}
+
+    # Determine how many pairs to correlate vs leave uncorrelated
+    num_uncorrelated = int(total_pairs * uncorrelated_ratio)
+    num_correlated = total_pairs - num_uncorrelated
+
+    # Randomly select which pairs to correlate
+    correlated_pairs = rng.sample(all_pairs, num_correlated)
+
+    # For correlated pairs, determine positive vs negative
+    num_positive = int(num_correlated * positive_ratio)
+    num_negative = num_correlated - num_positive
+
+    # Assign signs
+    signs = [1] * num_positive + [-1] * num_negative
+    rng.shuffle(signs)
+
+    # Generate correlation strengths
+    correlations = {}
+    for pair, sign in zip(correlated_pairs, signs):
+        # Sample correlation strength uniformly from range
+        strength = rng.uniform(min_correlation_strength, max_correlation_strength)
+        correlations[pair] = sign * strength
+
+    return correlations
+
+
+def generate_random_correlation_matrix(
+    num_features: int,
+    positive_ratio: float = 0.5,
+    uncorrelated_ratio: float = 0.3,
+    min_correlation_strength: float = 0.1,
+    max_correlation_strength: float = 0.8,
+    seed: int | None = None,
+) -> torch.Tensor:
+    """
+    Generate a random correlation matrix with specified constraints.
+
+    This is a convenience function that combines generate_random_correlations()
+    and create_correlation_matrix_from_correlations() into a single call.
+
+    Args:
+        num_features: Number of features
+        positive_ratio: Fraction of correlations that should be positive (0.0 to 1.0)
+        uncorrelated_ratio: Fraction of feature pairs that should remain uncorrelated (0.0 to 1.0)
+        min_correlation_strength: Minimum absolute correlation strength
+        max_correlation_strength: Maximum absolute correlation strength
+        seed: Random seed for reproducibility
+
+    Returns:
+        Random correlation matrix of shape [num_features, num_features]
+    """
+    # Generate random correlations
+    correlations = generate_random_correlations(
+        num_features=num_features,
+        positive_ratio=positive_ratio,
+        uncorrelated_ratio=uncorrelated_ratio,
+        min_correlation_strength=min_correlation_strength,
+        max_correlation_strength=max_correlation_strength,
+        seed=seed,
+    )
+
+    # Create and return correlation matrix
+    return create_correlation_matrix_from_correlations(
+        num_features=num_features, correlations=correlations
+    )

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,138 @@
+"""
+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,
+    target_cos_sim: float = 0,
+    num_steps: int = 200,
+    lr: float = 0.01,
+    show_progress: bool = False,
+) -> torch.Tensor:
+    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]
+
+    # Create a mask to zero out diagonal elements (avoid in-place operations)
+    off_diagonal_mask = ~torch.eye(
+        num_vectors, dtype=torch.bool, device=embeddings.device
+    )
+
+    pbar = tqdm(
+        range(num_steps), desc="Orthogonalizing vectors", disable=not show_progress
+    )
+    for _ in pbar:
+        optimizer.zero_grad()
+
+        dot_products = embeddings @ embeddings.T
+        diff = dot_products - target_cos_sim
+        # Use masking instead of in-place fill_diagonal_
+        off_diagonal_diff = diff * off_diagonal_mask.float()
+        loss = off_diagonal_diff.pow(2).sum()
+        loss = loss + num_vectors * (dot_products.diag() - 1).pow(2).sum()
+
+        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
+) -> 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,
+        )
+
+    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(),
+    ):
+        """
+        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.
+        """
+        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)
+        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), 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