drrik

 45class ActivationExtractor:
 46    """
 47    Extract MLP activations from language models using nnsight.
 48
 49    Handles model and dataset loading from HuggingFace Hub, runs
 50    inference with nnsight's tracing context to capture MLP outputs,
 51    and provides persistence via ``.npy`` / ``.pkl`` files.
 52
 53    The HuggingFace auth token is read automatically from the
 54    ``HUGGINGFACE_HUB_TOKEN`` environment variable (or ``.env`` file)
 55    via :func:`~drrik.settings.get_settings`.
 56
 57    Attributes:
 58        config: The resolved :class:`~drrik.config.ActivationExtractorConfig`.
 59        model: The loaded nnsight ``LanguageModel`` (``None`` until
 60            :meth:`load_model` is called).
 61        tokenizer: The HuggingFace tokenizer (``None`` until
 62            :meth:`load_model` is called).
 63        dataset: The loaded HuggingFace ``Dataset`` (``None`` until
 64            :meth:`load_dataset` is called).
 65
 66    Example:
 67        ```python
 68        extractor = ActivationExtractor(
 69            model_name="google/gemma-2b",
 70            dataset_name="wikitext",
 71            dataset_config="wikitext-2-raw-v1",
 72            mlp_layers=[0],
 73            num_samples=1000,
 74        )
 75        activations, metadata = extractor.extract()
 76        extractor.save_activations(activations, metadata, "./output")
 77        ```
 78    """
 79
 80    def __init__(self, config: Optional[ActivationExtractorConfig] = None, **kwargs):
 81        """
 82        Initialize the ActivationExtractor.
 83
 84        Accepts either a pre-built config object or flat keyword
 85        arguments that are automatically sorted into
 86        :class:`~drrik.config.ModelConfig`,
 87        :class:`~drrik.config.DatasetConfig`, and the remaining
 88        extraction-level settings.
 89
 90        Args:
 91            config: A fully constructed
 92                :class:`~drrik.config.ActivationExtractorConfig`.
 93                If ``None``, the config is built from ``**kwargs``.
 94            **kwargs: Flat config overrides. Recognised keys are
 95                distributed to the appropriate sub-config:
 96
 97                - *Model*: ``model_name``, ``revision``,
 98                  ``torch_dtype``, ``device_map``,
 99                  ``trust_remote_code``
100                - *Dataset*: ``dataset_name``, ``dataset_config``,
101                  ``split``, ``num_samples``, ``text_column``,
102                  ``max_length``
103                - *Extraction*: ``mlp_layers``, ``batch_size``,
104                  ``output_dir``
105        """
106        if config is None:
107            model_kwargs = {
108                k: v
109                for k, v in kwargs.items()
110                if k
111                in (
112                    "model_name",
113                    "revision",
114                    "torch_dtype",
115                    "device_map",
116                    "trust_remote_code",
117                )
118            }
119            dataset_kwargs = {
120                k: v
121                for k, v in kwargs.items()
122                if k
123                in (
124                    "dataset_name",
125                    "dataset_config",
126                    "split",
127                    "num_samples",
128                    "text_column",
129                    "max_length",
130                )
131            }
132            remaining_kwargs = {
133                k: v
134                for k, v in kwargs.items()
135                if k not in model_kwargs and k not in dataset_kwargs
136            }
137
138            model = ModelConfig(**model_kwargs) if model_kwargs else None
139            dataset = DatasetConfig(**dataset_kwargs) if dataset_kwargs else None
140
141            config = ActivationExtractorConfig(
142                model=model,
143                dataset=dataset,
144                **remaining_kwargs,
145            )
146
147        self.config = config
148        self.model = None
149        self.tokenizer = None
150        self.dataset = None
151        self._activations = []
152        self._metadata = []
153
154    def load_model(self) -> LanguageModel:
155        """
156        Load the model from HuggingFace Hub using nnsight.
157
158        If the model is already loaded, returns the cached instance.
159        The HuggingFace token is read from ``HUGGINGFACE_HUB_TOKEN``
160        via :func:`~drrik.settings.get_settings`.
161
162        Returns:
163            The loaded nnsight :class:`~nnsight.LanguageModel` wrapper.
164
165        Raises:
166            RuntimeError: If model loading fails.
167        """
168        if self.model is not None:
169            return self.model
170
171        try:
172            logger.info(f"Loading model: {self.config.model.model_name}")
173
174            # Get HF token from settings
175            settings = get_settings()
176            hf_token = settings.huggingface_hub_token
177
178            if hf_token:
179                logger.info("Using HuggingFace Hub token for authentication")
180
181            # Load tokenizer with token
182            tokenizer_kwargs = {
183                "revision": self.config.model.revision,
184                "trust_remote_code": self.config.model.trust_remote_code,
185            }
186            if hf_token:
187                tokenizer_kwargs["token"] = hf_token
188
189            self.tokenizer = AutoTokenizer.from_pretrained(
190                self.config.model.model_name, **tokenizer_kwargs
191            )
192
193            if self.tokenizer.pad_token is None:
194                self.tokenizer.pad_token = self.tokenizer.eos_token
195
196            # Determine dtype
197            dtype_map = {
198                "float16": torch.float16,
199                "bfloat16": torch.bfloat16,
200                "float32": torch.float32,
201            }
202            torch_dtype = dtype_map.get(self.config.model.torch_dtype, torch.float16)
203
204            # Load model with nnsight (with token if available)
205            nnsight_kwargs = {
206                "revision": self.config.model.revision,
207                "torch_dtype": torch_dtype,
208                "trust_remote_code": self.config.model.trust_remote_code,
209                "device_map": self.config.model.device_map,
210            }
211            if hf_token:
212                nnsight_kwargs["token"] = hf_token
213
214            self.model = LanguageModel(self.config.model.model_name, **nnsight_kwargs)
215
216            logger.info(f"Model loaded successfully on {self.model.device}")
217            return self.model
218
219        except Exception as e:
220            logger.error(f"Failed to load model: {e}")
221            raise RuntimeError(f"Model loading failed: {e}") from e
222
223    def load_dataset(self) -> Dataset:
224        """
225        Load the dataset from HuggingFace Hub.
226
227        If the dataset is already loaded, returns the cached instance.
228        The HuggingFace token is forwarded for gated datasets when
229        available.
230
231        Returns:
232            The loaded HuggingFace :class:`~datasets.Dataset`.
233
234        Raises:
235            RuntimeError: If dataset loading fails.
236        """
237        if self.dataset is not None:
238            return self.dataset
239
240        try:
241            logger.info(
242                f"Loading dataset: {self.config.dataset.dataset_name} "
243                f"({self.config.dataset.split} split)"
244            )
245
246            # Get HF token from settings (for gated datasets)
247            settings = get_settings()
248            hf_token = settings.huggingface_hub_token
249
250            # Load dataset with token if available
251            load_kwargs = {
252                "path": self.config.dataset.dataset_name,
253                "name": self.config.dataset.dataset_config,
254                "split": self.config.dataset.split,
255            }
256            if hf_token:
257                load_kwargs["token"] = hf_token
258
259            self.dataset = load_dataset(**load_kwargs)
260
261            logger.info(f"Dataset loaded with {len(self.dataset)} examples")
262            return self.dataset
263
264        except Exception as e:
265            logger.error(f"Failed to load dataset: {e}")
266            raise RuntimeError(f"Dataset loading failed: {e}") from e
267
268    def _get_mlp_layer_name(self, layer_idx: int) -> str:
269        """
270        Get the nnsight module path to an MLP layer.
271
272        Auto-detects the correct path pattern based on the model name
273        in the config.  For unknown architectures a warning is logged
274        and the default Gemma/Llama pattern is returned.
275
276        Args:
277            layer_idx: Zero-indexed transformer layer number.
278
279        Returns:
280            A dotted/bracket module path string, e.g.
281            ``"model.layers[3].mlp"``.
282        """
283        model_name_lower = self.config.model.model_name.lower()
284
285        # Gemma/GPT-style: model.layers.N.mlp
286        if any(name in model_name_lower for name in ["gemma", "gpt-2", "pythia"]):
287            return f"model.layers[{layer_idx}].mlp"
288
289        # Llama-style: model.layers.N.mlp
290        if "llama" in model_name_lower:
291            return f"model.layers[{layer_idx}].mlp"
292
293        # Phi-style: model.layers.N.mlp
294        if "phi" in model_name_lower:
295            return f"model.layers[{layer_idx}].mllp"  # Note: mlp can be 'mlp' or 'mlp'
296
297        # BERT-style: bert.encoder.layer.N.output.dense
298        if "bert" in model_name_lower:
299            return f"bert.encoder.layer[{layer_idx}].output"
300
301        # Default: try common patterns
302        common_patterns = [
303            f"model.layers[{layer_idx}].mlp",
304            f"model.layers[{layer_idx}].ffn",
305            f"transformer.h[{layer_idx}].mlp",
306            f"layers[{layer_idx}].mlp",
307        ]
308
309        logger.warning(
310            f"Unknown model architecture '{self.config.model.model_name}', "
311            f"trying common patterns"
312        )
313        return common_patterns[0]
314
315    def _resolve_layer_path(self, path: str):
316        """Resolve a dotted/bracket path to the actual nnsight module.
317
318        Delegates to the shared :func:`~drrik.steering.resolve_module_path`
319        utility so that path-resolution logic is maintained in one place.
320
321        Args:
322            path: A dot-separated module path with optional bracket
323                indexing, e.g., ``model.layers[2].mlp``.
324
325        Returns:
326            The resolved nnsight module object corresponding to the
327            given path.
328        """
329        return resolve_module_path(self.model, path)
330
331    def extract(
332        self,
333        num_samples: Optional[int] = None,
334    ) -> Tuple[np.ndarray, Dict[str, Any]]:
335        """
336        Extract MLP activations from the model.
337
338        Loads the model and dataset (if not already loaded), tokenizes
339        the data, and runs batched forward passes using nnsight's
340        tracing context.  For each sample, the **last-token** activation
341        from each requested MLP layer is collected and concatenated.
342
343        Args:
344            num_samples: Override for the number of samples to process.
345                If ``None``, uses ``config.dataset.num_samples``.
346
347        Returns:
348            A tuple of:
349
350            - **activations** (:class:`numpy.ndarray`) — shape
351              ``(n_samples, activation_dim * n_layers)``.  When a
352              single layer is requested the shape simplifies to
353              ``(n_samples, activation_dim)``.
354            - **metadata** (:class:`dict`) — contains the serialised
355              config, sample count, activation dimension, layer paths,
356              and per-sample metadata (index, truncated text, input
357              ids).
358
359        Raises:
360            RuntimeError: If extraction fails at any stage.
361        """
362        try:
363            # Load model and dataset
364            self.load_model()
365            self.load_dataset()
366
367            n_samples = num_samples or self.config.dataset.num_samples
368            logger.info(
369                f"Extracting activations from {len(self.config.mlp_layers)} MLP layers "
370                f"for {n_samples} samples"
371            )
372
373            # Prepare dataset
374            dataset = self.dataset.select(range(min(n_samples, len(self.dataset))))
375
376            # Tokenize dataset
377            def tokenize_function(examples):
378                return self.tokenizer(
379                    examples[self.config.dataset.text_column],
380                    padding="max_length",
381                    truncation=True,
382                    max_length=self.config.dataset.max_length,
383                    return_tensors="pt",
384                )
385
386            tokenized = dataset.map(
387                tokenize_function,
388                batched=True,
389                remove_columns=dataset.column_names,
390                desc="Tokenizing",
391            )
392
393            # Get MLP layer names
394            layer_paths = [
395                self._get_mlp_layer_name(layer_idx)
396                for layer_idx in self.config.mlp_layers
397            ]
398
399            logger.info(f"Extracting from layers: {layer_paths}")
400
401            # Collect activations using nnsight
402            self._activations = []
403            self._metadata = []
404
405            batch_size = self.config.batch_size
406            n_batches = (len(tokenized) + batch_size - 1) // batch_size
407
408            with torch.no_grad():
409                for batch_idx in tqdm(range(n_batches), desc="Extracting activations"):
410                    start_idx = batch_idx * batch_size
411                    end_idx = min(start_idx + batch_size, len(tokenized))
412
413                    batch = tokenized[start_idx:end_idx]
414                    input_ids = torch.tensor(batch["input_ids"])
415                    attention_mask = torch.tensor(batch["attention_mask"])
416
417                    layer_outputs = []
418
419                    # Use nnsight to extract activations
420                    with self.model.trace(input_ids, attention_mask=attention_mask):
421                        for layer_path in layer_paths:
422                            module = self._resolve_layer_path(layer_path)
423                            output = module.output.save()
424                            layer_outputs.append(output)
425
426                    # Process outputs
427                    batch_input_ids = input_ids.cpu().numpy()
428                    for sample_idx in range(len(input_ids)):
429                        sample_activations = []
430                        for layer_output in layer_outputs:
431                            # Get activation after non-linearity (if applicable)
432                            # Shape: (seq_len, hidden_dim) or (batch, seq_len, hidden_dim)
433                            act = layer_output
434
435                            if act.dim() == 3:
436                                act = act[sample_idx]  # (seq_len, hidden_dim)
437                            elif act.dim() == 2:
438                                act = act  # Already (seq_len, hidden_dim)
439
440                            # Use the last token's activation (common practice)
441                            act = act[-1]  # (hidden_dim,)
442
443                            sample_activations.append(act.cpu().numpy())
444
445                        # Concatenate activations from all layers
446                        self._activations.append(np.concatenate(sample_activations))
447
448                        # Store metadata
449                        self._metadata.append(
450                            {
451                                "sample_idx": start_idx + sample_idx,
452                                "text": dataset[start_idx + sample_idx][
453                                    self.config.dataset.text_column
454                                ][:200],
455                                "input_ids": batch_input_ids[sample_idx],
456                            }
457                        )
458
459            activations = np.array(self._activations)
460            logger.info(f"Extracted activations shape: {activations.shape}")
461
462            metadata = {
463                "config": self.config.model_dump(),
464                "n_samples": len(self._activations),
465                "activation_dim": activations.shape[-1],
466                "layer_paths": layer_paths,
467                "samples_metadata": self._metadata,
468            }
469
470            return activations, metadata
471
472        except Exception as e:
473            logger.error(f"Failed to extract activations: {e}")
474            raise RuntimeError(f"Activation extraction failed: {e}") from e
475
476    def save_activations(
477        self,
478        activations: np.ndarray,
479        metadata: Dict[str, Any],
480        output_dir: Optional[Union[str, Path]] = None,
481    ) -> Path:
482        """
483        Save extracted activations to disk.
484
485        Writes two files to *output_dir*:
486
487        - ``activations.npy`` — the activation array.
488        - ``metadata.pkl`` — the metadata dictionary.
489
490        Args:
491            activations: Activation array of shape
492                ``(n_samples, dim)``.
493            metadata: Metadata dictionary (as returned by
494                :meth:`extract`).
495            output_dir: Target directory.  If ``None``, uses
496                ``config.output_dir``.
497
498        Returns:
499            Path to the saved ``activations.npy`` file.
500
501        Raises:
502            ValueError: If *output_dir* is ``None`` and no output
503                directory is configured.
504        """
505        if output_dir is None:
506            output_dir = self.config.output_dir
507
508        output_dir = Path(output_dir)
509        output_dir.mkdir(parents=True, exist_ok=True)
510
511        activations_path = output_dir / "activations.npy"
512        np.save(str(activations_path), activations)
513
514        metadata_path = output_dir / "metadata.pkl"
515        with open(metadata_path, "wb") as f:
516            pickle.dump(metadata, f)
517
518        logger.info(f"Saved activations to {activations_path}")
519        return activations_path
520
521    def load_activations(
522        self,
523        filepath: Union[str, Path],
524    ) -> Tuple[np.ndarray, Dict[str, Any]]:
525        """
526        Load saved activations from disk.
527
528        Supports two formats:
529
530        - ``.npy`` — loads the array directly and looks for a
531          companion ``metadata.pkl`` in the same directory.
532        - ``.pkl`` (legacy) — loads a dict with ``"activations"``
533          and ``"metadata"`` keys.
534
535        Args:
536            filepath: Path to an ``.npy`` or ``.pkl`` file.
537
538        Returns:
539            A tuple of (activations array, metadata dict).
540
541        Raises:
542            FileNotFoundError: If the file (or companion metadata) does
543                not exist.
544        """
545        filepath = Path(filepath)
546
547        if filepath.suffix == ".npy":
548            activations = np.load(str(filepath))
549            metadata_path = filepath.parent / "metadata.pkl"
550            with open(metadata_path, "rb") as f:
551                metadata = pickle.load(f)
552        else:
553            with open(filepath, "rb") as f:
554                data = pickle.load(f)
555            activations = data["activations"]
556            metadata = data["metadata"]
557
558        logger.info(f"Loaded activations from {filepath}")
559        return activations, metadata

 47class SparseAutoencoder(nn.Module):
 48    """
 49    Sparse Autoencoder for extracting interpretable features from MLP activations.
 50
 51    Architecture:
 52        Input (activation_dim) -> [bias] -> Encoder (hidden_dim) -> ReLU -> [bias]
 53        -> Decoder (activation_dim) -> Output
 54
 55    The decoder weights are normalized to unit norm to prevent scaling collapse.
 56    L1 regularization on hidden activations encourages sparsity.
 57
 58    Based on the architecture described in:
 59    "Towards Monosemanticity: Decomposing Language Models With Dictionary Learning"
 60    https://transformer-circuits.pub/2023/monosemantic-features/index.html
 61
 62    Attributes:
 63        encoder: Linear layer from activation_dim to hidden_dim
 64        decoder: Linear layer from hidden_dim to activation_dim
 65        pre_encoder_bias: Learnable bias subtracted from input
 66        post_decoder_bias: Learnable bias added to output (tied to pre_encoder_bias)
 67
 68    Example:
 69        ```python
 70        sae = SparseAutoencoder(
 71            activation_dim=2048,
 72            hidden_dim=4096,  # 2x expansion
 73            l1_coefficient=0.01
 74        )
 75        sae.fit(activations, num_epochs=100)
 76        features = sae.encode(activations)
 77        reconstructed = sae.decode(features)
 78        ```
 79    """
 80
 81    def __init__(
 82        self,
 83        activation_dim: int,
 84        hidden_dim: int,
 85        l1_coefficient: float = 0.01,
 86        normalize_decoder: bool = True,
 87        pre_encoder_bias: bool = True,
 88    ):
 89        """
 90        Initialize the Sparse Autoencoder.
 91
 92        Args:
 93            activation_dim: Dimension of input MLP activations
 94            hidden_dim: Dimension of sparse hidden layer (overcomplete)
 95            l1_coefficient: L1 regularization strength
 96            normalize_decoder: Whether to normalize decoder columns to unit norm
 97            pre_encoder_bias: Whether to use pre-encoder bias (as in the paper)
 98        """
 99        super().__init__()
100
101        self.activation_dim = activation_dim
102        self.hidden_dim = hidden_dim
103        self.l1_coefficient = l1_coefficient
104        self.normalize_decoder = normalize_decoder
105        self.pre_encoder_bias = pre_encoder_bias
106
107        # Encoder: input -> hidden
108        self.encoder = nn.Linear(activation_dim, hidden_dim, bias=False)
109
110        # Decoder: hidden -> output
111        # Initialize with small random weights
112        self.decoder = nn.Linear(hidden_dim, activation_dim, bias=False)
113
114        # Initialize decoder weights to unit norm (Kaiming uniform)
115        nn.init.kaiming_uniform_(self.decoder.weight, a=np.sqrt(5))
116        with torch.no_grad():
117            self.decoder.weight.copy_(
118                self.decoder.weight / self.decoder.weight.norm(dim=0, keepdim=True)
119            )
120
121        # Initialize encoder weights
122        nn.init.kaiming_uniform_(self.encoder.weight, a=np.sqrt(5))
123
124        # Biases
125        if pre_encoder_bias:
126            # Pre-encoder bias (subtracted from input, added to output)
127            # Initialize to geometric median of data (will be set during fit)
128            self.bias = nn.Parameter(torch.zeros(activation_dim))
129        else:
130            self.register_parameter("bias", None)
131
132        # Encoder bias
133        self.encoder_bias = nn.Parameter(torch.zeros(hidden_dim))
134
135        # Training metrics
136        self.training_losses = []
137        self.training_l0_norms = []
138
139    def encode(self, x: torch.Tensor) -> torch.Tensor:
140        """
141        Encode input activations to sparse features.
142
143        Args:
144            x: Input tensor of shape (batch_size, activation_dim)
145
146        Returns:
147            Sparse feature activations of shape (batch_size, hidden_dim)
148        """
149        # Subtract bias if using pre-encoder bias
150        if self.bias is not None:
151            x = x - self.bias
152
153        # Apply encoder
154        hidden = F.linear(x, self.encoder.weight, self.encoder_bias)
155
156        # ReLU activation for sparsity
157        features = F.relu(hidden)
158
159        return features
160
161    def decode(self, features: torch.Tensor) -> torch.Tensor:
162        """
163        Decode sparse features back to activation space.
164
165        Args:
166            features: Sparse feature tensor of shape (batch_size, hidden_dim)
167
168        Returns:
169            Reconstructed activations of shape (batch_size, activation_dim)
170        """
171        # Apply decoder (weights are normalized to unit norm)
172        output = F.linear(features, self.decoder.weight)
173
174        # Add bias if using
175        if self.bias is not None:
176            output = output + self.bias
177
178        return output
179
180    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
181        """
182        Forward pass through the autoencoder.
183
184        Args:
185            x: Input tensor of shape (batch_size, activation_dim)
186
187        Returns:
188            Tuple of (reconstructed, features)
189        """
190        features = self.encode(x)
191        reconstructed = self.decode(features)
192        return reconstructed, features
193
194    def loss(
195        self, x: torch.Tensor, reconstructed: torch.Tensor, features: torch.Tensor
196    ) -> torch.Tensor:
197        """
198        Compute the loss with L1 regularization.
199
200        Loss = MSE(reconstructed, x) + lambda * L1(features)
201
202        Args:
203            x: Input tensor
204            reconstructed: Reconstructed tensor
205            features: Sparse feature activations
206
207        Returns:
208            Total loss
209        """
210        mse_loss = F.mse_loss(reconstructed, x)
211        l1_loss = self.l1_coefficient * features.abs().sum(dim=1).mean()
212        return mse_loss + l1_loss
213
214    def normalize_decoder_weights(self) -> None:
215        """Normalize decoder weight columns to unit L2 norm.
216
217        This prevents scaling collapse where the optimizer could
218        trivially reduce the loss by shrinking decoder weights and
219        scaling up encoder weights. Applied after each optimizer step
220        when ``normalize_decoder`` is ``True``.
221        """
222        if self.normalize_decoder:
223            with torch.no_grad():
224                self.decoder.weight.copy_(
225                    self.decoder.weight
226                    / self.decoder.weight.norm(dim=0, keepdim=True).clamp(min=1e-8)
227                )
228
229    def resample_dead_neurons(
230        self,
231        activations: torch.Tensor,
232        dead_threshold: float = 1e-8,
233        dead_mask: Optional[torch.Tensor] = None,
234    ) -> int:
235        """
236        Resample dead neurons that haven't fired recently.
237
238        This follows the procedure from the paper:
239        1. Identify neurons that haven't fired above threshold
240        2. Compute loss on a batch of data
241        3. Resample dead neurons to fit poorly reconstructed examples
242
243        Args:
244            activations: Batch of activations to use for resampling
245            dead_threshold: Activation threshold below which a neuron is considered dead
246            dead_mask: Pre-computed boolean mask of dead neurons. If None, computed
247                       from the current batch. When provided, uses a sliding window
248                       approach for more robust dead neuron detection.
249
250        Returns:
251            Number of neurons resampled
252        """
253        with torch.no_grad():
254            # Get feature activations
255            features = self.encode(activations)
256
257            # Find dead neurons
258            if dead_mask is None:
259                neuron_activity = features.abs().sum(dim=0)
260                dead_mask = neuron_activity < dead_threshold
261            n_dead = dead_mask.sum().item()
262
263            if n_dead == 0:
264                return 0
265
266            logger.info(f"Resampling {n_dead} dead neurons")
267
268            # Compute reconstruction loss for each sample
269            reconstructed, _ = self.forward(activations)
270            sample_losses = F.mse_loss(
271                reconstructed, activations, reduction="none"
272            ).sum(dim=1)
273
274            # Sample based on loss (higher loss = more likely to be resampled)
275            probs = sample_losses**2
276            probs = probs / probs.sum()
277
278            # For each dead neuron, resample from a high-loss example
279            for dead_idx in torch.where(dead_mask)[0]:
280                # Sample an example
281                sample_idx = torch.multinomial(probs, 1).item()
282                example = activations[sample_idx]
283
284                # Normalize example
285                example_normalized = example / (example.norm() + 1e-8)
286
287                # Set decoder weight
288                self.decoder.weight[:, dead_idx] = example_normalized
289
290                # Set encoder weight (smaller scale to prevent immediate firing)
291                avg_encoder_norm = (
292                    self.encoder.weight[~dead_mask].norm(dim=1).mean().item()
293                    if (~dead_mask).any() > 0
294                    else 0.1
295                )
296                self.encoder.weight[dead_idx, :] = (
297                    example_normalized * avg_encoder_norm * 0.2
298                )
299
300                # Reset encoder bias
301                self.encoder_bias[dead_idx] = 0
302
303            # Re-normalize decoder
304            self.normalize_decoder_weights()
305
306        return n_dead
307
308    def fit(
309        self,
310        activations: np.ndarray,
311        batch_size: int = 256,
312        num_epochs: int = 100,
313        learning_rate: float = 1e-4,
314        validation_split: float = 0.1,
315        resample_dead_neurons: bool = True,
316        resample_interval: int = 10000,
317        dead_threshold: float = 1e-8,
318        window_size: int = 100,
319        device: Optional[str] = None,
320        verbose: bool = True,
321        wandb_config: Optional[WandbConfig] = None,
322        wandb_enabled: bool = True,
323    ) -> "SparseAutoencoder":
324        """
325        Train the sparse autoencoder on MLP activations.
326
327        Args:
328            activations: Training data of shape (n_samples, activation_dim)
329            batch_size: Batch size for training
330            num_epochs: Number of training epochs
331            learning_rate: Learning rate for Adam optimizer
332            validation_split: Fraction of data to use for validation
333            resample_dead_neurons: Whether to resample dead neurons during training
334            resample_interval: Steps between resampling checks
335            dead_threshold: Activation threshold below which a neuron is considered dead
336            window_size: Number of recent batches to track for dead neuron detection
337            device: Device to train on (cuda/cpu). If None, auto-detect
338            verbose: Whether to show progress bars
339            wandb_config: Optional WandbConfig for experiment tracking
340            wandb_enabled: If True and wandb_config is None, create default WandbConfig
341
342        Returns:
343            self (trained model)
344        """
345        # Setup wandb if enabled
346        wandb_logger = None
347        _owns_wandb = False
348        if wandb_enabled:
349            if wandb_config is None:
350                wandb_logger = WandbConfig(
351                    config={
352                        "activation_dim": self.activation_dim,
353                        "hidden_dim": self.hidden_dim,
354                        "expansion_factor": self.hidden_dim / self.activation_dim,
355                        "l1_coefficient": self.l1_coefficient,
356                        "batch_size": batch_size,
357                        "learning_rate": learning_rate,
358                        "num_epochs": num_epochs,
359                    }
360                )
361                wandb_logger.initialize()
362                _owns_wandb = True
363            elif not wandb_config._initialized:
364                wandb_logger = wandb_config
365                wandb_logger.initialize()
366                _owns_wandb = True
367            else:
368                wandb_logger = wandb_config
369
370        # Determine device
371        if device is None:
372            device = "cuda" if torch.cuda.is_available() else "cpu"
373
374        self.to(device)
375
376        # Convert to tensor
377        if isinstance(activations, np.ndarray):
378            activations = torch.from_numpy(activations).float()
379
380        n_samples = len(activations)
381        n_val = int(n_samples * validation_split)
382        n_train = n_samples - n_val
383
384        # Split data
385        train_data = activations[:n_train]
386        val_data = activations[n_train:]
387
388        # Create data loaders
389        train_dataset = TensorDataset(train_data)
390        train_loader = DataLoader(
391            train_dataset,
392            batch_size=batch_size,
393            shuffle=True,
394        )
395
396        # Initialize bias to median if using
397        if self.bias is not None and self.pre_encoder_bias:
398            with torch.no_grad():
399                # Approximate geometric median using median for simplicity
400                median = train_data.to(device).median(dim=0).values
401                self.bias.data = median
402
403        # Optimizer
404        optimizer = torch.optim.Adam(
405            self.parameters(),
406            lr=learning_rate,
407        )
408
409        # Training loop
410        self.training_losses = []
411        self.training_l0_norms = []
412
413        n_steps_since_resample = 0
414        global_step = 0
415
416        # Sliding window for tracking neuron activity across batches
417        neuron_activity_window = []
418
419        if verbose:
420            pbar = tqdm(total=num_epochs, desc="Training SAE")
421        else:
422            pbar = None
423
424        for epoch in range(num_epochs):
425            epoch_loss = 0.0
426            epoch_l0 = 0.0
427
428            for batch in train_loader:
429                batch = batch[0].to(device)
430
431                # Forward pass
432                reconstructed, features = self.forward(batch)
433
434                # Track neuron activity for sliding window dead detection
435                neuron_activity = features.abs().sum(dim=0)
436                neuron_activity_window.append(neuron_activity)
437                if len(neuron_activity_window) > window_size:
438                    neuron_activity_window.pop(0)
439
440                # Compute loss
441                loss = self.loss(batch, reconstructed, features)
442
443                # Backward pass
444                optimizer.zero_grad()
445                loss.backward()
446
447                # Remove gradients parallel to decoder weights (Adam optimization trick from paper)
448                if self.normalize_decoder:
449                    with torch.no_grad():
450                        # Project gradients to be orthogonal to decoder weights
451                        decoder_w = self.decoder.weight
452                        decoder_grad = self.decoder.weight.grad
453
454                        # Compute projection
455                        projection = (decoder_grad * decoder_w).sum(
456                            dim=0, keepdim=True
457                        ) * decoder_w
458                        self.decoder.weight.grad = decoder_grad - projection
459
460                optimizer.step()
461
462                # Normalize decoder weights
463                self.normalize_decoder_weights()
464
465                # Track metrics
466                epoch_loss += loss.item()
467                epoch_l0 += (features > 0).sum(dim=1).float().mean().item()
468
469                n_steps_since_resample += 1
470                global_step += 1
471
472                # Resample dead neurons using sliding window
473                if (
474                    resample_dead_neurons
475                    and n_steps_since_resample >= resample_interval
476                ):
477                    if len(neuron_activity_window) > 0:
478                        avg_activity = torch.stack(neuron_activity_window).mean(dim=0)
479                        window_dead_mask = avg_activity < dead_threshold
480                    else:
481                        window_dead_mask = None
482
483                    n_resampled = self.resample_dead_neurons(
484                        batch, dead_threshold=dead_threshold, dead_mask=window_dead_mask
485                    )
486                    if n_resampled > 0:
487                        neuron_activity_window.clear()
488                    n_steps_since_resample = 0
489
490            # Average metrics
491            avg_loss = epoch_loss / len(train_loader)
492            avg_l0 = epoch_l0 / len(train_loader)
493
494            self.training_losses.append(avg_loss)
495            self.training_l0_norms.append(avg_l0)
496
497            # Validation
498            with torch.no_grad():
499                val_data_tensor = val_data.to(device)
500                val_reconstructed, val_features = self.forward(val_data_tensor)
501                val_loss = F.mse_loss(val_reconstructed, val_data_tensor).item()
502                val_l0 = (val_features > 0).sum(dim=1).float().mean().item()
503
504            # Log to wandb
505            if wandb_logger:
506                wandb_logger.log_metrics(
507                    {
508                        "epoch": epoch + 1,
509                        "train_loss": avg_loss,
510                        "train_l0_norm": avg_l0,
511                        "val_loss": val_loss,
512                        "val_l0_norm": val_l0,
513                        "learning_rate": learning_rate,
514                    },
515                    step=epoch + 1,
516                )
517
518            if verbose:
519                if pbar:
520                    pbar.update(1)
521                    pbar.set_postfix(
522                        {
523                            "loss": f"{avg_loss:.6f}",
524                            "val_loss": f"{val_loss:.6f}",
525                            "L0": f"{avg_l0:.2f}",
526                            "val_L0": f"{val_l0:.2f}",
527                        }
528                    )
529                else:
530                    logger.info(
531                        f"Epoch {epoch + 1}/{num_epochs}: "
532                        f"loss={avg_loss:.6f}, val_loss={val_loss:.6f}, "
533                        f"L0={avg_l0:.2f}, val_L0={val_l0:.2f}"
534                    )
535
536        if pbar:
537            pbar.close()
538
539        # Finalize wandb and log summary metrics
540        if wandb_logger:
541            # Log final feature densities
542            densities = self.get_feature_density(activations.cpu().numpy())
543            n_dead = (densities == 0).sum()
544            n_active = (densities > 0).sum()
545
546            wandb_logger.log_metrics(
547                {
548                    "final_train_loss": self.training_losses[-1],
549                    "final_val_loss": val_loss,
550                    "final_l0_norm": self.training_l0_norms[-1],
551                    "final_val_l0_norm": val_l0,
552                    "n_dead_features": int(n_dead),
553                    "n_active_features": int(n_active),
554                    "feature_sparsity": float(n_active / len(densities)),
555                }
556            )
557
558            # Log feature density histogram
559            wandb_logger.log_histogram(densities, "feature_density_histogram")
560
561            logger.info(f"wandb run URL: {wandb_logger.get_run_url()}")
562
563            # Finalize wandb (close the run) only if we created it
564            if _owns_wandb:
565                wandb_logger.finalize()
566
567        logger.info("Training complete!")
568        return self
569
570    def get_feature_density(self, activations: np.ndarray) -> np.ndarray:
571        """
572        Compute the density (fraction of non-zero activations) for each feature.
573
574        Args:
575            activations: Data to compute density on
576
577        Returns:
578            Array of feature densities of shape (hidden_dim,)
579        """
580        self.eval()
581        device = next(self.parameters()).device
582        with torch.no_grad():
583            if isinstance(activations, np.ndarray):
584                activations = torch.from_numpy(activations).float().to(device)
585            else:
586                activations = activations.to(device)
587
588            features = self.encode(activations)
589            density = (features > 0).float().mean(dim=0).cpu().numpy()
590
591        return density
592
593    def get_top_activating_examples(
594        self,
595        activations: np.ndarray,
596        feature_idx: int,
597        k: int = 10,
598    ) -> Tuple[np.ndarray, np.ndarray]:
599        """
600        Get the top k examples that activate a given feature the most.
601
602        Args:
603            activations: Data to search through
604            feature_idx: Index of the feature to analyze
605            k: Number of top examples to return
606
607        Returns:
608            Tuple of (top activations values, top indices)
609        """
610        self.eval()
611        device = next(self.parameters()).device
612        with torch.no_grad():
613            if isinstance(activations, np.ndarray):
614                activations = torch.from_numpy(activations).float().to(device)
615            else:
616                activations = activations.to(device)
617
618            features = self.encode(activations)
619            feature_activations = features[:, feature_idx].cpu().numpy()
620
621            top_k_indices = np.argsort(feature_activations)[-k:][::-1]
622            top_k_values = feature_activations[top_k_indices]
623
624        return top_k_values, top_k_indices
625
626    def save(self, filepath: Union[str, Path]) -> None:
627        """Save model state, hyperparameters, and training history to disk.
628
629        Saves a dictionary containing the model architecture parameters,
630        state dict, and training loss/L0 history as a PyTorch ``.pt`` file.
631        The file can be loaded with ``SparseAutoencoder.load()``.
632
633        Args:
634            filepath: Path where the model file will be saved. Parent
635                directories are created automatically if they don't exist.
636        """
637        filepath = Path(filepath)
638        filepath.parent.mkdir(parents=True, exist_ok=True)
639
640        state = {
641            "activation_dim": self.activation_dim,
642            "hidden_dim": self.hidden_dim,
643            "l1_coefficient": self.l1_coefficient,
644            "state_dict": self.state_dict(),
645            "training_losses": self.training_losses,
646            "training_l0_norms": self.training_l0_norms,
647        }
648
649        torch.save(state, filepath)
650        logger.info(f"Saved model to {filepath}")
651
652    @classmethod
653    def load(cls, filepath: Union[str, Path]) -> "SparseAutoencoder":
654        """Load a saved SparseAutoencoder from disk.
655
656        Reconstructs the model from a ``.pt`` file saved by ``save()``,
657        restoring the architecture, learned weights, and training
658        history (losses and L0 norms).
659
660        Args:
661            filepath: Path to the saved ``.pt`` model file.
662
663        Returns:
664            A ``SparseAutoencoder`` instance with restored weights and
665            training metrics.
666        """
667        filepath = Path(filepath)
668        state = torch.load(filepath, weights_only=False)
669
670        model = cls(
671            activation_dim=state["activation_dim"],
672            hidden_dim=state["hidden_dim"],
673            l1_coefficient=state["l1_coefficient"],
674        )
675        model.load_state_dict(state["state_dict"])
676        model.training_losses = state.get("training_losses", [])
677        model.training_l0_norms = state.get("training_l0_norms", [])
678
679        logger.info(f"Loaded model from {filepath}")
680        return model