drrik

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