feat: implement SAE manager for latent decomposition and steering library for contrastive activation addition

README.md

@@ -1,44 +1,89 @@
-# DT-Circuits
+# DT-Circuits: Mechanistic Interpretability for Decision Transformers
 
-A research-grade platform for the mechanistic interpretability of Decision Transformers.
+DT-Circuits is a research-grade framework designed for the rigorous mechanistic interpretability of Decision Transformers (DT). By leveraging the TransformerLens paradigm, this platform enables researchers to map internal neural circuits, decompose activations using Sparse Autoencoders, and perform causal interventions on agent decision-making.
 
-## Architecture
-- **Data**: PPO Trajectory Harvester for high-quality teacher data.
-- **Model**: `HookedDT` - A custom Decision Transformer wrapped in `TransformerLens` for full activation visibility.
-- **Interpretability**: Tools for Direct Logit Attribution (DLA), Activation Patching, and Induction Head detection.
-- **Dashboard**: Streamlit-based UI for real-time causal interventions.
+The primary objective is to move beyond behavioral observation and saliency maps toward a quantitative understanding of how Reward-to-Go, State, and Action tokens are processed within the residual stream.
 
-## Quick Start
+## Core Capabilities
 
-### 1. Install Dependencies
+### 1. Circuit Foundation
+- **Hooked-DT Architecture**: A custom Decision Transformer implementation wrapped in TransformerLens, providing full access to internal activations, weights, and the residual stream.
+- **Direct Logit Attribution (DLA)**: Quantitative mapping of individual attention heads and MLP layers to the final action logits.
+- **Induction Head Discovery**: Automated scanning tools to identify heads responsible for temporal pattern recognition and "memory" in RL tasks.
+
+### 2. Causal Interventions
+- **Activation Patching**: Surgical replacement of activations between "clean" and "corrupted" runs to identify bottleneck features and causal paths.
+- **Contrastive Activation Addition (CAA)**: Generation of steering vectors by calculating the mean difference between positive and negative activation sets.
+- **Steering Library**: A persistent library of pre-calculated vectors (e.g., success_vector, exploration_vector) that can be injected at inference time to manipulate agent behavior without retraining.
+
+### 3. Deep Discovery & Safety
+- **Sparse Autoencoder (SAE) Integration**: Tools to train and deploy SAEs on the residual stream, decomposing polysemantic neurons into monosemantic latents.
+- **Mechanistic Anomaly Detection**: Utilizing SAE reconstruction error as a high-fidelity proxy for detecting out-of-distribution (OOD) states.
+
+## Technical Architecture
+
+The platform is divided into four primary layers:
+- **Data Layer**: PPO Trajectory Harvester for generating high-quality expert demonstrations in Gymnasium environments (e.g., MiniGrid).
+- **Model Layer**: The HookedDT implementation which maintains compatibility with standard DT architectures while adding hook-based visibility.
+- **Interpretability Layer**: A suite of modules for attribution, patching, SAE management, and steering.
+- **Visualization Layer**: A Streamlit-based dashboard for real-time activation monitoring and interactive steering.
+
+## Getting Started
+
+### Prerequisites
+- Python 3.9+
+- PyTorch 2.x
+- TransformerLens
+- SAE-Lens
+
+### Installation
 ```bash
 pip install -r requirements.txt
 ```
 
-### 2. Collect Data & Train Mini-DT
-```bash
-python scripts/train_dt.py
-```
+### Basic Workflow
+1. **Generate Trajectories**:
+   Use the harvester to collect teacher data for model training or SAE feature extraction.
+   ```bash
+   python scripts/train_dt.py
+   ```
 
-### 3. Run Interpretation Dashboard
-```bash
-streamlit run src/dashboard/app.py
-```
+2. **Run Interpretability Dashboard**:
+   Launch the interactive UI to perform real-time patching and steering interventions.
+   ```bash
+   streamlit run src/dashboard/app.py
+   ```
 
 ## Testing
-Run the test suite to ensure system integrity:
+
 ```bash
-pytest tests/
+PYTHONPATH=. pytest tests/
 ```
 
-## Components
-- `src/data/harvester.py`: Collects trajectories from MiniGrid.
-- `src/models/hooked_dt.py`: Hookable transformer implementation.
-- `src/interpretability/`:
-    - `attribution.py`: Direct Logit Attribution logic.
-    - `patching.py`: Activation patching interface.
-    - `induction_scan.py`: Automated circuit discovery.
-
-## License
+## Project Structure
 
-MIT
+```text
+DT-Circuits/
+├── scripts/                # Training and harvesting entry points
+│   ├── train_dt.py         # Decision Transformer training pipeline
+│   └── train_sae.py        # Sparse Autoencoder (SAE) training script
+├── src/                    
+│   ├── dashboard/          
+│   │   └── app.py          # Streamlit-based visualization UI
+│   ├── data/               
+│   │   └── harvester.py    # PPO-based expert trajectory harvester
+│   ├── interpretability/   
+│   │   ├── attribution.py  # Direct Logit Attribution (DLA)
+│   │   ├── induction_scan.py # Induction head detection logic
+���   │   ├── patching.py     # Causal activation patching tools
+│   │   ├── sae_manager.py  # SAE deployment and anomaly detection
+│   │   └── steering.py     # Steering vector generation and injection
+│   ├── models/             
+│   │   └── hooked_dt.py    # TransformerLens-wrapped Decision Transformer
+│   └── utils/              
+├── tests/                  # Unit and integration test suite
+│   ├── test_components.py  
+│   └── test_sae_and_steering.py 
+├── config.yaml             # Experiment and environment configuration
+└── requirements.txt        # Environment dependencies
+```

src/interpretability/sae_manager.py

@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import os
+from typing import Dict, List, Optional, Tuple, Union
+from sae_lens import StandardSAE, StandardSAEConfig
+from jaxtyping import Float
+
+class SAEManager:
+    """
+    Research-grade manager for Sparse Autoencoders (SAEs) integrated with Decision Transformers.
+    Handles training, decomposition into monosemantic latents, and mechanistic anomaly detection.
+    """
+    def __init__(self, model: nn.Module, sae_dir: str = "artifacts/saes"):
+        self.model = model
+        self.sae_dir = sae_dir
+        self.saes: Dict[str, StandardSAE] = {}
+        os.makedirs(sae_dir, exist_ok=True)
+
+    def setup_sae(
+        self,
+        hook_point: str,
+        d_model: int,
+        expansion_factor: int = 8,
+    ) -> StandardSAE:
+        """
+        Initializes an SAE for a specific hook point in the transformer.
+        """
+        cfg = StandardSAEConfig(
+            d_in=d_model,
+            d_sae=d_model * expansion_factor,
+            device=str(next(self.model.parameters()).device)
+        )
+        sae = StandardSAE(cfg)
+        self.saes[hook_point] = sae
+        return sae
+
+    def train_on_trajectories(
+        self,
+        hook_point: str,
+        activations: Float[torch.Tensor, "n_samples d_model"],
+        l1_coefficient: float = 0.0001,
+        batch_size: int = 1024,
+        epochs: int = 10,
+    ):
+        """
+        Trains the SAE on collected trajectory activations.
+        """
+        if hook_point not in self.saes:
+            self.setup_sae(hook_point, activations.shape[-1])
+        
+        sae = self.saes[hook_point]
+        optimizer = torch.optim.Adam(sae.parameters(), lr=0.0004)
+        
+        sae.train()
+        n_samples = activations.shape[0]
+        
+        for epoch in range(epochs):
+            permutation = torch.randperm(n_samples)
+            epoch_loss = 0
+            for i in range(0, n_samples, batch_size):
+                indices = permutation[i:i+batch_size]
+                batch_acts = activations[indices].to(sae.device)
+                
+                optimizer.zero_grad()
+                
+                # Manual forward pass for training
+                feature_acts = sae.encode(batch_acts)
+                sae_out = sae.decode(feature_acts)
+                
+                mse_loss = torch.nn.functional.mse_loss(sae_out, batch_acts)
+                l1_loss = l1_coefficient * feature_acts.abs().sum()
+                loss = mse_loss + l1_loss
+                
+                loss.backward()
+                optimizer.step()
+                epoch_loss += loss.item()
+            
+            print(f"Epoch {epoch+1}/{epochs} - Loss: {epoch_loss / (n_samples / batch_size):.4f}")
+
+    def get_feature_activations(
+        self,
+        hook_point: str,
+        activations: Float[torch.Tensor, "... d_model"]
+    ) -> Float[torch.Tensor, "... d_sae"]:
+        """
+        Decomposes activations into monosemantic features.
+        """
+        if hook_point not in self.saes:
+            raise ValueError(f"SAE for {hook_point} not found. Train or load it first.")
+        
+        sae = self.saes[hook_point]
+        sae.eval()
+        with torch.no_grad():
+            feature_acts = sae.encode(activations.to(sae.device))
+        return feature_acts
+
+    def reconstruct(
+        self,
+        hook_point: str,
+        activations: Float[torch.Tensor, "... d_model"]
+    ) -> Float[torch.Tensor, "... d_model"]:
+        """
+        Reconstructs the original activations using the SAE.
+        """
+        if hook_point not in self.saes:
+            raise ValueError(f"SAE for {hook_point} not found.")
+        
+        sae = self.saes[hook_point]
+        sae.eval()
+        with torch.no_grad():
+            feature_acts = sae.encode(activations.to(sae.device))
+            sae_out = sae.decode(feature_acts)
+        return sae_out
+
+    def compute_anomaly_score(
+        self,
+        hook_point: str,
+        activations: Float[torch.Tensor, "... d_model"]
+    ) -> Float[torch.Tensor, "..."]:
+        """
+        Calculates reconstruction error as a proxy for mechanistic anomaly detection.
+        Formula: ||x - x_hat|| / ||x||
+        """
+        if hook_point not in self.saes:
+            raise ValueError(f"SAE for {hook_point} not found.")
+        
+        sae = self.saes[hook_point]
+        sae.eval()
+        with torch.no_grad():
+            x = activations.to(sae.device)
+            feature_acts = sae.encode(x)
+            x_hat = sae.decode(feature_acts)
+            
+            error = torch.norm(x - x_hat, dim=-1) / (torch.norm(x, dim=-1) + 1e-8)
+        return error
+
+    def save_all_saes(self):
+        for hook, sae in self.saes.items():
+            path = os.path.join(self.sae_dir, f"{hook.replace('.', '_')}_sae.pt")
+            torch.save({
+                'state_dict': sae.state_dict(),
+                'cfg': sae.cfg
+            }, path)
+            print(f"Saved SAE for {hook} to {path}")
+
+    def load_sae(self, hook_point: str):
+        path = os.path.join(self.sae_dir, f"{hook_point.replace('.', '_')}_sae.pt")
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"No saved SAE found at {path}")
+        
+        checkpoint = torch.load(path, map_location=str(next(self.model.parameters()).device))
+        sae = StandardSAE(checkpoint['cfg'])
+        sae.load_state_dict(checkpoint['state_dict'])
+        self.saes[hook_point] = sae
+        return sae

src/interpretability/steering.py

@@ -1,43 +1,78 @@
 import torch
 import torch.nn as nn
-from typing import Optional
+from typing import Dict, List, Optional
+class SteeringLibrary:
+    """
+    A persistent library of pre-calculated steering vectors (CAA).
+    Includes vectors for exploration, safety, and goal-directedness.
+    """
+    def __init__(self, d_model: int):
+        self.d_model = d_model
+        self.vectors: Dict[str, torch.Tensor] = {}
+
+    def add_vector(self, name: str, vector: torch.Tensor):
+        if vector.shape[-1] != self.d_model:
+            raise ValueError(f"Vector dimension {vector.shape[-1]} does not match d_model {self.d_model}")
+        self.vectors[name] = vector
+
+    def get_vector(self, name: str) -> torch.Tensor:
+        if name not in self.vectors:
+            raise KeyError(f"Vector '{name}' not found in library.")
+        return self.vectors[name]
+
+    def list_vectors(self) -> List[str]:
+        return list(self.vectors.keys())
 
 class RTGSteerer:
     """
-    Enables 'Behavioral Steering' by manipulating Reward-to-Go (RTG) tokens.
+    Enables 'Behavioral Steering' by manipulating Reward-to-Go (RTG) tokens or internal activations.
+    Supports Contrastive Activation Addition (CAA).
     """
-    def __init__(self, model):
+    def __init__(self, model, library: Optional[SteeringLibrary] = None):
         self.model = model
+        self.library = library or SteeringLibrary(model.cfg.d_model)
 
-    def steer(
+    def steer_rtg(
         self,
-        states: torch.Tensor,
-        actions: torch.Tensor,
         base_rtg: torch.Tensor,
-        steering_vector: torch.Tensor,
+        vector_name: Optional[str] = None,
+        custom_vector: Optional[torch.Tensor] = None,
         alpha: float = 1.0
-    ):
+    ) -> torch.Tensor:
         """
         Adds a steering vector to the RTG embeddings.
-        RTG_new = RTG_base + alpha * steering_vector
         """
-        # Embed base RTG
+        vector = custom_vector if custom_vector is not None else self.library.get_vector(vector_name)
+        
         with torch.no_grad():
             rtg_emb = self.model.embed_return(base_rtg)
-            
-            # Apply steering
-            steered_rtg_emb = rtg_emb + alpha * steering_vector
-            
-            # Hook the model to use the steered RTG
-            # This requires a slightly more complex hook in HookedDT
-            # For now, we returns the steered embedding to be used in a custom forward pass
-            return steered_rtg_emb
+            return rtg_emb + alpha * vector
+
+    def generate_caa_vector(
+        self,
+        positive_activations: torch.Tensor,
+        negative_activations: torch.Tensor,
+        method: str = "mean_diff"
+    ) -> torch.Tensor:
+        """
+        Generates a steering vector using Contrastive Activation Addition.
+        'mean_diff' calculates the difference between the means of positive and negative sets.
+        """
+        if method == "mean_diff":
+            pos_mean = positive_activations.mean(dim=0)
+            neg_mean = negative_activations.mean(dim=0)
+            return pos_mean - neg_mean
+        else:
+            raise NotImplementedError(f"Method {method} not implemented.")
 
-    def find_success_vector(self, high_reward_cache, low_reward_cache):
+    def apply_steering_hook(self, hook_point: str, vector_name: str, alpha: float = 1.0):
         """
-        Identifies the 'Success Vector' by comparing high vs low reward activations.
-        Vector = Mean(High Reward Residual) - Mean(Low Reward Residual)
+        Returns a HookedTransformer compatible hook function that applies steering.
         """
-        high_res = high_reward_cache["blocks.0.hook_resid_post"].mean(dim=(0, 1))
-        low_res = low_reward_cache["blocks.0.hook_resid_post"].mean(dim=(0, 1))
-        return high_res - low_res
+        vector = self.library.get_vector(vector_name)
+        
+        def steering_hook(activations, hook):
+            # activations: [batch, pos, d_model]
+            return activations + alpha * vector
+            
+        return steering_hook

tests/test_sae_and_steering.py

@@ -0,0 +1,75 @@
+import pytest
+import torch
+from src.models.hooked_dt import HookedDT
+from src.interpretability.sae_manager import SAEManager
+from src.interpretability.steering import RTGSteerer, SteeringLibrary
+
+@pytest.fixture
+def model():
+    return HookedDT.from_config(state_dim=10, action_dim=5, n_layers=1, n_heads=2, d_model=32)
+
+@pytest.fixture
+def sae_manager(model):
+    return SAEManager(model)
+
+def test_sae_lifecycle(sae_manager):
+    hook_point = "blocks.0.hook_resid_post"
+    d_model = 32
+    
+    # 1. Setup SAE
+    sae = sae_manager.setup_sae(hook_point, d_model, expansion_factor=2)
+    assert hook_point in sae_manager.saes
+    assert sae.cfg.d_sae == 64
+    
+    # 2. Mock activations
+    activations = torch.randn(100, d_model)
+    
+    # 3. Test reconstruction shape
+    reconstructed = sae_manager.reconstruct(hook_point, activations)
+    assert reconstructed.shape == activations.shape
+    
+    # 4. Test feature activations shape
+    latents = sae_manager.get_feature_activations(hook_point, activations)
+    assert latents.shape == (100, 64)
+    
+    # 5. Test anomaly score
+    scores = sae_manager.compute_anomaly_score(hook_point, activations)
+    assert scores.shape == (100,)
+    assert (scores >= 0).all()
+
+def test_steering_library():
+    lib = SteeringLibrary(d_model=32)
+    vec = torch.randn(32)
+    lib.add_vector("test_vec", vec)
+    
+    assert "test_vec" in lib.list_vectors()
+    assert torch.equal(lib.get_vector("test_vec"), vec)
+    
+    with pytest.raises(ValueError):
+        lib.add_vector("wrong_dim", torch.randn(16))
+
+def test_caa_generation(model):
+    steerer = RTGSteerer(model)
+    pos_acts = torch.randn(10, 32) + 1.0
+    neg_acts = torch.randn(10, 32) - 1.0
+    
+    vector = steerer.generate_caa_vector(pos_acts, neg_acts)
+    assert vector.shape == (32,)
+    # Mean difference should be around 2.0 for each dimension
+    assert vector.mean() > 0.0
+
+def test_steering_hook(model):
+    lib = SteeringLibrary(d_model=32)
+    vec = torch.ones(32)
+    lib.add_vector("boost", vec)
+    
+    steerer = RTGSteerer(model, library=lib)
+    hook = steerer.apply_steering_hook("blocks.0.hook_resid_post", "boost", alpha=2.0)
+    
+    input_acts = torch.zeros(1, 5, 32)
+    output_acts = hook(input_acts, None)
+    
+    assert torch.allclose(output_acts, torch.ones(1, 5, 32) * 2.0)
+
+if __name__ == "__main__":
+    pytest.main([__file__])