Fixed bug with BatchTopK SAEs and normalization in loss computation.

Author: jkminder

dictionary_learning/cache.py

@@ -503,15 +503,13 @@ def collect(
         ), "Shuffling shards and storing tokens is not supported yet"
 
         # Check if we need to store sequence ranges
-        has_bos_token = model.tokenizer.bos_token_id is not None
+        has_bos_token = model.tokenizer.bos_token is not None
         store_sequence_ranges = (
             store_tokens and 
             not shuffle_shards and 
             not has_bos_token
         )
-        if store_sequence_ranges:
-            print("No BOS token found. Will store sequence ranges.")
-        
+  
         dataloader = DataLoader(data, batch_size=batch_size, num_workers=num_workers)
 
         activation_cache = [[] for _ in submodules]

dictionary_learning/trainers/batch_top_k.py

@@ -173,7 +173,7 @@ def loss(
         if step > self.threshold_start_step:
             self.update_threshold(f)
 
-        x_hat = self.ae.decode(f, denormalize_activations=normalize_activations)
+        x_hat = self.ae.decode(f, denormalize_activations=False)
 
         e = x - x_hat
 

dictionary_learning/training.py

@@ -11,6 +11,7 @@
 import wandb
 from typing import List, Optional
 
+from .trainers.batch_top_k import BatchTopKTrainer
 from .trainers.crosscoder import CrossCoderTrainer, BatchTopKCrossCoderTrainer
 
 
@@ -300,7 +301,7 @@ def trainSAE(
                     use_threshold=False,
                     epoch_idx_per_step=epoch_idx_per_step,
                 )
-                if isinstance(trainer, BatchTopKCrossCoderTrainer):
+                if isinstance(trainer, BatchTopKCrossCoderTrainer) or isinstance(trainer, BatchTopKTrainer):
                     log_stats(
                         trainer,
                         step,