Add exposure-only input ablation

models.py

@@ -30,6 +30,7 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         dropout: float = 0.0,
         use_exposure_embeddings: bool = False,
+        input_ablation: str = "none",
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -38,12 +39,19 @@ class DeepHealth(nn.Module):
         if dist_mode not in ["exponential", "weibull", "mixed"]:
             raise ValueError(
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
+        if input_ablation not in {"none", "exposure_only"}:
+            raise ValueError("input_ablation must be 'none' or 'exposure_only'")
+        if input_ablation == "exposure_only" and not use_exposure_embeddings:
+            raise ValueError(
+                "input_ablation='exposure_only' requires exposure embeddings"
+            )
         self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
         self.gender_embedding = nn.Embedding(
             2, n_embd)  # Assuming binary gender
         self.target_mode = target_mode
         self.dist_mode = dist_mode
         self.use_exposure_embeddings = bool(use_exposure_embeddings)
+        self.input_ablation = input_ablation
         self.n_embd = n_embd
         self.vocab_size = vocab_size
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
@@ -61,6 +69,10 @@ class DeepHealth(nn.Module):
                 if isinstance(module, nn.Linear):
                     nn.init.normal_(module.weight, mean=0.0, std=0.02)
                     nn.init.zeros_(module.bias)
+            if self.input_ablation == "exposure_only":
+                # The additive gate is bypassed in this ablation. Excluding it
+                # from optimization also avoids an unused parameter under DDP.
+                self.exposure_gate.requires_grad_(False)
         else:
             self.exposure_adapter = None
             self.register_parameter("exposure_gate", None)
@@ -153,7 +165,10 @@ class DeepHealth(nn.Module):
             if self.exposure_adapter is None or self.exposure_gate is None:
                 raise RuntimeError("Exposure adapter is not initialized")
             exposure = self.exposure_adapter(exposure)
-            h_disease = h_disease + torch.sigmoid(self.exposure_gate) * exposure
+            if self.input_ablation == "exposure_only":
+                h_disease = exposure
+            else:
+                h_disease = h_disease + torch.sigmoid(self.exposure_gate) * exposure
         elif exposure_embedding is not None:
             raise ValueError(
                 "exposure_embedding provided but use_exposure_embeddings=False"