Unify FFN and TrajMixer model architectures

MODEL_ARCHITECTURES.md

@@ -0,0 +1,64 @@
+# 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.

backbones.py

@@ -4,6 +4,12 @@ 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):
@@ -213,7 +219,133 @@ class SwiGLU(nn.Module):
         return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
 
 
-class GPTBlock(nn.Module):
+class TrajMixer(nn.Module):
+    """PreNorm gated mixing within and across latent trajectory 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,
+        n_head: int = 10,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        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}"
+            )
+        self.n_embd = n_embd
+        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
+
+        self.norm = nn.LayerNorm(self.n_embd)
+
+        self.intra_gate_proj = nn.Parameter(
+            torch.empty(self.n_group, self.d_group, self.intra_hidden)
+        )
+        self.intra_value_proj = nn.Parameter(
+            torch.empty(self.n_group, self.d_group, self.intra_hidden)
+        )
+        self.intra_output_proj = nn.Parameter(
+            torch.empty(self.n_group, self.intra_hidden, self.d_group)
+        )
+        self.intra_gate_logits = nn.Parameter(
+            torch.empty(self.n_group, self.d_group)
+        )
+
+        self.gate_proj = nn.Parameter(
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
+        )
+        self.value_proj = nn.Parameter(
+            torch.empty(self.d_group, self.n_group, self.hidden_group)
+        )
+        self.output_proj = nn.Parameter(
+            torch.empty(self.d_group, self.hidden_group, self.n_group)
+        )
+        self.drop = nn.Dropout(dropout)
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        for group_idx in range(self.n_group):
+            nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
+            nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
+            nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
+        nn.init.constant_(
+            self.intra_gate_logits,
+            math.log(0.1 / 0.9),
+        )
+
+        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 _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
+        """Mix features independently inside each residual-space group."""
+        intra_gate = torch.einsum(
+            "blgd,gdh->blgh", grouped, self.intra_gate_proj
+        )
+        intra_value = torch.einsum(
+            "blgd,gdh->blgh", grouped, self.intra_value_proj
+        )
+        intra_hidden = F.silu(intra_gate) * intra_value
+        return torch.einsum(
+            "blgh,ghd->blgd", intra_hidden, self.intra_output_proj
+        )
+
+    def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
+        """Mix groups independently for each within-group coordinate."""
+        gate = torch.einsum(
+            "blgr,rgh->blhr", grouped, self.gate_proj
+        )
+        value = torch.einsum(
+            "blgr,rgh->blhr", grouped, self.value_proj
+        )
+        hidden = F.silu(gate) * value
+        return torch.einsum(
+            "blhr,rhg->blgr", hidden, self.output_proj
+        )
+
+    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)}"
+            )
+        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 = self.norm(x).reshape(
+            batch_size, seq_len, self.n_group, self.d_group
+        )
+
+        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
+
+        update = self._cross_mix(mixed_input).reshape(
+            batch_size, seq_len, self.n_embd
+        )
+        return x + self.drop(update)
+
+
+class TransformerFFNBlock(nn.Module):
     def __init__(
         self,
         n_embd: int,
@@ -250,6 +382,75 @@ class GPTBlock(nn.Module):
         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,
+        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 = TrajMixer(
+            n_embd=n_embd,
+            n_head=n_head,
+            dropout=mlp_dropout,
+        )
+        self.ln1 = 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)
+        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,6 +45,7 @@ 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 readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -313,19 +314,24 @@ 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) -> DeepHealth:
+def build_model_from_dataset(
+    args: argparse.Namespace,
+    cfg: Dict[str, Any],
+    dataset: HealthDataset,
+    state_dict: Optional[Dict[str, Any]] = None,
+) -> DeepHealth:
     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_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
+        n_layer=int(cfg["n_layer"]),
-        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,
@@ -336,6 +342,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)),
+        model_architecture=model_architecture,
     )
 
 
@@ -383,7 +390,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 def load_model_state(
-    model: torch.nn.Module,
+    model: DeepHealth,
     checkpoint_path: str,
     device: torch.device,
     state_dict: Optional[Dict[str, Any]] = None,
@@ -391,6 +398,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)
     model.load_state_dict(state, strict=True)
 
 
@@ -1371,14 +1379,19 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, str(model_ckpt_path),
                      device, state_dict=state_dict)
     model.eval()

evaluate_auc_v2.py

@@ -36,6 +36,7 @@ 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 readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -184,19 +185,24 @@ 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) -> DeepHealth:
+def build_model_from_dataset(
+    args: argparse.Namespace,
+    cfg: Dict[str, Any],
+    dataset: HealthDataset,
+    state_dict: Optional[Dict[str, Any]] = None,
+) -> DeepHealth:
     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_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
+        n_layer=int(cfg["n_layer"]),
-        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,
@@ -207,10 +213,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)),
+        model_architecture=model_architecture,
     )
 
 
-def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
+def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
+    resolve_model_architecture(model.model_architecture, state_dict)
     model.load_state_dict(state_dict, strict=True)
 
 
@@ -1392,12 +1400,17 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
 
     if (
         model_target_mode == "next_token"

evaluate_event_free_survival.py

@@ -649,7 +649,9 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

evaluate_extra_info_attribution.py

@@ -758,7 +758,9 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

evaluate_single_disease_mortality_attribution.py

@@ -553,7 +553,9 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

export_tquery_logits_hidden.py

@@ -180,7 +180,9 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

export_weibull_death_parameter_stats.py

@@ -381,7 +381,9 @@ 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).to(device)
+    model = build_model_from_dataset(
+        args, cfg_model, dataset, state_dict=state_dict
+    ).to(device)
     load_model_state(model, state_dict)
     model.eval()
 

model_architectures.py

@@ -0,0 +1,131 @@
+"""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

@@ -6,11 +6,12 @@ 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
 
 
@@ -147,8 +148,7 @@ class DeepHealth(nn.Module):
         vocab_size: int,
         n_embd: int,
         n_head: int,
-        n_hist_layer: int,
+        n_layer: int,
-        n_tab_layer: int,
         n_types: int,
         n_cont_types: int,
         n_categories: int,
@@ -160,6 +160,7 @@ 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"]:
@@ -173,6 +174,9 @@ 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
@@ -189,6 +193,8 @@ 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)
@@ -208,26 +214,28 @@ class DeepHealth(nn.Module):
         if time_mode == "absolute":
             self.age_encoding = AgeSinusoidalEncoding(n_embd)
             self.blocks = nn.ModuleList([
-                GPTBlock(
+                build_backbone_block(
+                    model_architecture,
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=False,
                     use_rbf_bias=False,
                     mlp_dropout=dropout,
-                ) for _ in range(n_hist_layer)
+                ) for _ in range(n_layer)
             ])
             self.rope = None
             self.rbf = None
         elif time_mode == "relative":
             self.age_encoding = None
             self.blocks = nn.ModuleList([
-                GPTBlock(
+                build_backbone_block(
+                    model_architecture,
                     n_embd=n_embd,
                     n_head=n_head,
                     use_time_rope=True,
                     use_rbf_bias=True,
                     mlp_dropout=dropout,
-                ) for _ in range(n_hist_layer)
+                ) for _ in range(n_layer)
             ])
             self.rope = TimeRoPE(n_embd // n_head)
             self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)

run_missing_evaluations.sh

@@ -1,10 +1,11 @@
 #!/usr/bin/env bash
 set -euo pipefail
 
-# Run all non-wrapper evaluation scripts for every completed experiment under
+# Run all non-wrapper evaluation scripts for every completed current-format
-# runs/.  The script is written for Linux servers with bash 4.2.
+# 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}"
@@ -81,6 +82,20 @@ 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
@@ -115,8 +130,30 @@ raise SystemExit(0 if mode == "all_future" else 1)
 PY
 }
 
-for run_path in runs/*; do
+run_has_current_model_config() {
-  [[ -d "${run_path}" ]] || continue
+  "${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}"
 
   echo "==> ${run_path}"
   if [[ ! -f "${run_path}/train_config.json" ]]; then
@@ -127,6 +164,10 @@ for run_path in runs/*; 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}")
@@ -137,18 +178,16 @@ for run_path in runs/*; do
   cpu_reduce_extra=()
   while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
 
-  run_dir_result_if_missing \
+  run_file_result_if_missing \
     "evaluate_auc.py" \
     "${run_path}" \
-    "df_both.csv" \
+    "df_auc_delphi2m_report.csv" \
-    "df_auc_unpooled.csv" \
     "${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
 
-  run_dir_result_if_missing \
+  run_file_result_if_missing \
     "evaluate_auc_v2.py" \
     "${run_path}" \
-    "df_auc_landmark.csv" \
+    "df_auc_landmark_delphi2m_report.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,7 +10,8 @@ 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.
+# not call evaluate_*.py and does not add extra random seeds. Set
+# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
 
 cd "$(dirname "${BASH_SOURCE[0]}")"
 
@@ -18,6 +19,8 @@ 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"
@@ -36,8 +39,8 @@ COMMON_ARGS=(
   --min_future_events 1
   --n_embd 120
   --n_head 10
-  --n_hist_layer 12
+  --n_layer "${N_LAYER}"
-  --n_tab_layer 4
+  --model_architecture "${MODEL_ARCHITECTURE}"
   --n_bins 16
   --extra_pool_reduce mean
   --dropout 0.0
@@ -57,15 +60,23 @@ COMMON_ARGS=(
 
 already_trained() {
   local extra_file="$1"
-  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
+  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
 import json
 import sys
 from pathlib import Path
 
-time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
+(
+    time_mode,
+    dist_mode,
+    extra_file,
+    seed,
+    validation_query_seed,
+    model_architecture,
+    n_layer,
+) = sys.argv[1:8]
 extra_name = Path(extra_file).name
 
-for config_path in Path("runs").glob("*/train_config.json"):
+for config_path in Path("runs").rglob("train_config.json"):
     try:
         cfg = json.loads(config_path.read_text(encoding="utf-8"))
     except Exception:
@@ -78,6 +89,8 @@ for config_path in Path("runs").glob("*/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
@@ -100,7 +113,7 @@ train_if_missing() {
     return 2
   fi
 
-  echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
+  echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${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

@@ -0,0 +1,247 @@
+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_traj_mixer.py

@@ -0,0 +1,166 @@
+import unittest
+
+import torch
+
+from backbones import TrajMixer
+
+
+class TrajMixerTest(unittest.TestCase):
+    def test_default_shape_parameters_and_initialization(self) -> None:
+        mixer = TrajMixer(
+            n_embd=120,
+            n_head=10,
+            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()), 32_040)
+        self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
+        self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
+        torch.testing.assert_close(
+            torch.sigmoid(mixer.intra_gate_logits.detach()),
+            torch.full((10, 12), 0.1),
+        )
+        self.assertEqual(mixer.intra_hidden, 48)
+        self.assertEqual(
+            tuple(mixer.intra_gate_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_value_proj.shape),
+            (10, 12, 48),
+        )
+        self.assertEqual(
+            tuple(mixer.intra_output_proj.shape),
+            (10, 48, 12),
+        )
+        self.assertEqual(mixer.hidden_group, 40)
+        self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
+        self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
+        self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
+
+    def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
+        torch.manual_seed(0)
+        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        with torch.no_grad():
+            mixer.output_proj.zero_()
+        x = torch.randn(2, 5, 12)
+        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.eval()
+        x = torch.randn(2, 5, 12)
+
+        grouped = mixer.norm(x).reshape(2, 5, 3, 4)
+        intra_output = mixer._intra_mix(grouped)
+        static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
+            1, 1, 3, 4
+        )
+        mixed_input = grouped + static_gate * intra_output
+        update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)
+
+        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.eval()
+
+        grouped = torch.randn(2, 4, 3, 4)
+        changed = grouped.clone()
+        changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])
+
+        original_out = mixer._intra_mix(grouped)
+        changed_out = mixer._intra_mix(changed)
+        unchanged_groups = torch.tensor([0, 2])
+        torch.testing.assert_close(
+            original_out.index_select(2, unchanged_groups),
+            changed_out.index_select(2, unchanged_groups),
+        )
+
+    def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
+        mixer = TrajMixer(6, n_head=3, dropout=0.0)
+        mixer.eval()
+        with torch.no_grad():
+            mixer.gate_proj.zero_()
+            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.
+            mixer.gate_proj[0, 0, 0] = 1.0
+            mixer.value_proj[0, 0, 0] = 1.0
+            mixer.output_proj[0, 0, 1] = 1.0
+
+        grouped = torch.tensor(
+            [[[
+                [-1.0, 4.0],
+                [0.0, 5.0],
+                [1.0, 6.0],
+            ]]]
+        )
+        changed = grouped.clone()
+        changed[0, 0, 0, 0] = 2.0
+
+        original_out = mixer._cross_mix(grouped)
+        changed_out = mixer._cross_mix(changed)
+
+        self.assertNotEqual(
+            original_out[0, 0, 1, 0].item(),
+            changed_out[0, 0, 1, 0].item(),
+        )
+        torch.testing.assert_close(
+            original_out[..., 1],
+            changed_out[..., 1],
+        )
+
+    def test_mixer_does_not_mix_sequence_positions(self) -> None:
+        torch.manual_seed(0)
+        mixer = TrajMixer(12, n_head=3, dropout=0.0)
+        mixer.eval()
+        x = torch.randn(2, 5, 12)
+        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_every_projection_family(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(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)
+            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_invalid_group_partition_is_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "divisible"):
+            TrajMixer(n_embd=121, n_head=10)
+
+
+if __name__ == "__main__":
+    unittest.main()

train_all_future.py

@@ -27,12 +27,17 @@ 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 targets import CHECKUP_IDX, PAD_IDX
 from train_util import (
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    get_model_parameter_counts,
     load_extra_info_types_file,
     resolve_device,
     save_checkpoint,
@@ -64,6 +69,7 @@ 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)
 
@@ -79,8 +85,7 @@ 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_hist_layer", type=int, default=12)
+    parser.add_argument("--n_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"])
@@ -89,6 +94,12 @@ 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)
@@ -147,8 +158,7 @@ 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_hist_layer=args.n_hist_layer,
+        n_layer=args.n_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,
@@ -159,6 +169,7 @@ 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,
     )
 
 
@@ -298,6 +309,7 @@ def build_metadata(
         "dataset_class": "AllFutureHealthDataset",
         "collate_fn": "all_future_collate_fn",
         "model_class": "DeepHealth",
+        "model_architecture": args.model_architecture,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
         "dist_mode": args.dist_mode,
@@ -335,12 +347,14 @@ 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}"
+        lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}",
+        runs_root=Path(args.runs_root) / args.model_architecture,
     )
     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...")
@@ -434,6 +448,12 @@ def main() -> None:
     )
 
     model = build_model(args, train_dataset).to(device)
+    parameter_counts = get_model_parameter_counts(model)
+    logger.info(
+        "Model parameters: "
+        f"total={parameter_counts['model_parameter_count']:,}, "
+        f"trainable={parameter_counts['trainable_parameter_count']:,}"
+    )
     optimizer = AdamW(
         model.parameters(),
         lr=args.base_lr,
@@ -443,10 +463,14 @@ def main() -> None:
     criterion = build_criterion(args, train_dataset)
     adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
 
+    train_metadata = build_metadata(
+        args, train_dataset, run_name, train_subset, val_subset, test_subset
+    )
+    train_metadata.update(parameter_counts)
     save_config(
         args,
         run_dir / "train_config.json",
-        extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset),
+        extra=train_metadata,
     )
 
     best_val = float("inf")

train_next_step.py

@@ -24,6 +24,10 @@ 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 readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -31,6 +35,7 @@ from train_util import (
     configure_torch_for_training,
     create_unique_run_dir,
     format_extra_info_types,
+    get_model_parameter_counts,
     load_extra_info_types_file,
     resolve_device,
     save_checkpoint,
@@ -61,6 +66,7 @@ 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)
@@ -75,14 +81,19 @@ 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_hist_layer", type=int, default=12)
+    parser.add_argument("--n_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"])
@@ -152,8 +163,7 @@ 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_hist_layer=args.n_hist_layer,
+        n_layer=args.n_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,
@@ -164,6 +174,7 @@ 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,
     )
 
 
@@ -484,6 +495,7 @@ def build_metadata(
         "dataset_class": "NextStepHealthDataset",
         "collate_fn": "next_step_collate_fn",
         "model_class": "DeepHealth",
+        "model_architecture": args.model_architecture,
         "model_target_mode": "next_token",
         "target_mode": args.target_mode,
         "dist_mode": "exponential",
@@ -521,12 +533,14 @@ 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}")
 
@@ -596,6 +610,12 @@ def main() -> None:
     )
 
     model = build_model(args, dataset).to(device)
+    parameter_counts = get_model_parameter_counts(model)
+    logger.info(
+        "Model parameters: "
+        f"total={parameter_counts['model_parameter_count']:,}, "
+        f"trainable={parameter_counts['trainable_parameter_count']:,}"
+    )
     readout = build_next_step_readout(args).to(device)
     criterion = build_next_step_loss(args)
     optimizer = AdamW(
@@ -606,10 +626,14 @@ def main() -> None:
     )
     adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
 
+    train_metadata = build_metadata(
+        args, dataset, run_name, train_subset, val_subset, test_subset
+    )
+    train_metadata.update(parameter_counts)
     save_config(
         args,
         run_dir / "train_config.json",
-        extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
+        extra=train_metadata,
     )
 
     best_val = float("inf")

train_util.py

@@ -300,6 +300,20 @@ 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 {
+        "model_parameter_count": sum(
+            parameter.numel() for parameter in model.parameters()
+        ),
+        "trainable_parameter_count": sum(
+            parameter.numel()
+            for parameter in model.parameters()
+            if parameter.requires_grad
+        ),
+    }
+
+
 def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
     for param_group in optimizer.param_groups:
         param_group["lr"] = lr