From db0947ce9d6c5f9324c71305282bddf43c63a5df Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 22 Jul 2026 11:52:44 +0800 Subject: [PATCH] Implement TrajMixer block --- TrajMixer_设计方案.md | 330 ++++++++++++++++++++++++++++++++++++++++++ backbones.py | 114 ++++++++++++--- evaluate_auc.py | 9 +- evaluate_auc_v2.py | 9 +- models.py | 31 ++++ test_traj_mixer.py | 96 ++++++++++++ train_all_future.py | 5 +- train_next_step.py | 5 +- 8 files changed, 575 insertions(+), 24 deletions(-) create mode 100644 TrajMixer_设计方案.md create mode 100644 test_traj_mixer.py diff --git a/TrajMixer_设计方案.md b/TrajMixer_设计方案.md new file mode 100644 index 0000000..eb3d10b --- /dev/null +++ b/TrajMixer_设计方案.md @@ -0,0 +1,330 @@ +# TrajMixer Block 最终设计方案 + +> 状态：**Frozen implementation baseline** +> +> 版本：**v1.0** +> +> 固化日期：**2026-07-22** + +本文档是 TrajMixer 后续实现与实验的唯一结构基线。除显式标记为消融项的配置外，所有实现均应遵循本文档；若结构发生变化，应先更新版本和实验记录。 + +## 1. 目标 + +在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的轨迹交互模块替换 FFN。 + +保持不变的组件包括： + +- 原始 causal mask； +- 原始 TimeRoPE / Relative Time Attention Bias； +- 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)； +- 原始序列建模与训练目标。 + +TrajMixer 不修改 Attention，只替换每个 Transformer block 中的 FFN residual branch。 + +## 2. Block 总体结构 + +概念结构： + +```text +PreNorm Causal Multi-Head Attention +→ Residual +→ Standard Mixer PreNorm +→ Group-wise Feature Alignment +→ SwiGLU Cross-Group Mixer +→ Residual +``` + +完整计算为： + +\[ +U = X^{(l)} + \operatorname{Dropout}\!\left( +\operatorname{CausalMHA}\left( +\operatorname{LN}_{\mathrm{attn}}(X^{(l)}), +\text{time information} +\right)\right), +\] + +\[ +N = \operatorname{LN}_{\mathrm{mixer}}(U), +\] + +\[ +\Delta = \operatorname{TrajMixer}(N), +\] + +\[ +X^{(l+1)} = U + \operatorname{Dropout}(\Delta). +\] + +首版中的 \(\operatorname{LN}_{\mathrm{mixer}}\) 是作用于完整 \(d=120\) 维 residual representation 的标准 LayerNorm。 + +## 3. Latent Trajectory Group 定义 + +Attention 输出经过 \(W_O\) 后仍是标准 residual representation： + +\[ +N\in\mathbb{R}^{B\times L\times d},\qquad d=120. +\] + +将 hidden dimension 划分为与 Attention head 数量相同的 group 数量： + +\[ +n_{\mathrm{group}}:=n_{\mathrm{head}}=10, +\qquad d_{\mathrm{group}}=\frac{d}{n_{\mathrm{head}}}=12, +\] + +`n_group` 不再是独立超参数，代码统一使用 `n_head` 确定 residual group 数量。二者只共享数量；这些 residual groups 在语义和张量来源上仍不等同于原始 Attention heads。 + +并 reshape 为： + +\[ +N_{\mathrm{group}}in +\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}. +\] + +这些 group 是 residual space 中的 **latent trajectory groups**，不等同于原始 Attention heads。本文中的 group、trajectory group 均指这一 residual-channel partition。 + +## 4. Group-wise Feature Alignment + +为缓解不同 group 内部坐标不对齐的问题，每个 group 使用独立的小矩阵： + +\[ +B_i\in\mathbb{R}^{d_{\mathrm{group}}\times d_{\mathrm{group}}}, +\qquad i=1,\ldots,n_{\mathrm{group}}. +\] + +对每个 group 内的特征进行可学习对齐： + +\[ +Z_{b,t,i,:}=N_{\mathrm{group},b,t,i,:}B_i. +\] + +因此： + +\[ +Z\in +\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}. +\] + +首版实现约定： + +- \(B_i\) 不带 bias； +- \(B_i\) 使用单位矩阵初始化； +- Alignment 只作用于 Mixer residual branch，不改变 Attention residual stream； +- 首版不增加逆变换或额外的 group 内输出投影。 + +Alignment 每层权重参数量为： + +\[ +n_{\mathrm{group}}d_{\mathrm{group}}^2 +=10\times12^2 +=1{,}440. +\] + +## 5. SwiGLU Cross-Group Mixer + +Mixer 只沿 group 维度交互，不沿序列维度交互，因此不会引入时间递归或未来信息泄漏。 + +对于每个 group 内特征维度： + +\[ +r=1,\ldots,d_{\mathrm{group}}, +\] + +定义： + +\[ +A_g^{(r)},A_v^{(r)} +\in\mathbb{R}^{n_{\mathrm{group}}\times h_{\mathrm{group}}}, +\] + +\[ +A_o^{(r)} +\in\mathbb{R}^{h_{\mathrm{group}}\times n_{\mathrm{group}}}. +\] + +其中默认： + +\[ +h_{\mathrm{group}}=20. +\] + +对固定的 batch、时间位置和内部特征维度 \(r\)，将： + +\[ +Z_{b,t,:,r}\in\mathbb{R}^{n_{\mathrm{group}}} +\] + +视为 row vector，计算： + +\[ +G_{b,t,:,r}=Z_{b,t,:,r}A_g^{(r)}, +\] + +\[ +V_{b,t,:,r}=Z_{b,t,:,r}A_v^{(r)}, +\] + +\[ +M_{b,t,:,r}=\operatorname{SiLU}(G_{b,t,:,r})\odot V_{b,t,:,r}, +\] + +\[ +Y_{b,t,:,r}=M_{b,t,:,r}A_o^{(r)}. +\] + +其中： + +- gate 分支控制信息写入； +- value 分支提供交互内容； +- output matrix 将隐藏 group 表示投影回原始 group 数量； +- 当 \(h_{\mathrm{group}}=n_{\mathrm{group}}=10\) 时，三类矩阵退化为原始的 \(10\times10\) 方阵形式。 + +所有 \(r\) 的输出组合为： + +\[ +Y\in +\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}, +\] + +再 reshape 为： + +\[ +\Delta\in\mathbb{R}^{B\times L\times d}. +\] + +## 6. 参数张量与无歧义索引 + +建议的实现存储形状为： + +```text +group_align: [n_group, d_group, d_group] +gate_proj: [d_group, n_group, hidden_group] +value_proj: [d_group, n_group, hidden_group] +output_proj: [d_group, hidden_group, n_group] +``` + +对应的索引公式为： + +\[ +G_{b,t,q,r} +=\sum_i Z_{b,t,i,r}\,A_{g,r,i,q}, +\] + +\[ +V_{b,t,q,r} +=\sum_i Z_{b,t,i,r}\,A_{v,r,i,q}, +\] + +\[ +Y_{b,t,i,r} +=\sum_q +\left[\operatorname{SiLU}(G_{b,t,q,r})V_{b,t,q,r}\right] +A_{o,r,q,i}. +\] + +首版的三个 Mixer projection 均不带 bias。 + +## 7. Mixer Hidden Width 与参数量 + +Mixer 的 group 维变换为： + +\[ +n_{\mathrm{group}} +\rightarrow +h_{\mathrm{group}} +\rightarrow +n_{\mathrm{group}}. +\] + +默认 \(h_{\mathrm{group}}=20\) 时，Mixer 每层权重参数量为： + +\[ +3d_{\mathrm{group}}n_{\mathrm{group}}h_{\mathrm{group}} +=3\times12\times10\times20 +=7{,}200. +\] + +加上 Group Feature Alignment 后，TrajMixer residual branch 每层共有： + +\[ +7{,}200+1{,}440=8{,}640 +\] + +个主要权重参数。作为对照，原始 \(120\rightarrow480\rightarrow120\) FFN 每层约有 115,800 个参数。 + +参数对照口径说明：上面的 115,800 对应结构方案中的标准两层 FFN。当前代码库在 TrajMixer 替换前实际使用的是隐藏宽度 300 的全维度 SwiGLU（gate/value/output 三个线性层），每层共有 108,720 个参数（含 bias）。代码实验和 checkpoint 参数量比较必须以 108,720 作为历史实现基线，不能与概念方案中的标准 FFN 参数量混用。 + +## 8. LayerNorm 基线与消融 + +为保持与原始 Transformer 的可比性，首版固定使用： + +```text +原始 FFN baseline：FFN + 标准 LayerNorm +TrajMixer baseline：Mixer + 标准 LayerNorm +``` + +以下配置不属于首版主实验，只作为独立消融： + +```text +Mixer + Group-wise LayerNorm +``` + +不得将 Group-wise LayerNorm 的结果直接作为“仅替换 FFN”的对照结果。 + +## 9. 初始化 + +首版初始化约定： + +- Group Alignment \(B_i\)：单位矩阵初始化； +- \(A_g/A_v\)：Xavier uniform 初始化； +- \(A_o\)：均值为 0、标准差为 \(10^{-3}\) 的正态初始化； +- Dropout 概率沿用原始 FFN residual branch 的配置。 + +小初始化的 \(A_o\) 使新增分支在训练初期接近恒等残差更新，同时允许模型逐步学习轨迹交互。 + +## 10. 核心设计思想 + +**Attention**：负责从历史疾病序列中选择并整合相关信息。 + +**Group Feature Alignment**：负责学习不同 latent trajectory groups 的内部特征对齐。 + +**Cross-Group Mixer**：负责不同潜在疾病轨迹之间的非线性门控交互。 + +整个模块保持： + +- 无时间递归； +- 序列维度完全并行； +- 参数量远低于原始 FFN； +- 保留 Transformer 的因果历史建模能力； +- 不把 residual groups 误解释为原始 Attention heads。 + +## 11. 首版固定配置 + +```yaml +model_architecture: traj_mixer_v1 +d_model: 120 +n_head: 10 # 同时决定 residual group 数量 +d_group: 12 +hidden_group: 20 +attention: unchanged +attention_output_projection: unchanged +mixer_norm: standard_layer_norm +group_alignment: per_group_12x12 +group_alignment_bias: false +group_alignment_init: identity +mixer_bias: false +gate_value_init: xavier_uniform +output_init_std: 0.001 +group_wise_layer_norm: false +``` + +训练时必须将 `model_architecture: traj_mixer_v1` 写入 `train_config.json`。本分支的评估和导出入口只接受带有该标识、且 checkpoint 中包含 TrajMixer 参数张量的模型；其他分支生成的模型应直接拒绝加载。 + +必须满足： + +\[ +d=n_{\mathrm{group}}d_{\mathrm{group}}. +\] + +后续实现、单元测试、参数量核验和主实验均以以上配置为默认基线。 diff --git a/backbones.py b/backbones.py index ea61077..6d6e6ab 100644 --- a/backbones.py +++ b/backbones.py @@ -176,38 +176,106 @@ class TemporalAttention(nn.Module): return self.resid_drop(self.out_proj(out)) -class SwiGLU(nn.Module): +class TrajMixer(nn.Module): + """Lightweight gated interaction across latent residual-space groups. + + The groups are contiguous partitions of the post-``W_O`` residual + representation. They are deliberately not treated as attention heads. + All operations are position-wise, so the sequence dimension remains fully + parallel and no temporal information can leak between positions here. + """ + def __init__( self, n_embd: int, - hidden_dim: int | None = None, + n_head: int = 10, + hidden_group: int = 20, 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) + if n_embd <= 0: + raise ValueError(f"n_embd must be > 0, got {n_embd}") + if n_head <= 0: + raise ValueError(f"n_head must be > 0, got {n_head}") + if n_embd % n_head != 0: + raise ValueError( + f"n_embd must be divisible by n_head, got {n_embd} and {n_head}" + ) + if hidden_group <= 0: + raise ValueError( + f"hidden_group must be > 0, got {hidden_group}" + ) - 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.n_embd = n_embd + # The residual-group count is tied to n_head, but the resulting groups + # are still residual-space partitions rather than attention heads. + self.n_group = n_head + self.d_group = n_embd // n_head + self.hidden_group = hidden_group + + # Per-group feature alignment: [group, input feature, output feature]. + self.group_align = nn.Parameter( + torch.empty(self.n_group, self.d_group, self.d_group) + ) + + # Per-feature cross-group projections. The feature index is kept + # independent, exactly as specified by the TrajMixer baseline. + self.gate_proj = nn.Parameter( + torch.empty(self.d_group, self.n_group, hidden_group) + ) + self.value_proj = nn.Parameter( + torch.empty(self.d_group, self.n_group, hidden_group) + ) + self.output_proj = nn.Parameter( + torch.empty(self.d_group, hidden_group, self.n_group) + ) 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) + with torch.no_grad(): + identity = torch.eye( + self.d_group, + dtype=self.group_align.dtype, + device=self.group_align.device, + ) + self.group_align.copy_(identity.unsqueeze(0).expand_as(self.group_align)) + + # Initialise each feature-specific matrix independently so Xavier's + # fan-in/fan-out calculation sees a two-dimensional matrix. + for feature_idx in range(self.d_group): + nn.init.xavier_uniform_(self.gate_proj[feature_idx]) + nn.init.xavier_uniform_(self.value_proj[feature_idx]) + nn.init.normal_(self.output_proj, mean=0.0, std=1e-3) 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))) + """Map ``(B, L, n_embd)`` to an equally shaped residual update.""" + if x.ndim != 3: + raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}") + if x.size(-1) != self.n_embd: + raise ValueError( + f"Expected hidden size {self.n_embd}, got {x.size(-1)}" + ) + + batch_size, seq_len, _ = x.shape + grouped = x.reshape( + batch_size, seq_len, self.n_group, self.d_group + ) + aligned = torch.einsum( + "blgd,gde->blge", grouped, self.group_align + ) + + gate = torch.einsum( + "blgr,rgh->blhr", aligned, self.gate_proj + ) + value = torch.einsum( + "blgr,rgh->blhr", aligned, self.value_proj + ) + hidden = F.silu(gate) * value + mixed = torch.einsum( + "blhr,rhg->blgr", hidden, self.output_proj + ) + return self.drop(mixed.reshape(batch_size, seq_len, self.n_embd)) class GPTBlock(nn.Module): @@ -218,6 +286,7 @@ class GPTBlock(nn.Module): attn_dropout: float = 0.0, mlp_dropout: float = 0.0, + hidden_group: int = 20, use_time_rope: bool = False, use_rbf_bias: bool = False, n_rbf_bases: int = 16, @@ -231,7 +300,12 @@ class GPTBlock(nn.Module): use_time_rope=use_time_rope, use_rbf_bias=use_rbf_bias, ) - self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout) + self.mlp = TrajMixer( + n_embd=n_embd, + n_head=n_head, + hidden_group=hidden_group, + dropout=mlp_dropout, + ) self.ln1 = nn.LayerNorm(n_embd) self.ln2 = nn.LayerNorm(n_embd) diff --git a/evaluate_auc.py b/evaluate_auc.py index 96d05b7..dbc1de7 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -40,7 +40,11 @@ from tqdm.auto import tqdm from dataset import HealthDataset from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn -from models import DeepHealth +from models import ( + DeepHealth, + validate_traj_mixer_config, + validate_traj_mixer_state_dict, +) from readouts import build_readout from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX @@ -309,6 +313,7 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth: + validate_traj_mixer_config(cfg) model_target_mode = str(cfg_get( args, cfg, "model_target_mode", "next_token")).lower() if model_target_mode not in {"next_token", "all_future"}: @@ -331,6 +336,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), + hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)), ) @@ -386,6 +392,7 @@ def load_model_state( state = state_dict if state_dict is not None else load_checkpoint_state_dict( checkpoint_path, map_location=device) + validate_traj_mixer_state_dict(state) model.load_state_dict(state, strict=True) diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 952d5f1..f7abae3 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -29,7 +29,11 @@ from tqdm.auto import tqdm from dataset import HealthDataset from eval_data import load_sequence_eval_dataset -from models import DeepHealth +from models import ( + DeepHealth, + validate_traj_mixer_config, + validate_traj_mixer_state_dict, +) from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX @@ -178,6 +182,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth: + validate_traj_mixer_config(cfg) model_target_mode = str(cfg_get( args, cfg, "model_target_mode", "next_token")).lower() if model_target_mode not in {"next_token", "all_future"}: @@ -200,10 +205,12 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), + hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)), ) def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None: + validate_traj_mixer_state_dict(state_dict) model.load_state_dict(state_dict, strict=True) diff --git a/models.py b/models.py index d8b6e4e..9359ee7 100644 --- a/models.py +++ b/models.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from dataclasses import dataclass import torch @@ -14,6 +15,33 @@ from backbones import ( from targets import PAD_IDX +TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v1" + + +def validate_traj_mixer_config(config: Mapping[str, object]) -> None: + actual = config.get("model_architecture") + if actual != TRAJ_MIXER_ARCHITECTURE: + raise ValueError( + "This branch only accepts models trained with the TrajMixer " + f"architecture marker {TRAJ_MIXER_ARCHITECTURE!r}; got {actual!r}." + ) + + +def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None: + required_keys = { + "blocks.0.mlp.group_align", + "blocks.0.mlp.gate_proj", + "blocks.0.mlp.value_proj", + "blocks.0.mlp.output_proj", + } + missing = sorted(required_keys.difference(state_dict)) + if missing: + raise ValueError( + "Checkpoint is not a TrajMixer checkpoint; missing required " + f"parameters: {', '.join(missing)}" + ) + + @dataclass class DeepHealthOutput: hidden: torch.Tensor @@ -160,6 +188,7 @@ class DeepHealth(nn.Module): dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" extra_pool_reduce: str = "mean", dropout: float = 0.0, + hidden_group: int = 20, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -214,6 +243,7 @@ class DeepHealth(nn.Module): use_time_rope=False, use_rbf_bias=False, mlp_dropout=dropout, + hidden_group=hidden_group, ) for _ in range(n_hist_layer) ]) self.rope = None @@ -227,6 +257,7 @@ class DeepHealth(nn.Module): use_time_rope=True, use_rbf_bias=True, mlp_dropout=dropout, + hidden_group=hidden_group, ) for _ in range(n_hist_layer) ]) self.rope = TimeRoPE(n_embd // n_head) diff --git a/test_traj_mixer.py b/test_traj_mixer.py new file mode 100644 index 0000000..c52ad6a --- /dev/null +++ b/test_traj_mixer.py @@ -0,0 +1,96 @@ +import unittest + +import torch + +from backbones import GPTBlock, TrajMixer +from models import ( + TRAJ_MIXER_ARCHITECTURE, + validate_traj_mixer_config, + validate_traj_mixer_state_dict, +) + + +class TrajMixerTest(unittest.TestCase): + def test_default_shape_parameters_and_initialization(self) -> None: + mixer = TrajMixer( + n_embd=120, + n_head=10, + hidden_group=20, + dropout=0.0, + ) + + x = torch.randn(2, 7, 120) + self.assertEqual(mixer(x).shape, x.shape) + self.assertEqual(sum(p.numel() for p in mixer.parameters()), 8_640) + + expected = torch.eye(12).expand(10, 12, 12) + torch.testing.assert_close(mixer.group_align.detach(), expected) + self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 20)) + self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 20)) + self.assertEqual(tuple(mixer.output_proj.shape), (12, 20, 10)) + + def test_mixer_does_not_mix_sequence_positions(self) -> None: + torch.manual_seed(0) + mixer = TrajMixer(120, n_head=10, hidden_group=20, dropout=0.0) + mixer.eval() + x = torch.randn(2, 5, 120) + changed = x.clone() + changed[:, 3, :] += torch.randn_like(changed[:, 3, :]) + + original_out = mixer(x) + changed_out = mixer(changed) + unchanged_positions = torch.tensor([0, 1, 2, 4]) + torch.testing.assert_close( + original_out.index_select(1, unchanged_positions), + changed_out.index_select(1, unchanged_positions), + ) + + def test_gradients_reach_all_projection_families(self) -> None: + torch.manual_seed(1) + mixer = TrajMixer(120, n_head=10, hidden_group=20, dropout=0.0) + x = torch.randn(2, 4, 120, requires_grad=True) + + mixer(x).square().mean().backward() + + self.assertIsNotNone(x.grad) + for name, parameter in mixer.named_parameters(): + self.assertIsNotNone(parameter.grad, name) + self.assertTrue(torch.isfinite(parameter.grad).all(), name) + + def test_gpt_block_defaults_to_traj_mixer_and_standard_layer_norm(self) -> None: + block = GPTBlock(n_embd=120, n_head=10) + self.assertIsInstance(block.mlp, TrajMixer) + self.assertIsInstance(block.ln2, torch.nn.LayerNorm) + self.assertEqual(tuple(block.ln2.normalized_shape), (120,)) + + x = torch.randn(2, 6, 120) + self.assertEqual(block(x).shape, x.shape) + + def test_architecture_marker_is_required(self) -> None: + validate_traj_mixer_config( + {"model_architecture": TRAJ_MIXER_ARCHITECTURE} + ) + with self.assertRaisesRegex(ValueError, "only accepts models trained"): + validate_traj_mixer_config({}) + with self.assertRaisesRegex(ValueError, "only accepts models trained"): + validate_traj_mixer_config({"model_architecture": "delphi_swiglu"}) + + def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None: + block = GPTBlock(n_embd=120, n_head=10) + state_dict = { + f"blocks.0.{key}": value + for key, value in block.state_dict().items() + } + validate_traj_mixer_state_dict(state_dict) + + state_dict.pop("blocks.0.mlp.group_align") + with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"): + validate_traj_mixer_state_dict(state_dict) + + def test_invalid_group_partition_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "divisible"): + TrajMixer(n_embd=121, n_head=10, hidden_group=20) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_all_future.py b/train_all_future.py index 9e1068f..0fadc99 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -27,7 +27,7 @@ from tqdm.auto import tqdm from dataset import AllFutureHealthDataset, all_future_collate_fn from losses import build_loss -from models import DeepHealth +from models import TRAJ_MIXER_ARCHITECTURE, DeepHealth from targets import CHECKUP_IDX, PAD_IDX from train_util import ( configure_torch_for_training, @@ -89,6 +89,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dist_mode", type=str, default="exponential", choices=["exponential", "weibull", "mixed"]) parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument("--hidden_group", type=int, default=20) parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) @@ -159,6 +160,7 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De time_mode=args.time_mode, dist_mode=args.dist_mode, dropout=args.dropout, + hidden_group=args.hidden_group, ) @@ -298,6 +300,7 @@ def build_metadata( "dataset_class": "AllFutureHealthDataset", "collate_fn": "all_future_collate_fn", "model_class": "DeepHealth", + "model_architecture": TRAJ_MIXER_ARCHITECTURE, "model_target_mode": "all_future", "target_mode": "all_future", "dist_mode": args.dist_mode, diff --git a/train_next_step.py b/train_next_step.py index 35f86ff..22e5f53 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -24,7 +24,7 @@ from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn from losses import build_loss -from models import DeepHealth, DeepHealthOutput +from models import TRAJ_MIXER_ARCHITECTURE, DeepHealth, DeepHealthOutput from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( @@ -83,6 +83,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dropout", type=float, default=0.0) + parser.add_argument("--hidden_group", type=int, default=20) parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) @@ -164,6 +165,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: time_mode=args.time_mode, dist_mode="exponential", dropout=args.dropout, + hidden_group=args.hidden_group, ) @@ -484,6 +486,7 @@ def build_metadata( "dataset_class": "NextStepHealthDataset", "collate_fn": "next_step_collate_fn", "model_class": "DeepHealth", + "model_architecture": TRAJ_MIXER_ARCHITECTURE, "model_target_mode": "next_token", "target_mode": args.target_mode, "dist_mode": "exponential",