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

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]

sae_lens/synthetic/plotting.py

@@ -0,0 +1,230 @@
+"""
+Plotting utilities for visualizing SAE training on synthetic data.
+
+This module provides functions for:
+
+- Plotting cosine similarities between SAE features and true features
+- Automatically reordering features for better visualization
+- Creating comparison plots between encoder and decoder
+"""
+
+from collections.abc import Iterable
+from pathlib import Path
+from typing import Any
+
+import plotly.graph_objects as go
+import torch
+from plotly.subplots import make_subplots
+
+from sae_lens.saes import SAE
+from sae_lens.synthetic.feature_dictionary import FeatureDictionary
+from sae_lens.util import cosine_similarities
+
+
+def find_best_feature_ordering(
+    sae_features: torch.Tensor,
+    true_features: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Find the best ordering of SAE features to match true features.
+
+    Reorders SAE features so that each SAE latent aligns with its best-matching
+    true feature in order. This makes cosine similarity plots more interpretable.
+
+    Args:
+        sae_features: SAE decoder weights of shape [d_sae, hidden_dim]
+        true_features: True feature vectors of shape [num_features, hidden_dim]
+
+    Returns:
+        Tensor of indices that reorders sae_features for best alignment
+    """
+    cos_sims = cosine_similarities(sae_features, true_features)
+    best_matches = torch.argmax(torch.abs(cos_sims), dim=1)
+    return torch.argsort(best_matches)
+
+
+def find_best_feature_ordering_from_sae(
+    sae: torch.nn.Module,
+    feature_dict: FeatureDictionary,
+) -> torch.Tensor:
+    """
+    Find the best feature ordering for an SAE given a feature dictionary.
+
+    Args:
+        sae: SAE with W_dec attribute of shape [d_sae, hidden_dim]
+        feature_dict: The feature dictionary containing true features
+
+    Returns:
+        Tensor of indices that reorders SAE latents for best alignment
+    """
+    sae_features = sae.W_dec.detach()  # type: ignore[attr-defined]
+    true_features = feature_dict.feature_vectors.detach()
+    return find_best_feature_ordering(sae_features, true_features)
+
+
+def find_best_feature_ordering_across_saes(
+    saes: Iterable[torch.nn.Module],
+    feature_dict: FeatureDictionary,
+) -> torch.Tensor:
+    """
+    Find the best feature ordering that works across multiple SAEs.
+
+    Useful for creating consistent orderings across training snapshots.
+
+    Args:
+        saes: Iterable of SAEs to consider
+        feature_dict: The feature dictionary containing true features
+
+    Returns:
+        The best ordering tensor found across all SAEs
+    """
+    best_score = float("-inf")
+    best_ordering: torch.Tensor | None = None
+
+    true_features = feature_dict.feature_vectors.detach()
+
+    for sae in saes:
+        sae_features = sae.W_dec.detach()  # type: ignore[attr-defined]
+        cos_sims = cosine_similarities(sae_features, true_features)
+        cos_sims = torch.round(cos_sims * 100) / 100  # Reduce numerical noise
+
+        ordering = find_best_feature_ordering(sae_features, true_features)
+        score = cos_sims[ordering, torch.arange(cos_sims.shape[1])].mean().item()
+
+        if score > best_score:
+            best_score = score
+            best_ordering = ordering
+
+    if best_ordering is None:
+        raise ValueError("No SAEs provided")
+
+    return best_ordering
+
+
+def plot_sae_feature_similarity(
+    sae: SAE[Any],
+    feature_dict: FeatureDictionary,
+    title: str | None = None,
+    reorder_features: bool | torch.Tensor = False,
+    decoder_only: bool = False,
+    show_values: bool = False,
+    height: int = 400,
+    width: int = 800,
+    save_path: str | Path | None = None,
+    show_plot: bool = True,
+    dtick: int | None = 1,
+    scale: float = 1.0,
+):
+    """
+    Plot cosine similarities between SAE features and true features.
+
+    Creates a heatmap showing how well each SAE latent aligns with each
+    true feature. Useful for understanding what the SAE has learned.
+
+    Args:
+        sae: The SAE to visualize. Must have W_enc and W_dec attributes.
+        feature_dict: The feature dictionary containing true features
+        title: Plot title. If None, a default title is used.
+        reorder_features: If True, automatically reorders features for best alignment.
+            If a tensor, uses that as the ordering.
+        decoder_only: If True, only plots the decoder (not encoder and decoder side-by-side)
+        show_values: If True, shows numeric values on the heatmap
+        height: Height of the figure in pixels
+        width: Width of the figure in pixels
+        save_path: If provided, saves the figure to this path
+        show_plot: If True, displays the plot
+        dtick: Tick spacing for axes
+        scale: Scale factor for image resolution when saving
+    """
+    # Get cosine similarities
+    true_features = feature_dict.feature_vectors.detach()
+    dec_cos_sims = cosine_similarities(sae.W_dec.detach(), true_features)  # type: ignore[attr-defined]
+    enc_cos_sims = cosine_similarities(sae.W_enc.T.detach(), true_features)  # type: ignore[attr-defined]
+
+    # Round to reduce numerical noise
+    dec_cos_sims = torch.round(dec_cos_sims * 100) / 100
+    enc_cos_sims = torch.round(enc_cos_sims * 100) / 100
+
+    # Apply feature reordering if requested
+    if reorder_features is not False:
+        if isinstance(reorder_features, bool):
+            sorted_indices = find_best_feature_ordering(
+                sae.W_dec.detach(),
+                true_features,  # type: ignore[attr-defined]
+            )
+        else:
+            sorted_indices = reorder_features
+        dec_cos_sims = dec_cos_sims[sorted_indices]
+        enc_cos_sims = enc_cos_sims[sorted_indices]
+
+    hovertemplate = "True feature: %{x}<br>SAE Latent: %{y}<br>Cosine Similarity: %{z:.3f}<extra></extra>"
+
+    if decoder_only:
+        fig = make_subplots(rows=1, cols=1)
+
+        decoder_args: dict[str, Any] = {
+            "z": dec_cos_sims.cpu().numpy(),
+            "zmin": -1,
+            "zmax": 1,
+            "colorscale": "RdBu",
+            "colorbar": dict(title="cos sim", x=1.0, dtick=1, tickvals=[-1, 0, 1]),
+            "hovertemplate": hovertemplate,
+        }
+        if show_values:
+            decoder_args["texttemplate"] = "%{z:.2f}"
+            decoder_args["textfont"] = {"size": 10}
+
+        fig.add_trace(go.Heatmap(**decoder_args), row=1, col=1)
+        fig.update_xaxes(title_text="True feature", row=1, col=1, dtick=dtick)
+        fig.update_yaxes(title_text="SAE Latent", row=1, col=1, dtick=dtick)
+    else:
+        fig = make_subplots(
+            rows=1, cols=2, subplot_titles=("SAE encoder", "SAE decoder")
+        )
+
+        # Encoder heatmap
+        encoder_args: dict[str, Any] = {
+            "z": enc_cos_sims.cpu().numpy(),
+            "zmin": -1,
+            "zmax": 1,
+            "colorscale": "RdBu",
+            "showscale": False,
+            "hovertemplate": hovertemplate,
+        }
+        if show_values:
+            encoder_args["texttemplate"] = "%{z:.2f}"
+            encoder_args["textfont"] = {"size": 10}
+
+        fig.add_trace(go.Heatmap(**encoder_args), row=1, col=1)
+
+        # Decoder heatmap
+        decoder_args = {
+            "z": dec_cos_sims.cpu().numpy(),
+            "zmin": -1,
+            "zmax": 1,
+            "colorscale": "RdBu",
+            "colorbar": dict(title="cos sim", x=1.0, dtick=1, tickvals=[-1, 0, 1]),
+            "hovertemplate": hovertemplate,
+        }
+        if show_values:
+            decoder_args["texttemplate"] = "%{z:.2f}"
+            decoder_args["textfont"] = {"size": 10}
+
+        fig.add_trace(go.Heatmap(**decoder_args), row=1, col=2)
+
+        fig.update_xaxes(title_text="True feature", row=1, col=1, dtick=dtick)
+        fig.update_xaxes(title_text="True feature", row=1, col=2, dtick=dtick)
+        fig.update_yaxes(title_text="SAE Latent", row=1, col=1, dtick=dtick)
+        fig.update_yaxes(title_text="SAE Latent", row=1, col=2, dtick=dtick)
+
+    # Set main title
+    if title is None:
+        title = "Cosine similarity with true features"
+    fig.update_layout(height=height, width=width, title_text=title)
+
+    if save_path:
+        Path(save_path).parent.mkdir(parents=True, exist_ok=True)
+        fig.write_image(save_path, scale=scale)
+
+    if show_plot:
+        fig.show()