Revert cross-attention integration into GPTBlock

backbones.py

@@ -220,7 +220,6 @@ 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__()
@@ -232,17 +231,6 @@ 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)
@@ -253,31 +241,8 @@ 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,6 +5,7 @@ import torch.nn.functional as F
 from backbones import (
     AgeSinusoidalEncoding,
     BaselineEncoder,
+    CrossAttention,
     GPTBlock,
     GaussianRBFTimeBasis,
     TimeRoPE,
@@ -55,6 +56,12 @@ 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
@@ -82,8 +89,6 @@ 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)
             ])
@@ -97,8 +102,6 @@ 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)
             ])
@@ -153,8 +156,7 @@ class DeepHealth(nn.Module):
     ) -> torch.Tensor:
         if unused_kwargs:
             unknown = ", ".join(sorted(unused_kwargs))
-            raise TypeError(
+            raise TypeError(f"Unexpected DeepHealth forward arguments: {unknown}")
-                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:
@@ -173,8 +175,7 @@ class DeepHealth(nn.Module):
         if padding_mask is None:
             padding_mask = event_seq > 0
         else:
-            padding_mask = padding_mask.to(
+            padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
-                device=event_seq.device, dtype=torch.bool)
         h_disease = self.token_embedding(event_seq)
         t_disease = time_seq
 
@@ -199,8 +200,7 @@ 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 * \
+            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-                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,6 +210,18 @@ 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,
@@ -221,21 +233,14 @@ 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 = block(
+        h_disease = self.cross_attention(
-                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":