Implement TrajMixer block

2026-07-22 11:52:44 +08:00
parent f6bde7e167
commit db0947ce9d
8 changed files with 575 additions and 24 deletions
--- a/TrajMixer_设计方案.md
+++ 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}}.
+\]
+
+后续实现、单元测试、参数量核验和主实验均以以上配置为默认基线。
--- 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)

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


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


--- 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)
--- a/test_traj_mixer.py
+++ 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()
--- 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,
--- 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",