Upload BatchTopKSAE

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,28 @@
+{
+  "act_size": 512,
+  "architectures": [
+    "BatchTopKSAE"
+  ],
+  "auto_map": {
+    "AutoConfig": "config.SAEConfig",
+    "AutoModel": "model.BatchTopKSAE"
+  },
+  "aux_penalty": 0.03125,
+  "bandwidth": 0.001,
+  "dict_size": 16384,
+  "dtype": "float32",
+  "input_mean": -3.768720489460975e-06,
+  "input_std": 0.3794997036457062,
+  "input_unit_norm": true,
+  "l1_coeff": 0.001,
+  "model_type": "sae",
+  "n_batches_to_dead": 10,
+  "parent_hook_point": "resid_post",
+  "parent_layer": 1,
+  "parent_model_name": "EleutherAI/pythia-70m",
+  "sae_dtype": "float32",
+  "sae_type": "batchtopk",
+  "top_k_aux": 512,
+  "torch_dtype": "float32",
+  "transformers_version": "4.47.1"
+}

config.py

@@ -0,0 +1,203 @@
+from dataclasses import dataclass, field
+from typing import Optional
+import torch
+import pyrallis
+from transformers import PretrainedConfig
+from typing import Optional
+
+
+@dataclass
+class TrainingConfig:
+    # Model settings
+    model_name: str = "unsloth/Meta-Llama-3.1-8B"
+    layer: int = 12
+    hook_point: str = "resid_post"
+    act_size: Optional[int] = None  # Will be set after model initialization
+    
+    # SAE settings
+    sae_type: str = "batchtopk"
+    dict_size: int = 2**15
+    aux_penalty: float = 1/32
+    input_unit_norm: bool = True
+    
+    # TopK specific settings
+    top_k: int = 50
+    top_k_warmup_steps_fraction: float = 0.1
+    start_top_k: int = 4096
+    top_k_aux: int = 512
+
+    n_batches_to_dead: int = 10
+    
+    # Training settings
+    lr: float = 3e-4
+    bandwidth: float = 0.001
+    l1_coeff: float = 0.0018
+    num_tokens: int = int(1e9)
+    seq_len: int = 1024
+    model_batch_size: int = 16
+    num_batches_in_buffer: int = 5
+    max_grad_norm: float = 1.0
+    batch_size: int = 8192
+
+    # scheduler
+    warmup_fraction: float = 0.1
+    scheduler_type: str = 'linear'
+    
+    # Hardware settings
+    device: str = "cuda"
+    dtype: torch.dtype = field(default=torch.float32)
+    sae_dtype: torch.dtype = field(default=torch.float32)
+    
+    # Dataset settings
+    dataset_path: str = "cerebras/SlimPajama-627B"
+    
+    # Logging settings
+    wandb_project: str = "turbo-llama-lens"
+
+    performance_log_steps: int = 100
+    save_checkpoint_steps: int = 10_000
+    def __post_init__(self):
+        if self.device == "cuda" and not torch.cuda.is_available():
+            print("CUDA not available, falling back to CPU")
+            self.device = "cpu"
+        
+        # Convert string dtype to torch.dtype if needed
+        if isinstance(self.dtype, str):
+            self.dtype = getattr(torch, self.dtype)
+
+
+class SAEConfig(PretrainedConfig):
+    model_type = "sae"
+    
+    def __init__(
+        self,
+        # SAE architecture
+        act_size: int = None,
+        dict_size: int = 2**15,
+        sae_type: str = "batchtopk",
+        input_unit_norm: bool = True,
+        
+        # TopK specific settings
+        top_k: int = 50,
+        top_k_aux: int = 512,
+        n_batches_to_dead: int = 10,
+        
+        # Training hyperparameters
+        aux_penalty: float = 1/32,
+        l1_coeff: float = 0.0018,
+        bandwidth: float = 0.001,
+        
+        # Hardware settings
+        dtype: str = "float32",
+        sae_dtype: str = "float32",
+        
+        # Optional parent model info
+        parent_model_name: Optional[str] = None,
+        parent_layer: Optional[int] = None,
+        parent_hook_point: Optional[str] = None,
+        
+        # Input normalization settings
+        input_mean: Optional[float] = None,
+        input_std: Optional[float] = None,
+        
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.act_size = act_size
+        self.dict_size = dict_size
+        self.sae_type = sae_type
+        self.input_unit_norm = input_unit_norm
+        
+        self.top_k = top_k
+        self.top_k_aux = top_k_aux
+        self.n_batches_to_dead = n_batches_to_dead
+        
+        self.aux_penalty = aux_penalty
+        self.l1_coeff = l1_coeff
+        self.bandwidth = bandwidth
+        
+        self.dtype = dtype
+        self.sae_dtype = sae_dtype
+        
+        self.parent_model_name = parent_model_name
+        self.parent_layer = parent_layer
+        self.parent_hook_point = parent_hook_point
+        
+        self.input_mean = input_mean
+        self.input_std = input_std
+    
+    def get_torch_dtype(self, dtype_str: str) -> torch.dtype:
+        dtype_map = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+        }
+        return dtype_map.get(dtype_str, torch.float32)
+    
+    @classmethod
+    def from_training_config(cls, cfg: TrainingConfig):
+        """Convert TrainingConfig to SAEConfig"""
+        return cls(
+            act_size=cfg.act_size,
+            dict_size=cfg.dict_size,
+            sae_type=cfg.sae_type,
+            input_unit_norm=cfg.input_unit_norm,
+            top_k=cfg.top_k,
+            top_k_aux=cfg.top_k_aux,
+            n_batches_to_dead=cfg.n_batches_to_dead,
+            aux_penalty=cfg.aux_penalty,
+            l1_coeff=cfg.l1_coeff,
+            bandwidth=cfg.bandwidth,
+            dtype=str(cfg.dtype).split('.')[-1],
+            sae_dtype=str(cfg.sae_dtype).split('.')[-1],
+            parent_model_name=cfg.model_name,
+            parent_layer=cfg.layer,
+            parent_hook_point=cfg.hook_point,
+            input_mean=cfg.input_mean if hasattr(cfg, 'input_mean') else None,
+            input_std=cfg.input_std if hasattr(cfg, 'input_std') else None,
+        )
+    
+    def to_training_config(self) -> TrainingConfig:
+        """Convert SAEConfig back to TrainingConfig"""
+        return TrainingConfig(
+            dtype=self.get_torch_dtype(self.dtype),
+            sae_dtype=self.get_torch_dtype(self.sae_dtype),
+            model_name=self.parent_model_name,
+            layer=self.parent_layer,
+            hook_point=self.parent_hook_point,
+            act_size=self.act_size,
+            dict_size=self.dict_size,
+            sae_type=self.sae_type,
+            input_unit_norm=self.input_unit_norm,
+            top_k=self.top_k,
+            top_k_aux=self.top_k_aux,
+            n_batches_to_dead=self.n_batches_to_dead,
+            aux_penalty=self.aux_penalty,
+            l1_coeff=self.l1_coeff,
+            bandwidth=self.bandwidth,
+        )
+
+
+@pyrallis.wrap()
+def get_config() -> TrainingConfig:
+    return TrainingConfig()
+
+
+# For backward compatibility
+def get_default_cfg() -> TrainingConfig:
+    return get_config()
+
+
+def post_init_cfg(cfg: TrainingConfig, activation_store = None) -> TrainingConfig:
+    """
+    Any additional configuration setup that needs to happen after model initialization
+    Args:
+        cfg: Training configuration
+        activation_store: Optional activation store to get input statistics
+    """
+    if activation_store is not None:
+        cfg.input_mean = activation_store.mean
+        cfg.input_std = activation_store.std
+        print(f"Setting input statistics from activation store - Mean: {cfg.input_mean:.4f}, Std: {cfg.input_std:.4f}")
+    
+    return cfg

model.py

@@ -0,0 +1,562 @@
+from transformers import PreTrainedModel
+from typing import Optional, Dict, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.autograd as autograd
+from .config import SAEConfig
+
+
+class BaseSAE(PreTrainedModel):
+    """Base class for autoencoder models."""
+    config_class = SAEConfig
+    base_model_prefix = "sae"
+
+    def __init__(self, config: SAEConfig):
+        super().__init__(config)
+        print(config)
+        self.config = config
+        torch.manual_seed(42)
+
+        self.b_dec = nn.Parameter(torch.zeros(self.config.act_size))
+        self.b_enc = nn.Parameter(torch.zeros(self.config.dict_size))
+        self.W_enc = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(self.config.act_size, self.config.dict_size)
+            )
+        )
+        self.W_dec = nn.Parameter(
+            torch.nn.init.kaiming_uniform_(
+                torch.empty(self.config.dict_size, self.config.act_size)
+            )
+        )
+        self.W_dec.data[:] = self.W_enc.t().data
+        self.W_dec.data[:] = self.W_dec / self.W_dec.norm(dim=-1, keepdim=True)
+        self.register_buffer('num_batches_not_active', torch.zeros((self.config.dict_size,)))
+
+        self.to(self.config.get_torch_dtype(self.config.dtype))
+
+        # Initialize input normalization parameters if provided
+        if config.input_mean is not None and config.input_std is not None:
+            self.register_buffer('input_mean', torch.tensor(config.input_mean))
+            self.register_buffer('input_std', torch.tensor(config.input_std))
+        else:
+            self.input_mean = None
+            self.input_std = None
+
+        # Initialize CUDA events for timing
+        if torch.cuda.is_available():
+            self.start_event = torch.cuda.Event(enable_timing=True)
+            self.end_event = torch.cuda.Event(enable_timing=True)
+        else:
+            self.start_event = None
+            self.end_event = None
+
+    def preprocess_input(self, x):
+        x = x.to(self.config.get_torch_dtype(self.config.sae_dtype))
+        if self.config.input_unit_norm:
+            if self.input_mean is not None and self.input_std is not None:
+                # Use pre-computed statistics
+                x = (x - self.input_mean) / (self.input_std + 1e-5)
+                return x, self.input_mean, self.input_std
+            else:
+                # Compute statistics on the fly
+                x_mean = x.mean(dim=-1, keepdim=True)
+                x = x - x_mean
+                x_std = x.std(dim=-1, keepdim=True)
+                x = x / (x_std + 1e-5)
+                return x, x_mean, x_std
+        else:
+            return x, None, None
+
+    def postprocess_output(self, x_reconstruct, x_mean, x_std):
+        if self.config.input_unit_norm:
+            x_reconstruct = x_reconstruct * x_std + x_mean
+        return x_reconstruct
+
+    @torch.no_grad()
+    def make_decoder_weights_and_grad_unit_norm(self):
+        W_dec_normed = self.W_dec / self.W_dec.norm(dim=-1, keepdim=True)
+        W_dec_grad_proj = (self.W_dec.grad * W_dec_normed).sum(
+            -1, keepdim=True
+        ) * W_dec_normed
+        self.W_dec.grad -= W_dec_grad_proj
+        self.W_dec.data = W_dec_normed
+
+    def update_inactive_features(self, acts):
+        self.num_batches_not_active += (acts.sum(0) == 0).float()
+        self.num_batches_not_active[acts.sum(0) > 0] = 0
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode input tensor to sparse features
+        Args:
+            x: Input tensor of shape (batch_size, act_size)
+        Returns:
+            Encoded features of shape (batch_size, dict_size)
+        """
+        if self.config.input_unit_norm:
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            x_std = x.std(dim=-1, keepdim=True)
+            x = x / (x_std + 1e-5)
+        
+        return F.relu(x @ self.W_enc + self.b_enc)
+
+    def decode(self, h: torch.Tensor) -> torch.Tensor:
+        """
+        Decode features back to input space
+        Args:
+            h: Encoded features of shape (batch_size, dict_size)
+        Returns:
+            Reconstructed input of shape (batch_size, act_size)
+        """
+        return h @ self.W_dec + self.b_dec
+
+    def forward(self, x):
+        # Start timing if CUDA is available
+        if self.start_event is not None:
+            self.start_event.record()
+            
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        x_reconstruct = acts @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts)
+        output = self.get_loss_dict(x, x_reconstruct, acts, x_mean, x_std)
+        
+        # End timing if CUDA is available
+        if self.end_event is not None:
+            self.end_event.record()
+            torch.cuda.synchronize()
+            output["forward_time_ms"] = self.start_event.elapsed_time(self.end_event)
+        
+        return output
+
+    @torch.no_grad()
+    def fold_stats_into_weights(self, mean: torch.Tensor = None, std: torch.Tensor = None) -> "BaseSAE":
+        """
+        Fold normalization statistics into the encoder weights and biases
+        """
+        print("Folding statistics into encoder...")
+        
+        if mean is not None and std is not None:
+            mean = mean.to(self.W_enc.device)
+            std = std.to(self.W_enc.device)
+        else:
+            mean = self.input_mean
+            std = self.input_std
+        
+        # Original forward pass:
+        # x_norm = (x - mean) / std
+        # acts = relu(x_norm @ W_enc + b_enc)
+        # x_hat = acts @ W_dec + b_dec
+        
+        # Folding steps:
+        # 1. x_norm = (x - mean) / std
+        # 2. acts = relu(x_norm @ W_enc + b_enc)
+        #        = relu((x/std - mean/std) @ W_enc + b_enc)
+        #        = relu(x @ (W_enc/std) - mean @ (W_enc/std) + b_enc)
+        
+        # First scale encoder weights
+        self.W_enc.data = self.W_enc / std
+        
+        # Then adjust encoder bias
+        self.b_enc.data = self.b_enc - mean * (self.W_enc.sum(0))
+        
+        # Scale decoder to preserve reconstruction
+        self.W_dec.data = self.W_dec * std
+        self.b_dec.data = self.b_dec * std + mean
+        
+        # Turn off input normalization
+        self.config.input_unit_norm = False
+        
+        return self
+    
+    @torch.no_grad()
+    def fold_W_dec_norm(self):
+        """
+        Make decoder weights unit norm and adjust encoder accordingly
+        """
+        # Get current decoder norms
+        W_dec_norm = self.W_dec.norm(dim=-1, keepdim=True)
+        
+        # Original: acts @ W_dec + b_dec
+        # After: acts @ (W_dec/norm) + b_dec
+        # Need: (acts * norm) @ (W_dec/norm) + b_dec
+        # So: scale W_enc by norm.T
+        
+        # Scale encoder weights first
+        self.W_enc.data = self.W_enc * W_dec_norm.t()
+        
+        # Then normalize decoder weights
+        self.W_dec.data = self.W_dec / W_dec_norm
+        
+        # Scale encoder bias to compensate for the scaling of activations
+        self.b_enc.data = self.b_enc * W_dec_norm.squeeze(-1)
+        
+        return self
+
+    def set_mean_std(self, mean: float, std: float):
+        """
+        Set input normalization statistics after model initialization
+        
+        Args:
+            mean: Mean scalar value for input normalization
+            std: Standard deviation scalar value for input normalization
+        """
+        self.register_buffer('input_mean', torch.tensor(mean, device=self.device))
+        self.register_buffer('input_std', torch.tensor(std, device=self.device))
+        self.config.input_unit_norm = True
+        return self
+
+
+
+class BatchTopKSAE(BaseSAE):
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode input tensor to sparse features with batch-wise top-k
+        """
+        if self.config.input_unit_norm:
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            x_std = x.std(dim=-1, keepdim=True)
+            x = x / (x_std + 1e-5)
+            
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        acts_topk = torch.topk(acts.flatten(), self.config.top_k * x.shape[0], dim=-1)
+        return (
+            torch.zeros_like(acts.flatten())
+            .scatter(-1, acts_topk.indices, acts_topk.values)
+            .reshape(acts.shape)
+        )
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        acts_topk = torch.topk(acts.flatten(), self.config.top_k * x.shape[0], dim=-1)
+        acts_topk = (
+            torch.zeros_like(acts.flatten())
+            .scatter(-1, acts_topk.indices, acts_topk.values)
+            .reshape(acts.shape)
+        )
+        x_reconstruct = acts_topk @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts_topk)
+        output = self.get_loss_dict(x, x_reconstruct, acts, acts_topk, x_mean, x_std)
+        return output
+
+    def get_loss_dict(self, x, x_reconstruct, acts, acts_topk, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts_topk.float().abs().sum(-1).mean()
+        l1_loss = self.config.l1_coeff * l1_norm
+        l0_norm = (acts_topk > 0).float().sum(-1).mean()
+        aux_loss = self.get_auxiliary_loss(x, x_reconstruct, acts)
+        loss = l2_loss + aux_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.config.n_batches_to_dead
+        ).sum()
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        per_token_l2_loss_A = (x_reconstruct.float() - x.float()).pow(2).sum(-1).squeeze()
+        total_variance_A = (x.float() - x.float().mean(0)).pow(2).sum(-1).squeeze()
+        explained_variance = (1 - per_token_l2_loss_A / total_variance_A).mean()
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts_topk,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+            "aux_loss": aux_loss,
+            "explained_variance": explained_variance,
+            "top_k": self.config.top_k
+        }
+        return output
+
+    def get_auxiliary_loss(self, x, x_reconstruct, acts):
+        dead_features = self.num_batches_not_active >= self.config.n_batches_to_dead
+        if dead_features.sum() > 0:
+            residual = x.float() - x_reconstruct.float()
+            acts_topk_aux = torch.topk(
+                acts[:, dead_features],
+                min(self.config.top_k_aux, dead_features.sum()),
+                dim=-1,
+            )
+            acts_aux = torch.zeros_like(acts[:, dead_features]).scatter(
+                -1, acts_topk_aux.indices, acts_topk_aux.values
+            )
+            x_reconstruct_aux = acts_aux @ self.W_dec[dead_features]
+            l2_loss_aux = (
+                self.config.aux_penalty
+                * (x_reconstruct_aux.float() - residual.float()).pow(2).mean()
+            )
+            return l2_loss_aux
+        else:
+            return torch.tensor(0, dtype=x.dtype, device=x.device)
+
+
+class TopKSAE(BaseSAE):
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode input tensor to sparse features with per-sample top-k
+        """
+        if self.config.input_unit_norm:
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            x_std = x.std(dim=-1, keepdim=True)
+            x = x / (x_std + 1e-5)
+            
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        acts_topk = torch.topk(acts, self.config.top_k, dim=-1)
+        return torch.zeros_like(acts).scatter(
+            -1, acts_topk.indices, acts_topk.values
+        )
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        acts_topk = torch.topk(acts, self.config.top_k, dim=-1)
+        acts_topk = torch.zeros_like(acts).scatter(
+            -1, acts_topk.indices, acts_topk.values
+        )
+        x_reconstruct = acts_topk @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts_topk)
+        output = self.get_loss_dict(x, x_reconstruct, acts, acts_topk, x_mean, x_std)
+        return output
+
+    def get_loss_dict(self, x, x_reconstruct, acts, acts_topk, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts_topk.float().abs().sum(-1).mean()
+        l1_loss = self.config.l1_coeff * l1_norm
+        l0_norm = (acts_topk > 0).float().sum(-1).mean()
+        aux_loss = self.get_auxiliary_loss(x, x_reconstruct, acts)
+        loss = l2_loss + l1_loss + aux_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.config.n_batches_to_dead
+        ).sum()
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        per_token_l2_loss_A = (x_reconstruct.float() - x.float()).pow(2).sum(-1).squeeze()
+        total_variance_A = (x.float() - x.float().mean(0)).pow(2).sum(-1).squeeze()
+        explained_variance = (1 - per_token_l2_loss_A / total_variance_A).mean()
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts_topk,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+            "explained_variance": explained_variance,
+            "aux_loss": aux_loss,
+        }
+        return output
+
+    def get_auxiliary_loss(self, x, x_reconstruct, acts):
+        dead_features = self.num_batches_not_active >= self.config.n_batches_to_dead
+        if dead_features.sum() > 0:
+            residual = x.float() - x_reconstruct.float()
+            acts_topk_aux = torch.topk(
+                acts[:, dead_features],
+                min(self.config.top_k_aux, dead_features.sum()),
+                dim=-1,
+            )
+            acts_aux = torch.zeros_like(acts[:, dead_features]).scatter(
+                -1, acts_topk_aux.indices, acts_topk_aux.values
+            )
+            x_reconstruct_aux = acts_aux @ self.W_dec[dead_features]
+            l2_loss_aux = (
+                self.config.aux_penalty
+                * (x_reconstruct_aux.float() - residual.float()).pow(2).mean()
+            )
+            return l2_loss_aux
+        else:
+            return torch.tensor(0, dtype=x.dtype, device=x.device)
+
+
+class VanillaSAE(BaseSAE):
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        acts = F.relu(x @ self.W_enc + self.b_enc)
+        x_reconstruct = acts @ self.W_dec + self.b_dec
+        self.update_inactive_features(acts)
+        output = self.get_loss_dict(x, x_reconstruct, acts, x_mean, x_std)
+        return output
+
+    def get_loss_dict(self, x, x_reconstruct, acts, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+        l1_norm = acts.float().abs().sum(-1).mean()
+        l1_loss = self.config.l1_coeff * l1_norm
+        l0_norm = (acts > 0).float().sum(-1).mean()
+        loss = l2_loss + l1_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.config.n_batches_to_dead
+        ).sum()
+
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        per_token_l2_loss_A = (x_reconstruct.float() - x.float()).pow(2).sum(-1).squeeze()
+        total_variance_A = (x.float() - x.float().mean(0)).pow(2).sum(-1).squeeze()
+        explained_variance = (1 - per_token_l2_loss_A / total_variance_A).mean()
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0_norm,
+            "l1_norm": l1_norm,
+            "explained_variance": explained_variance,
+        }
+        
+        # Add timing if available
+        if hasattr(self, "start_event") and self.start_event is not None:
+            output["forward_time_ms"] = self.start_event.elapsed_time(self.end_event)
+        
+        return output
+
+
+import torch
+import torch.nn as nn
+
+class RectangleFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x):
+        ctx.save_for_backward(x)
+        return ((x > -0.5) & (x < 0.5)).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        grad_input = grad_output.clone()
+        grad_input[(x <= -0.5) | (x >= 0.5)] = 0
+        return grad_input
+
+class JumpReLUFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return x * (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = (x > threshold).float() * grad_output
+        threshold_grad = (
+            -(threshold / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None  # None for bandwidth
+
+class JumpReLU(nn.Module):
+    def __init__(self, feature_size, bandwidth, device='cpu'):
+        super(JumpReLU, self).__init__()
+        self.log_threshold = nn.Parameter(torch.zeros(feature_size, device=device))
+        self.bandwidth = bandwidth
+
+    def forward(self, x):
+        return JumpReLUFunction.apply(x, self.log_threshold, self.bandwidth)
+
+class StepFunction(autograd.Function):
+    @staticmethod
+    def forward(ctx, x, log_threshold, bandwidth):
+        ctx.save_for_backward(x, log_threshold, torch.tensor(bandwidth))
+        threshold = torch.exp(log_threshold)
+        return (x > threshold).float()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        x, log_threshold, bandwidth_tensor = ctx.saved_tensors
+        bandwidth = bandwidth_tensor.item()
+        threshold = torch.exp(log_threshold)
+        x_grad = torch.zeros_like(x)
+        threshold_grad = (
+            -(1.0 / bandwidth)
+            * RectangleFunction.apply((x - threshold) / bandwidth)
+            * grad_output
+        )
+        return x_grad, threshold_grad, None  # None for bandwidth
+
+class JumpReLUSAE(BaseSAE):
+    def __init__(self, config: SAEConfig):
+        super().__init__(config)
+        self.jumprelu = JumpReLU(
+            feature_size=config.dict_size,
+            bandwidth=config.bandwidth,
+            device=config.device if hasattr(config, 'device') else 'cpu'
+        )
+
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Encode input tensor to sparse features using JumpReLU
+        """
+        if self.config.input_unit_norm:
+            x_mean = x.mean(dim=-1, keepdim=True)
+            x = x - x_mean
+            x_std = x.std(dim=-1, keepdim=True)
+            x = x / (x_std + 1e-5)
+        
+        pre_activations = F.relu(x @ self.W_enc + self.b_enc)
+        return self.jumprelu(pre_activations)
+
+    def forward(self, x):
+        x, x_mean, x_std = self.preprocess_input(x)
+        pre_activations = torch.relu(x @ self.W_enc + self.b_enc)
+        feature_magnitudes = self.jumprelu(pre_activations)
+        x_reconstructed = feature_magnitudes @ self.W_dec + self.b_dec
+        return self.get_loss_dict(x, x_reconstructed, feature_magnitudes, x_mean, x_std)
+
+    def get_loss_dict(self, x, x_reconstruct, acts, x_mean, x_std):
+        l2_loss = (x_reconstruct.float() - x.float()).pow(2).mean()
+
+        l0 = StepFunction.apply(acts, self.jumprelu.log_threshold, self.config.bandwidth).sum(dim=-1).mean()
+        l0_loss = self.config.l1_coeff * l0
+        l1_loss = l0_loss
+
+        loss = l2_loss + l1_loss
+        num_dead_features = (
+            self.num_batches_not_active > self.config.n_batches_to_dead
+        ).sum()
+
+        sae_out = self.postprocess_output(x_reconstruct, x_mean, x_std)
+        per_token_l2_loss_A = (x_reconstruct.float() - x.float()).pow(2).sum(-1).squeeze()
+        total_variance_A = (x.float() - x.float().mean(0)).pow(2).sum(-1).squeeze()
+        explained_variance = (1 - per_token_l2_loss_A / total_variance_A).mean()
+        output = {
+            "sae_out": sae_out,
+            "feature_acts": acts,
+            "num_dead_features": num_dead_features,
+            "loss": loss,
+            "l1_loss": l1_loss,
+            "l2_loss": l2_loss,
+            "l0_norm": l0,
+            "l1_norm": l0,
+            "explained_variance": explained_variance,
+        }
+        return output
+    
+    @torch.no_grad()
+    def fold_W_dec_norm(self):
+        """
+        Make decoder weights unit norm and adjust encoder and thresholds accordingly
+        """
+        # Get current decoder norms
+        W_dec_norm = self.W_dec.norm(dim=-1, keepdim=True)
+        
+        # Call parent method to handle weights
+        super().fold_W_dec_norm()
+        
+        # Scale thresholds to preserve sparsity (remove keepdim to match threshold shape)
+        self.jumprelu.log_threshold.data = self.jumprelu.log_threshold + torch.log(W_dec_norm.squeeze(-1))
+        
+        return self
+    
+
+SAEConfig.register_for_auto_class("AutoConfig")
+BatchTopKSAE.register_for_auto_class()
+JumpReLUSAE.register_for_auto_class()
+VanillaSAE.register_for_auto_class()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f618af78706dddf6d13d50b575599f3d86ea34b7e39613d188daaf5dcdfd1b3
+size 67242576