Add target construction and training script for DeepHealth model
- Implemented target construction in `targets.py` for next-token and unique-time set supervision. - Added validation functions and utility methods for target building. - Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging. - Integrated loss functions and readout mechanisms based on target modes. - Established dataset splitting and DataLoader configurations for training, validation, and testing.
This commit is contained in:
630
backbones.py
Normal file
630
backbones.py
Normal file
@@ -0,0 +1,630 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class TimeRoPE(nn.Module):
|
||||
def __init__(self, dim: int, base: float = 10000.0):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0, "RoPE dim must be even"
|
||||
self.dim = dim
|
||||
# inv_freq: (dim // 2,) — not trainable, but should move with device
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def precompute_cache(self, tau: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
t = tau.unsqueeze(-1) # (B, L, 1)
|
||||
angles = t * self.inv_freq # (B, L, dim//2)
|
||||
# Pre-expand for heads and interleave once (avoids N_layers repeats)
|
||||
cos = angles.cos().unsqueeze(1).repeat_interleave(2, dim=-1)
|
||||
sin = angles.sin().unsqueeze(1).repeat_interleave(2, dim=-1)
|
||||
return cos, sin # (B, 1, L, dim)
|
||||
|
||||
@staticmethod
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Rotate pairs: ``[-x2, x1, -x4, x3, ...]``."""
|
||||
x1 = x[..., 0::2]
|
||||
x2 = x[..., 1::2]
|
||||
return torch.stack((-x2, x1), dim=-1).flatten(-2)
|
||||
|
||||
@staticmethod
|
||||
def apply_from_cache(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos, sin = rope_cache # each (B, 1, L, dim)
|
||||
q_rot = q * cos + TimeRoPE._rotate_half(q) * sin
|
||||
k_rot = k * cos + TimeRoPE._rotate_half(k) * sin
|
||||
return q_rot, k_rot
|
||||
|
||||
def forward(
|
||||
self,
|
||||
tau: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cache = self.precompute_cache(tau)
|
||||
return self.apply_from_cache(q, k, cache)
|
||||
|
||||
|
||||
class GaussianRBFTimeBasis(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_bases: int = 16,
|
||||
max_time_diff: float = 40.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.n_bases = n_bases
|
||||
|
||||
# Evenly spaced RBF centres for non-negative linear time differences.
|
||||
# Causal masking enforces query_time >= key_time, so diff is >= 0.
|
||||
centers = torch.linspace(0.0, max_time_diff, n_bases)
|
||||
self.register_buffer("centers", centers,
|
||||
persistent=False) # (n_bases,)
|
||||
|
||||
# Learnable log-widths (initialized to center spacing on linear scale).
|
||||
init_width = max(max_time_diff / max(n_bases - 1, 1), 1e-3)
|
||||
init_log_width = math.log(init_width)
|
||||
self.log_widths = nn.Parameter(torch.full((n_bases,), init_log_width))
|
||||
|
||||
def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
time_coord = tau.float() # (B, L)
|
||||
# Pairwise signed difference: query_i − key_j
|
||||
diff = time_coord.unsqueeze(
|
||||
2) - time_coord.unsqueeze(1) # (B, L_q, L_k)
|
||||
# Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2)
|
||||
diff = diff.unsqueeze(-1) # (B, L, L, 1)
|
||||
widths = self.log_widths.exp() # (n_bases,)
|
||||
rbf_acts = torch.exp(
|
||||
-0.5 * ((diff - self.centers) / widths).square()
|
||||
# (B, L, L, n_bases)
|
||||
)
|
||||
return rbf_acts
|
||||
|
||||
|
||||
class TemporalAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
n_rbf_bases: int = 16,
|
||||
dropout: float = 0.0,
|
||||
use_time_rope: bool = True,
|
||||
use_rbf_bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
||||
self.n_head = n_head
|
||||
self.d_head = n_embd // n_head
|
||||
self.scale = 1.0 / math.sqrt(self.d_head)
|
||||
self.use_time_rope = use_time_rope
|
||||
self.use_rbf_bias = use_rbf_bias
|
||||
|
||||
# QKV projection (fused for efficiency)
|
||||
self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
|
||||
# Output projection
|
||||
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
|
||||
# Layer-specific projection from shared RBF basis activations to per-head attention bias.
|
||||
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
|
||||
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
|
||||
|
||||
self.resid_drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""Match the previous version's GPT-style weight initialization."""
|
||||
nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.rbf_proj.weight)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.use_time_rope:
|
||||
assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
|
||||
if self.use_rbf_bias:
|
||||
assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"
|
||||
|
||||
B, L, _ = x.shape
|
||||
H, D = self.n_head, self.d_head
|
||||
|
||||
# --- QKV ----------------------------------------------------------
|
||||
qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv.unbind(0) # each (B, H, L, D)
|
||||
|
||||
# --- Apply RoPE (from shared cache) --------------------------------
|
||||
if self.use_time_rope:
|
||||
q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
|
||||
|
||||
# Build additive attention bias mask: time bias + causal/padding mask.
|
||||
time_bias = None
|
||||
if self.use_rbf_bias:
|
||||
time_bias = self.rbf_proj(rbf_cache).permute(
|
||||
0, 3, 1, 2) # (B, H, L, L)
|
||||
time_bias = self.time_bias_scale.tanh() * time_bias
|
||||
|
||||
if time_bias is not None and attn_mask is not None:
|
||||
attn_bias = time_bias + attn_mask.to(time_bias.dtype)
|
||||
elif time_bias is not None:
|
||||
attn_bias = time_bias
|
||||
elif attn_mask is not None:
|
||||
attn_bias = attn_mask
|
||||
else:
|
||||
attn_bias = None
|
||||
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_bias,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=self.scale,
|
||||
)
|
||||
|
||||
# --- Aggregate & project out --------------------------------------
|
||||
out = out.transpose(1, 2).reshape(B, L, H * D)
|
||||
return self.resid_drop(self.out_proj(out))
|
||||
|
||||
|
||||
class SwiGLU(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
hidden_dim: int | None = None,
|
||||
dropout: float = 0.0,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
hidden_dim = hidden_dim if hidden_dim is not None else int(
|
||||
n_embd * 2.5)
|
||||
|
||||
self.w1 = nn.Linear(n_embd, hidden_dim, bias=bias) # gate path
|
||||
self.w2 = nn.Linear(n_embd, hidden_dim, bias=bias) # value path
|
||||
# output projection
|
||||
self.w3 = nn.Linear(hidden_dim, n_embd, bias=bias)
|
||||
self.drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
"""GPT-style parameter initialization for MLP paths."""
|
||||
nn.init.normal_(self.w1.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.w2.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.w3.weight, mean=0.0, std=0.02)
|
||||
if self.w1.bias is not None:
|
||||
nn.init.zeros_(self.w1.bias)
|
||||
nn.init.zeros_(self.w2.bias)
|
||||
nn.init.zeros_(self.w3.bias)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""``(B, L, n_embd) -> (B, L, n_embd)``."""
|
||||
return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
|
||||
|
||||
|
||||
class GPTBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
self.attn = TemporalAttention(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
dropout=attn_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
)
|
||||
self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
|
||||
self.ln1 = nn.LayerNorm(n_embd)
|
||||
self.ln2 = nn.LayerNorm(n_embd)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
|
||||
x = x + self.mlp(self.ln2(x))
|
||||
return x
|
||||
|
||||
|
||||
class TokenAutoDiscretization(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_cont_types: int,
|
||||
n_bins: int,
|
||||
n_embd: int,
|
||||
):
|
||||
super().__init__()
|
||||
if n_cont_types <= 0:
|
||||
raise ValueError(f"n_cont_types must be > 0, got {n_cont_types}")
|
||||
if n_bins <= 1:
|
||||
raise ValueError(f"n_bins must be > 1, got {n_bins}")
|
||||
if n_embd <= 0:
|
||||
raise ValueError(f"n_embd must be > 0, got {n_embd}")
|
||||
|
||||
self.n_cont_types = n_cont_types
|
||||
self.n_bins = n_bins
|
||||
self.n_embd = n_embd
|
||||
self.weight = nn.Parameter(torch.empty(n_cont_types, n_bins))
|
||||
self.bias = nn.Parameter(torch.empty(n_cont_types, n_bins))
|
||||
self.bin_emb = nn.Parameter(torch.empty(n_cont_types, n_bins, n_embd))
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.normal_(self.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.bias)
|
||||
nn.init.normal_(self.bin_emb, mean=0.0, std=0.02)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
cont_type_idx: torch.LongTensor, # (N,)
|
||||
value: torch.Tensor, # (N,)
|
||||
) -> torch.Tensor:
|
||||
if cont_type_idx.dim() != 1:
|
||||
raise ValueError(
|
||||
f"cont_type_idx must be 1D, got {tuple(cont_type_idx.shape)}"
|
||||
)
|
||||
if value.dim() != 1:
|
||||
raise ValueError(f"value must be 1D, got {tuple(value.shape)}")
|
||||
if cont_type_idx.numel() != value.numel():
|
||||
raise ValueError("cont_type_idx and value must have the same length")
|
||||
|
||||
w = self.weight[cont_type_idx] # (N, n_bins)
|
||||
b = self.bias[cont_type_idx] # (N, n_bins)
|
||||
e = self.bin_emb[cont_type_idx] # (N, n_bins, D)
|
||||
logits = value.unsqueeze(-1) * w + b
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
return torch.einsum("nb,nbd->nd", probs, e)
|
||||
|
||||
|
||||
class BaselineEncoder(nn.Module):
|
||||
PAD_KIND = 0
|
||||
CONT_KIND = 1
|
||||
CATE_KIND = 2
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
n_types: int,
|
||||
n_cont_types: int,
|
||||
n_categories: int,
|
||||
cont_type_ids: list[int],
|
||||
n_value_kinds: int = 3,
|
||||
n_bins: int = 16,
|
||||
n_tab_layer: int = 2,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
if len(cont_type_ids) != n_cont_types:
|
||||
raise ValueError(
|
||||
"cont_type_ids length must match n_cont_types, got "
|
||||
f"{len(cont_type_ids)} vs {n_cont_types}"
|
||||
)
|
||||
if n_types <= 0:
|
||||
raise ValueError(f"n_types must include PAD and be > 0, got {n_types}")
|
||||
if n_categories <= 0:
|
||||
raise ValueError(
|
||||
f"n_categories must include PAD and be > 0, got {n_categories}"
|
||||
)
|
||||
if n_value_kinds <= self.CATE_KIND:
|
||||
raise ValueError(
|
||||
f"n_value_kinds must be > {self.CATE_KIND}, got {n_value_kinds}"
|
||||
)
|
||||
|
||||
self.n_embd = n_embd
|
||||
self.type_emb = nn.Embedding(n_types, n_embd, padding_idx=0)
|
||||
self.kind_emb = nn.Embedding(n_value_kinds, n_embd, padding_idx=0)
|
||||
self.cont_value_encoder = (
|
||||
TokenAutoDiscretization(
|
||||
n_cont_types=n_cont_types,
|
||||
n_bins=n_bins,
|
||||
n_embd=n_embd,
|
||||
)
|
||||
if n_cont_types > 0
|
||||
else None
|
||||
)
|
||||
self.cate_value_emb = nn.Embedding(
|
||||
n_categories,
|
||||
n_embd,
|
||||
padding_idx=0,
|
||||
)
|
||||
|
||||
cont_type_index = torch.full((n_types,), -1, dtype=torch.long)
|
||||
for idx, type_id in enumerate(cont_type_ids):
|
||||
if type_id <= 0 or type_id >= n_types:
|
||||
raise ValueError(
|
||||
f"continuous type id {type_id} must be in [1, {n_types})"
|
||||
)
|
||||
cont_type_index[type_id] = idx
|
||||
self.register_buffer(
|
||||
"cont_type_index",
|
||||
cont_type_index,
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
self.blocks = nn.ModuleList([
|
||||
GPTBlock(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
use_time_rope=False,
|
||||
use_rbf_bias=False,
|
||||
mlp_dropout=dropout,
|
||||
) for _ in range(n_tab_layer)
|
||||
])
|
||||
self.ln = nn.LayerNorm(n_embd)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.normal_(self.type_emb.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.type_emb.weight[0])
|
||||
nn.init.normal_(self.kind_emb.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.kind_emb.weight[0])
|
||||
nn.init.normal_(self.cate_value_emb.weight, mean=0.0, std=0.02)
|
||||
nn.init.zeros_(self.cate_value_emb.weight[0])
|
||||
|
||||
def _make_attn_mask(self, mask: torch.Tensor, dtype: torch.dtype):
|
||||
return torch.zeros(
|
||||
mask.size(0),
|
||||
1,
|
||||
1,
|
||||
mask.size(1),
|
||||
device=mask.device,
|
||||
dtype=dtype,
|
||||
).masked_fill(~mask[:, None, None, :], -1e4)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
other_type: torch.LongTensor, # (B, K), 0 = padding
|
||||
other_value: torch.Tensor, # (B, K), cate stores global id
|
||||
other_value_kind: torch.LongTensor, # (B, K), 0=PAD, 1=CONT, 2=CATE
|
||||
):
|
||||
if other_type.shape != other_value.shape:
|
||||
raise ValueError(
|
||||
"other_type and other_value must have the same shape, got "
|
||||
f"{tuple(other_type.shape)} vs {tuple(other_value.shape)}"
|
||||
)
|
||||
if other_type.shape != other_value_kind.shape:
|
||||
raise ValueError(
|
||||
"other_type and other_value_kind must have the same shape, got "
|
||||
f"{tuple(other_type.shape)} vs {tuple(other_value_kind.shape)}"
|
||||
)
|
||||
|
||||
other_valid = other_type > 0
|
||||
type_emb = self.type_emb(other_type)
|
||||
kind_emb = self.kind_emb(other_value_kind)
|
||||
value_emb = torch.zeros_like(type_emb)
|
||||
|
||||
cont_pos = other_valid & (other_value_kind == self.CONT_KIND)
|
||||
if cont_pos.any():
|
||||
if self.cont_value_encoder is None:
|
||||
raise ValueError("continuous tokens found but n_cont_types is 0")
|
||||
cont_idx = self.cont_type_index[other_type[cont_pos]]
|
||||
if (cont_idx < 0).any():
|
||||
bad_type = other_type[cont_pos][cont_idx < 0][0].item()
|
||||
raise ValueError(
|
||||
f"type_id={bad_type} is marked continuous but is not in "
|
||||
"cont_type_ids"
|
||||
)
|
||||
value_emb[cont_pos] = self.cont_value_encoder(
|
||||
cont_type_idx=cont_idx,
|
||||
value=other_value[cont_pos].to(type_emb.dtype),
|
||||
)
|
||||
|
||||
cate_pos = other_valid & (other_value_kind == self.CATE_KIND)
|
||||
if cate_pos.any():
|
||||
cate_id = other_value[cate_pos].long()
|
||||
value_emb[cate_pos] = self.cate_value_emb(cate_id)
|
||||
|
||||
f = type_emb + kind_emb + value_emb
|
||||
f = f * other_valid.unsqueeze(-1).to(f.dtype)
|
||||
|
||||
if f.size(1) == 0:
|
||||
return f, other_valid
|
||||
|
||||
attn_mask = self._make_attn_mask(other_valid, f.dtype)
|
||||
for block in self.blocks:
|
||||
f = block(f, attn_mask=attn_mask)
|
||||
f = f * other_valid.unsqueeze(-1).to(f.dtype)
|
||||
|
||||
h = self.ln(f)
|
||||
h = h * other_valid.unsqueeze(-1).to(h.dtype)
|
||||
return h, other_valid
|
||||
|
||||
|
||||
class CrossAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
dropout: float = 0.0,
|
||||
n_rbf_bases: int = 16,
|
||||
max_time_diff: float = 40.0,
|
||||
):
|
||||
super().__init__()
|
||||
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
||||
self.n_head = n_head
|
||||
self.d_head = n_embd // n_head
|
||||
self.scale = 1.0 / math.sqrt(self.d_head)
|
||||
self.dropout = dropout
|
||||
self.mask_value = -1e4
|
||||
|
||||
self.q_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
self.k_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
self.v_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
|
||||
self.time_rope = TimeRoPE(self.d_head)
|
||||
self.rbf_time_basis = GaussianRBFTimeBasis(
|
||||
n_bases=n_rbf_bases,
|
||||
max_time_diff=max_time_diff,
|
||||
)
|
||||
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
|
||||
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
|
||||
self.resid_drop = nn.Dropout(dropout)
|
||||
self.ln = nn.LayerNorm(n_embd)
|
||||
self.out_ln = nn.LayerNorm(n_embd)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
|
||||
nn.init.normal_(self.rbf_proj.weight, mean=0.0, std=0.02)
|
||||
|
||||
def _make_attn_mask(
|
||||
self,
|
||||
token_mask: torch.Tensor, # (B, K), True = valid
|
||||
t_disease: torch.Tensor, # (B, L)
|
||||
t_token: torch.Tensor, # (B, K)
|
||||
dtype: torch.dtype,
|
||||
):
|
||||
valid_token = token_mask[:, None, :] # (B, 1, K)
|
||||
visible_by_time = t_token[:, None, :] <= t_disease[:, :, None]
|
||||
valid = visible_by_time & valid_token # (B, L, K)
|
||||
|
||||
has_any_valid = valid.any(dim=-1, keepdim=True)
|
||||
safe_valid = valid.clone()
|
||||
safe_valid[..., 0:1] = torch.where(
|
||||
has_any_valid,
|
||||
safe_valid[..., 0:1],
|
||||
torch.ones_like(safe_valid[..., 0:1]),
|
||||
)
|
||||
|
||||
attn_mask = torch.zeros(
|
||||
safe_valid.shape,
|
||||
device=safe_valid.device,
|
||||
dtype=dtype,
|
||||
).masked_fill(~safe_valid, self.mask_value)
|
||||
return attn_mask[:, None, :, :], valid
|
||||
|
||||
def _cross_rbf_cache(
|
||||
self,
|
||||
t_disease: torch.Tensor,
|
||||
t_token: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
tau = torch.cat([t_disease, t_token], dim=1)
|
||||
rbf_cache = self.rbf_time_basis.precompute_cache(tau)
|
||||
n_disease = t_disease.size(1)
|
||||
return rbf_cache[:, :n_disease, n_disease:, :]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
h_disease: torch.Tensor, # (B, L, D)
|
||||
t_disease: torch.Tensor, # (B, L)
|
||||
h_token: torch.Tensor, # (B, K, D)
|
||||
t_token: torch.Tensor, # (B, K)
|
||||
token_mask: torch.Tensor, # (B, K), True = valid
|
||||
need_weights: bool = False,
|
||||
):
|
||||
B, L, _ = h_disease.shape
|
||||
K = h_token.size(1)
|
||||
H, Dh = self.n_head, self.d_head
|
||||
if K == 0:
|
||||
empty_weights = h_disease.new_zeros(B, L, 0)
|
||||
if need_weights:
|
||||
return h_disease, empty_weights
|
||||
return h_disease
|
||||
|
||||
attn_mask, valid = self._make_attn_mask(
|
||||
token_mask=token_mask.to(device=h_disease.device, dtype=torch.bool),
|
||||
t_disease=t_disease,
|
||||
t_token=t_token,
|
||||
dtype=h_disease.dtype,
|
||||
)
|
||||
|
||||
q = self.q_proj(h_disease).reshape(B, L, H, Dh).transpose(1, 2)
|
||||
k = self.k_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
|
||||
v = self.v_proj(h_token).reshape(B, K, H, Dh).transpose(1, 2)
|
||||
|
||||
q_cos, q_sin = self.time_rope.precompute_cache(t_disease)
|
||||
k_cos, k_sin = self.time_rope.precompute_cache(t_token)
|
||||
q = q * q_cos + TimeRoPE._rotate_half(q) * q_sin
|
||||
k = k * k_cos + TimeRoPE._rotate_half(k) * k_sin
|
||||
|
||||
rbf_cache = self._cross_rbf_cache(t_disease, t_token) # (B, L, K, R)
|
||||
time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2)
|
||||
time_bias = self.time_bias_scale.tanh() * time_bias
|
||||
|
||||
attn_bias = attn_mask.to(time_bias.dtype) + time_bias
|
||||
attn_out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_bias,
|
||||
dropout_p=self.dropout if self.training else 0.0,
|
||||
is_causal=False,
|
||||
scale=self.scale,
|
||||
)
|
||||
attn_out = attn_out.transpose(1, 2).reshape(B, L, H * Dh)
|
||||
attn_out = self.resid_drop(self.out_proj(attn_out))
|
||||
|
||||
has_any_valid = valid.any(dim=-1) # (B, L)
|
||||
attn_out = attn_out * has_any_valid.unsqueeze(-1).to(attn_out.dtype)
|
||||
|
||||
attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale
|
||||
attn_scores = attn_scores + attn_bias
|
||||
attn_weights = torch.softmax(attn_scores, dim=-1).mean(dim=1)
|
||||
attn_weights = attn_weights * valid.to(attn_weights.dtype)
|
||||
weight_sum = attn_weights.sum(dim=-1, keepdim=True).clamp_min(1e-12)
|
||||
norm_weights = attn_weights / weight_sum
|
||||
norm_weights = norm_weights * has_any_valid.unsqueeze(-1).to(
|
||||
norm_weights.dtype
|
||||
)
|
||||
out = self.out_ln(h_disease + attn_out)
|
||||
out = torch.where(has_any_valid.unsqueeze(-1), out, h_disease)
|
||||
|
||||
if need_weights:
|
||||
return out, norm_weights
|
||||
return out
|
||||
|
||||
|
||||
class AgeSinusoidalEncoding(nn.Module):
|
||||
|
||||
def __init__(self, embedding_dim: int):
|
||||
|
||||
super().__init__()
|
||||
if embedding_dim % 2 != 0:
|
||||
raise ValueError(
|
||||
f"Embedding dimension must be an even number, but got {embedding_dim}")
|
||||
|
||||
self.embedding_dim = embedding_dim
|
||||
|
||||
i = torch.arange(0, self.embedding_dim, 2, dtype=torch.float32)
|
||||
divisor = torch.pow(10000, i / self.embedding_dim)
|
||||
self.register_buffer('divisor', divisor)
|
||||
self.linear = nn.Linear(embedding_dim, embedding_dim, bias=False)
|
||||
|
||||
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
t_years = t
|
||||
# Broadcast (B, L, 1) against (1, 1, D/2) to get (B, L, D/2)
|
||||
args = t_years.unsqueeze(-1) / self.divisor.view(1, 1, -1)
|
||||
# Interleave cos and sin along the last dimension
|
||||
output = torch.zeros(t.shape[0], t.shape[1],
|
||||
self.embedding_dim, device=t.device)
|
||||
output[:, :, 0::2] = torch.cos(args)
|
||||
output[:, :, 1::2] = torch.sin(args)
|
||||
output = self.linear(output)
|
||||
return output
|
||||
Reference in New Issue
Block a user