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

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

sae_lens/synthetic/hierarchy.py

@@ -0,0 +1,335 @@
+"""
+Hierarchical feature modifier for activation generators.
+
+This module provides HierarchyNode, which enforces hierarchical dependencies
+on feature activations. Child features are deactivated when their parent is inactive,
+and children can optionally be mutually exclusive.
+
+Based on Noa Nabeshima's Matryoshka SAEs:
+https://github.com/noanabeshima/matryoshka-saes/blob/main/toy_model.py
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import torch
+
+ActivationsModifier = Callable[[torch.Tensor], torch.Tensor]
+
+
+def _validate_hierarchy(roots: Sequence[HierarchyNode]) -> None:
+    """
+    Validate a forest of hierarchy trees.
+
+    Treats the input as children of a virtual root node and validates the
+    entire structure.
+
+    Checks that:
+    1. There are no loops (no node is its own ancestor)
+    2. Each node has at most one parent (no node appears in multiple children lists)
+    3. No feature index appears in multiple trees
+
+    Args:
+        roots: Root nodes of the hierarchy trees to validate
+
+    Raises:
+        ValueError: If the hierarchy is invalid
+    """
+    if not roots:
+        return
+
+    # Collect all nodes and check for loops, treating roots as children of virtual root
+    all_nodes: list[HierarchyNode] = []
+    virtual_root_id = id(roots)  # Use the list itself as virtual root identity
+
+    for root in roots:
+        all_nodes.append(root)
+        _collect_nodes_and_check_loops(root, all_nodes, ancestors={virtual_root_id})
+
+    # Check for multiple parents (same node appearing multiple times)
+    seen_ids: set[int] = set()
+    for node in all_nodes:
+        node_id = id(node)
+        if node_id in seen_ids:
+            node_desc = _node_description(node)
+            raise ValueError(
+                f"Node ({node_desc}) has multiple parents. "
+                "Each node must have at most one parent."
+            )
+        seen_ids.add(node_id)
+
+    # Check for overlapping feature indices across trees
+    if len(roots) > 1:
+        all_indices: set[int] = set()
+        for root in roots:
+            tree_indices = root.get_all_feature_indices()
+            overlap = all_indices & set(tree_indices)
+            if overlap:
+                raise ValueError(
+                    f"Feature indices {overlap} appear in multiple hierarchy trees. "
+                    "Each feature should belong to at most one hierarchy."
+                )
+            all_indices.update(tree_indices)
+
+
+def _collect_nodes_and_check_loops(
+    node: HierarchyNode,
+    all_nodes: list[HierarchyNode],
+    ancestors: set[int],
+) -> None:
+    """Recursively collect nodes and check for loops."""
+    node_id = id(node)
+
+    if node_id in ancestors:
+        node_desc = _node_description(node)
+        raise ValueError(f"Loop detected: node ({node_desc}) is its own ancestor.")
+
+    # Add to ancestors for children traversal
+    new_ancestors = ancestors | {node_id}
+
+    for child in node.children:
+        # Collect child (before recursing, so we can detect multiple parents)
+        all_nodes.append(child)
+        _collect_nodes_and_check_loops(child, all_nodes, new_ancestors)
+
+
+def _node_description(node: HierarchyNode) -> str:
+    """Get a human-readable description of a node for error messages."""
+    if node.feature_index is not None:
+        return f"feature_index={node.feature_index}"
+    if node.feature_id:
+        return f"id={node.feature_id}"
+    return "unnamed node"
+
+
+def hierarchy_modifier(
+    roots: Sequence[HierarchyNode] | HierarchyNode,
+) -> ActivationsModifier:
+    """
+    Create an activations modifier from one or more hierarchy trees.
+
+    This is the recommended way to use hierarchies with ActivationGenerator.
+    It validates the hierarchy structure and returns a modifier function that
+    applies all hierarchy constraints.
+
+    Args:
+        roots: One or more root HierarchyNode objects. Each root defines an
+            independent hierarchy tree. All trees are validated and applied.
+
+    Returns:
+        An ActivationsModifier function that can be passed to ActivationGenerator.
+
+    Raises:
+        ValueError: If validate=True and any hierarchy contains loops or
+            nodes with multiple parents.
+    """
+    if not roots:
+        # No hierarchies - return identity function
+        def identity(activations: torch.Tensor) -> torch.Tensor:
+            return activations
+
+        return identity
+
+    if isinstance(roots, HierarchyNode):
+        roots = [roots]
+    _validate_hierarchy(roots)
+
+    # Create modifier function that applies all hierarchies
+    def modifier(activations: torch.Tensor) -> torch.Tensor:
+        result = activations.clone()
+        for root in roots:
+            root._apply_hierarchy(result, parent_active_mask=None)
+        return result
+
+    return modifier
+
+
+class HierarchyNode:
+    """
+    Represents a node in a feature hierarchy tree.
+
+    Used to define hierarchical dependencies between features. Children are
+    deactivated when their parent is inactive, and children can optionally
+    be mutually exclusive.
+
+    Use `hierarchy_modifier()` to create an ActivationsModifier from one or
+    more HierarchyNode trees.
+
+
+    Attributes:
+        feature_index: Index of this feature in the activation tensor
+        children: Child HierarchyNode nodes
+        mutually_exclusive_children: If True, at most one child is active per sample
+        feature_id: Optional identifier for debugging
+    """
+
+    children: Sequence[HierarchyNode]
+    feature_index: int | None
+
+    @classmethod
+    def from_dict(cls, tree_dict: dict[str, Any]) -> HierarchyNode:
+        """
+        Create a HierarchyNode from a dictionary specification.
+
+        Args:
+            tree_dict: Dictionary with keys:
+
+                - feature_index (optional): Index in the activation tensor
+                - children (optional): List of child tree dictionaries
+                - mutually_exclusive_children (optional): Whether children are exclusive
+                - id (optional): Identifier for this node
+
+        Returns:
+            HierarchyNode instance
+        """
+        children = [
+            HierarchyNode.from_dict(child_dict)
+            for child_dict in tree_dict.get("children", [])
+        ]
+        return cls(
+            feature_index=tree_dict.get("feature_index"),
+            children=children,
+            mutually_exclusive_children=tree_dict.get(
+                "mutually_exclusive_children", False
+            ),
+            feature_id=tree_dict.get("id"),
+        )
+
+    def __init__(
+        self,
+        feature_index: int | None = None,
+        children: Sequence[HierarchyNode] | None = None,
+        mutually_exclusive_children: bool = False,
+        feature_id: str | None = None,
+    ):
+        """
+        Create a new HierarchyNode.
+
+        Args:
+            feature_index: Index of this feature in the activation tensor.
+                Use None for organizational nodes that don't correspond to a feature.
+            children: Child nodes that depend on this feature
+            mutually_exclusive_children: If True, only one child can be active per sample
+            feature_id: Optional identifier for debugging
+        """
+        self.feature_index = feature_index
+        self.children = children or []
+        self.mutually_exclusive_children = mutually_exclusive_children
+        self.feature_id = feature_id
+
+        if self.mutually_exclusive_children and len(self.children) < 2:
+            raise ValueError("Need at least 2 children for mutual exclusion")
+
+    def _apply_hierarchy(
+        self,
+        activations: torch.Tensor,
+        parent_active_mask: torch.Tensor | None,
+    ) -> None:
+        """Recursively apply hierarchical constraints."""
+        batch_size = activations.shape[0]
+
+        # Determine which samples have this node active
+        if self.feature_index is not None:
+            is_active = activations[:, self.feature_index] > 0
+        else:
+            # Non-readout node: active if parent is active (or always if root)
+            is_active = (
+                parent_active_mask
+                if parent_active_mask is not None
+                else torch.ones(batch_size, dtype=torch.bool, device=activations.device)
+            )
+
+        # Deactivate this node if parent is inactive
+        if parent_active_mask is not None and self.feature_index is not None:
+            activations[~parent_active_mask, self.feature_index] = 0
+            # Update is_active after deactivation
+            is_active = activations[:, self.feature_index] > 0
+
+        # Handle mutually exclusive children
+        if self.mutually_exclusive_children and len(self.children) >= 2:
+            self._enforce_mutual_exclusion(activations, is_active)
+
+        # Recursively process children
+        for child in self.children:
+            child._apply_hierarchy(activations, parent_active_mask=is_active)
+
+    def _enforce_mutual_exclusion(
+        self,
+        activations: torch.Tensor,
+        parent_active_mask: torch.Tensor,
+    ) -> None:
+        """Ensure at most one child is active per sample."""
+        batch_size = activations.shape[0]
+
+        # Get indices of children that have feature indices
+        child_indices = [
+            child.feature_index
+            for child in self.children
+            if child.feature_index is not None
+        ]
+
+        if len(child_indices) < 2:
+            return
+
+        # For each sample where parent is active, enforce mutual exclusion.
+        # Note: This loop is not vectorized because we need to randomly select
+        # which child to keep active per sample. Vectorizing would require either
+        # a deterministic selection (losing randomness) or complex gather/scatter
+        # operations that aren't more efficient for typical batch sizes.
+        for batch_idx in range(batch_size):
+            if not parent_active_mask[batch_idx]:
+                continue
+
+            # Find which children are active
+            active_children = [
+                i
+                for i, feat_idx in enumerate(child_indices)
+                if activations[batch_idx, feat_idx] > 0
+            ]
+
+            if len(active_children) <= 1:
+                continue
+
+            # Randomly select one to keep active
+            random_idx = int(torch.randint(len(active_children), (1,)).item())
+            keep_idx = active_children[random_idx]
+
+            # Deactivate all others
+            for i, feat_idx in enumerate(child_indices):
+                if i != keep_idx and i in active_children:
+                    activations[batch_idx, feat_idx] = 0
+
+    def get_all_feature_indices(self) -> list[int]:
+        """Get all feature indices in this subtree."""
+        indices = []
+        if self.feature_index is not None:
+            indices.append(self.feature_index)
+        for child in self.children:
+            indices.extend(child.get_all_feature_indices())
+        return indices
+
+    def validate(self) -> None:
+        """
+        Validate the hierarchy structure.
+
+        Checks that:
+        1. There are no loops (no node is its own ancestor)
+        2. Each node has at most one parent (no node appears in multiple children lists)
+
+        Raises:
+            ValueError: If the hierarchy is invalid
+        """
+        _validate_hierarchy([self])
+
+    def __repr__(self, indent: int = 0) -> str:
+        s = " " * (indent * 2)
+        s += str(self.feature_index) if self.feature_index is not None else "-"
+        s += "x" if self.mutually_exclusive_children else " "
+        if self.feature_id:
+            s += f" ({self.feature_id})"
+
+        for child in self.children:
+            s += "\n" + child.__repr__(indent + 2)
+        return s

sae_lens/synthetic/initialization.py

@@ -0,0 +1,40 @@
+import torch
+
+from sae_lens.synthetic import FeatureDictionary
+
+
+@torch.no_grad()
+def init_sae_to_match_feature_dict(
+    sae: torch.nn.Module,
+    feature_dict: FeatureDictionary,
+    noise_level: float = 0.0,
+    feature_ordering: torch.Tensor | None = None,
+) -> None:
+    """
+    Initialize an SAE's weights to match a feature dictionary.
+
+    This can be useful for:
+
+    - Starting training from a known good initialization
+    - Testing SAE evaluation code with ground truth
+    - Ablation studies on initialization
+
+    Args:
+        sae: The SAE to initialize. Must have W_enc and W_dec attributes.
+        feature_dict: The feature dictionary to match
+        noise_level: Standard deviation of Gaussian noise to add (0 = exact match)
+        feature_ordering: Optional permutation of feature indices
+    """
+    features = feature_dict.feature_vectors  # [num_features, hidden_dim]
+    min_dim = min(sae.W_enc.shape[1], features.shape[0])  # type: ignore[attr-defined]
+
+    if feature_ordering is not None:
+        features = features[feature_ordering]
+
+    features = features[:min_dim]
+
+    # W_enc is [hidden_dim, d_sae], feature vectors are [num_features, hidden_dim]
+    sae.W_enc.data[:, :min_dim] = (  # type: ignore[index]
+        features.T + torch.randn_like(features.T) * noise_level
+    )
+    sae.W_dec.data = sae.W_enc.data.T.clone()  # type: ignore[union-attr]