Add SAE model definition

sae_model.py

@@ -0,0 +1,213 @@
+"""
+SAE Model - TopK Sparse Autoencoder
+"""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class SAEConfig:
+    """SAE 配置"""
+    input_dim: int = 4096
+    dict_size: int = 32768
+    k: int = 128
+    device: str = "cuda"
+    dtype: str = "bfloat16"
+    
+    def get_torch_dtype(self) -> torch.dtype:
+        dtype_map = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+        }
+        return dtype_map.get(self.dtype, torch.bfloat16)       
+
+
+class TopKSAE(nn.Module):
+    """
+    TopK Sparse Autoencoder    
+    """
+    def __init__(self, config: SAEConfig):
+        super().__init__()
+        self.config = config
+        
+        self.decoder = nn.Linear(config.dict_size, config.input_dim, bias=False)
+        self._normalize_decoder() 
+        self.encoder = nn.Linear(config.input_dim, config.dict_size, bias=False)
+        self.encoder.weight.data = self.decoder.weight.data.T.clone()
+        self.pre_bias = nn.Parameter(torch.zeros(config.input_dim))
+
+        self.to(config.device)
+        self.to(config.get_torch_dtype())
+    
+    def _normalize_decoder(self):
+        """归一化 decoder 权重（每列归一化）"""
+        with torch.no_grad():
+            self.decoder.weight.data = F.normalize(
+                self.decoder.weight.data, dim=0
+            )
+    
+    def encode(self, x: torch.Tensor) -> torch.Tensor:
+        """编码：input -> latent activations
+        Args:
+            x: [batch, input_dim]
+        Returns:
+            latent: [batch, dict_size] 稀疏激活
+        """
+        centered_x = x - self.pre_bias
+        pre_activation = self.encoder(centered_x)
+        
+        topk_values, topk_indices = torch.topk(
+            pre_activation, k=self.config.k, dim=-1
+        )
+        
+        latents = torch.zeros_like(pre_activation)
+        latents.scatter_(-1, topk_indices, F.relu(topk_values))
+        
+        return latents
+    
+    def decode(self, latents: torch.Tensor) -> torch.Tensor:
+        """解码：latent -> reconstruction
+        Args:
+            latent: [batch, dict_size]
+        Returns:
+            reconstruction: [batch, input_dim]
+        """
+        return self.decoder(latents) + self.pre_bias 
+    
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        latents = self.encode(x)
+        x_hat = self.decode(latents)
+        return x_hat, latents
+    
+    def compute_loss(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, dict]:
+        """计算损失
+        Args:
+            x: [batch, input_dim]
+        Returns:
+            (loss, loss_dict)
+        """
+        x_hat, latent = self.forward(x)
+        loss = (((x_hat - x) ** 2).mean(dim=-1) / (x**2).mean(dim=-1)).mean()
+        loss_dict = {
+            "mean_activation": latent[latent > 0].mean().item() if (latent > 0).any() else 0,
+        }
+        return loss, loss_dict
+        
+    def get_feature_activations(
+        self,
+        x: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """获取特征激活及其索引
+        Args:
+            x: [batch, input_dim]
+        Returns:
+            values: [batch, k] TopK 激活值
+            indices: [batch, k] TopK 特征索引
+        """
+        centered_x = x - self.pre_bias
+        pre_activation = self.encoder(centered_x)
+        topk_values, topk_indices = torch.topk(
+            pre_activation, k=self.config.k, dim=-1
+        )
+        return F.relu(topk_values), topk_indices
+    
+    def get_decoder_vectors(self, indices: torch.Tensor) -> torch.Tensor:
+        """获取指定特征的 decoder 向量
+        Args:
+            indices: [batch, k] 特征索引
+        Returns:
+            vectors: [batch, k, input_dim] decoder 向量
+        """
+        return self.decoder.weight[:, indices].permute(1, 2, 0)
+    
+    def steer(
+        self,
+        x: torch.Tensor,
+        feature_idx: int,
+        strength: float,
+    ) -> torch.Tensor:
+        """在原始 activation 上直接加减 decoder 方向来 steer 指定特征。
+
+        h_new = h + (strength - 1) * z_i * d_i
+
+        其中 z_i 是 feature_idx 的 TopK 编码激活值，d_i 是其 decoder 列向量。
+        不走完整 encode→decode，避免 TopK 重建误差污染其他特征。
+        strength=1.0 时变化量为 0（真正的 baseline）。
+
+        Args:
+            x: [batch, input_dim] 输入激活
+            feature_idx: 要 steering 的特征索引
+            strength: 目标特征缩放系数（1.0 不变，>1 增强，<1 抑制，0 完全消除）
+        Returns:
+            steered_x: [batch, input_dim] 调整后的激活
+        """
+        with torch.no_grad():
+            latents = self.encode(x)                         # [batch, dict_size]
+            z_i = latents[:, feature_idx]                   # [batch]
+            d_i = self.decoder.weight[:, feature_idx]       # [input_dim]
+            delta = (strength - 1) * z_i.unsqueeze(-1) * d_i.unsqueeze(0)
+            return x + delta
+
+    def steer_multi(
+        self,
+        x: torch.Tensor,
+        feature_indices: list,
+        strengths: list,
+    ) -> torch.Tensor:
+        """在原始 activation 上直接加减 decoder 方向来同时 steer 多个特征。
+
+        h_new = h + sum_i (strength_i - 1) * z_i * d_i
+
+        不走完整 encode→decode，避免 TopK 重建误差污染其他特征。
+        strength=1.0 时对应特征的变化量为 0（真正的 baseline）。
+
+        Args:
+            x: [batch, input_dim] 输入激活
+            feature_indices: 要 steering 的特征索引列表
+            strengths: 对应的缩放系数列表
+        Returns:
+            steered_x: [batch, input_dim] 调整后的激活
+        """
+        with torch.no_grad():
+            latents = self.encode(x)                         # [batch, dict_size]
+            delta = torch.zeros_like(x)
+            for feat_idx, strength in zip(feature_indices, strengths):
+                z_i = latents[:, feat_idx]                  # [batch]
+                d_i = self.decoder.weight[:, feat_idx]      # [input_dim]
+                delta += (strength - 1) * z_i.unsqueeze(-1) * d_i.unsqueeze(0)
+            return x + delta
+    
+    def save(self, path: str):
+        """保存模型"""
+        torch.save({
+            "config": self.config,
+            "state_dict": self.state_dict(),
+        }, path)
+    
+    @classmethod
+    def load(cls, path: str, device: str = "cuda") -> "TopKSAE":
+        """加载模型"""
+        checkpoint = torch.load(path, map_location=device, weights_only=False)
+        config = checkpoint["config"]
+        config.device = device
+        
+        model = cls(config)
+        model.load_state_dict(checkpoint["state_dict"])
+        model.to(config.get_torch_dtype())
+        
+        return model
+
+