feat: Add top-k and groupmax activation functions with bias initialization

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/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

src/clt_forge/training/clt_trainer.py

@@ -74,7 +74,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 +101,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):

src/clt_forge/training/optim.py

@@ -1,6 +1,7 @@
 import math
 import torch 
-from typing import Any
+import torch.nn.functional as F
+from typing import Any, Literal
 
 class LearningRateScheduler:
     def __init__(
@@ -135,3 +136,112 @@ def backward(  # type: ignore[override]
             dim=0,
         )
         return x_grad, threshold_grad, None
+
+
+class FusedEncoder(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx, input, weight, bias, k: int, activation: Literal["groupmax", "topk"]
+    ):
+        """
+        input:  (B, L, D)
+        weight: (L, D, M)
+        bias:   (L, M)
+        k:      int (number of top elements to select along dim=-1)
+        """
+        preacts = torch.einsum("bld,ldm->blm", input, weight)
+        if bias is not None:
+            preacts = preacts + bias.unsqueeze(0)
+        preacts = F.relu(preacts)
+
+        # Get top-k values and indices for each row
+        if activation == "topk":
+            values, indices = torch.topk(preacts, k, dim=-1, sorted=False)
+        elif activation == "groupmax":
+            values, indices = preacts.unflatten(-1, (k, -1)).max(dim=-1)
+
+            # torch.max gives us indices into each group, but we want indices into the
+            # flattened tensor. Add the offsets to get the correct indices.
+            num_latents = preacts.shape[-1]
+            offsets = torch.arange(
+                0, num_latents, num_latents // k, device=preacts.device
+            )
+            indices = offsets + indices
+        else:
+            raise ValueError(f"Unknown activation: {activation}")
+
+        ctx.save_for_backward(input, weight, bias, indices)
+        ctx.k = k
+        ctx.activation = activation
+        return values, indices, preacts
+
+    @staticmethod
+    def backward(ctx, grad_values, grad_indices, grad_preacts):
+        input, weight, bias, indices = ctx.saved_tensors
+
+        B, L, D = input.shape
+        _, _, M = weight.shape
+
+        grad_input = grad_weight = grad_bias = None
+
+        # --- Grad w.r.t. input ---
+        if ctx.needs_input_grad[0]:
+            grad_input = torch.zeros_like(input)
+            for l in range(L):
+                embedding_weight = weight[l].T  # Shape: (M, D)
+                grad_input[:, l, :] = F.embedding_bag(
+                    indices[:, l, :],
+                    embedding_weight,
+                    mode="sum",
+                    per_sample_weights=grad_values[:, l, :].type_as(embedding_weight),
+                )
+
+        # --- Grad w.r.t. weight ---
+        if ctx.needs_input_grad[1]:
+            grad_weight = torch.zeros_like(weight)
+            for l in range(L):
+                # Compute contributions from each top-k element for layer l
+                contributions_l = grad_values[:, l, :].unsqueeze(2) * input[:, l, :].unsqueeze(1) # Shape: (B, k, D)
+                contributions_l = contributions_l.reshape(-1, D) # Shape: (B*k, D)
+
+                grad_weight_l_T = torch.zeros(M, D, device=weight.device, dtype=weight.dtype)
+                grad_weight_l_T.index_add_(0, indices[:, l, :].flatten(), contributions_l.type_as(weight))
+                grad_weight[l] = grad_weight_l_T.T
+
+        # --- Grad w.r.t. bias ---
+        if bias is not None and ctx.needs_input_grad[2]:
+            grad_bias = torch.zeros_like(bias)
+            for l in range(L):
+                grad_bias[l].index_add_(
+                    0, indices[:, l, :].flatten(), grad_values[:, l, :].flatten().type_as(bias)
+                )
+
+        return grad_input, grad_weight, grad_bias, None, None
+
+
+def fused_encoder(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    k: int,
+    activation: Literal["groupmax", "topk"],
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Convenience wrapper that performs a multi-layer projection followed by `activation`
+    with a backward pass optimized using index_add and embedding_bag.
+    """
+    is_2d = (input.dim() == 2)
+    if is_2d:
+        input = input.unsqueeze(1)    # (B, 1, D)
+        weight = weight.unsqueeze(0)  # (1, D, M)
+        if bias is not None:
+            bias = bias.unsqueeze(0)  # (1, M)
+
+    values, indices, preacts = FusedEncoder.apply(input, weight, bias, k, activation)
+
+    if is_2d:
+        values = values.squeeze(1)
+        indices = indices.squeeze(1)
+        preacts = preacts.squeeze(1)
+
+    return values, indices, preacts