JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Revert cross-attention integration into GPTBlock

Browse Source
This commit is contained in:
Jiarui Li
2026-06-15 10:05:13 +08:00
parent c87c3b9f7c
commit 593ecd2e71
2 changed files with 29 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
models.py
View File

@@ -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":