Add cross-attention support to GPTBlock and update DeepHealth model integration

backbones.py

@@ -220,6 +220,7 @@ class GPTBlock(nn.Module):
         mlp_dropout: float = 0.0,
         use_time_rope: bool = False,
         use_rbf_bias: bool = False,
+        use_cross_attention: bool = False,
         n_rbf_bases: int = 16,
     ):
         super().__init__()
@@ -231,6 +232,17 @@ class GPTBlock(nn.Module):
             use_time_rope=use_time_rope,
             use_rbf_bias=use_rbf_bias,
         )
+        self.use_cross_attention = bool(use_cross_attention)
+        self.cross_attn = (
+            CrossAttention(
+                n_embd=n_embd,
+                n_head=n_head,
+                dropout=attn_dropout,
+                n_rbf_bases=n_rbf_bases,
+            )
+            if self.use_cross_attention
+            else None
+        )
         self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
         self.ln1 = nn.LayerNorm(n_embd)
         self.ln2 = nn.LayerNorm(n_embd)
@@ -241,8 +253,31 @@ class GPTBlock(nn.Module):
         rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
         rbf_cache: torch.Tensor | None = None,
         attn_mask: torch.Tensor | None = None,
+        t_disease: torch.Tensor | None = None,
+        h_token: torch.Tensor | None = None,
+        t_token: torch.Tensor | None = None,
+        token_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
+        if self.use_cross_attention:
+            if (
+                self.cross_attn is None
+                or t_disease is None
+                or h_token is None
+                or t_token is None
+                or token_mask is None
+            ):
+                raise ValueError(
+                    "GPTBlock cross-attention requires t_disease, h_token, "
+                    "t_token, and token_mask."
+                )
+            x = self.cross_attn(
+                h_disease=x,
+                t_disease=t_disease,
+                h_token=h_token,
+                t_token=t_token,
+                token_mask=token_mask,
+            )
         x = x + self.mlp(self.ln2(x))
         return x
 

models.py

@@ -5,7 +5,6 @@ import torch.nn.functional as F
 from backbones import (
     AgeSinusoidalEncoding,
     BaselineEncoder,
-    CrossAttention,
     GPTBlock,
     GaussianRBFTimeBasis,
     TimeRoPE,
@@ -56,12 +55,6 @@ class DeepHealth(nn.Module):
             n_tab_layer=n_tab_layer,
             dropout=dropout,
         )
-        self.cross_attention = CrossAttention(
-            n_embd=n_embd,
-            n_head=n_head,
-            dropout=dropout,
-            n_rbf_bases=16,
-        )
         self.target_mode = target_mode
         self.time_mode = time_mode
         self.dist_mode = dist_mode
@@ -89,6 +82,8 @@ class DeepHealth(nn.Module):
                     n_head=n_head,
                     use_time_rope=False,
                     use_rbf_bias=False,
+                    use_cross_attention=True,
+                    attn_dropout=dropout,
                     mlp_dropout=dropout,
                 ) for _ in range(n_hist_layer)
             ])
@@ -102,6 +97,8 @@ class DeepHealth(nn.Module):
                     n_head=n_head,
                     use_time_rope=True,
                     use_rbf_bias=True,
+                    use_cross_attention=True,
+                    attn_dropout=dropout,
                     mlp_dropout=dropout,
                 ) for _ in range(n_hist_layer)
             ])
@@ -156,7 +153,8 @@ class DeepHealth(nn.Module):
     ) -> torch.Tensor:
         if unused_kwargs:
             unknown = ", ".join(sorted(unused_kwargs))
-            raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
+            raise TypeError(
+                f"Unexpected DeepHealth forward arguments: {unknown}")
         if mode not in {"next_token", "all_future"}:
             raise ValueError("mode must be either 'next_token' or 'all_future'")
         if mode == "all_future" and t_query is None:
@@ -175,7 +173,8 @@ class DeepHealth(nn.Module):
         if padding_mask is None:
             padding_mask = event_seq > 0
         else:
-            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
+            padding_mask = padding_mask.to(
+                device=event_seq.device, dtype=torch.bool)
         h_disease = self.token_embedding(event_seq)
         t_disease = time_seq
 
@@ -200,7 +199,8 @@ class DeepHealth(nn.Module):
         rbf_cache = None
         if self.time_mode == "absolute":
             h_disease = h_disease + self.age_encoding(t_disease)
-            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+            h_disease = h_disease * \
+                padding_mask.unsqueeze(-1).to(h_disease.dtype)
         elif self.time_mode == "relative":
             rope_cache = self.rope.precompute_cache(t_disease)
             rbf_cache = self.rbf.precompute_cache(t_disease)
@@ -210,18 +210,6 @@ class DeepHealth(nn.Module):
             time_seq=t_disease,
             dtype=h_disease.dtype,
         )
-        for block in self.blocks:
-            h_disease = block(
-                h_disease,
-                rope_cache=rope_cache,
-                rbf_cache=rbf_cache,
-                attn_mask=attn_mask,
-            )
-            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-
-        h_disease = self.final_ln(h_disease)
-        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-
         h_token, token_mask = self._encode_other_tokens(
             other_type=other_type,
             other_value=other_value,
@@ -233,14 +221,21 @@ class DeepHealth(nn.Module):
                 f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
             )
         token_time = other_time.to(device=h_token.device, dtype=t_disease.dtype)
-
+        for block in self.blocks:
-        h_disease = self.cross_attention(
+            h_disease = block(
-            h_disease=h_disease,
+                h_disease,
+                rope_cache=rope_cache,
+                rbf_cache=rbf_cache,
+                attn_mask=attn_mask,
                 t_disease=t_disease,
                 h_token=h_token,
                 t_token=token_time,
                 token_mask=token_mask,
             )
+            h_disease = h_disease * \
+                padding_mask.unsqueeze(-1).to(h_disease.dtype)
+
+        h_disease = self.final_ln(h_disease)
         h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
 
         if mode == "all_future":