Feat: Add Evaluation Metrics, W&B Logging, and Interleaved Evaluation

src/clt_forge/clt.py

@@ -13,7 +13,7 @@
 
 from clt_forge.config import CLTConfig
 from clt_forge.utils import DTYPE_MAP, CLT_WEIGHTS_FILENAME, CLT_CFG_FILENAME
-from clt_forge.training.optim import JumpReLU
+from clt_forge.training.optim import JumpReLU, fused_encoder
 from clt_forge import logger
 
 C_l0_COEF = 4
@@ -97,6 +97,8 @@ def __init__(self, cfg: CLTConfig, rank: int = 0, world_size: int = 1) -> None:
         self.log_threshold = nn.Parameter(
             torch.full((self.N_layers, self.local_d_latent), math.log(cfg.jumprelu_init_threshold), dtype=self.dtype, device=init_device)
         )
+        if cfg.activation_fn != "jumprelu":
+            self.log_threshold.requires_grad = False
         self.bandwidth = cfg.jumprelu_bandwidth
 
         self.register_buffer('feature_count', 
@@ -153,7 +155,7 @@ def _initialize_b_enc(self, hidden_pre: Float[torch.Tensor, "..."]) -> None:
                 for feature in range(self.local_d_latent):
                     feature_pre_acts = hidden_pre[:, layer, feature]  # [B]
                     sorted_acts, _ = torch.sort(feature_pre_acts, descending=True)
-                    target_idx = int(target_activation_rate * B) + 1
+                    target_idx = min(int(target_activation_rate * B) + 1, B - 1)
                     threshold_value = sorted_acts[target_idx]
                     required_bias = thresh[layer, feature] - threshold_value
 
@@ -182,24 +184,40 @@ def encode(
         output: tuple([B, N_layers, local_d_latent], [B, N_layers, local_d_latent]) if layer is None, else [B, local_d_latent]
         """
 
-        if layer is None: 
-            hidden_pre = (torch.einsum( # double check einsum and autocast
-                "bnd,ndk->bnk",
-                x,
-                self.W_enc,
-            ) + self.b_enc)
-
-            thresh = torch.exp(self.log_threshold) #shape [N_layers, d_latent]
-        else: 
-            assert 0 <= layer < self.N_layers, f"Layer {layer} out of range"
-            hidden_pre = F.linear(
-                x,
-                self.W_enc[layer].T,
-                self.b_enc[layer]
-            )            
-            thresh = torch.exp(self.log_threshold[layer]) 
-
-        feat_act = JumpReLU.apply(hidden_pre, thresh, self.bandwidth)
+        if self.cfg.activation_fn in ["topk", "groupmax"]:
+            assert self.cfg.k is not None, f"k must be specified in config for {self.cfg.activation_fn} activation"
+            if layer is None:
+                weight = self.W_enc
+                bias = self.b_enc
+            else:
+                assert 0 <= layer < self.N_layers, f"Layer {layer} out of range"
+                weight = self.W_enc[layer]
+                bias = self.b_enc[layer]
+
+            values, indices, preacts = fused_encoder(x, weight, bias, self.cfg.k, self.cfg.activation_fn)
+            feat_act = torch.zeros_like(preacts)
+            feat_act.scatter_(-1, indices, values)
+            hidden_pre = preacts
+        elif self.cfg.activation_fn == "jumprelu":
+            if layer is None: 
+                hidden_pre = (torch.einsum(
+                    "bnd,ndk->bnk",
+                    x,
+                    self.W_enc,
+                ) + self.b_enc)
+                thresh = torch.exp(self.log_threshold)
+            else: 
+                assert 0 <= layer < self.N_layers, f"Layer {layer} out of range"
+                hidden_pre = F.linear(
+                    x,
+                    self.W_enc[layer].T,
+                    self.b_enc[layer]
+                )            
+                thresh = torch.exp(self.log_threshold[layer])
+
+            feat_act = JumpReLU.apply(hidden_pre, thresh, self.bandwidth)
+        else:
+            raise ValueError(f"Unsupported activation_fn: {self.cfg.activation_fn}")
         return feat_act, hidden_pre
 
     def decode(
@@ -305,36 +323,41 @@ def loss(self, act_in: torch.Tensor, act_out: torch.Tensor, l0_coef: float, df_c
         mse_loss_accross_layers = mse_loss_tensor.sum(dim=-1).mean(dim=0)
         mse_loss = mse_loss_accross_layers.sum()
 
-        if self.cfg.cross_layer_decoders:
-            squared_norms = (self.W_dec.float()**2).sum(dim=2)
-            feature_norms_local = torch.sqrt(torch.matmul(self.layer_mask.float(), squared_norms)) 
-        else: 
-            feature_norms_local = self.W_dec.float().norm(dim=2)
-
-        # Compute L0 loss local
-        weighted_activations = feat_act.float() * feature_norms_local
-        tanh_weighted_activations = torch.tanh(C_l0_COEF * weighted_activations)
-        l0_loss_accross_layers = l0_coef * tanh_weighted_activations.sum(dim=-1).mean(dim=0)
-        l0_loss = l0_loss_accross_layers.sum().float()
-
-        # SUM losses across ranks using autograd-aware all_reduce
-        if self.cfg.is_sharded:
-            l0_loss = all_reduce(l0_loss, op=dist.ReduceOp.SUM)
-            # l0_loss /= self.world_size
-            l0_loss_accross_layers = all_reduce(l0_loss_accross_layers, op=dist.ReduceOp.SUM)
-            # l0_loss_accross_layers /= self.world_size
-
-        if self.cfg.debug: 
-            self.log_loss_debug(feat_act, feature_norms_local, l0_loss)
-
-        ### Dead feature penalty 
-        dead_feature_loss = df_coef * torch.relu(torch.exp(self.log_threshold.float()) - hidden_pre.float()) * feature_norms_local
-        dead_feature_loss = dead_feature_loss.sum(dim=-1).mean(dim=0).sum()
-
-        # SUM losses across ranks using autograd-aware all_reduce
-        if self.cfg.is_sharded:
-            dead_feature_loss = all_reduce(dead_feature_loss, op=dist.ReduceOp.SUM)
-            # dead_feature_loss /= self.world_size
+        if self.cfg.activation_fn in ["topk", "groupmax"]:
+            l0_loss = torch.tensor(0.0, device=act_in.device, dtype=torch.float32)
+            l0_loss_accross_layers = torch.zeros(self.N_layers, device=act_in.device, dtype=torch.float32)
+            dead_feature_loss = torch.tensor(0.0, device=act_in.device, dtype=torch.float32)
+        else:
+            if self.cfg.cross_layer_decoders:
+                squared_norms = (self.W_dec.float()**2).sum(dim=2)
+                feature_norms_local = torch.sqrt(torch.matmul(self.layer_mask.float(), squared_norms)) 
+            else: 
+                feature_norms_local = self.W_dec.float().norm(dim=2)
+
+            # Compute L0 loss local
+            weighted_activations = feat_act.float() * feature_norms_local
+            tanh_weighted_activations = torch.tanh(C_l0_COEF * weighted_activations)
+            l0_loss_accross_layers = l0_coef * tanh_weighted_activations.sum(dim=-1).mean(dim=0)
+            l0_loss = l0_loss_accross_layers.sum().float()
+
+            # SUM losses across ranks using autograd-aware all_reduce
+            if self.cfg.is_sharded:
+                l0_loss = all_reduce(l0_loss, op=dist.ReduceOp.SUM)
+                # l0_loss /= self.world_size
+                l0_loss_accross_layers = all_reduce(l0_loss_accross_layers, op=dist.ReduceOp.SUM)
+                # l0_loss_accross_layers /= self.world_size
+
+            if self.cfg.debug: 
+                self.log_loss_debug(feat_act, feature_norms_local, l0_loss)
+
+            ### Dead feature penalty 
+            dead_feature_loss = df_coef * torch.relu(torch.exp(self.log_threshold.float()) - hidden_pre.float()) * feature_norms_local
+            dead_feature_loss = dead_feature_loss.sum(dim=-1).mean(dim=0).sum()
+
+            # SUM losses across ranks using autograd-aware all_reduce
+            if self.cfg.is_sharded:
+                dead_feature_loss = all_reduce(dead_feature_loss, op=dist.ReduceOp.SUM)
+                # dead_feature_loss /= self.world_size
 
         ### Dead feature count local
         with torch.no_grad():

src/clt_forge/clt_training_runner.py

@@ -164,6 +164,7 @@ def run(self):
         trainer = CLTTrainer(
             clt=self.clt,
             activations_store=self.activations_store,
+            val_activations_store=getattr(self, "val_activations_store", None),
             save_checkpoint_fn=self.save_checkpoint,
             cfg=self.cfg,
             rank=self.rank, 

src/clt_forge/config/clt_config.py

@@ -24,6 +24,8 @@ class CLTConfig(BaseModel):
     cross_layer_decoders: bool
     context_size: int
     functional_loss: Optional[str] = None
+    activation_fn: str = "jumprelu"
+    k: Optional[int] = None
 
     # -----Sparsity---------------------------
     l0_coefficient: float

src/clt_forge/config/clt_training_runner_config.py

@@ -38,6 +38,8 @@ class CLTTrainingRunnerConfig(BaseModel):
     jumprelu_init_threshold: float = 0.03
     jumprelu_bandwidth: float = 1.
     normalize_decoder: bool = False
+    activation_fn: str = "jumprelu"
+    k: Optional[int] = None
 
     # -----ActivationStore Parameters---------
     context_size: int = 32
@@ -81,6 +83,9 @@ class CLTTrainingRunnerConfig(BaseModel):
     # -----Metrics----------------------------
     dead_feature_window: int = 250
     # n_eval_batches: int = 10
+    eval_step_size: Optional[int] = None
+    val_step_size: Optional[int] = None
+    tau_sweep_steps: int = 10
 
     # -----WANDB------------------------------
     log_to_wandb: bool = True

src/clt_forge/training/clt_trainer.py

@@ -30,6 +30,7 @@ def __init__(
         activations_store: ActivationsStore,
         cfg: CLTTrainingRunnerConfig,
         save_checkpoint_fn: Callable[["CLTTrainer", str], None],
+        val_activations_store: Optional[ActivationsStore] = None,
         rank: int = 0,
         world_size: int = 1
     ) -> None:
@@ -38,9 +39,13 @@ def __init__(
         self.is_main_process = rank == 0
         self.clt = clt
         self.activations_store = activations_store
+        self.val_activations_store = val_activations_store
         self.cfg = cfg
         self.save_checkpoint_fn = save_checkpoint_fn
         self.n_training_steps: int = 0
+
+        # Evaluator instantiated dynamically to save memory if not evaluating
+        self.evaluator = None
 
         self.checkpoint_thresholds = []
         if self.cfg.n_checkpoints > 0:
@@ -74,7 +79,7 @@ def __init__(
         print(f"warm up {cfg.l0_warm_up_steps}")
 
         self.optimizer = Adam(
-            self.clt.parameters(),
+            [p for p in self.clt.parameters() if p.requires_grad],
             lr=cfg.lr,
             betas=(
                 cfg.adam_beta1,
@@ -101,7 +106,6 @@ def __init__(
         self.accumulation_step: int = 0
 
     def _initialize_b_enc(self, n_batches: int = 5):
-
         model = self._get_clt()
 
         def get_hidden_pre(acts_in):
@@ -195,6 +199,14 @@ def fit(self):
             if self.accumulation_step == 0: 
                 self._log_train_step(loss_metrics)
                 self._checkpoint_if_needed()
+
+                # Check for evaluation
+                if self.is_main_process and self.cfg.eval_step_size is not None and self.n_training_steps > 0 and self.n_training_steps % self.cfg.eval_step_size == 0:
+                    self._evaluate_and_log_metrics(prefix="train")
+
+                # Check for validation
+                if self.is_main_process and self.cfg.val_step_size is not None and self.n_training_steps > 0 and self.n_training_steps % self.cfg.val_step_size == 0:
+                    self._evaluate_and_log_metrics(prefix="val")
 
             # if self.cfg.functional_loss is not None and self.fc_scheduler.get_lr() > 0 and start_func_finetuning:
             #     self._enable_functional_training()
@@ -222,6 +234,9 @@ def fit(self):
             if self.accumulation_step == 0:
                 self.n_training_steps += 1
 
+        if self.is_main_process and self.cfg.eval_step_size is not None:
+            self._evaluate_and_log_metrics(prefix="train")
+
         self.save_checkpoint_fn(
             trainer=self,
             checkpoint_name=f"final_{self.n_tokens}",
@@ -604,6 +619,81 @@ def _update_layer_metrics_history(
 
         return log_dict
 
+    def _evaluate_and_log_metrics(self, prefix: str = "train"):
+        """Evaluates metrics via CLTEvaluator and logs to W&B."""
+        if not self.cfg.log_to_wandb:
+            return
+
+        import numpy as np
+
+        if self.evaluator is None:
+            from clt_forge.training.eval_metrics import CLTEvaluator
+            self.evaluator = CLTEvaluator(self._get_clt(), self.cfg)
+
+        store = self.val_activations_store if prefix == "val" and self.val_activations_store else self.activations_store
+
+        # Get a batch of tokens
+        # The activations store __iter__ yields: `*tokens, acts_in, acts_out`
+        try:
+            batch = next(store.__iter__())
+            tokens = batch[0] if len(batch) > 2 else None
+        except StopIteration:
+            logger.warning(f"Failed to get evaluation batch for {prefix}.")
+            return
+
+        if tokens is None:
+            logger.warning(f"No tokens available in ActivationsStore for evaluation.")
+            return
+
+        logger.info(f"Running evaluation ({prefix}) at step {self.n_training_steps}...")
+
+        # Evaluate Replacement Score and KL Divergence
+        metrics_kl = self.evaluator.evaluate_replacement_and_kl(tokens)
+
+        # Evaluate Pruning-time Sparsity Sweeps
+        tau_values = np.linspace(0, 1.0, self.cfg.tau_sweep_steps)
+        pruning_table = wandb.Table(columns=["tau", "pruned_l0_norm"])
+
+        base_sparsity_metrics = None
+        for tau in tau_values:
+            sparsity_metrics = self.evaluator.evaluate_sparsity_and_l0(tokens, tau=tau)
+            pruning_table.add_data(tau, sparsity_metrics["pruned_l0_norm"])
+            if tau == 0.0:
+                base_sparsity_metrics = sparsity_metrics
+
+        # Evaluate Activation Density
+        density_metrics = self.evaluator.evaluate_activation_density(tokens)
+
+        # Log basic metrics
+        log_dict = {
+            f"{prefix}_eval/replacement_score": metrics_kl["replacement_score"],
+            f"{prefix}_eval/kl_divergence": metrics_kl["kl_divergence"],
+            f"{prefix}_eval/l0_norm": base_sparsity_metrics["l0_norm"],
+            f"{prefix}_eval/mean_activation_density": density_metrics["mean_activation_density"],
+            f"{prefix}_eval/dead_features_ratio": density_metrics["dead_features_ratio"],
+        }
+
+        # 1. Pareto Frontier
+        pareto_table = wandb.Table(columns=["l0_norm", "replacement_score", "step"])
+        pareto_table.add_data(base_sparsity_metrics["l0_norm"], metrics_kl["replacement_score"], self.n_training_steps)
+        log_dict[f"{prefix}_eval/Pareto_Frontier"] = wandb.plot.scatter(pareto_table, "l0_norm", "replacement_score", title=f"Pareto Frontier ({prefix})")
+
+        # 3. Pruning Sweep
+        log_dict[f"{prefix}_eval/Pruning_Sweep"] = wandb.plot.line(pruning_table, "tau", "pruned_l0_norm", title=f"Pruning-time Sparsity Sweep ({prefix})")
+
+        # 4. Dead Feature Histogram
+        density_vector = density_metrics["activation_density_per_feature"].flatten().cpu().numpy().tolist()
+        log_dict[f"{prefix}_eval/Activation_Density_Hist"] = wandb.Histogram(density_vector)
+
+        # 5. Per-Layer Density Table for W&B Native Custom Charts
+        layer_density_table = wandb.Table(columns=["Layer", "Step", "Mean_Density"])
+        density_per_layer = density_metrics["activation_density_per_feature"].mean(dim=1).cpu().numpy()
+        for l_idx, density in enumerate(density_per_layer):
+            layer_density_table.add_data(l_idx, self.n_training_steps, float(density))
+        log_dict[f"{prefix}_eval/Layer_Density_Table"] = layer_density_table
+
+        wandb.log(log_dict)
+
     # def _enable_functional_training(self):
     #     """Enable functional training by configuring activations store for token return."""