sae-lens 6.32.1__py3-none-any.whl → 6.33.0__py3-none-any.whl

sae_lens/__init__.py

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

sae_lens/loading/pretrained_saes_directory.py

@@ -57,28 +57,6 @@ def get_pretrained_saes_directory() -> dict[str, PretrainedSAELookup]:
     return directory
 
 
-def get_norm_scaling_factor(release: str, sae_id: str) -> float | None:
-    """
-    Retrieve the norm_scaling_factor for a specific SAE if it exists.
-
-    Args:
-        release (str): The release name of the SAE.
-        sae_id (str): The ID of the specific SAE.
-
-    Returns:
-        float | None: The norm_scaling_factor if it exists, None otherwise.
-    """
-    package = "sae_lens"
-    yaml_file = files(package).joinpath("pretrained_saes.yaml")
-    with yaml_file.open("r") as file:
-        data = yaml.safe_load(file)
-        if release in data:
-            for sae_info in data[release]["saes"]:
-                if sae_info["id"] == sae_id:
-                    return sae_info.get("norm_scaling_factor")
-    return None
-
-
 def get_repo_id_and_folder_name(release: str, sae_id: str) -> tuple[str, str]:
     saes_directory = get_pretrained_saes_directory()
     sae_info = saes_directory.get(release, None)

sae_lens/saes/sae.py

@@ -45,7 +45,6 @@ from sae_lens.loading.pretrained_sae_loaders import (
 )
 from sae_lens.loading.pretrained_saes_directory import (
     get_config_overrides,
-    get_norm_scaling_factor,
     get_pretrained_saes_directory,
     get_releases_for_repo_id,
     get_repo_id_and_folder_name,
@@ -638,24 +637,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
                     stacklevel=2,
                 )
         elif sae_id not in sae_directory[release].saes_map:
-            # Handle special cases like Gemma Scope
-            if (
-                "gemma-scope" in release
-                and "canonical" not in release
-                and f"{release}-canonical" in sae_directory
-            ):
-                canonical_ids = list(
-                    sae_directory[release + "-canonical"].saes_map.keys()
-                )
-                # Shorten the lengthy string of valid IDs
-                if len(canonical_ids) > 5:
-                    str_canonical_ids = str(canonical_ids[:5])[:-1] + ", ...]"
-                else:
-                    str_canonical_ids = str(canonical_ids)
-                value_suffix = f" If you don't want to specify an L0 value, consider using release {release}-canonical which has valid IDs {str_canonical_ids}"
-            else:
-                value_suffix = ""
-
             valid_ids = list(sae_directory[release].saes_map.keys())
             # Shorten the lengthy string of valid IDs
             if len(valid_ids) > 5:
@@ -665,7 +646,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
 
             raise ValueError(
                 f"ID {sae_id} not found in release {release}. Valid IDs are {str_valid_ids}."
-                + value_suffix
             )
 
         conversion_loader = (
@@ -702,17 +682,6 @@ class SAE(HookedRootModule, Generic[T_SAE_CONFIG], ABC):
         sae.process_state_dict_for_loading(state_dict)
         sae.load_state_dict(state_dict, assign=True)
 
-        # Apply normalization if needed
-        if cfg_dict.get("normalize_activations") == "expected_average_only_in":
-            norm_scaling_factor = get_norm_scaling_factor(release, sae_id)
-            if norm_scaling_factor is not None:
-                sae.fold_activation_norm_scaling_factor(norm_scaling_factor)
-                cfg_dict["normalize_activations"] = "none"
-            else:
-                warnings.warn(
-                    f"norm_scaling_factor not found for {release} and {sae_id}, but normalize_activations is 'expected_average_only_in'. Skipping normalization folding."
-                )
-
         # the loaders should already handle the dtype / device conversion
         # but this is a fallback to guarantee the SAE is on the correct device and dtype
         return (

sae_lens/synthetic/__init__.py

@@ -50,6 +50,13 @@ from sae_lens.synthetic.plotting import (
     find_best_feature_ordering_from_sae,
     plot_sae_feature_similarity,
 )
+from sae_lens.synthetic.stats import (
+    CorrelationMatrixStats,
+    SuperpositionStats,
+    compute_correlation_matrix_stats,
+    compute_low_rank_correlation_matrix_stats,
+    compute_superposition_stats,
+)
 from sae_lens.synthetic.training import (
     SyntheticActivationIterator,
     train_toy_sae,
@@ -80,6 +87,12 @@ __all__ = [
     "orthogonal_initializer",
     "FeatureDictionaryInitializer",
     "cosine_similarities",
+    # Statistics
+    "compute_correlation_matrix_stats",
+    "compute_low_rank_correlation_matrix_stats",
+    "compute_superposition_stats",
+    "CorrelationMatrixStats",
+    "SuperpositionStats",
     # Training utilities
     "SyntheticActivationIterator",
     "SyntheticDataEvalResult",

sae_lens/synthetic/correlation.py

@@ -3,6 +3,7 @@ from typing import NamedTuple
 
 import torch
 
+from sae_lens import logger
 from sae_lens.util import str_to_dtype
 
 
@@ -268,7 +269,7 @@ def generate_random_correlation_matrix(
 def generate_random_low_rank_correlation_matrix(
     num_features: int,
     rank: int,
-    correlation_scale: float = 0.1,
+    correlation_scale: float = 0.075,
     seed: int | None = None,
     device: torch.device | str = "cpu",
     dtype: torch.dtype | str = torch.float32,
@@ -331,20 +332,17 @@ def generate_random_low_rank_correlation_matrix(
     factor_sq_sum = (factor**2).sum(dim=1)
     diag_term = 1 - factor_sq_sum
 
-    # Ensure diagonal terms are at least _MIN_DIAG for numerical stability
-    # If any diagonal term is too small, scale down the factor matrix
-    if torch.any(diag_term < _MIN_DIAG):
-        # Scale factor so max row norm squared is at most (1 - _MIN_DIAG)
-        # This ensures all diagonal terms are >= _MIN_DIAG
-        max_factor_contribution = 1 - _MIN_DIAG
-        max_sq_sum = factor_sq_sum.max()
-        scale = torch.sqrt(
-            torch.tensor(max_factor_contribution, device=device, dtype=dtype)
-            / max_sq_sum
+    # alternatively, we can rescale each row independently to ensure the diagonal is 1
+    mask = diag_term < _MIN_DIAG
+    factor[mask, :] *= torch.sqrt((1 - _MIN_DIAG) / factor_sq_sum[mask].unsqueeze(1))
+    factor_sq_sum = (factor**2).sum(dim=1)
+    diag_term = 1 - factor_sq_sum
+
+    total_rescaled = mask.sum().item()
+    if total_rescaled > 0:
+        logger.warning(
+            f"{total_rescaled} / {num_features} rows were capped. Either reduce the rank or reduce the correlation_scale to avoid rescaling."
         )
-        factor = factor * scale
-        factor_sq_sum = (factor**2).sum(dim=1)
-        diag_term = 1 - factor_sq_sum
 
     return LowRankCorrelationMatrix(
         correlation_factor=factor, correlation_diag=diag_term

sae_lens/synthetic/stats.py

@@ -0,0 +1,205 @@
+from dataclasses import dataclass
+
+import torch
+
+from sae_lens.synthetic.correlation import LowRankCorrelationMatrix
+from sae_lens.synthetic.feature_dictionary import FeatureDictionary
+
+
+@dataclass
+class CorrelationMatrixStats:
+    """Statistics computed from a correlation matrix."""
+
+    rms_correlation: float  # Root mean square of off-diagonal correlations
+    mean_correlation: float  # Mean of off-diagonal correlations (not absolute)
+    num_features: int
+
+
+@torch.no_grad()
+def compute_correlation_matrix_stats(
+    correlation_matrix: torch.Tensor,
+) -> CorrelationMatrixStats:
+    """Compute correlation statistics from a dense correlation matrix.
+
+    Args:
+        correlation_matrix: Dense correlation matrix of shape (n, n)
+
+    Returns:
+        CorrelationMatrixStats with correlation statistics
+    """
+    num_features = correlation_matrix.shape[0]
+
+    # Extract off-diagonal elements
+    mask = ~torch.eye(num_features, dtype=torch.bool, device=correlation_matrix.device)
+    off_diag = correlation_matrix[mask]
+
+    rms_correlation = (off_diag**2).mean().sqrt().item()
+    mean_correlation = off_diag.mean().item()
+
+    return CorrelationMatrixStats(
+        rms_correlation=rms_correlation,
+        mean_correlation=mean_correlation,
+        num_features=num_features,
+    )
+
+
+@torch.no_grad()
+def compute_low_rank_correlation_matrix_stats(
+    correlation_matrix: LowRankCorrelationMatrix,
+) -> CorrelationMatrixStats:
+    """Compute correlation statistics from a LowRankCorrelationMatrix.
+
+    The correlation matrix is represented as:
+        correlation = factor @ factor.T + diag(diag_term)
+
+    The off-diagonal elements are simply factor @ factor.T (the diagonal term
+    only affects the diagonal).
+
+    All statistics are computed efficiently in O(n*r²) time and O(r²) memory
+    without materializing the full n×n correlation matrix.
+
+    Args:
+        correlation_matrix: Low-rank correlation matrix
+
+    Returns:
+        CorrelationMatrixStats with correlation statistics
+    """
+
+    factor = correlation_matrix.correlation_factor
+    num_features = factor.shape[0]
+    num_off_diag = num_features * (num_features - 1)
+
+    # RMS correlation: uses ||F @ F.T||_F² = ||F.T @ F||_F²
+    # This avoids computing the (num_features, num_features) matrix
+    G = factor.T @ factor  # (rank, rank) - small!
+    frobenius_sq = (G**2).sum()
+    row_norms_sq = (factor**2).sum(dim=1)  # ||F[i]||² for each row
+    diag_sq_sum = (row_norms_sq**2).sum()  # Σᵢ ||F[i]||⁴
+    off_diag_sq_sum = frobenius_sq - diag_sq_sum
+    rms_correlation = (off_diag_sq_sum / num_off_diag).sqrt().item()
+
+    # Mean correlation (not absolute): sum(C) = ||col_sums(F)||², trace(C) = Σ||F[i]||²
+    col_sums = factor.sum(dim=0)  # (rank,)
+    sum_all = (col_sums**2).sum()  # 1ᵀ C 1
+    trace_C = row_norms_sq.sum()
+    mean_correlation = ((sum_all - trace_C) / num_off_diag).item()
+
+    return CorrelationMatrixStats(
+        rms_correlation=rms_correlation,
+        mean_correlation=mean_correlation,
+        num_features=num_features,
+    )
+
+
+@dataclass
+class SuperpositionStats:
+    """Statistics measuring superposition in a feature dictionary."""
+
+    # Per-latent statistics: for each latent, max and percentile of |cos_sim| with others
+    max_abs_cos_sims: torch.Tensor  # Shape: (num_features,)
+    percentile_abs_cos_sims: dict[int, torch.Tensor]  # {percentile: (num_features,)}
+
+    # Summary statistics (means of the per-latent values)
+    mean_max_abs_cos_sim: float
+    mean_percentile_abs_cos_sim: dict[int, float]
+    mean_abs_cos_sim: float  # Mean |cos_sim| across all pairs
+
+    # Metadata
+    num_features: int
+    hidden_dim: int
+
+
+@torch.no_grad()
+def compute_superposition_stats(
+    feature_dictionary: FeatureDictionary,
+    batch_size: int = 1024,
+    device: str | torch.device | None = None,
+    percentiles: list[int] | None = None,
+) -> SuperpositionStats:
+    """Compute superposition statistics for a feature dictionary.
+
+    Computes pairwise cosine similarities in batches to handle large dictionaries.
+
+    For each latent i, computes:
+
+    - max |cos_sim(i, j)| over all j != i
+    - kth percentile of |cos_sim(i, j)| over all j != i (for each k in percentiles)
+
+    Args:
+        feature_dictionary: FeatureDictionary containing the feature vectors
+        batch_size: Number of features to process per batch
+        device: Device for computation (defaults to feature dictionary's device)
+        percentiles: List of percentiles to compute per latent (default: [95, 99])
+
+    Returns:
+        SuperpositionStats with superposition metrics
+    """
+    if percentiles is None:
+        percentiles = [95, 99]
+
+    feature_vectors = feature_dictionary.feature_vectors
+    num_features, hidden_dim = feature_vectors.shape
+
+    if num_features < 2:
+        raise ValueError("Need at least 2 features to compute superposition stats")
+    if device is None:
+        device = feature_vectors.device
+
+    # Normalize features to unit norm for cosine similarity
+    features_normalized = feature_vectors.to(device).float()
+    norms = torch.linalg.norm(features_normalized, dim=1, keepdim=True)
+    features_normalized = features_normalized / norms.clamp(min=1e-8)
+
+    # Track per-latent statistics
+    max_abs_cos_sims = torch.zeros(num_features, device=device)
+    percentile_abs_cos_sims = {
+        p: torch.zeros(num_features, device=device) for p in percentiles
+    }
+    sum_abs_cos_sim = 0.0
+    n_pairs = 0
+
+    # Process in batches: for each batch of features, compute similarities with all others
+    for i in range(0, num_features, batch_size):
+        batch_end = min(i + batch_size, num_features)
+        batch = features_normalized[i:batch_end]  # (batch_size, hidden_dim)
+
+        # Compute cosine similarities with all features: (batch_size, num_features)
+        cos_sims = batch @ features_normalized.T
+
+        # Absolute cosine similarities
+        abs_cos_sims = cos_sims.abs()
+
+        # Process each latent in the batch
+        for j, idx in enumerate(range(i, batch_end)):
+            # Get similarities with all other features (exclude self)
+            row = abs_cos_sims[j].clone()
+            row[idx] = 0.0  # Exclude self for max
+            max_abs_cos_sims[idx] = row.max()
+
+            # For percentiles, exclude self and compute
+            other_sims = torch.cat([abs_cos_sims[j, :idx], abs_cos_sims[j, idx + 1 :]])
+            for p in percentiles:
+                percentile_abs_cos_sims[p][idx] = torch.quantile(other_sims, p / 100.0)
+
+            # Sum for mean computation (only count pairs once - with features after this one)
+            sum_abs_cos_sim += abs_cos_sims[j, idx + 1 :].sum().item()
+            n_pairs += num_features - idx - 1
+
+    # Compute summary statistics
+    mean_max_abs_cos_sim = max_abs_cos_sims.mean().item()
+    mean_percentile_abs_cos_sim = {
+        p: percentile_abs_cos_sims[p].mean().item() for p in percentiles
+    }
+    mean_abs_cos_sim = sum_abs_cos_sim / n_pairs if n_pairs > 0 else 0.0
+
+    return SuperpositionStats(
+        max_abs_cos_sims=max_abs_cos_sims.cpu(),
+        percentile_abs_cos_sims={
+            p: v.cpu() for p, v in percentile_abs_cos_sims.items()
+        },
+        mean_max_abs_cos_sim=mean_max_abs_cos_sim,
+        mean_percentile_abs_cos_sim=mean_percentile_abs_cos_sim,
+        mean_abs_cos_sim=mean_abs_cos_sim,
+        num_features=num_features,
+        hidden_dim=hidden_dim,
+    )

{sae_lens-6.32.1.dist-info → sae_lens-6.33.0.dist-info}/METADATA

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

{sae_lens-6.32.1.dist-info → sae_lens-6.33.0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-sae_lens/__init__.py,sha256=Y_TVKGehpnTvQw8tvIn0fjo8uAw-XAYi7carZS_cRjQ,5168
+sae_lens/__init__.py,sha256=gHaxlySzLskrAUg2oUZ3aOpnI3U_AVIHce-agGJL9rI,5168
 sae_lens/analysis/__init__.py,sha256=FZExlMviNwWR7OGUSGRbd0l-yUDGSp80gglI_ivILrY,412
 sae_lens/analysis/compat.py,sha256=cgE3nhFcJTcuhppxbL71VanJS7YqVEOefuneB5eOaPw,538
 sae_lens/analysis/hooked_sae_transformer.py,sha256=LpnjxSAcItqqXA4SJyZuxY4Ki0UOuWV683wg9laYAsY,14050
@@ -12,7 +12,7 @@ sae_lens/llm_sae_training_runner.py,sha256=M7BK55gSFYu2qFQKABHX3c8i46P1LfODCeyHF
 sae_lens/load_model.py,sha256=C8AMykctj6H7tz_xRwB06-EXj6TfW64PtSJZR5Jxn1Y,8649
 sae_lens/loading/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sae_lens/loading/pretrained_sae_loaders.py,sha256=kshvA0NivOc7B3sL19lHr_zrC_DDfW2T6YWb5j0hgAk,63930
-sae_lens/loading/pretrained_saes_directory.py,sha256=1at_aQbD8WFywchQCKuwfP-yvCq_Z2aUYrpKDnSN5Nc,4283
+sae_lens/loading/pretrained_saes_directory.py,sha256=lSnHl77IO5dd7iO21ynCzZNMrzuJAT8Za4W5THNq0qw,3554
 sae_lens/pretokenize_runner.py,sha256=amJwIz3CKi2s2wNQn-10E7eAV7VFhNqtFDNTeTkwEI8,7133
 sae_lens/pretrained_saes.yaml,sha256=IVBLLR8_XNllJ1O-kVv9ED4u0u44Yn8UOL9R-f8Idp4,1511936
 sae_lens/registry.py,sha256=nhy7BPSudSATqW4lo9H_k3Na7sfGHmAf9v-3wpnLL_o,1490
@@ -22,20 +22,21 @@ sae_lens/saes/gated_sae.py,sha256=V_2ZNlV4gRD-rX5JSx1xqY7idT8ChfdQ5yxWDdu_6hg,88
 sae_lens/saes/jumprelu_sae.py,sha256=miiF-xI_yXdV9EkKjwAbU9zSMsx9KtKCz5YdXEzkN8g,13313
 sae_lens/saes/matching_pursuit_sae.py,sha256=08_G9p1YMLnE5qZVCPp6gll-iG6nHRbMMASf4_bkFt8,13207
 sae_lens/saes/matryoshka_batchtopk_sae.py,sha256=Qr6htt1HHOuO9FXI9hyaPSnGFIiJG-v7y1t1CEmkFzM,5995
-sae_lens/saes/sae.py,sha256=xRmgiLuaFlDCv8SyLbL-5TwdrWHpNLqSGe8mC1L6WcI,40942
+sae_lens/saes/sae.py,sha256=wkwqzNragj-1189cV52S3_XeRtEgBd2ZNwvL2EsKkWw,39429
 sae_lens/saes/standard_sae.py,sha256=_hldNZkFPAf9VGrxouR1-tN8T2OEk8IkWBcXoatrC1o,5749
 sae_lens/saes/temporal_sae.py,sha256=S44sPddVj2xujA02CC8gT1tG0in7c_CSAhspu9FHbaA,13273
 sae_lens/saes/topk_sae.py,sha256=vrMRPrCQR1o8G_kXqY_EAoGZARupkQNFB2dNZVLsusE,21073
 sae_lens/saes/transcoder.py,sha256=CTpJs8ASOK06npih7gZHygZuxqTR7HICWlOYfTiKjI4,13501
-sae_lens/synthetic/__init__.py,sha256=MtTnGkTfHV2WjkIgs7zZyx10EK9U5fjOHXy69Aq3uKw,3095
+sae_lens/synthetic/__init__.py,sha256=hRRA3xhEQUacGyFbJXkLVYg_8A1bbSYYWlVovb0g4KU,3503
 sae_lens/synthetic/activation_generator.py,sha256=8L9nwC4jFRv_wg3QN-n1sFwX8w1NqwJMysWaJ41lLlY,15197
-sae_lens/synthetic/correlation.py,sha256=tMTLo9fBfDpeXwqhyUgFqnTipj9x2W0t4oEtNxB7AG0,13256
+sae_lens/synthetic/correlation.py,sha256=tD8J9abWfuFtGZrEbbFn4P8FeTcNKF2V5JhBLwDUmkg,13146
 sae_lens/synthetic/evals.py,sha256=Nhi314ZnRgLfhBj-3tm_zzI-pGyFTcwllDXbIpPFXeU,4584
 sae_lens/synthetic/feature_dictionary.py,sha256=Nd4xjSTxKMnKilZ3uYi8Gv5SS5D4bv4wHiSL1uGB69E,6933
 sae_lens/synthetic/firing_probabilities.py,sha256=yclz1pWl5gE1r8LAxFvzQS88Lxwk5-3r8BCX9HLVejA,3370
 sae_lens/synthetic/hierarchy.py,sha256=nm7nwnTswktVJeKUsRZ0hLOdXcFWGbxnA1b6lefHm-4,33592
 sae_lens/synthetic/initialization.py,sha256=orMGW-786wRDHIS2W7bEH0HmlVFQ4g2z4bnnwdv5w4s,1386
 sae_lens/synthetic/plotting.py,sha256=5lFrej1QOkGAcImFNo5-o-8mI_rUVqvEI57KzUQPPtQ,8208
+sae_lens/synthetic/stats.py,sha256=BoDPKDx8pgFF5Ko_IaBRZTczm7-ANUIRjjF5W5Qh3Lk,7441
 sae_lens/synthetic/training.py,sha256=fHcX2cZ6nDupr71GX0Gk17f1NvQ0SKIVXIA6IuAb2dw,5692
 sae_lens/tokenization_and_batching.py,sha256=uoHtAs9z3XqG0Fh-iQVYVlrbyB_E3kFFhrKU30BosCo,5438
 sae_lens/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -48,7 +49,7 @@ sae_lens/training/types.py,sha256=1FpLx_Doda9vZpmfm-x1e8wGBYpyhe9Kpb_JuM5nIFM,90
 sae_lens/training/upload_saes_to_huggingface.py,sha256=r_WzI1zLtGZ5TzAxuG3xa_8T09j3zXJrWd_vzPsPGkQ,4469
 sae_lens/tutorial/tsea.py,sha256=fd1am_XXsf2KMbByDapJo-2qlxduKaa62Z2qcQZ3QKU,18145
 sae_lens/util.py,sha256=oIMoeyEP2IzcPFmRbKUzOAycgEyMcOasGeO_BGVZbc4,4846
-sae_lens-6.32.1.dist-info/METADATA,sha256=TcO6hFEXKdbLp32UTiVluHcMXFetfYJDqTHNCsx9PRw,6566
-sae_lens-6.32.1.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
-sae_lens-6.32.1.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
-sae_lens-6.32.1.dist-info/RECORD,,
+sae_lens-6.33.0.dist-info/METADATA,sha256=X6XqngWTNEsfdaPPWXxtF8Kvdp8fAk8i68sfRtDb2xo,6566
+sae_lens-6.33.0.dist-info/WHEEL,sha256=3ny-bZhpXrU6vSQ1UPG34FoxZBp3lVcvK0LkgUz6VLk,88
+sae_lens-6.33.0.dist-info/licenses/LICENSE,sha256=DW6e-hDosiu4CfW0-imI57sV1I5f9UEslpviNQcOAKs,1069
+sae_lens-6.33.0.dist-info/RECORD,,