diff --git a/backbones.py b/backbones.py index 48cfc24..edf0dcb 100644 --- a/backbones.py +++ b/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 diff --git a/models.py b/models.py index 094d90c..e74db4b 100644 --- a/models.py +++ b/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( - 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: @@ -173,8 +175,7 @@ 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 @@ -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 * \ - 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,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, - 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 = self.cross_attention( + h_disease=h_disease, + 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) if mode == "all_future":