Improve all-future first-onset training

models.py

@@ -221,6 +221,8 @@ class DeepHealth(nn.Module):
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
         model_architecture: str | None = None,
+        risk_head_bias: bool = False,
+        risk_head_bias_init: torch.Tensor | list[float] | None = None,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -302,7 +304,33 @@ class DeepHealth(nn.Module):
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
 
         self.final_ln = nn.LayerNorm(n_embd)
-        self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
+        self.risk_head = nn.Linear(
+            n_embd,
+            vocab_size,
+            bias=bool(risk_head_bias),
+        )
+        if risk_head_bias_init is not None:
+            if self.risk_head.bias is None:
+                raise ValueError(
+                    "risk_head_bias_init requires risk_head_bias=True"
+                )
+            initial_bias = torch.as_tensor(
+                risk_head_bias_init,
+                dtype=self.risk_head.bias.dtype,
+                device=self.risk_head.bias.device,
+            )
+            if initial_bias.shape != (vocab_size,):
+                raise ValueError(
+                    "risk_head_bias_init must have shape "
+                    f"({vocab_size},), got {tuple(initial_bias.shape)}"
+                )
+            if not torch.isfinite(initial_bias).all():
+                raise ValueError("risk_head_bias_init must contain only finite values")
+            with torch.no_grad():
+                # Start exactly at the fitted marginal baseline; covariate and
+                # history effects are learned away from zero during training.
+                self.risk_head.weight.zero_()
+                self.risk_head.bias.copy_(initial_bias)
         if target_mode == "next_token":
             self.risk_head.weight = self.token_embedding.weight
         self.query_token = nn.Parameter(torch.zeros(n_embd))