Merge pull request #50 from saprmarks/topk_annealing

Topk annealing

Author: adamkarvonen

dictionary_learning/buffer.py

@@ -56,8 +56,14 @@ def __init__(self,
         self.refresh_batch_size = refresh_batch_size
         self.out_batch_size = out_batch_size
         self.device = device
-        self.remove_bos = remove_bos and (self.model.tokenizer.bos_token_id is not None)
         self.add_special_tokens = add_special_tokens
+        self.remove_bos = remove_bos
+
+        if remove_bos and self.model.tokenizer.bos_token_id is None:
+            print(
+                "\n\n\nWARNING: remove_bos is True but tokenizer does not have a bos token. We are removing the first non-pad token instead. Don't use sequence packing.\n\n\n"
+            )
+            
 
     def __iter__(self):
         return self
@@ -138,9 +144,17 @@ def refresh(self):
             hidden_states = hidden_states.value
             if isinstance(hidden_states, tuple):
                 hidden_states = hidden_states[0]
+                
             if self.remove_bos:
-                bos_mask = (input.value[1]["input_ids"] == self.model.tokenizer.bos_token_id)
-                mask = mask & ~bos_mask
+                if self.model.tokenizer.bos_token_id is not None:
+                    bos_mask = input.value[1]["input_ids"] == self.model.tokenizer.bos_token_id
+                    mask = mask & ~bos_mask
+                else:
+                    # some models (like Qwen) don't have a bos token, so we need to remove the first non-pad token
+                    assert mask.dim() == 2, "expected shape (batch_size, seq_len)"
+                    first_one = (mask.to(t.int64).cumsum(dim=1) == 1) & mask
+                    mask = mask & ~first_one
+
             hidden_states = hidden_states[mask]
 
             remaining_space = self.activation_buffer_size - current_idx

dictionary_learning/pytorch_buffer.py

@@ -119,7 +119,12 @@ def __init__(
         self.device = device
         self.add_special_tokens = add_special_tokens
         self.tokenizer = AutoTokenizer.from_pretrained(model.name_or_path)
-        self.remove_bos = remove_bos and (self.tokenizer.bos_token_id is not None)
+        self.remove_bos = remove_bos
+
+        if remove_bos and self.tokenizer.bos_token_id is None:
+            print(
+                "\n\n\nWARNING: remove_bos is True but tokenizer does not have a bos token. We are removing the first non-pad token instead. Don't use sequence packing.\n\n\n"
+            )
 
         if not self.tokenizer.pad_token:
             self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -192,10 +197,17 @@ def refresh(self):
             with t.no_grad():
                 input = self.tokenized_batch()
                 hidden_states = collect_activations(self.model, self.submodule, input)
-            mask = (input["attention_mask"] != 0)
+            mask = input["attention_mask"] != 0
             if self.remove_bos:
-                bos_mask = (input["input_ids"] == self.tokenizer.bos_token_id)
-                mask = mask & ~bos_mask
+                if self.tokenizer.bos_token_id is not None:
+                    bos_mask = input["input_ids"] == self.tokenizer.bos_token_id
+                    mask = mask & ~bos_mask
+                else:
+                    # some models (like Qwen) don't have a bos token, so we need to remove the first non-pad token
+                    assert mask.dim() == 2, "expected shape (batch_size, seq_len)"
+                    first_one = (mask.to(t.int64).cumsum(dim=1) == 1) & mask
+                    mask = mask & ~first_one
+
             hidden_states = hidden_states[mask]
 
             remaining_space = self.activation_buffer_size - current_idx

dictionary_learning/trainers/batch_top_k.py

@@ -34,11 +34,15 @@ def __init__(self, activation_dim: int, dict_size: int, k: int):
         self.encoder.bias.data.zero_()
         self.b_dec = nn.Parameter(t.zeros(activation_dim))
 
-    def encode(self, x: t.Tensor, return_active: bool = False, use_threshold: bool = True):
+    def encode(
+        self, x: t.Tensor, return_active: bool = False, use_threshold: bool = True
+    ):
         post_relu_feat_acts_BF = nn.functional.relu(self.encoder(x - self.b_dec))
 
         if use_threshold:
-            encoded_acts_BF = post_relu_feat_acts_BF * (post_relu_feat_acts_BF > self.threshold)
+            encoded_acts_BF = post_relu_feat_acts_BF * (
+                post_relu_feat_acts_BF > self.threshold
+            )
         else:
             # Flatten and perform batch top-k
             flattened_acts = post_relu_feat_acts_BF.flatten()
@@ -105,6 +109,7 @@ def __init__(
         decay_start: Optional[int] = None,  # when does the lr decay start
         threshold_beta: float = 0.999,
         threshold_start_step: int = 1000,
+        k_anneal_steps: Optional[int] = None,
         seed: Optional[int] = None,
         device: Optional[str] = None,
         wandb_name: str = "BatchTopKSAE",
@@ -122,6 +127,7 @@ def __init__(
         self.k = k
         self.threshold_beta = threshold_beta
         self.threshold_start_step = threshold_start_step
+        self.k_anneal_steps = k_anneal_steps
 
         if seed is not None:
             t.manual_seed(seed)
@@ -146,17 +152,43 @@ def __init__(
         self.dead_feature_threshold = 10_000_000
         self.top_k_aux = activation_dim // 2  # Heuristic from B.1 of the paper
         self.num_tokens_since_fired = t.zeros(dict_size, dtype=t.long, device=device)
-        self.logging_parameters = ["effective_l0", "dead_features", "pre_norm_auxk_loss"]
+        self.logging_parameters = [
+            "effective_l0",
+            "dead_features",
+            "pre_norm_auxk_loss",
+        ]
         self.effective_l0 = -1
         self.dead_features = -1
         self.pre_norm_auxk_loss = -1
 
-        self.optimizer = t.optim.Adam(self.ae.parameters(), lr=self.lr, betas=(0.9, 0.999))
+        self.optimizer = t.optim.Adam(
+            self.ae.parameters(), lr=self.lr, betas=(0.9, 0.999)
+        )
 
         lr_fn = get_lr_schedule(steps, warmup_steps, decay_start=decay_start)
 
         self.scheduler = t.optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=lr_fn)
 
+    def update_annealed_k(
+        self, step: int, activation_dim: int, k_anneal_steps: Optional[int] = None
+    ) -> None:
+        """Update k buffer in-place with annealed value"""
+        if k_anneal_steps is None:
+            return
+
+        assert 0 <= k_anneal_steps < self.steps, (
+            "k_anneal_steps must be >= 0 and < steps."
+        )
+        # self.k is the target k set for the trainer, not the dictionary's current k
+        assert activation_dim > self.k, "activation_dim must be greater than k"
+
+        step = min(step, k_anneal_steps)
+        ratio = step / k_anneal_steps
+        annealed_value = activation_dim * (1 - ratio) + self.k * ratio
+
+        # Update in-place
+        self.ae.k.fill_(int(annealed_value))
+
     def get_auxiliary_loss(self, residual_BD: t.Tensor, post_relu_acts_BF: t.Tensor):
         dead_features = self.num_tokens_since_fired >= self.dead_feature_threshold
         self.dead_features = int(dead_features.sum())
@@ -170,19 +202,28 @@ def get_auxiliary_loss(self, residual_BD: t.Tensor, post_relu_acts_BF: t.Tensor)
             auxk_acts, auxk_indices = auxk_latents.topk(k_aux, sorted=False)
 
             auxk_buffer_BF = t.zeros_like(post_relu_acts_BF)
-            auxk_acts_BF = auxk_buffer_BF.scatter_(dim=-1, index=auxk_indices, src=auxk_acts)
+            auxk_acts_BF = auxk_buffer_BF.scatter_(
+                dim=-1, index=auxk_indices, src=auxk_acts
+            )
 
             # Note: decoder(), not decode(), as we don't want to apply the bias
             x_reconstruct_aux = self.ae.decoder(auxk_acts_BF)
             l2_loss_aux = (
-                (residual_BD.float() - x_reconstruct_aux.float()).pow(2).sum(dim=-1).mean()
+                (residual_BD.float() - x_reconstruct_aux.float())
+                .pow(2)
+                .sum(dim=-1)
+                .mean()
             )
 
             self.pre_norm_auxk_loss = l2_loss_aux
 
             # normalization from OpenAI implementation: https://github.com/openai/sparse_autoencoder/blob/main/sparse_autoencoder/kernels.py#L614
-            residual_mu = residual_BD.mean(dim=0)[None, :].broadcast_to(residual_BD.shape)
-            loss_denom = (residual_BD.float() - residual_mu.float()).pow(2).sum(dim=-1).mean()
+            residual_mu = residual_BD.mean(dim=0)[None, :].broadcast_to(
+                residual_BD.shape
+            )
+            loss_denom = (
+                (residual_BD.float() - residual_mu.float()).pow(2).sum(dim=-1).mean()
+            )
             normalized_auxk_loss = l2_loss_aux / loss_denom
 
             return normalized_auxk_loss.nan_to_num(0.0)
@@ -220,7 +261,7 @@ def loss(self, x, step=None, logging=False):
 
         e = x - x_hat
 
-        self.effective_l0 = self.k
+        self.effective_l0 = self.ae.k.item()
 
         num_tokens_in_step = x.size(0)
         did_fire = t.zeros_like(self.num_tokens_since_fired, dtype=t.bool)
@@ -239,7 +280,11 @@ def loss(self, x, step=None, logging=False):
                 x,
                 x_hat,
                 f,
-                {"l2_loss": l2_loss.item(), "auxk_loss": auxk_loss.item(), "loss": loss.item()},
+                {
+                    "l2_loss": l2_loss.item(),
+                    "auxk_loss": auxk_loss.item(),
+                    "loss": loss.item(),
+                },
             )
 
     def update(self, step, x):
@@ -263,6 +308,7 @@ def update(self, step, x):
         self.optimizer.step()
         self.optimizer.zero_grad()
         self.scheduler.step()
+        self.update_annealed_k(step, self.ae.activation_dim, self.k_anneal_steps)
 
         # Make sure the decoder is still unit-norm
         self.ae.decoder.weight.data = set_decoder_norm_to_unit_norm(