From 6dfeb5a6969cc080d11eb1a047ac4f7b32ff03c9 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Tue, 7 Jul 2026 16:40:43 +0800 Subject: [PATCH] Add TimesNet exposure encoder --- backbones.py | 325 +++++++++++++++++++++++++++++++++++++++++++++++++++ models.py | 131 +++++++++++++++++++++ 2 files changed, 456 insertions(+) diff --git a/backbones.py b/backbones.py index ea61077..4a95694 100644 --- a/backbones.py +++ b/backbones.py @@ -326,3 +326,328 @@ class AgeSinusoidalEncoding(nn.Module): output[:, :, 1::2] = torch.sin(args) output = self.linear(output) return output + + +class GRN2d(nn.Module): + """Global Response Normalization from ConvNeXt V2 for NCHW tensors. + + Reference: + Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with + Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808 + """ + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1)) + self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True) + nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps) + return x + self.gamma * (x * nx) + self.beta + + +class LayerNorm2d(nn.Module): + """Channel-wise LayerNorm for NCHW tensors.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.bias = nn.Parameter(torch.zeros(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.permute(0, 2, 3, 1) + x = F.layer_norm(x, (self.weight.numel(),), self.weight, self.bias, self.eps) + return x.permute(0, 3, 1, 2) + + +class ConvNeXtV2Block2d(nn.Module): + """Lightweight ConvNeXt V2-style 2D block for TimesNet period images. + + This is intentionally a block, not the full ConvNeXt V2 image backbone: + TimesNet's 2D tensors are reshaped time-series period maps rather than + natural images, so aggressive visual downsampling would destroy axis + semantics. + """ + + def __init__( + self, + dim: int, + kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + ): + super().__init__() + padding = kernel_size // 2 + hidden_dim = int(dim * mlp_ratio) + self.dwconv = nn.Conv2d( + dim, + dim, + kernel_size=kernel_size, + padding=padding, + groups=dim, + ) + self.norm = LayerNorm2d(dim) + self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1) + self.act = nn.GELU() + self.grn = GRN2d(hidden_dim) + self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1) + self.drop = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.dwconv.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.dwconv.bias) + nn.init.normal_(self.pwconv1.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.pwconv1.bias) + nn.init.normal_(self.pwconv2.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.pwconv2.bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x = self.dwconv(x) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.grn(x) + x = self.pwconv2(x) + x = self.drop(x) + return residual + x + + +class TimesNetBlock(nn.Module): + """TimesNet block with ConvNeXt V2-style 2D extraction. + + The block follows TimesNet's idea: discover dominant periods with FFT, + reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional + extractor, then fuse the top-k period branches. + + Reference: + Wu et al., "TimesNet: Temporal 2D-Variation Modeling for General Time + Series Analysis", ICLR 2023. https://arxiv.org/abs/2210.02186 + """ + + def __init__( + self, + d_model: int, + top_k: int = 3, + n_convnext_blocks: int = 2, + conv_kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + ): + super().__init__() + if top_k <= 0: + raise ValueError(f"top_k must be > 0, got {top_k}") + self.top_k = top_k + self.norm = nn.LayerNorm(d_model) + self.extractor = nn.Sequential(*[ + ConvNeXtV2Block2d( + dim=d_model, + kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + ) + for _ in range(n_convnext_blocks) + ]) + + def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]: + B, T, C = x.shape + spectrum = torch.fft.rfft(x.float(), dim=1) + amplitude = spectrum.abs().mean(dim=(0, 2)) + if amplitude.numel() <= 1: + return [max(T, 1)], x.new_ones(1) + + amplitude = amplitude.clone() + amplitude[0] = 0.0 + k = min(self.top_k, amplitude.numel() - 1) + weights, indices = torch.topk(amplitude, k=k) + periods = [max(1, T // int(idx.item())) for idx in indices] + return periods, weights.to(dtype=x.dtype, device=x.device) + + def _period_branch(self, x: torch.Tensor, period: int) -> torch.Tensor: + B, T, C = x.shape + if T % period != 0: + padded_len = ((T // period) + 1) * period + pad = x.new_zeros(B, padded_len - T, C) + x_pad = torch.cat([x, pad], dim=1) + else: + padded_len = T + x_pad = x + + n_periods = padded_len // period + x_2d = x_pad.reshape(B, n_periods, period, C).permute(0, 3, 1, 2) + y = self.extractor(x_2d) + y = y.permute(0, 2, 3, 1).reshape(B, padded_len, C) + return y[:, :T, :] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + residual = x + x_norm = self.norm(x) + periods, weights = self._select_periods(x_norm) + branches = torch.stack( + [self._period_branch(x_norm, period) for period in periods], + dim=-1, + ) + weights = torch.softmax(weights, dim=0).view(1, 1, 1, -1) + return residual + (branches * weights).sum(dim=-1) + + +class TimesNetEncoder(nn.Module): + """Encode a multivariate time series into one fixed-size embedding.""" + + def __init__( + self, + input_dim: int, + d_model: int, + n_layers: int = 2, + top_k: int = 3, + n_convnext_blocks: int = 2, + conv_kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + append_observed_mask: bool = True, + ): + super().__init__() + self.input_dim = input_dim + self.append_observed_mask = append_observed_mask + in_dim = input_dim * 2 if append_observed_mask else input_dim + self.input_proj = nn.Linear(in_dim, d_model) + self.blocks = nn.ModuleList([ + TimesNetBlock( + d_model=d_model, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + ) + for _ in range(n_layers) + ]) + self.final_ln = nn.LayerNorm(d_model) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.input_proj.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.input_proj.bias) + + def forward( + self, + x: torch.Tensor, + observed_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if x.dim() != 3: + raise ValueError(f"x must have shape (B, T, C), got {tuple(x.shape)}") + if x.size(-1) != self.input_dim: + raise ValueError( + f"last dim must be input_dim={self.input_dim}, got {x.size(-1)}" + ) + + finite_mask = torch.isfinite(x) + x = torch.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + if observed_mask is None: + observed_mask = finite_mask + elif observed_mask.shape == x.shape[:2]: + observed_mask = observed_mask.unsqueeze(-1).expand_as(x) + elif observed_mask.shape != x.shape: + raise ValueError( + "observed_mask must have shape (B, T) or (B, T, C), got " + f"{tuple(observed_mask.shape)}" + ) + + observed_mask = observed_mask.to(device=x.device, dtype=x.dtype) + x = x * observed_mask + if self.append_observed_mask: + x = torch.cat([x, observed_mask], dim=-1) + + h = self.input_proj(x) + for block in self.blocks: + h = block(h) + h = self.final_ln(h) + + valid_time = observed_mask.amax(dim=-1) + pooled = (h * valid_time.unsqueeze(-1)).sum(dim=1) + denom = valid_time.sum(dim=1, keepdim=True).clamp_min(1.0) + return pooled / denom + + +class TimesNetExposureEncoder(nn.Module): + """Encode pre-onset environmental exposure into an event-level embedding. + + Expected inputs: + daily: (B, 1826, 4) for tmean, tmax, tmin, rhmean + monthly: (B, 241, 2) for tmean, rhmean + + Output: + (B, n_embd), suitable for adding to a disease event embedding in the + route-2 single-stream event-enhancement setup. + """ + + def __init__( + self, + n_embd: int, + daily_input_dim: int = 4, + monthly_input_dim: int = 2, + d_model: int | None = None, + n_layers: int = 2, + top_k: int = 3, + n_convnext_blocks: int = 2, + conv_kernel_size: int = 7, + mlp_ratio: float = 4.0, + dropout: float = 0.0, + use_gate: bool = True, + ): + super().__init__() + d_model = n_embd if d_model is None else d_model + self.daily_encoder = TimesNetEncoder( + input_dim=daily_input_dim, + d_model=d_model, + n_layers=n_layers, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + append_observed_mask=True, + ) + self.monthly_encoder = TimesNetEncoder( + input_dim=monthly_input_dim, + d_model=d_model, + n_layers=n_layers, + top_k=top_k, + n_convnext_blocks=n_convnext_blocks, + conv_kernel_size=conv_kernel_size, + mlp_ratio=mlp_ratio, + dropout=dropout, + append_observed_mask=True, + ) + self.out_proj = nn.Sequential( + nn.LayerNorm(2 * d_model), + nn.Linear(2 * d_model, n_embd), + nn.GELU(), + nn.Linear(n_embd, n_embd), + ) + self.gate = nn.Parameter(torch.tensor(-2.0)) if use_gate else None + self.reset_parameters() + + def reset_parameters(self) -> None: + for module in self.out_proj: + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=0.02) + nn.init.zeros_(module.bias) + + def forward( + self, + daily: torch.Tensor, + monthly: torch.Tensor, + daily_mask: torch.Tensor | None = None, + monthly_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + h_daily = self.daily_encoder(daily, observed_mask=daily_mask) + h_monthly = self.monthly_encoder(monthly, observed_mask=monthly_mask) + h = self.out_proj(torch.cat([h_daily, h_monthly], dim=-1)) + if self.gate is not None: + h = torch.sigmoid(self.gate) * h + return h diff --git a/models.py b/models.py index d8b6e4e..04e59af 100644 --- a/models.py +++ b/models.py @@ -8,6 +8,7 @@ from backbones import ( AgeSinusoidalEncoding, GPTBlock, GaussianRBFTimeBasis, + TimesNetExposureEncoder, TimeRoPE, TokenAutoDiscretization, ) @@ -160,6 +161,16 @@ class DeepHealth(nn.Module): dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" extra_pool_reduce: str = "mean", dropout: float = 0.0, + use_exposure_encoder: bool = False, + exposure_daily_input_dim: int = 4, + exposure_monthly_input_dim: int = 2, + exposure_d_model: int | None = None, + exposure_n_layers: int = 2, + exposure_top_k: int = 3, + exposure_n_convnext_blocks: int = 2, + exposure_conv_kernel_size: int = 7, + exposure_mlp_ratio: float = 4.0, + exposure_use_gate: bool = True, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -189,8 +200,26 @@ class DeepHealth(nn.Module): self.time_mode = time_mode self.dist_mode = dist_mode self.extra_pool_reduce = extra_pool_reduce + self.use_exposure_encoder = use_exposure_encoder self.n_embd = n_embd self.vocab_size = vocab_size + self.exposure_encoder = ( + TimesNetExposureEncoder( + n_embd=n_embd, + daily_input_dim=exposure_daily_input_dim, + monthly_input_dim=exposure_monthly_input_dim, + d_model=exposure_d_model, + n_layers=exposure_n_layers, + top_k=exposure_top_k, + n_convnext_blocks=exposure_n_convnext_blocks, + conv_kernel_size=exposure_conv_kernel_size, + mlp_ratio=exposure_mlp_ratio, + dropout=dropout, + use_gate=exposure_use_gate, + ) + if use_exposure_encoder + else None + ) nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02) nn.init.zeros_(self.token_embedding.weight[0]) nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02) @@ -317,6 +346,94 @@ class DeepHealth(nn.Module): pooled_mask = arange_groups.unsqueeze(0) < group_count.unsqueeze(1) return pooled_h, pooled_time, pooled_mask + def _encode_event_exposure( + self, + exposure_daily: torch.Tensor | None, + exposure_monthly: torch.Tensor | None, + exposure_daily_mask: torch.Tensor | None, + exposure_monthly_mask: torch.Tensor | None, + event_shape: tuple[int, int], + ) -> torch.Tensor | None: + if self.exposure_encoder is None: + if exposure_daily is not None or exposure_monthly is not None: + raise ValueError( + "Exposure tensors were provided but use_exposure_encoder=False" + ) + return None + if exposure_daily is None or exposure_monthly is None: + raise ValueError( + "exposure_daily and exposure_monthly are required when " + "use_exposure_encoder=True" + ) + + batch_size, event_len = event_shape + if exposure_daily.shape[:2] != event_shape: + raise ValueError( + "exposure_daily must have shape (B, L, T_daily, C_daily), got " + f"{tuple(exposure_daily.shape)} for event shape {event_shape}" + ) + if exposure_monthly.shape[:2] != event_shape: + raise ValueError( + "exposure_monthly must have shape (B, L, T_monthly, C_monthly), got " + f"{tuple(exposure_monthly.shape)} for event shape {event_shape}" + ) + if exposure_daily.dim() != 4 or exposure_monthly.dim() != 4: + raise ValueError( + "exposure_daily and exposure_monthly must both be 4D tensors" + ) + + def flatten_mask(mask: torch.Tensor | None, ref: torch.Tensor) -> torch.Tensor | None: + if mask is None: + return None + if mask.shape[:2] != event_shape: + raise ValueError( + "exposure mask must start with shape (B, L), got " + f"{tuple(mask.shape)}" + ) + if mask.dim() not in {3, 4}: + raise ValueError( + "exposure mask must have shape (B, L, T) or (B, L, T, C), " + f"got {tuple(mask.shape)}" + ) + if mask.shape[2] != ref.shape[2]: + raise ValueError( + "exposure mask time dimension does not match exposure tensor" + ) + if mask.dim() == 4 and mask.shape[3] != ref.shape[3]: + raise ValueError( + "exposure mask channel dimension does not match exposure tensor" + ) + return mask.reshape(batch_size * event_len, *mask.shape[2:]) + + daily = exposure_daily.reshape( + batch_size * event_len, + exposure_daily.size(2), + exposure_daily.size(3), + ) + monthly = exposure_monthly.reshape( + batch_size * event_len, + exposure_monthly.size(2), + exposure_monthly.size(3), + ) + daily_mask = flatten_mask(exposure_daily_mask, exposure_daily) + monthly_mask = flatten_mask(exposure_monthly_mask, exposure_monthly) + + param = next(self.exposure_encoder.parameters()) + daily = daily.to(device=param.device, dtype=param.dtype) + monthly = monthly.to(device=param.device, dtype=param.dtype) + if daily_mask is not None: + daily_mask = daily_mask.to(device=param.device) + if monthly_mask is not None: + monthly_mask = monthly_mask.to(device=param.device) + + exposure_emb = self.exposure_encoder( + daily=daily, + monthly=monthly, + daily_mask=daily_mask, + monthly_mask=monthly_mask, + ) + return exposure_emb.reshape(batch_size, event_len, self.n_embd) + def _forward_shared( self, event_seq: torch.LongTensor, @@ -329,6 +446,10 @@ class DeepHealth(nn.Module): other_value: torch.Tensor | None = None, other_value_kind: torch.LongTensor | None = None, other_time: torch.FloatTensor | None = None, + exposure_daily: torch.Tensor | None = None, + exposure_monthly: torch.Tensor | None = None, + exposure_daily_mask: torch.Tensor | None = None, + exposure_monthly_mask: torch.Tensor | None = None, return_output: bool = False, **unused_kwargs, ) -> torch.Tensor | DeepHealthOutput: @@ -357,6 +478,16 @@ class DeepHealth(nn.Module): event_len = event_seq.size(1) h_disease = self.token_embedding(event_seq) + h_exposure = self._encode_event_exposure( + exposure_daily=exposure_daily, + exposure_monthly=exposure_monthly, + exposure_daily_mask=exposure_daily_mask, + exposure_monthly_mask=exposure_monthly_mask, + event_shape=(event_seq.size(0), event_len), + ) + if h_exposure is not None: + h_exposure = h_exposure.to(device=event_seq.device, dtype=h_disease.dtype) + h_disease = h_disease + h_exposure t_disease = time_seq if other_time.shape != other_type.shape: