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

35
backbones.py
View File

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