Report AUCs in Delphi2M format

This reverts commit 06f29c0f0a.

MODEL_ARCHITECTURES.md

@@ -1,64 +0,0 @@
-# Model architectures
-
-DeepHealth uses one codebase for both supported history-block architectures.
-Select the architecture explicitly when starting a training run:
-
-| `model_architecture` | History block | Checkpoint fingerprint |
-| --- | --- | --- |
-| `transformer_ffn_v1` | Temporal attention + SwiGLU FFN | `blocks.*.mlp.w1/w2/w3` and `blocks.*.ln2` |
-| `traj_mixer_v5` | Temporal attention + TrajMixer | `blocks.*.mlp.intra_*`, `gate_proj`, and `output_proj` |
-
-`transformer_ffn_v1` is the CLI default; pass `traj_mixer_v5` explicitly for
-TrajMixer runs.
-
-## Training
-
-Next-step example:
-
-```powershell
-python train_next_step.py --model_architecture traj_mixer_v5 --n_layer 12
-```
-
-All-future example:
-
-```powershell
-python train_all_future.py --model_architecture transformer_ffn_v1 --n_layer 12
-```
-
-New runs are separated by architecture:
-
-```text
-runs/
-  transformer_ffn_v1/
-    <run_name>/
-  traj_mixer_v5/
-    <run_name>/
-```
-
-Use `--runs_root` to place this structure under a different root. Existing run
-directories are not moved or renamed.
-
-Each generated `train_config.json` records `model_architecture`, total parameter
-count, and trainable parameter count.
-
-Both training entry points use the single `--n_layer` option to set the number
-of history backbone blocks. The same value is passed to `DeepHealth.n_layer`
-and saved as `n_layer` in `train_config.json`; it must be at least 1.
-
-## Architecture validation
-
-Evaluation resolves the architecture before constructing the model and always
-loads weights with `strict=True`.
-
-- Every config must include an explicit `model_architecture` marker.
-- Checkpoint fingerprints are used to validate that the selected architecture
-  matches the stored weights.
-- A config marker that conflicts with the checkpoint fingerprint raises an
-  error instead of silently choosing one architecture.
-- Unsupported historical TrajMixer markers such as `traj_mixer_v2`,
-  `traj_mixer_v3`, and `traj_mixer_v4` are rejected.
-- Checkpoints and configs created before architecture markers were introduced
-  are intentionally unsupported.
-
-Project code should use the architecture factory rather than instantiate a
-history block directly.

TrajMixer_设计方案.md

@@ -0,0 +1,385 @@
+# TrajMixer Block 最终设计方案
+
+> 状态：**Frozen implementation baseline**
+>
+> 版本：**v3.0 / traj_mixer_v5**
+>
+> 固化日期：**2026-07-24**
+
+本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。
+
+## 1. 目标
+
+在不改变 Delphi Transformer Attention 的前提下，用轻量、完全并行的 TrajMixer 替换 FFN。
+
+保持不变：
+
+- causal mask；
+- TimeRoPE；
+- Relative Time Attention Bias；
+- Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
+- 序列建模和训练目标。
+
+TrajMixer 不沿序列维度混合，也不引入时间递归。
+
+## 2. Block 结构
+
+```text
+PreNorm Causal Multi-Head Attention
+→ Attention Residual
+→ Full-width TrajMixer PreNorm
+→ reshape [B, L, n_group, d_group]
+→ Per-Group SwiGLU: d_group → 4d_group → d_group
+→ Static Gated Fusion
+→ Cross-Group SwiGLU: n_group → 4n_group → n_group
+→ reshape [B, L, n_embd]
+→ Dropout
+→ One TrajMixer Residual
+```
+
+Attention 阶段：
+
+\[
+X
+=X^{(l)}
++\operatorname{CausalMHA}
+\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
+\]
+
+TrajMixer 阶段：
+
+\[
+N=\operatorname{LN}_{d}(X),
+\]
+
+\[
+G=\operatorname{reshape}(N)
+\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
+\]
+
+\[
+P=\operatorname{IntraMixer}(G),
+\]
+
+\[
+U=G+\sigma(\Theta)\odot P,
+\]
+
+\[
+\Delta=\operatorname{reshape}
+\left(\operatorname{CrossMixer}(U)\right),
+\]
+
+\[
+X^{(l+1)}=X+\operatorname{Dropout}(\Delta).
+\]
+
+整个 TrajMixer 只有最后一次 `X + update` 是 residual。`U=G+\sigma(\Theta)\odot P` 是 update 分支内部的静态门控特征融合，不是相对于主 residual stream 的独立 residual stage。
+
+## 3. Group 定义
+
+Attention 输出经过 \(W_O\) 后仍是标准 residual representation：
+
+\[
+X\in\mathbb{R}^{B\times L\times d}.
+\]
+
+定义：
+
+\[
+n_{\mathrm{group}}:=n_{\mathrm{head}},
+\qquad
+d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}},
+\qquad
+d=n_{\mathrm{group}}d_{\mathrm{group}}.
+\]
+
+默认：
+
+\[
+d=120,\qquad
+n_{\mathrm{group}}=10,\qquad
+d_{\mathrm{group}}=12.
+\]
+
+这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。
+
+## 4. 唯一的 Full-Width PreNorm
+
+TrajMixer 只使用一个：
+
+```text
+norm: LayerNorm(n_embd)
+```
+
+LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape：
+
+\[
+G=\operatorname{reshape}
+\left(\operatorname{LN}_{d}(X)\right).
+\]
+
+本版本明确删除：
+
+```text
+intra_norm
+cross_norm
+group_align
+```
+
+不得在组内或跨组阶段再增加额外 LayerNorm。
+
+## 5. 组内 SwiGLU
+
+每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行：
+
+\[
+d_{\mathrm{group}}
+\rightarrow
+4d_{\mathrm{group}}
+\rightarrow
+d_{\mathrm{group}}.
+\]
+
+对 group \(g\)：
+
+\[
+W_{g,\mathrm{intra}}^{(g)},
+W_{v,\mathrm{intra}}^{(g)}
+\in
+\mathbb{R}^{d_{\mathrm{group}}\times4d_{\mathrm{group}}},
+\]
+
+\[
+W_{o,\mathrm{intra}}^{(g)}
+\in
+\mathbb{R}^{4d_{\mathrm{group}}\times d_{\mathrm{group}}}.
+\]
+
+计算：
+
+\[
+H_g
+=
+\operatorname{SiLU}
+\left(G_gW_{g,\mathrm{intra}}^{(g)}\right)
+\odot
+\left(G_gW_{v,\mathrm{intra}}^{(g)}\right),
+\]
+
+\[
+P_g=H_gW_{o,\mathrm{intra}}^{(g)}.
+\]
+
+实现形状：
+
+```text
+intra_gate_proj:   [n_group, d_group, 4 * d_group]
+intra_value_proj:  [n_group, d_group, 4 * d_group]
+intra_output_proj: [n_group, 4 * d_group, d_group]
+```
+
+三个 projection 均不带 bias。
+
+## 6. 静态门控融合
+
+定义可学习 gate logits：
+
+\[
+\Theta\in
+\mathbb{R}^{n_{\mathrm{group}}\times d_{\mathrm{group}}}.
+\]
+
+实际门值为：
+
+\[
+\Gamma=\sigma(\Theta).
+\]
+
+初始化：
+
+\[
+\Theta_{g,r}
+=\operatorname{logit}(0.1)
+=\log\frac{0.1}{0.9}
+\approx-2.1972,
+\]
+
+因此：
+
+\[
+\Gamma_{g,r}\approx0.1.
+\]
+
+融合：
+
+\[
+U=G+\Gamma\odot P.
+\]
+
+\(\Gamma\) 对 batch 和序列位置共享，但每个 group、每个内部坐标拥有独立可学习值。
+
+## 7. 跨 Group SwiGLU
+
+对于每个内部坐标 \(r\)，独立沿 group 维度执行：
+
+\[
+n_{\mathrm{group}}
+\rightarrow
+4n_{\mathrm{group}}
+\rightarrow
+n_{\mathrm{group}}.
+\]
+
+定义：
+
+\[
+A_g^{(r)},A_v^{(r)}
+\in
+\mathbb{R}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}},
+\]
+
+\[
+A_o^{(r)}
+\in
+\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}.
+\]
+
+计算：
+
+\[
+Q_{:,r}
+=
+\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right)
+\odot
+\left(U_{:,r}A_v^{(r)}\right),
+\]
+
+\[
+\Delta_{:,r}=Q_{:,r}A_o^{(r)}.
+\]
+
+实现形状：
+
+```text
+gate_proj:   [d_group, n_group, 4 * n_group]
+value_proj:  [d_group, n_group, 4 * n_group]
+output_proj: [d_group, 4 * n_group, n_group]
+```
+
+三个 projection 均不带 bias。不同内部坐标拥有独立的跨 group 参数，且不沿序列维度交互。
+
+## 8. 唯一的外层 Residual
+
+跨 group 输出 reshape 回：
+
+\[
+\Delta\in\mathbb{R}^{B\times L\times d}.
+\]
+
+最终：
+
+\[
+\operatorname{TrajMixer}(X)
+=X+\operatorname{Dropout}(\Delta).
+\]
+
+固定约束：
+
+- 组内阶段后不执行独立 residual；
+- 跨组阶段后不执行独立 residual；
+- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
+- 整个 TrajMixer 只有一次主 residual。
+
+## 9. 初始化
+
+固定初始化：
+
+- `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform；
+- `intra_output_proj`：每个 group 独立 Xavier uniform；
+- `intra_gate_logits`：初始化为 \(\operatorname{logit}(0.1)\)；
+- 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform；
+- 最终跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布；
+- Full-width LayerNorm：PyTorch 默认 affine 初始化；
+- Dropout：沿用 `mlp_dropout`。
+
+组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。
+
+Relative Time Attention Bias 初始化固定为：
+
+- `rbf_proj.weight`：零初始化；
+- `time_bias_scale`：初始化为 \(1.0\)；
+- 初始 RBF attention bias 严格为零；
+- `rbf_proj.weight` 从第一个优化步骤即可获得梯度。
+
+## 10. 参数量
+
+默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。
+
+Full-width LayerNorm：
+
+\[
+2d=240.
+\]
+
+组内 projections：
+
+\[
+3n_{\mathrm{group}}d_{\mathrm{group}}
+\left(4d_{\mathrm{group}}\right)
+=17{,}280.
+\]
+
+静态门控：
+
+\[
+n_{\mathrm{group}}d_{\mathrm{group}}
+=120.
+\]
+
+跨 group projections：
+
+\[
+3d_{\mathrm{group}}n_{\mathrm{group}}
+\left(4n_{\mathrm{group}}\right)
+=14{,}400.
+\]
+
+每层 TrajMixer 合计：
+
+\[
+240+17{,}280+120+14{,}400
+=\boxed{32{,}040}.
+\]
+
+默认 relative-time、12 层、`vocab_size=1256`、无额外信息类型时，完整模型参数量为：
+
+\[
+\boxed{1{,}232{,}428}.
+\]
+
+## 11. 固定配置与 checkpoint 约束
+
+```yaml
+model_architecture: traj_mixer_v5
+d_model: 120
+n_head: 10
+n_group_rule: n_head
+d_group_rule: d_model / n_group
+traj_mixer_norm: layer_norm_over_n_embd
+intra_hidden_rule: 4 * d_group
+intra_gate_shape: [n_group, d_group]
+intra_gate_initial_sigmoid: 0.1
+cross_hidden_rule: 4 * n_group
+group_alignment: false
+intra_residual: false
+cross_residual: false
+traj_mixer_outer_residual: true
+projection_bias: false
+intra_output_init: xavier_uniform
+cross_output_init_std: 0.001
+```
+
+训练时必须将 `model_architecture: traj_mixer_v5`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在日志中打印参数量。
+
+评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。

backbones.py

@@ -4,12 +4,6 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from model_architectures import (
-    TRAJ_MIXER_ARCHITECTURE,
-    TRANSFORMER_FFN_ARCHITECTURE,
-    resolve_model_architecture,
-)
-
 
 class TimeRoPE(nn.Module):
     def __init__(self, dim: int, base: float = 10000.0):
@@ -185,40 +179,6 @@ class TemporalAttention(nn.Module):
         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 TrajMixer(nn.Module):
     """PreNorm gated mixing within and across latent trajectory groups.
 
@@ -244,13 +204,17 @@ class TrajMixer(nn.Module):
                 f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
             )
         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.intra_hidden = 4 * self.d_group
         self.hidden_group = 4 * n_head
 
+        # A single full-width PreNorm serves the entire TrajMixer branch.
         self.norm = nn.LayerNorm(self.n_embd)
 
+        # Stage 1: each group independently mixes its internal features.
         self.intra_gate_proj = nn.Parameter(
             torch.empty(self.n_group, self.d_group, self.intra_hidden)
         )
@@ -264,6 +228,8 @@ class TrajMixer(nn.Module):
             torch.empty(self.n_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, self.hidden_group)
         )
@@ -286,6 +252,8 @@ class TrajMixer(nn.Module):
             math.log(0.1 / 0.9),
         )
 
+        # 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])
@@ -320,9 +288,7 @@ class TrajMixer(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Apply one full-width PreNorm and one outer residual update."""
         if x.ndim != 3:
-            raise ValueError(
-                f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}"
-            )
+            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)}"
@@ -333,60 +299,26 @@ class TrajMixer(nn.Module):
             batch_size, seq_len, self.n_group, self.d_group
         )
 
+        # The static per-channel gate starts at sigmoid(logit) ~= 0.1.
         intra_output = self._intra_mix(grouped)
         intra_gate = torch.sigmoid(self.intra_gate_logits).view(
             1, 1, self.n_group, self.d_group
         )
         mixed_input = grouped + intra_gate * intra_output
 
+        # Stage 2: n_group -> 4*n_group -> n_group for each coordinate.
         update = self._cross_mix(mixed_input).reshape(
             batch_size, seq_len, self.n_embd
         )
         return x + self.drop(update)
 
 
-class TransformerFFNBlock(nn.Module):
+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 TrajMixerBlock(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,
@@ -420,37 +352,6 @@ class TrajMixerBlock(nn.Module):
         return self.mlp(x)
 
 
-def build_backbone_block(
-    model_architecture: str,
-    *,
-    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,
-) -> nn.Module:
-    """Build one history block for a supported model architecture."""
-    architecture = resolve_model_architecture(model_architecture)
-    block_class: type[nn.Module]
-    if architecture == TRANSFORMER_FFN_ARCHITECTURE:
-        block_class = TransformerFFNBlock
-    elif architecture == TRAJ_MIXER_ARCHITECTURE:
-        block_class = TrajMixerBlock
-    else:  # pragma: no cover - guarded by resolve_model_architecture.
-        raise ValueError(f"Unsupported model architecture: {architecture!r}")
-    return block_class(
-        n_embd=n_embd,
-        n_head=n_head,
-        attn_dropout=attn_dropout,
-        mlp_dropout=mlp_dropout,
-        use_time_rope=use_time_rope,
-        use_rbf_bias=use_rbf_bias,
-        n_rbf_bases=n_rbf_bases,
-    )
-
-
 class TokenAutoDiscretization(nn.Module):
     def __init__(
         self,

evaluate_auc.py

@@ -45,8 +45,11 @@ from delphi2m_auc_report import (
     build_delphi2m_auc_report,
 )
 from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
-from model_architectures import resolve_model_architecture
-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
 
@@ -314,24 +317,20 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
     return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
 
 
-def build_model_from_dataset(
-    args: argparse.Namespace,
-    cfg: Dict[str, Any],
-    dataset: HealthDataset,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
+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"}:
         raise ValueError(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
         n_head=int(cfg_get(args, cfg, "n_head", 10)),
-        n_layer=int(cfg["n_layer"]),
+        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
+        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -342,7 +341,6 @@ def build_model_from_dataset(
         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)),
-        model_architecture=model_architecture,
     )
 
 
@@ -390,7 +388,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 def load_model_state(
-    model: DeepHealth,
+    model: torch.nn.Module,
     checkpoint_path: str,
     device: torch.device,
     state_dict: Optional[Dict[str, Any]] = None,
@@ -398,7 +396,7 @@ def load_model_state(
     state = state_dict if state_dict is not None else load_checkpoint_state_dict(
         checkpoint_path, map_location=device)
 
-    resolve_model_architecture(model.model_architecture, state)
+    validate_traj_mixer_state_dict(state)
     model.load_state_dict(state, strict=True)
 
 
@@ -1379,19 +1377,14 @@ def main() -> None:
     cfg = dict(cfg)
     cfg["dist_mode"] = dist_mode
     cfg["model_target_mode"] = model_target_mode
-    model_architecture = resolve_model_architecture(cfg, state_dict)
-    cfg["model_architecture"] = model_architecture
     print(f"Resolved dist_mode for evaluation: {dist_mode}")
-    print(f"Resolved model architecture: {model_architecture}")
     print(f"Model target mode for AUC: {model_target_mode}")
     print(
         "AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
         "dist_mode affects model loading but is not converted to horizon-specific risk probability."
     )
 
-    model = build_model_from_dataset(
-        args, cfg, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg, dataset).to(device)
     load_model_state(model, str(model_ckpt_path),
                      device, state_dict=state_dict)
     model.eval()

evaluate_auc_v2.py

@@ -36,8 +36,11 @@ from delphi2m_auc_report import (
     build_delphi2m_auc_report,
 )
 from eval_data import load_sequence_eval_dataset
-from model_architectures import resolve_model_architecture
-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
 
@@ -185,24 +188,20 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
     return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
 
 
-def build_model_from_dataset(
-    args: argparse.Namespace,
-    cfg: Dict[str, Any],
-    dataset: HealthDataset,
-    state_dict: Optional[Dict[str, Any]] = None,
-) -> DeepHealth:
+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"}:
         raise ValueError(
             f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
         )
-    model_architecture = resolve_model_architecture(cfg, state_dict)
     return DeepHealth(
         vocab_size=dataset.vocab_size,
         n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
         n_head=int(cfg_get(args, cfg, "n_head", 10)),
-        n_layer=int(cfg["n_layer"]),
+        n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
+        n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -213,12 +212,11 @@ def build_model_from_dataset(
         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)),
-        model_architecture=model_architecture,
     )
 
 
-def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
-    resolve_model_architecture(model.model_architecture, state_dict)
+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)
 
 
@@ -1400,17 +1398,12 @@ def main() -> None:
 
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
-    model_architecture = resolve_model_architecture(cfg_model, state_dict)
-    cfg_model["model_architecture"] = model_architecture
-    print(f"Resolved model architecture: {model_architecture}")
 
     device = resolve_eval_device(args.device)
     if device.type == "cuda":
         torch.backends.cudnn.benchmark = True
 
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
 
     if (
         model_target_mode == "next_token"

evaluate_event_free_survival.py

@@ -649,9 +649,7 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
     device = resolve_eval_device(args.device)
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

evaluate_extra_info_attribution.py

@@ -758,9 +758,7 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
     device = resolve_eval_device(args.device)
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

evaluate_single_disease_mortality_attribution.py

@@ -553,9 +553,7 @@ def main() -> None:
     device = resolve_eval_device(args.device)
     selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
     selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

export_tquery_logits_hidden.py

@@ -180,9 +180,7 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
     device = resolve_eval_device(args.device)
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

export_weibull_death_parameter_stats.py

@@ -381,9 +381,7 @@ def main() -> None:
     cfg_model = dict(cfg)
     cfg_model["dist_mode"] = dist_mode
     device = resolve_eval_device(args.device)
-    model = build_model_from_dataset(
-        args, cfg_model, dataset, state_dict=state_dict
-    ).to(device)
+    model = build_model_from_dataset(args, cfg_model, dataset).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

model_architectures.py

@@ -1,131 +0,0 @@
-"""Model-architecture identifiers and checkpoint validation helpers."""
-
-from __future__ import annotations
-
-import re
-from collections.abc import Mapping
-
-
-TRANSFORMER_FFN_ARCHITECTURE = "transformer_ffn_v1"
-TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
-DEFAULT_MODEL_ARCHITECTURE = TRANSFORMER_FFN_ARCHITECTURE
-SUPPORTED_MODEL_ARCHITECTURES = (
-    TRANSFORMER_FFN_ARCHITECTURE,
-    TRAJ_MIXER_ARCHITECTURE,
-)
-
-
-_FFN_STATE_KEY = re.compile(
-    r"(?:^|\.)blocks\.\d+\.mlp\.w[123]\.(?:weight|bias)$"
-)
-_TRAJ_MIXER_STATE_KEY = re.compile(
-    r"(?:^|\.)blocks\.\d+\.mlp\.(?:"
-    r"norm\.(?:weight|bias)|"
-    r"intra_gate_proj|"
-    r"intra_value_proj|"
-    r"intra_output_proj|"
-    r"intra_gate_logits|"
-    r"gate_proj|"
-    r"value_proj|"
-    r"output_proj"
-    r")$"
-)
-
-
-def _validate_model_architecture(model_architecture: object) -> str:
-    if not isinstance(model_architecture, str):
-        raise ValueError(
-            "model_architecture must be one of "
-            f"{SUPPORTED_MODEL_ARCHITECTURES}, got {model_architecture!r}"
-        )
-    if model_architecture not in SUPPORTED_MODEL_ARCHITECTURES:
-        raise ValueError(
-            f"Unsupported model_architecture={model_architecture!r}; "
-            f"expected one of {SUPPORTED_MODEL_ARCHITECTURES}."
-        )
-    return model_architecture
-
-
-def detect_model_architecture_from_state_dict(
-    state_dict: Mapping[str, object],
-) -> str:
-    """Infer the architecture from block parameter names.
-
-    Detection deliberately accepts any ``blocks.<index>`` prefix rather than
-    assuming that block zero is present.
-    """
-
-    if not isinstance(state_dict, Mapping):
-        raise TypeError(
-            "state_dict must be a mapping, got "
-            f"{type(state_dict).__name__}"
-        )
-
-    has_ffn = False
-    has_traj_mixer = False
-    for raw_key in state_dict:
-        key = str(raw_key)
-        has_ffn = has_ffn or _FFN_STATE_KEY.search(key) is not None
-        has_traj_mixer = (
-            has_traj_mixer
-            or _TRAJ_MIXER_STATE_KEY.search(key) is not None
-        )
-        if has_ffn and has_traj_mixer:
-            raise ValueError(
-                "Checkpoint contains both Transformer FFN and TrajMixer "
-                "block parameters; its model architecture is ambiguous."
-            )
-
-    if has_ffn:
-        return TRANSFORMER_FFN_ARCHITECTURE
-    if has_traj_mixer:
-        return TRAJ_MIXER_ARCHITECTURE
-    raise ValueError(
-        "Could not detect model architecture from checkpoint parameters. "
-        "Expected a blocks.<index>.mlp FFN or TrajMixer parameter."
-    )
-
-
-def resolve_model_architecture(
-    config_or_marker: Mapping[str, object] | str | None = None,
-    state_dict: Mapping[str, object] | None = None,
-) -> str:
-    """Resolve and cross-check a configured and checkpoint architecture.
-
-    Every saved run must provide an explicit architecture marker. Checkpoint
-    parameter names are used only to verify that the marker describes the
-    weights being loaded.
-    """
-
-    if isinstance(config_or_marker, Mapping):
-        configured = config_or_marker.get("model_architecture")
-    elif isinstance(config_or_marker, str) or config_or_marker is None:
-        configured = config_or_marker
-    else:
-        raise TypeError(
-            "config_or_marker must be a config mapping, string, or None, got "
-            f"{type(config_or_marker).__name__}"
-        )
-
-    resolved_config = (
-        _validate_model_architecture(configured)
-        if configured is not None
-        else None
-    )
-    detected = (
-        detect_model_architecture_from_state_dict(state_dict)
-        if state_dict is not None
-        else None
-    )
-
-    if resolved_config is None:
-        raise ValueError(
-            "model_architecture is required; expected one of "
-            f"{SUPPORTED_MODEL_ARCHITECTURES}."
-        )
-    if detected is not None and resolved_config != detected:
-        raise ValueError(
-            "Configured model architecture conflicts with checkpoint: "
-            f"config={resolved_config!r}, checkpoint={detected!r}."
-        )
-    return resolved_config

models.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 import torch
@@ -6,15 +7,46 @@ import torch.nn.functional as F
 
 from backbones import (
     AgeSinusoidalEncoding,
+    GPTBlock,
     GaussianRBFTimeBasis,
     TimeRoPE,
     TokenAutoDiscretization,
-    build_backbone_block,
 )
-from model_architectures import resolve_model_architecture
 from targets import PAD_IDX
 
 
+TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
+
+
+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.norm.weight",
+        "blocks.0.mlp.norm.bias",
+        "blocks.0.mlp.intra_gate_proj",
+        "blocks.0.mlp.intra_value_proj",
+        "blocks.0.mlp.intra_output_proj",
+        "blocks.0.mlp.intra_gate_logits",
+        "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
@@ -148,7 +180,8 @@ class DeepHealth(nn.Module):
         vocab_size: int,
         n_embd: int,
         n_head: int,
-        n_layer: int,
+        n_hist_layer: int,
+        n_tab_layer: int,
         n_types: int,
         n_cont_types: int,
         n_categories: int,
@@ -160,7 +193,6 @@ class DeepHealth(nn.Module):
         dist_mode: str = "exponential",  # "exponential", "weibull" or "mixed"
         extra_pool_reduce: str = "mean",
         dropout: float = 0.0,
-        model_architecture: str | None = None,
     ):
         super().__init__()
         if target_mode not in ["next_token", "all_future"]:
@@ -174,9 +206,6 @@ class DeepHealth(nn.Module):
                 "dist_mode must be either 'exponential', 'weibull' or 'mixed'")
         if extra_pool_reduce not in {"mean", "sum"}:
             raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
-        if n_layer < 1:
-            raise ValueError(f"n_layer must be >= 1, got {n_layer}")
-        model_architecture = resolve_model_architecture(model_architecture)
         self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
         self.gender_embedding = nn.Embedding(
             2, n_embd)  # Assuming binary gender
@@ -193,8 +222,6 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
-        self.model_architecture = model_architecture
-        self.n_layer = n_layer
         self.n_embd = n_embd
         self.vocab_size = vocab_size
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
@@ -214,28 +241,26 @@ class DeepHealth(nn.Module):
         if time_mode == "absolute":
             self.age_encoding = AgeSinusoidalEncoding(n_embd)
             self.blocks = nn.ModuleList([
-                build_backbone_block(
-                    model_architecture,
+                GPTBlock(
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
-                ) for _ in range(n_layer)
+                ) for _ in range(n_hist_layer)
             ])
             self.rope = None
             self.rbf = None
         elif time_mode == "relative":
             self.age_encoding = None
             self.blocks = nn.ModuleList([
-                build_backbone_block(
-                    model_architecture,
+                GPTBlock(
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
-                ) for _ in range(n_layer)
+                ) for _ in range(n_hist_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)

run_missing_evaluations.sh

@@ -1,11 +1,10 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-# Run all non-wrapper evaluation scripts for every completed current-format
-# experiment under runs/. The script is written for Linux servers with bash 4.2.
+# Run all non-wrapper evaluation scripts for every completed experiment under
+# runs/.  The script is written for Linux servers with bash 4.2.
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
-shopt -s globstar nullglob
 
 PYTHON_BIN="${PYTHON_BIN:-python}"
 DEVICE="${DEVICE:-cuda}"
@@ -82,20 +81,6 @@ run_dir_result_if_missing() {
   run_command "$@"
 }
 
-run_file_result_if_missing() {
-  local label="$1"
-  local result_dir="$2"
-  local required="$3"
-  shift 3
-
-  if [[ -s "${result_dir}/${required}" ]]; then
-    echo "    skip ${label}: found ${result_dir}/${required}"
-    return 0
-  fi
-
-  run_command "$@"
-}
-
 run_has_extra_info() {
   "${PYTHON_BIN}" - "$1" <<'PY'
 import json
@@ -130,30 +115,8 @@ raise SystemExit(0 if mode == "all_future" else 1)
 PY
 }
 
-run_has_current_model_config() {
-  "${PYTHON_BIN}" - "$1" <<'PY'
-import json
-import sys
-from pathlib import Path
-
-cfg_path = Path(sys.argv[1]) / "train_config.json"
-try:
-    cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
-    n_layer = int(cfg.get("n_layer", 0))
-except Exception:
-    raise SystemExit(1)
-
-supported = {"transformer_ffn_v1", "traj_mixer_v5"}
-raise SystemExit(
-    0
-    if cfg.get("model_architecture") in supported and n_layer >= 1
-    else 1
-)
-PY
-}
-
-for config_path in runs/**/train_config.json; do
-  run_path="${config_path%/train_config.json}"
+for run_path in runs/*; do
+  [[ -d "${run_path}" ]] || continue
 
   echo "==> ${run_path}"
   if [[ ! -f "${run_path}/train_config.json" ]]; then
@@ -164,10 +127,6 @@ for config_path in runs/**/train_config.json; do
     echo "    skip run: missing best_model.pt"
     continue
   fi
-  if ! run_has_current_model_config "${run_path}"; then
-    echo "    skip run: config lacks current model_architecture/n_layer fields"
-    continue
-  fi
 
   common=()
   while IFS= read -r arg; do common+=("${arg}"); done < <(common_args_with_device "${run_path}")
@@ -178,16 +137,18 @@ for config_path in runs/**/train_config.json; do
   cpu_reduce_extra=()
   while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
 
-  run_file_result_if_missing \
+  run_dir_result_if_missing \
     "evaluate_auc.py" \
     "${run_path}" \
-    "df_auc_delphi2m_report.csv" \
+    "df_both.csv" \
+    "df_auc_unpooled.csv" \
     "${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
 
-  run_file_result_if_missing \
+  run_dir_result_if_missing \
     "evaluate_auc_v2.py" \
     "${run_path}" \
-    "df_auc_landmark_delphi2m_report.csv" \
+    "df_auc_landmark.csv" \
+    "df_auc_landmark_unpooled.csv" \
     "${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"
 
   if ! run_is_all_future "${run_path}"; then

run_missing_training_runs.sh

@@ -10,8 +10,7 @@ set -euo pipefail
 #   all_future + relative time + mixed death/risk head
 #
 # This script only launches those missing training jobs. It intentionally does
-# not call evaluate_*.py and does not add extra random seeds. Set
-# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
+# not call evaluate_*.py and does not add extra random seeds.
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
@@ -19,8 +18,6 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
 DEVICE="${DEVICE:-cuda}"
 NUM_WORKERS="${NUM_WORKERS:-4}"
 PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
-MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
-N_LAYER="${N_LAYER:-12}"
 
 TIME_MODE="relative"
 DIST_MODE="mixed"
@@ -39,8 +36,8 @@ COMMON_ARGS=(
   --min_future_events 1
   --n_embd 120
   --n_head 10
-  --n_layer "${N_LAYER}"
-  --model_architecture "${MODEL_ARCHITECTURE}"
+  --n_hist_layer 12
+  --n_tab_layer 4
   --n_bins 16
   --extra_pool_reduce mean
   --dropout 0.0
@@ -60,23 +57,15 @@ COMMON_ARGS=(
 
 already_trained() {
   local extra_file="$1"
-  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
+  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
 import json
 import sys
 from pathlib import Path
 
-(
-    time_mode,
-    dist_mode,
-    extra_file,
-    seed,
-    validation_query_seed,
-    model_architecture,
-    n_layer,
-) = sys.argv[1:8]
+time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
 extra_name = Path(extra_file).name
 
-for config_path in Path("runs").rglob("train_config.json"):
+for config_path in Path("runs").glob("*/train_config.json"):
     try:
         cfg = json.loads(config_path.read_text(encoding="utf-8"))
     except Exception:
@@ -89,8 +78,6 @@ for config_path in Path("runs").rglob("train_config.json"):
 
     if (
         cfg.get("model_target_mode") == "all_future"
-        and cfg.get("model_architecture") == model_architecture
-        and int(cfg.get("n_layer", -1)) == int(n_layer)
         and cfg.get("time_mode") == time_mode
         and cfg.get("dist_mode") == dist_mode
         and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
@@ -113,7 +100,7 @@ train_if_missing() {
     return 2
   fi
 
-  echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
+  echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
   if existing_run="$(already_trained "$extra_file")"; then
     echo "    skip: already trained at ${existing_run}"
     return 0

test_model_architectures.py

@@ -1,247 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import (
-    SwiGLU,
-    TrajMixer,
-    TrajMixerBlock,
-    TransformerFFNBlock,
-    build_backbone_block,
-)
-from model_architectures import (
-    TRAJ_MIXER_ARCHITECTURE,
-    TRANSFORMER_FFN_ARCHITECTURE,
-    detect_model_architecture_from_state_dict,
-    resolve_model_architecture,
-)
-from models import DeepHealth
-
-
-def _build_block(model_architecture: str):
-    return build_backbone_block(
-        model_architecture,
-        n_embd=12,
-        n_head=3,
-        use_time_rope=False,
-        use_rbf_bias=False,
-        mlp_dropout=0.0,
-    )
-
-
-def _as_model_state_dict(block: torch.nn.Module) -> dict[str, torch.Tensor]:
-    return {
-        f"blocks.0.{name}": value.detach().clone()
-        for name, value in block.state_dict().items()
-    }
-
-
-def _build_model(
-    model_architecture: str | None,
-    *,
-    n_layer: int = 1,
-) -> DeepHealth:
-    return DeepHealth(
-        vocab_size=8,
-        n_embd=12,
-        n_head=3,
-        n_layer=n_layer,
-        n_types=2,
-        n_cont_types=0,
-        n_categories=2,
-        cont_type_ids=[],
-        time_mode="absolute",
-        model_architecture=model_architecture,
-    )
-
-
-class ModelArchitectureFactoryTest(unittest.TestCase):
-    def test_factory_builds_both_architectures_with_expected_topology(self) -> None:
-        ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.assertIsInstance(ffn_block, TransformerFFNBlock)
-        self.assertIsInstance(ffn_block.mlp, SwiGLU)
-        self.assertTrue(hasattr(ffn_block, "ln1"))
-        self.assertTrue(hasattr(ffn_block, "ln2"))
-
-        traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.assertIsInstance(traj_block, TrajMixerBlock)
-        self.assertIsInstance(traj_block.mlp, TrajMixer)
-        self.assertTrue(hasattr(traj_block, "ln1"))
-        self.assertFalse(hasattr(traj_block, "ln2"))
-
-    def test_both_architectures_forward_and_backward(self) -> None:
-        for architecture in (
-            TRANSFORMER_FFN_ARCHITECTURE,
-            TRAJ_MIXER_ARCHITECTURE,
-        ):
-            with self.subTest(architecture=architecture):
-                torch.manual_seed(0)
-                block = _build_block(architecture)
-                x = torch.randn(2, 5, 12, requires_grad=True)
-
-                output = block(x)
-                self.assertEqual(output.shape, x.shape)
-                output.square().mean().backward()
-
-                self.assertIsNotNone(x.grad)
-                self.assertTrue(torch.isfinite(x.grad).all())
-                self.assertGreater(x.grad.abs().sum().item(), 0.0)
-                self.assertIsNotNone(block.attn.qkv.weight.grad)
-                self.assertGreater(
-                    block.attn.qkv.weight.grad.abs().sum().item(),
-                    0.0,
-                )
-
-                if architecture == TRANSFORMER_FFN_ARCHITECTURE:
-                    branch_parameters = (
-                        block.mlp.w1.weight,
-                        block.mlp.w2.weight,
-                        block.mlp.w3.weight,
-                    )
-                else:
-                    branch_parameters = (
-                        block.mlp.intra_gate_proj,
-                        block.mlp.intra_value_proj,
-                        block.mlp.output_proj,
-                    )
-                for parameter in branch_parameters:
-                    self.assertIsNotNone(parameter.grad)
-                    self.assertTrue(torch.isfinite(parameter.grad).all())
-                    self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
-
-    def test_unknown_architecture_is_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            _build_block("unknown_architecture")
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            _build_model(None)
-
-    def test_deephealth_rejects_fewer_than_one_layer(self) -> None:
-        for n_layer in (0, -1):
-            with self.subTest(n_layer=n_layer):
-                with self.assertRaisesRegex(ValueError, "n_layer must be >= 1"):
-                    _build_model(
-                        TRANSFORMER_FFN_ARCHITECTURE,
-                        n_layer=n_layer,
-                    )
-
-    def test_deephealth_uses_factory_and_strictly_reloads_both_models(self) -> None:
-        for architecture, block_class in (
-            (TRANSFORMER_FFN_ARCHITECTURE, TransformerFFNBlock),
-            (TRAJ_MIXER_ARCHITECTURE, TrajMixerBlock),
-        ):
-            with self.subTest(architecture=architecture):
-                model = _build_model(architecture)
-                self.assertEqual(model.model_architecture, architecture)
-                self.assertIsInstance(model.blocks[0], block_class)
-                self.assertEqual(
-                    detect_model_architecture_from_state_dict(
-                        model.state_dict()
-                    ),
-                    architecture,
-                )
-
-                reloaded = _build_model(architecture)
-                incompatible = reloaded.load_state_dict(
-                    model.state_dict(),
-                    strict=True,
-                )
-                self.assertEqual(incompatible.missing_keys, [])
-                self.assertEqual(incompatible.unexpected_keys, [])
-
-
-class ModelArchitectureResolutionTest(unittest.TestCase):
-    def setUp(self) -> None:
-        self.ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        self.traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
-        self.ffn_state = _as_model_state_dict(self.ffn_block)
-        self.traj_state = _as_model_state_dict(self.traj_block)
-
-    def test_state_dict_detection_recognizes_both_architectures(self) -> None:
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.ffn_state),
-            TRANSFORMER_FFN_ARCHITECTURE,
-        )
-        self.assertEqual(
-            detect_model_architecture_from_state_dict(self.traj_state),
-            TRAJ_MIXER_ARCHITECTURE,
-        )
-
-    def test_explicit_markers_resolve_when_checkpoint_matches(self) -> None:
-        for architecture, state_dict in (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.ffn_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.traj_state),
-        ):
-            with self.subTest(architecture=architecture):
-                self.assertEqual(
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    ),
-                    architecture,
-                )
-
-    def test_architecture_marker_is_required_for_checkpoint_loading(self) -> None:
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture({}, self.ffn_state)
-        with self.assertRaisesRegex(ValueError, "model_architecture is required"):
-            resolve_model_architecture(None, self.traj_state)
-
-    def test_explicit_marker_conflicting_with_state_dict_is_rejected(self) -> None:
-        conflicts = (
-            (TRANSFORMER_FFN_ARCHITECTURE, self.traj_state),
-            (TRAJ_MIXER_ARCHITECTURE, self.ffn_state),
-        )
-        for architecture, state_dict in conflicts:
-            with self.subTest(architecture=architecture):
-                with self.assertRaises(ValueError):
-                    resolve_model_architecture(
-                        {"model_architecture": architecture},
-                        state_dict,
-                    )
-
-    def test_unknown_marker_and_ambiguous_state_dict_are_rejected(self) -> None:
-        with self.assertRaises(ValueError):
-            resolve_model_architecture(
-                {"model_architecture": "traj_mixer_v4"}
-            )
-
-        ambiguous_state = dict(self.ffn_state)
-        ambiguous_state.update(self.traj_state)
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(ambiguous_state)
-
-        with self.assertRaises(ValueError):
-            detect_model_architecture_from_state_dict(
-                {"token_embedding.weight": torch.empty(2, 2)}
-            )
-
-    def test_ffn_block_schema_is_stable_and_strictly_loadable(self) -> None:
-        expected_keys = {
-            "attn.time_bias_scale",
-            "attn.qkv.weight",
-            "attn.out_proj.weight",
-            "attn.rbf_proj.weight",
-            "mlp.w1.weight",
-            "mlp.w1.bias",
-            "mlp.w2.weight",
-            "mlp.w2.bias",
-            "mlp.w3.weight",
-            "mlp.w3.bias",
-            "ln1.weight",
-            "ln1.bias",
-            "ln2.weight",
-            "ln2.bias",
-        }
-        state = self.ffn_block.state_dict()
-        self.assertSetEqual(set(state), expected_keys)
-        self.assertEqual(tuple(state["mlp.w1.weight"].shape), (30, 12))
-        self.assertEqual(tuple(state["mlp.w3.weight"].shape), (12, 30))
-
-        reloaded = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
-        incompatible = reloaded.load_state_dict(state, strict=True)
-        self.assertEqual(incompatible.missing_keys, [])
-        self.assertEqual(incompatible.unexpected_keys, [])
-
-
-if __name__ == "__main__":
-    unittest.main()

test_temporal_attention.py

@@ -1,46 +0,0 @@
-import unittest
-
-import torch
-
-from backbones import TemporalAttention
-
-
-class TemporalAttentionTest(unittest.TestCase):
-    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
-        torch.manual_seed(0)
-        attention = TemporalAttention(
-            n_embd=12,
-            n_head=3,
-            use_time_rope=False,
-            use_rbf_bias=True,
-        )
-        features = torch.randn(2, 4, 4, 16)
-        target = torch.randn(2, 4, 4, 3)
-
-        initial_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
-
-        (initial_bias * target).sum().backward()
-        projection_grad = attention.rbf_proj.weight.grad
-        self.assertIsNotNone(projection_grad)
-        self.assertGreater(projection_grad.abs().sum().item(), 0.0)
-
-        with torch.no_grad():
-            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
-        attention.zero_grad(set_to_none=True)
-        updated_bias = (
-            attention.time_bias_scale.tanh()
-            * attention.rbf_proj(features)
-        )
-        (updated_bias * target).sum().backward()
-
-        scale_grad = attention.time_bias_scale.grad
-        self.assertIsNotNone(scale_grad)
-        self.assertGreater(scale_grad.abs().item(), 0.0)
-
-
-if __name__ == "__main__":
-    unittest.main()

test_traj_mixer.py

@@ -2,10 +2,52 @@ import unittest
 
 import torch
 
-from backbones import TrajMixer
+from backbones import GPTBlock, TemporalAttention, TrajMixer
+from models import (
+    TRAJ_MIXER_ARCHITECTURE,
+    validate_traj_mixer_config,
+    validate_traj_mixer_state_dict,
+)
+from train_util import get_model_parameter_counts
 
 
 class TrajMixerTest(unittest.TestCase):
+    def test_zero_rbf_bias_has_live_projection_gradient(self) -> None:
+        attention = TemporalAttention(
+            n_embd=12,
+            n_head=3,
+            use_time_rope=False,
+            use_rbf_bias=True,
+        )
+        features = torch.randn(2, 4, 4, 16)
+        target = torch.randn(2, 4, 4, 3)
+
+        initial_bias = (
+            attention.time_bias_scale.tanh()
+            * attention.rbf_proj(features)
+        )
+        torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias))
+
+        loss = (initial_bias * target).sum()
+        loss.backward()
+
+        projection_grad = attention.rbf_proj.weight.grad
+        self.assertIsNotNone(projection_grad)
+        self.assertGreater(projection_grad.abs().sum().item(), 0.0)
+
+        with torch.no_grad():
+            attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3)
+        attention.zero_grad(set_to_none=True)
+        updated_bias = (
+            attention.time_bias_scale.tanh()
+            * attention.rbf_proj(features)
+        )
+        (updated_bias * target).sum().backward()
+
+        scale_grad = attention.time_bias_scale.grad
+        self.assertIsNotNone(scale_grad)
+        self.assertGreater(scale_grad.abs().item(), 0.0)
+
     def test_default_shape_parameters_and_initialization(self) -> None:
         mixer = TrajMixer(
             n_embd=120,
@@ -16,6 +58,9 @@ class TrajMixerTest(unittest.TestCase):
         x = torch.randn(2, 7, 120)
         self.assertEqual(mixer(x).shape, x.shape)
         self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
+        self.assertFalse(hasattr(mixer, "group_align"))
+        self.assertFalse(hasattr(mixer, "intra_norm"))
+        self.assertFalse(hasattr(mixer, "cross_norm"))
         self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
         self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
         torch.testing.assert_close(
@@ -42,40 +87,40 @@ class TrajMixerTest(unittest.TestCase):
 
     def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
         torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
         with torch.no_grad():
             mixer.output_proj.zero_()
-        x = torch.randn(2, 5, 12)
+        x = torch.randn(2, 5, 120)
         torch.testing.assert_close(mixer(x), x)
 
     def test_forward_matches_single_outer_residual_formula(self) -> None:
         torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
         mixer.eval()
-        x = torch.randn(2, 5, 12)
+        x = torch.randn(2, 5, 120)
 
-        grouped = mixer.norm(x).reshape(2, 5, 3, 4)
+        grouped = mixer.norm(x).reshape(2, 5, 10, 12)
         intra_output = mixer._intra_mix(grouped)
         static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
-            1, 1, 3, 4
+            1, 1, 10, 12
         )
         mixed_input = grouped + static_gate * intra_output
-        update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)
+        update = mixer._cross_mix(mixed_input).reshape(2, 5, 120)
 
         torch.testing.assert_close(mixer(x), x + update)
 
     def test_intra_stage_is_independent_across_groups(self) -> None:
         torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
         mixer.eval()
 
-        grouped = torch.randn(2, 4, 3, 4)
+        grouped = torch.randn(2, 4, 10, 12)
         changed = grouped.clone()
-        changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])
+        changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :])
 
         original_out = mixer._intra_mix(grouped)
         changed_out = mixer._intra_mix(changed)
-        unchanged_groups = torch.tensor([0, 2])
+        unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9])
         torch.testing.assert_close(
             original_out.index_select(2, unchanged_groups),
             changed_out.index_select(2, unchanged_groups),
@@ -89,8 +134,8 @@ class TrajMixerTest(unittest.TestCase):
             mixer.value_proj.zero_()
             mixer.output_proj.zero_()
 
-            # Coordinate 0 reads group 0 through hidden unit 0 and writes it
-            # into group 1. Coordinate 1 must remain independent.
+            # For coordinate 0 only, read group 0 through hidden unit 0 and
+            # write the resulting gated value into group 1.
             mixer.gate_proj[0, 0, 0] = 1.0
             mixer.value_proj[0, 0, 0] = 1.0
             mixer.output_proj[0, 0, 1] = 1.0
@@ -119,9 +164,9 @@ class TrajMixerTest(unittest.TestCase):
 
     def test_mixer_does_not_mix_sequence_positions(self) -> None:
         torch.manual_seed(0)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
         mixer.eval()
-        x = torch.randn(2, 5, 12)
+        x = torch.randn(2, 5, 120)
         changed = x.clone()
         changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
 
@@ -133,34 +178,76 @@ class TrajMixerTest(unittest.TestCase):
             changed_out.index_select(1, unchanged_positions),
         )
 
-    def test_gradients_reach_every_projection_family(self) -> None:
+    def test_gradients_reach_all_projection_families(self) -> None:
         torch.manual_seed(1)
-        mixer = TrajMixer(12, n_head=3, dropout=0.0)
-        x = torch.randn(2, 4, 12, requires_grad=True)
+        mixer = TrajMixer(120, n_head=10, dropout=0.0)
+        x = torch.randn(2, 4, 120, requires_grad=True)
 
         mixer(x).square().mean().backward()
 
         self.assertIsNotNone(x.grad)
-        self.assertTrue(torch.isfinite(x.grad).all())
-        self.assertGreater(x.grad.abs().sum().item(), 0.0)
-        projection_names = (
-            "intra_gate_proj",
-            "intra_value_proj",
-            "intra_output_proj",
-            "gate_proj",
-            "value_proj",
-            "output_proj",
-        )
-        for name in projection_names:
-            parameter = getattr(mixer, name)
+        for name, parameter in mixer.named_parameters():
             self.assertIsNotNone(parameter.grad, name)
             self.assertTrue(torch.isfinite(parameter.grad).all(), name)
-            self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name)
+
+    def test_gpt_block_delegates_single_mixer_residual_to_traj_mixer(self) -> None:
+        block = GPTBlock(n_embd=120, n_head=10)
+        self.assertIsInstance(block.mlp, TrajMixer)
+        self.assertFalse(hasattr(block, "ln2"))
+        self.assertIsInstance(block.mlp.norm, torch.nn.LayerNorm)
+        self.assertFalse(hasattr(block.mlp, "intra_norm"))
+        self.assertFalse(hasattr(block.mlp, "cross_norm"))
+
+        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"})
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v2"}
+            )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v3"}
+            )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_traj_mixer_config(
+                {"model_architecture": "traj_mixer_v4"}
+            )
+
+    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.intra_gate_proj")
+        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)
 
+    def test_parameter_counts_match_traj_mixer_parameters(self) -> None:
+        mixer = TrajMixer(n_embd=120, n_head=10)
+        self.assertEqual(
+            get_model_parameter_counts(mixer),
+            {
+                "model_parameter_count": 32_040,
+                "trainable_parameter_count": 32_040,
+            },
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

train_all_future.py

@@ -27,11 +27,7 @@ from tqdm.auto import tqdm
 
 from dataset import AllFutureHealthDataset, all_future_collate_fn
 from losses import build_loss
-from model_architectures import (
-    DEFAULT_MODEL_ARCHITECTURE,
-    SUPPORTED_MODEL_ARCHITECTURES,
-)
-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,
@@ -69,7 +65,6 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument("--data_prefix", type=str, default="ukb")
     parser.add_argument("--labels_file", type=str, default="labels.csv")
-    parser.add_argument("--runs_root", type=str, default="runs")
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--extra_info_types_file", type=str, default=None)
 
@@ -85,7 +80,8 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument("--n_embd", type=int, default=120)
     parser.add_argument("--n_head", type=int, default=10)
-    parser.add_argument("--n_layer", type=int, default=12)
+    parser.add_argument("--n_hist_layer", type=int, default=12)
+    parser.add_argument("--n_tab_layer", type=int, default=4)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
@@ -94,12 +90,6 @@ 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(
-        "--model_architecture",
-        type=str,
-        default=DEFAULT_MODEL_ARCHITECTURE,
-        choices=SUPPORTED_MODEL_ARCHITECTURES,
-    )
 
     parser.add_argument("--batch_size", type=int, default=128)
     parser.add_argument("--base_lr", type=float, default=3e-4)
@@ -158,7 +148,8 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         vocab_size=dataset.vocab_size,
         n_embd=args.n_embd,
         n_head=args.n_head,
-        n_layer=args.n_layer,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -169,7 +160,6 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
         time_mode=args.time_mode,
         dist_mode=args.dist_mode,
         dropout=args.dropout,
-        model_architecture=args.model_architecture,
     )
 
 
@@ -309,7 +299,7 @@ def build_metadata(
         "dataset_class": "AllFutureHealthDataset",
         "collate_fn": "all_future_collate_fn",
         "model_class": "DeepHealth",
-        "model_architecture": args.model_architecture,
+        "model_architecture": TRAJ_MIXER_ARCHITECTURE,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
         "dist_mode": args.dist_mode,
@@ -347,14 +337,12 @@ def main() -> None:
     configure_torch_for_training(device)
 
     run_dir, run_name = create_unique_run_dir(
-        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}",
-        runs_root=Path(args.runs_root) / args.model_architecture,
+        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
     )
     logger = setup_logging(run_dir)
 
     logger.info(f"Starting all-future training run: {run_name}")
     logger.info(f"Device: {device}")
-    logger.info(f"Model architecture: {args.model_architecture}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
 
     logger.info("Loading all-future datasets...")

train_next_step.py

@@ -24,11 +24,7 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset, collate_fn
 from losses import build_loss
-from model_architectures import (
-    DEFAULT_MODEL_ARCHITECTURE,
-    SUPPORTED_MODEL_ARCHITECTURES,
-)
-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 (
@@ -66,7 +62,6 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument("--data_prefix", type=str, default="ukb")
     parser.add_argument("--labels_file", type=str, default="labels.csv")
-    parser.add_argument("--runs_root", type=str, default="runs")
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--extra_info_types_file", type=str, default=None)
     parser.add_argument("--no_event_interval_years", type=float, default=5.0)
@@ -81,19 +76,14 @@ def parse_args() -> argparse.Namespace:
 
     parser.add_argument("--n_embd", type=int, default=120)
     parser.add_argument("--n_head", type=int, default=10)
-    parser.add_argument("--n_layer", type=int, default=12)
+    parser.add_argument("--n_hist_layer", type=int, default=12)
+    parser.add_argument("--n_tab_layer", type=int, default=4)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
     parser.add_argument("--time_mode", type=str, default="relative",
                         choices=["relative", "absolute"])
     parser.add_argument("--dropout", type=float, default=0.0)
-    parser.add_argument(
-        "--model_architecture",
-        type=str,
-        default=DEFAULT_MODEL_ARCHITECTURE,
-        choices=SUPPORTED_MODEL_ARCHITECTURES,
-    )
 
     parser.add_argument("--target_mode", type=str, default="uts",
                         choices=["delphi2m", "uts"])
@@ -163,7 +153,8 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         vocab_size=dataset.vocab_size,
         n_embd=args.n_embd,
         n_head=args.n_head,
-        n_layer=args.n_layer,
+        n_hist_layer=args.n_hist_layer,
+        n_tab_layer=args.n_tab_layer,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -174,7 +165,6 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
         time_mode=args.time_mode,
         dist_mode="exponential",
         dropout=args.dropout,
-        model_architecture=args.model_architecture,
     )
 
 
@@ -495,7 +485,7 @@ def build_metadata(
         "dataset_class": "NextStepHealthDataset",
         "collate_fn": "next_step_collate_fn",
         "model_class": "DeepHealth",
-        "model_architecture": args.model_architecture,
+        "model_architecture": TRAJ_MIXER_ARCHITECTURE,
         "model_target_mode": "next_token",
         "target_mode": args.target_mode,
         "dist_mode": "exponential",
@@ -533,14 +523,12 @@ def main() -> None:
         lambda timestamp: (
             f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
             f"gap_{args.no_event_interval_years:g}y_{timestamp}"
-        ),
-        runs_root=Path(args.runs_root) / args.model_architecture,
+        )
     )
     logger = setup_logging(run_dir)
 
     logger.info(f"Starting next-step training run: {run_name}")
     logger.info(f"Device: {device}")
-    logger.info(f"Model architecture: {args.model_architecture}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
     logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}")
 

train_util.py

@@ -301,7 +301,7 @@ def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
 
 
 def get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]:
-    """Return stable total and trainable parameter counts."""
+    """Return stable parameter-count fields for logs and train_config.json."""
     return {
         "model_parameter_count": sum(
             parameter.numel() for parameter in model.parameters()