Add nano and tiny model presets

Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md

@@ -0,0 +1,389 @@
+# Event–Trajectory Shared Reasoning Backbone
+
+> 状态：**Frozen implementation baseline**
+> 架构标识：`event_trajectory_shared_v2`
+> 固化日期：**2026-07-23**
+
+## 1. 核心定义
+
+使用一个共享的 Attention–TrajMixer 推理核心，对固定 Event Memory 进行多轮读取，并持续更新 Trajectory State。
+
+模型只实例化：
+
+```python
+self.reasoning_core = SharedEventTrajectoryCore(...)
+```
+
+禁止为不同推理轮创建独立 Transformer blocks。参数只保存一套，计算上顺序运行多轮。
+
+模型规模固定为五档：
+
+| model_size | d_model | n_trajectory | trajectory_dim | traj_hidden |
+|---|---:|---:|---:|---:|
+| nano | 120 | 6 | 20 | 24 |
+| tiny | 256 | 8 | 32 | 32 |
+| small | 512 | 8 | 64 | 32 |
+| medium | 768 | 12 | 64 | 48 |
+| huge | 1024 | 16 | 64 | 64 |
+
+默认使用 `model_size=nano`。`n_reasoning_rounds` 是独立参数，默认值为12，
+不属于模型规模预设；任意模型规模均可单独指定推理轮数。
+
+必须满足：
+
+\[
+d_{\mathrm{model}}
+=n_{\mathrm{trajectory}}d_{\mathrm{trajectory}}.
+\]
+
+## 2. 固定 Event Memory
+
+疾病事件与可选协变量首先组成事件序列：
+
+\[
+X_E\in\mathbb{R}^{B\times L\times d_{\mathrm{model}}}.
+\]
+
+事件特征由以下信息相加：
+
+```text
+disease / covariate embedding
++ age/time encoding
++ sex context
+```
+
+然后只编码一次：
+
+\[
+E=\operatorname{EventNorm}
+\left(\operatorname{EventProjection}(X_E)\right).
+\]
+
+进入 reasoning loop 后，\(E\) 的数值保持不变，但不执行 `detach`，梯度仍可回传到事件编码器。
+
+Key 和 Value 同样每次 forward 只投影一次：
+
+```python
+event_key_value = reasoning_core.project_event_memory(E)
+```
+
+12 轮共享并复用该结果。
+
+## 3. Trajectory State
+
+每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用6个：
+
+\[
+S\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
+d_{\mathrm{trajectory}}}.
+\]
+
+其中：
+
+- all-future：\(Q=1\)；
+- next-token：\(Q=L\)，所有查询位置并行计算。
+
+定义可学习原型：
+
+\[
+P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}.
+\]
+
+查询上下文经过投影并 reshape：
+
+\[
+C_Q
+=\operatorname{QueryProjection}(\text{query features})
+\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
+d_{\mathrm{trajectory}}},
+\]
+
+\[
+S^{(0)}=P+C_Q.
+\]
+
+all-future 的 query features 包含可学习 query token、查询年龄和性别；next-token 的 query features 使用当前位置的事件、时间和性别表示，以保留 token-level 预测语义。
+
+## 4. 共享 Trajectory-to-Event Attention
+
+Trajectory State 作为 Query，固定 Event Memory 作为 Key 和 Value：
+
+\[
+Q^{(r)}
+=W_Q\operatorname{LN}_{A}(S^{(r)}),
+\]
+
+\[
+K=W_KE,\qquad V=W_VE.
+\]
+
+形状为：
+
+```text
+Q: [B, query, trajectory, trajectory_dim]
+K: [B, trajectory, event, trajectory_dim]
+V: [B, trajectory, event, trajectory_dim]
+```
+
+Attention：
+
+\[
+\operatorname{score}_{b,q,h,l}
+=
+\frac{
+\left\langle Q_{b,q,h,:},K_{b,h,l,:}\right\rangle
+}{
+\sqrt{d_{\mathrm{trajectory}}}
+}.
+\]
+
+每个 trajectory slot 独立读取整段 Event Memory。Attention 不包含跨 trajectory 的完整输出投影；trajectory 之间的交互只由后续 TrajMixer 完成。
+
+外部 `padding_mask` 的语义固定为 `True = valid`。
+
+all-future 的内部 mask 必须满足：
+
+\[
+\operatorname{valid}_{b,q,l}
+=
+\operatorname{eventValid}_{b,l}
+\land
+(t_l\le t_q).
+\]
+
+next-token 还必须对相同时间戳加入位置因果约束：
+
+\[
+\operatorname{valid}_{b,q,l}
+=
+\operatorname{eventValid}_{b,l}
+\land
+\left[
+(t_l<t_q)
+\lor
+\left((t_l=t_q)\land(l\le q)\right)
+\right].
+\]
+
+这可以阻止前一个 token 直接读到同时间的后续目标 token；它只改变并行
+Attention 的可见性矩阵，不沿疾病时间轴递归。
+
+两种 mask 均不能读取未来事件或协变量。
+
+全 masked query 的 Attention readout 必须显式返回零，不能产生 NaN。
+
+## 5. 时间信息
+
+所有模式都在 Event Memory 和 query context 中加入 age/time encoding。
+
+当 `time_mode=relative` 时，shared cross-attention 额外使用：
+
+- query-time 对 event-time 的 Cross-TimeRoPE；
+- query–event 时间差的 Gaussian RBF bias。
+
+对应缓存形状为：
+
+\[
+\text{RBF cache}\in\mathbb{R}^{B\times Q\times L\times n_{\mathrm{rbf}}}.
+\]
+
+## 6. 共享 TrajMixer
+
+TrajMixer 输入：
+
+\[
+U\in\mathbb{R}^{B\times Q\times H\times D_h}.
+\]
+
+它只沿 trajectory 轴交互，不使用普通全维度 FFN，也不包含旧版 Group Alignment。
+
+对每个内部坐标 \(r\)：
+
+\[
+G_r=U_rW_g^{(r)},\qquad
+V_r=U_rW_v^{(r)},
+\]
+
+\[
+M_r
+=
+\operatorname{SiLU}(G_r)\odot V_r,
+\]
+
+\[
+Y_r=M_rW_o^{(r)}.
+\]
+
+参数形状：
+
+```text
+W_g: [trajectory_dim, n_trajectory, traj_hidden]
+W_v: [trajectory_dim, n_trajectory, traj_hidden]
+W_o: [trajectory_dim, traj_hidden, n_trajectory]
+```
+
+默认：
+
+```text
+n_trajectory -> 4 * n_trajectory -> n_trajectory
+```
+
+其中：
+
+\[
+d_{\mathrm{trajHidden}}=4n_{\mathrm{trajectory}}.
+\]
+
+## 7. 单轮共享核心
+
+单轮计算：
+
+\[
+R^{(r)}
+=
+A_\theta\left(
+\operatorname{LN}_A(S^{(r)}),E
+\right),
+\]
+
+\[
+U^{(r)}
+=
+S^{(r)}
++\alpha_A\operatorname{Dropout}(R^{(r)}),
+\]
+
+\[
+S^{(r+1)}
+=
+U^{(r)}
++\alpha_M\operatorname{Dropout}
+\left(
+M_\phi(\operatorname{LN}_M(U^{(r)}))
+\right).
+\]
+
+残差统一使用加法。
+
+## 8. 多轮参数共享
+
+同一个核心重复运行：
+
+```python
+for _ in range(n_reasoning_rounds):
+    S = self.reasoning_core(...)
+```
+
+所有轮次共享：
+
+```text
+q_proj / k_proj / v_proj
+relative-time projection
+TrajMixer parameters
+LayerNorm parameters
+attn_scale / mixer_scale
+```
+
+因此推理轮数不改变模型参数量：
+
+\[
+A^{(1)}=\cdots=A^{(12)}=A_\theta,
+\]
+
+\[
+M^{(1)}=\cdots=M^{(12)}=M_\phi.
+\]
+
+但每轮 state 不同，因此 Query 与 Attention weights 也不同。
+
+## 9. 稳定性设计
+
+共享残差缩放初始化为：
+
+\[
+\alpha_A=\alpha_M
+=
+\frac{1}{\sqrt{n_{\mathrm{reasoningRounds}}}}.
+\]
+
+两个标量可学习，并由全部轮次共享。
+
+第一版不加入：
+
+```text
+round-specific parameters
+round embedding
+每轮独立 LayerNorm
+每轮独立 residual scale
+GRU 或其他时间递归
+```
+
+## 10. 输出接口
+
+推理结束后按固定顺序 flatten trajectory slots：
+
+\[
+H
+=
+\operatorname{FinalNorm}
+\left(
+\operatorname{Flatten}(S^{(R)})
+\right).
+\]
+
+- all-future 输出：`[B, d_model]`；
+- next-token 输出：`[B, L, d_model]`；
+- next-token 的 risk-head weight tying 保持不变；
+- Weibull 与 mixed heads 继续使用同一最终 hidden。
+
+next-token 的 query 位置全部并行，只有 reasoning rounds 顺序执行，因此不存在沿疾病时间轴的状态递归。
+
+## 11. 配置与 checkpoint 约束
+
+训练配置必须写入：
+
+```yaml
+model_architecture: event_trajectory_shared_v2
+model_size: nano
+d_model: 120
+n_trajectory: 6
+trajectory_dim: 20
+traj_hidden: 24
+n_reasoning_rounds: 12
+model_parameter_count: <runtime count>
+trainable_parameter_count: <runtime count>
+```
+
+评估和导出入口必须同时验证：
+
+1. `model_architecture` 完全匹配；
+2. `model_size` 属于 `nano / tiny / small / medium / huge`；
+3. `d_model`、`n_trajectory`、`trajectory_dim` 和 `traj_hidden`
+   与对应规模预设完全匹配；
+4. checkpoint 包含一套且仅一套 `reasoning_core` 关键参数；
+5. checkpoint 内持久化的 `d_model`、`n_trajectory` 和
+   `n_reasoning_rounds` 架构指纹与训练配置完全一致；
+6. 不接受旧 `traj_mixer_v2` checkpoint。
+
+其中 `n_reasoning_rounds` 必须进入 checkpoint 架构指纹，因为改变轮数
+不会改变参数 shape，不能仅依赖 `load_state_dict(strict=True)` 检出错配。
+
+## 12. 信息流
+
+```text
+E ─────────────┬──────────────┬──────────────┬──────────────┐
+               │              │              │              │
+               ▼              ▼              ▼              ▼
+S0 -> Shared Core -> S1 -> Shared Core -> S2 -> ... -> Shared Core -> S12
+       同一套参数          同一套参数                   同一套参数
+```
+
+整体定义：
+
+\[
+\boxed{
+\text{一个共享 Event–Trajectory 推理核心}
+\times
+\text{多轮状态依赖推理}
+}
+\]

backbones.py

@@ -29,6 +29,14 @@ class TimeRoPE(nn.Module):
         x2 = x[..., 1::2]
         return torch.stack((-x2, x1), dim=-1).flatten(-2)
 
+    @staticmethod
+    def apply_single_from_cache(
+        x: torch.Tensor,
+        rope_cache: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        cos, sin = rope_cache
+        return x * cos + TimeRoPE._rotate_half(x) * sin
+
     @staticmethod
     def apply_from_cache(
         q: torch.Tensor,
@@ -85,166 +93,280 @@ class GaussianRBFTimeBasis(nn.Module):
         )
         return rbf_acts
 
+    def precompute_cross_cache(
+        self,
+        query_tau: torch.Tensor,
+        key_tau: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return RBF activations for query-time minus event-time."""
+        diff = query_tau.float().unsqueeze(2) - key_tau.float().unsqueeze(1)
+        widths = self.log_widths.exp()
+        return torch.exp(
+            -0.5
+            * (
+                (diff.unsqueeze(-1) - self.centers)
+                / widths
+            ).square()
+        )
+
+
+class TrajectoryCrossAttention(nn.Module):
+    """Shared trajectory-slot queries reading a fixed event memory."""
 
-class TemporalAttention(nn.Module):
     def __init__(
         self,
-        n_embd: int,
-        n_head: int,
+        d_model: int,
+        n_trajectory: int,
         n_rbf_bases: int = 16,
-        dropout: float = 0.0,
-        use_time_rope: bool = True,
-        use_rbf_bias: bool = True,
+        use_time_rope: bool = False,
+        use_rbf_bias: bool = False,
     ):
         super().__init__()
-        assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
-        self.n_head = n_head
-        self.d_head = n_embd // n_head
-        self.scale = 1.0 / math.sqrt(self.d_head)
+        if d_model <= 0 or n_trajectory <= 0:
+            raise ValueError("d_model and n_trajectory must be positive")
+        if d_model % n_trajectory != 0:
+            raise ValueError(
+                "d_model must be divisible by n_trajectory, got "
+                f"{d_model} and {n_trajectory}"
+            )
+        self.d_model = d_model
+        self.n_trajectory = n_trajectory
+        self.trajectory_dim = d_model // n_trajectory
+        self.scale = self.trajectory_dim ** -0.5
         self.use_time_rope = use_time_rope
         self.use_rbf_bias = use_rbf_bias
 
-        # QKV projection (fused for efficiency)
-        self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
-        # Output projection
-        self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
-
-        # Layer-specific projection from shared RBF basis activations to per-head attention bias.
-        self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
-        self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
-
-        self.resid_drop = nn.Dropout(dropout)
+        # q_proj acts on each slot independently and is shared across slots.
+        self.q_proj = nn.Linear(
+            self.trajectory_dim,
+            self.trajectory_dim,
+            bias=False,
+        )
+        self.k_proj = nn.Linear(d_model, d_model, bias=False)
+        self.v_proj = nn.Linear(d_model, d_model, bias=False)
+        if use_rbf_bias:
+            self.rbf_proj = nn.Linear(
+                n_rbf_bases,
+                n_trajectory,
+                bias=False,
+            )
+            self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
+        else:
+            self.rbf_proj = None
+            self.register_parameter("time_bias_scale", None)
         self.reset_parameters()
 
     def reset_parameters(self) -> None:
-        """Match the previous version's GPT-style weight initialization."""
-        nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02)
-        nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02)
-        nn.init.zeros_(self.rbf_proj.weight)
+        nn.init.normal_(self.q_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.k_proj.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.v_proj.weight, mean=0.0, std=0.02)
+        if self.rbf_proj is not None:
+            # The scalar gate starts at zero, so the relative-time bias still
+            # starts disabled. A nonzero projection is necessary for the gate
+            # itself to receive a gradient on the first optimization step.
+            nn.init.xavier_uniform_(self.rbf_proj.weight)
+
+    def project_event_memory(
+        self,
+        event_memory: torch.Tensor,
+        event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Project K/V once for reuse by every reasoning round."""
+        batch_size, memory_len, _ = event_memory.shape
+        key = self.k_proj(event_memory).reshape(
+            batch_size,
+            memory_len,
+            self.n_trajectory,
+            self.trajectory_dim,
+        ).transpose(1, 2)
+        value = self.v_proj(event_memory).reshape(
+            batch_size,
+            memory_len,
+            self.n_trajectory,
+            self.trajectory_dim,
+        ).transpose(1, 2)
+        if self.use_time_rope:
+            if event_rope_cache is None:
+                raise ValueError(
+                    "event_rope_cache is required when TimeRoPE is enabled"
+                )
+            key = TimeRoPE.apply_single_from_cache(key, event_rope_cache)
+        return key, value
 
     def forward(
         self,
-        x: torch.Tensor,
-        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        trajectory_state: torch.Tensor,
+        event_key_value: tuple[torch.Tensor, torch.Tensor],
+        event_invalid_mask: torch.Tensor,
+        query_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
         rbf_cache: torch.Tensor | None = None,
-        attn_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if self.use_time_rope:
-            assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True"
-        if self.use_rbf_bias:
-            assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True"
-
-        B, L, _ = x.shape
-        H, D = self.n_head, self.d_head
-
-        # --- QKV ----------------------------------------------------------
-        qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
-        q, k, v = qkv.unbind(0)  # each (B, H, L, D)
-
-        # --- Apply RoPE (from shared cache) --------------------------------
-        if self.use_time_rope:
-            q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
-
-        # Build additive attention bias mask: time bias + causal/padding mask.
-        time_bias = None
-        if self.use_rbf_bias:
-            time_bias = self.rbf_proj(rbf_cache).permute(
-                0, 3, 1, 2)  # (B, H, L, L)
-            time_bias = self.time_bias_scale.tanh() * time_bias
-
-        if time_bias is not None and attn_mask is not None:
-            attn_bias = time_bias + attn_mask.to(time_bias.dtype)
-        elif time_bias is not None:
-            attn_bias = time_bias
-        elif attn_mask is not None:
-            attn_bias = attn_mask
-        else:
-            attn_bias = None
-
-        out = F.scaled_dot_product_attention(
-            q,
-            k,
-            v,
-            attn_mask=attn_bias,
-            dropout_p=0.0,
-            is_causal=False,
-            scale=self.scale,
+        """Read memory for states shaped ``(B, Q, H, Dh)``."""
+        if trajectory_state.ndim != 4:
+            raise ValueError(
+                "trajectory_state must have shape (B, Q, H, Dh), got "
+                f"{tuple(trajectory_state.shape)}"
+            )
+        batch_size, n_query, n_trajectory, trajectory_dim = (
+            trajectory_state.shape
         )
+        if (n_trajectory, trajectory_dim) != (
+            self.n_trajectory,
+            self.trajectory_dim,
+        ):
+            raise ValueError(
+                "Unexpected trajectory shape: "
+                f"{(n_trajectory, trajectory_dim)}"
+            )
+        key, value = event_key_value
+        memory_len = key.size(2)
+        if event_invalid_mask.shape != (batch_size, n_query, memory_len):
+            raise ValueError(
+                "event_invalid_mask must have shape "
+                f"{(batch_size, n_query, memory_len)}, got "
+                f"{tuple(event_invalid_mask.shape)}"
+            )
 
-        # --- Aggregate & project out --------------------------------------
-        out = out.transpose(1, 2).reshape(B, L, H * D)
-        return self.resid_drop(self.out_proj(out))
+        query = self.q_proj(trajectory_state).transpose(1, 2)
+        if self.use_time_rope:
+            if query_rope_cache is None:
+                raise ValueError(
+                    "query_rope_cache is required when TimeRoPE is enabled"
+                )
+            query = TimeRoPE.apply_single_from_cache(query, query_rope_cache)
+        query = query.transpose(1, 2)
+
+        scores = torch.einsum("bqhd,bhld->bqhl", query, key) * self.scale
+        if self.use_rbf_bias:
+            if rbf_cache is None or self.rbf_proj is None:
+                raise ValueError(
+                    "rbf_cache is required when relative time bias is enabled"
+                )
+            time_bias = self.rbf_proj(rbf_cache).permute(0, 1, 3, 2)
+            scores = scores + self.time_bias_scale.tanh() * time_bias
+
+        mask = event_invalid_mask.unsqueeze(2)
+        min_value = torch.finfo(scores.dtype).min
+        masked_scores = scores.masked_fill(mask, min_value)
+        weights = torch.softmax(masked_scores.float(), dim=-1).to(scores.dtype)
+        weights = weights.masked_fill(mask, 0.0)
+        denominator = weights.sum(dim=-1, keepdim=True)
+        weights = weights / denominator.clamp_min(
+            torch.finfo(weights.dtype).eps
+        )
+        return torch.einsum("bqhl,bhld->bqhd", weights, value)
 
 
-class SwiGLU(nn.Module):
-    def __init__(
-        self,
-        n_embd: int,
-        hidden_dim: int | None = None,
-        dropout: float = 0.0,
-        bias: bool = True,
-    ):
+class SharedTrajectoryMixer(nn.Module):
+    """SwiGLU interaction along the trajectory axis only."""
+
+    def __init__(self, n_trajectory: int, trajectory_dim: int):
         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)
+        if n_trajectory <= 0 or trajectory_dim <= 0:
+            raise ValueError("trajectory dimensions must be positive")
+        self.n_trajectory = n_trajectory
+        self.trajectory_dim = trajectory_dim
+        self.traj_hidden = 4 * n_trajectory
+        self.gate_proj = nn.Parameter(
+            torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
+        )
+        self.value_proj = nn.Parameter(
+            torch.empty(trajectory_dim, n_trajectory, self.traj_hidden)
+        )
+        self.output_proj = nn.Parameter(
+            torch.empty(trajectory_dim, self.traj_hidden, n_trajectory)
+        )
         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)
+        for feature_idx in range(self.trajectory_dim):
+            nn.init.xavier_uniform_(self.gate_proj[feature_idx])
+            nn.init.xavier_uniform_(self.value_proj[feature_idx])
+        nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """``(B, L, n_embd) -> (B, L, n_embd)``."""
-        return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
+    def forward(self, state: torch.Tensor) -> torch.Tensor:
+        if state.shape[-2:] != (self.n_trajectory, self.trajectory_dim):
+            raise ValueError(
+                "Expected trailing trajectory shape "
+                f"{(self.n_trajectory, self.trajectory_dim)}, got "
+                f"{tuple(state.shape[-2:])}"
+            )
+        gate = torch.einsum("...hr,rhk->...kr", state, self.gate_proj)
+        value = torch.einsum("...hr,rhk->...kr", state, self.value_proj)
+        hidden = F.silu(gate) * value
+        return torch.einsum("...kr,rkh->...hr", hidden, self.output_proj)
 
 
-class GPTBlock(nn.Module):
+class SharedEventTrajectoryCore(nn.Module):
+    """One parameter-shared reasoning core reused across all rounds."""
+
     def __init__(
         self,
-        n_embd: int,
-        n_head: int,
-
-        attn_dropout: float = 0.0,
-        mlp_dropout: float = 0.0,
+        d_model: int,
+        n_trajectory: int,
+        n_reasoning_rounds: int,
+        dropout: float = 0.0,
+        n_rbf_bases: int = 16,
         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,
+        if n_reasoning_rounds <= 0:
+            raise ValueError("n_reasoning_rounds must be positive")
+        if d_model <= 0 or n_trajectory <= 0:
+            raise ValueError("d_model and n_trajectory must be positive")
+        if d_model % n_trajectory != 0:
+            raise ValueError("d_model must equal n_trajectory * trajectory_dim")
+        trajectory_dim = d_model // n_trajectory
+        self.norm_attn = nn.LayerNorm(trajectory_dim)
+        self.cross_attention = TrajectoryCrossAttention(
+            d_model=d_model,
+            n_trajectory=n_trajectory,
             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)
+        self.norm_mixer = nn.LayerNorm(trajectory_dim)
+        self.traj_mixer = SharedTrajectoryMixer(
+            n_trajectory=n_trajectory,
+            trajectory_dim=trajectory_dim,
+        )
+        initial_scale = 1.0 / math.sqrt(n_reasoning_rounds)
+        self.attn_scale = nn.Parameter(torch.tensor(initial_scale))
+        self.mixer_scale = nn.Parameter(torch.tensor(initial_scale))
+        self.dropout = nn.Dropout(dropout)
+
+    def project_event_memory(
+        self,
+        event_memory: torch.Tensor,
+        event_rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        return self.cross_attention.project_event_memory(
+            event_memory,
+            event_rope_cache=event_rope_cache,
+        )
 
     def forward(
         self,
-        x: torch.Tensor,
-        rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
+        trajectory_state: torch.Tensor,
+        event_key_value: tuple[torch.Tensor, torch.Tensor],
+        event_invalid_mask: torch.Tensor,
+        query_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
+        readout = self.cross_attention(
+            trajectory_state=self.norm_attn(trajectory_state),
+            event_key_value=event_key_value,
+            event_invalid_mask=event_invalid_mask,
+            query_rope_cache=query_rope_cache,
+            rbf_cache=rbf_cache,
+        )
+        updated = (
+            trajectory_state
+            + self.attn_scale * self.dropout(readout)
+        )
+        mixed = self.traj_mixer(self.norm_mixer(updated))
+        return updated + self.mixer_scale * self.dropout(mixed)
 
 
 class TokenAutoDiscretization(nn.Module):

evaluate_auc.py

@@ -40,7 +40,11 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset
 from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
-from models import DeepHealth
+from models import (
+    DeepHealth,
+    validate_event_trajectory_config,
+    validate_event_trajectory_state_dict,
+)
 from readouts import build_readout
 from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
 
@@ -309,6 +313,7 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
+    validate_event_trajectory_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"}:
@@ -317,10 +322,10 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         )
     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_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
+        model_size=str(cfg_get(args, cfg, "model_size", "nano")),
+        n_reasoning_rounds=int(
+            cfg_get(args, cfg, "n_reasoning_rounds", 12)
+        ),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -386,6 +391,12 @@ def load_model_state(
     state = state_dict if state_dict is not None else load_checkpoint_state_dict(
         checkpoint_path, map_location=device)
 
+    validate_event_trajectory_state_dict(
+        state,
+        expected_d_model=model.d_model,
+        expected_n_trajectory=model.n_trajectory,
+        expected_n_reasoning_rounds=model.n_reasoning_rounds,
+    )
     model.load_state_dict(state, strict=True)
 
 
@@ -522,7 +533,7 @@ def infer_readout_hidden(
                 hidden = torch.zeros(
                     batch_size,
                     seq_len,
-                    model.n_embd,
+                    model.d_model,
                     device=event_seq.device,
                     dtype=torch.float32,
                 )

evaluate_auc_v2.py

@@ -29,7 +29,11 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset
 from eval_data import load_sequence_eval_dataset
-from models import DeepHealth
+from models import (
+    DeepHealth,
+    validate_event_trajectory_config,
+    validate_event_trajectory_state_dict,
+)
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
 
@@ -178,6 +182,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
 
 
 def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
+    validate_event_trajectory_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"}:
@@ -186,10 +191,10 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
         )
     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_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
+        model_size=str(cfg_get(args, cfg, "model_size", "nano")),
+        n_reasoning_rounds=int(
+            cfg_get(args, cfg, "n_reasoning_rounds", 12)
+        ),
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -204,6 +209,12 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
 
 
 def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
+    validate_event_trajectory_state_dict(
+        state_dict,
+        expected_d_model=model.d_model,
+        expected_n_trajectory=model.n_trajectory,
+        expected_n_reasoning_rounds=model.n_reasoning_rounds,
+    )
     model.load_state_dict(state_dict, strict=True)
 
 

export_tquery_logits_hidden.py

@@ -205,7 +205,7 @@ def main() -> None:
 
     n_rows = len(landmark_dataset)
     vocab_size = int(dataset.vocab_size)
-    hidden_dim = int(getattr(model, "n_embd", cfg_get(args, cfg_model, "n_embd", 120)))
+    hidden_dim = int(model.d_model)
     logits_dtype = numpy_float_dtype(args.logits_dtype)
     hidden_dtype = numpy_float_dtype(args.hidden_dtype)
 

models.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from dataclasses import dataclass
 
 import torch
@@ -6,14 +7,200 @@ import torch.nn.functional as F
 
 from backbones import (
     AgeSinusoidalEncoding,
-    GPTBlock,
     GaussianRBFTimeBasis,
+    SharedEventTrajectoryCore,
     TimeRoPE,
     TokenAutoDiscretization,
 )
 from targets import PAD_IDX
 
 
+EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"
+
+
+@dataclass(frozen=True)
+class EventTrajectoryModelSize:
+    d_model: int
+    n_trajectory: int
+
+    @property
+    def trajectory_dim(self) -> int:
+        return self.d_model // self.n_trajectory
+
+    @property
+    def traj_hidden(self) -> int:
+        return 4 * self.n_trajectory
+
+
+MODEL_SIZE_PRESETS = {
+    "nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
+    "tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
+    "small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
+    "medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
+    "huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),
+}
+MODEL_SIZE_NAMES = tuple(MODEL_SIZE_PRESETS)
+
+
+def resolve_model_size(model_size: str) -> EventTrajectoryModelSize:
+    if not isinstance(model_size, str):
+        raise ValueError(
+            f"model_size must be a string, got {type(model_size).__name__}"
+        )
+    normalized = model_size.strip().lower()
+    try:
+        return MODEL_SIZE_PRESETS[normalized]
+    except KeyError as exc:
+        choices = ", ".join(MODEL_SIZE_NAMES)
+        raise ValueError(
+            f"Unknown model_size {model_size!r}; expected one of: {choices}"
+        ) from exc
+
+
+def _required_config_int(
+    config: Mapping[str, object],
+    key: str,
+) -> int:
+    raw_value = config.get(key)
+    if isinstance(raw_value, bool):
+        raise ValueError(f"Config field {key!r} must be an integer")
+    try:
+        value = int(raw_value)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(
+            f"Config field {key!r} must be present and integer-valued; "
+            f"got {raw_value!r}"
+        ) from exc
+    if isinstance(raw_value, float) and not raw_value.is_integer():
+        raise ValueError(f"Config field {key!r} must be an integer")
+    return value
+
+
+def validate_event_trajectory_config(config: Mapping[str, object]) -> None:
+    actual = config.get("model_architecture")
+    if actual != EVENT_TRAJECTORY_ARCHITECTURE:
+        raise ValueError(
+            "This branch only accepts models trained with the shared "
+            "event-trajectory architecture marker "
+            f"{EVENT_TRAJECTORY_ARCHITECTURE!r}; got {actual!r}."
+        )
+    raw_model_size = config.get("model_size")
+    if not isinstance(raw_model_size, str):
+        raise ValueError(
+            "Config field 'model_size' must be one of: "
+            + ", ".join(MODEL_SIZE_NAMES)
+        )
+    model_size = raw_model_size.strip().lower()
+    preset = resolve_model_size(model_size)
+    d_model = _required_config_int(config, "d_model")
+    n_trajectory = _required_config_int(config, "n_trajectory")
+    n_reasoning_rounds = _required_config_int(
+        config,
+        "n_reasoning_rounds",
+    )
+    trajectory_dim = _required_config_int(config, "trajectory_dim")
+    traj_hidden = _required_config_int(config, "traj_hidden")
+    if n_reasoning_rounds <= 0:
+        raise ValueError(
+            "n_reasoning_rounds must be positive"
+        )
+    expected_values = {
+        "d_model": preset.d_model,
+        "n_trajectory": preset.n_trajectory,
+        "trajectory_dim": preset.trajectory_dim,
+        "traj_hidden": preset.traj_hidden,
+    }
+    actual_values = {
+        "d_model": d_model,
+        "n_trajectory": n_trajectory,
+        "trajectory_dim": trajectory_dim,
+        "traj_hidden": traj_hidden,
+    }
+    mismatches = [
+        f"{key}: expected {expected}, got {actual_values[key]}"
+        for key, expected in expected_values.items()
+        if actual_values[key] != expected
+    ]
+    if mismatches:
+        raise ValueError(
+            f"Config does not match model_size={model_size!r}: "
+            + "; ".join(mismatches)
+        )
+
+
+def _checkpoint_scalar_int(
+    state_dict: Mapping[str, object],
+    key: str,
+) -> int:
+    value = state_dict[key]
+    if not isinstance(value, torch.Tensor) or value.numel() != 1:
+        raise ValueError(
+            f"Checkpoint architecture field {key!r} must be a scalar tensor"
+        )
+    return int(value.detach().cpu().item())
+
+
+def validate_event_trajectory_state_dict(
+    state_dict: Mapping[str, object],
+    *,
+    expected_d_model: int | None = None,
+    expected_n_trajectory: int | None = None,
+    expected_n_reasoning_rounds: int | None = None,
+) -> None:
+    required_keys = {
+        "architecture_d_model",
+        "architecture_n_trajectory",
+        "architecture_n_reasoning_rounds",
+        "event_projection.weight",
+        "trajectory_prototypes",
+        "query_projection.weight",
+        "reasoning_core.cross_attention.q_proj.weight",
+        "reasoning_core.cross_attention.k_proj.weight",
+        "reasoning_core.cross_attention.v_proj.weight",
+        "reasoning_core.traj_mixer.gate_proj",
+        "reasoning_core.traj_mixer.value_proj",
+        "reasoning_core.traj_mixer.output_proj",
+        "reasoning_core.attn_scale",
+        "reasoning_core.mixer_scale",
+    }
+    missing = sorted(required_keys.difference(state_dict))
+    if missing:
+        raise ValueError(
+            "Checkpoint is not a shared event-trajectory checkpoint; "
+            "missing required "
+            f"parameters: {', '.join(missing)}"
+        )
+    checkpoint_values = {
+        "d_model": _checkpoint_scalar_int(
+            state_dict,
+            "architecture_d_model",
+        ),
+        "n_trajectory": _checkpoint_scalar_int(
+            state_dict,
+            "architecture_n_trajectory",
+        ),
+        "n_reasoning_rounds": _checkpoint_scalar_int(
+            state_dict,
+            "architecture_n_reasoning_rounds",
+        ),
+    }
+    expected_values = {
+        "d_model": expected_d_model,
+        "n_trajectory": expected_n_trajectory,
+        "n_reasoning_rounds": expected_n_reasoning_rounds,
+    }
+    mismatches = [
+        f"{name}: checkpoint={checkpoint_values[name]}, expected={expected}"
+        for name, expected in expected_values.items()
+        if expected is not None and checkpoint_values[name] != expected
+    ]
+    if mismatches:
+        raise ValueError(
+            "Checkpoint architecture does not match the constructed model: "
+            + "; ".join(mismatches)
+        )
+
+
 @dataclass
 class DeepHealthOutput:
     hidden: torch.Tensor
@@ -145,10 +332,8 @@ class DeepHealth(nn.Module):
     def __init__(
         self,
         vocab_size: int,
-        n_embd: int,
-        n_head: int,
-        n_hist_layer: int,
-        n_tab_layer: int,
+        model_size: str,
+        n_reasoning_rounds: int,
         n_types: int,
         n_cont_types: int,
         n_categories: int,
@@ -173,11 +358,21 @@ 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'")
-        self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
+        if n_reasoning_rounds <= 0:
+            raise ValueError(
+                "n_reasoning_rounds must be positive, got "
+                f"{n_reasoning_rounds}"
+            )
+        size_config = resolve_model_size(model_size)
+        normalized_model_size = model_size.strip().lower()
+        d_model = size_config.d_model
+        n_trajectory = size_config.n_trajectory
+
+        self.token_embedding = nn.Embedding(vocab_size, d_model, padding_idx=0)
         self.gender_embedding = nn.Embedding(
-            2, n_embd)  # Assuming binary gender
+            2, d_model)  # Assuming binary gender
         self.tokenizer = OtherInfoTokenizer(
-            n_embd=n_embd,
+            n_embd=d_model,
             n_types=n_types,
             n_cont_types=n_cont_types,
             n_categories=n_categories,
@@ -189,70 +384,101 @@ class DeepHealth(nn.Module):
         self.time_mode = time_mode
         self.dist_mode = dist_mode
         self.extra_pool_reduce = extra_pool_reduce
-        self.n_embd = n_embd
+        self.model_size = normalized_model_size
+        self.d_model = d_model
+        self.n_trajectory = n_trajectory
+        self.trajectory_dim = d_model // n_trajectory
+        self.traj_hidden = 4 * n_trajectory
+        self.n_reasoning_rounds = n_reasoning_rounds
         self.vocab_size = vocab_size
+        self.register_buffer(
+            "architecture_d_model",
+            torch.tensor(d_model, dtype=torch.int64),
+        )
+        self.register_buffer(
+            "architecture_n_trajectory",
+            torch.tensor(n_trajectory, dtype=torch.int64),
+        )
+        self.register_buffer(
+            "architecture_n_reasoning_rounds",
+            torch.tensor(n_reasoning_rounds, dtype=torch.int64),
+        )
         nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
         nn.init.zeros_(self.token_embedding.weight[0])
         nn.init.normal_(self.gender_embedding.weight, mean=0.0, std=0.02)
         if dist_mode == "weibull":
-            self.rho_head = nn.Linear(n_embd, vocab_size)
+            self.rho_head = nn.Linear(d_model, vocab_size)
             nn.init.zeros_(self.rho_head.weight)
             nn.init.constant_(self.rho_head.bias, 0.5413)
 
         if dist_mode == "mixed":
             self.death_idx = vocab_size - 1
-            self.rho_death_head = nn.Linear(n_embd, 1)
+            self.rho_death_head = nn.Linear(d_model, 1)
             nn.init.zeros_(self.rho_death_head.weight)
             nn.init.constant_(self.rho_death_head.bias, 0.5413)
 
-        if time_mode == "absolute":
-            self.age_encoding = AgeSinusoidalEncoding(n_embd)
-            self.blocks = nn.ModuleList([
-                GPTBlock(
-                    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)
-            ])
-            self.rope = None
-            self.rbf = None
-        elif time_mode == "relative":
-            self.age_encoding = None
-            self.blocks = nn.ModuleList([
-                GPTBlock(
-                    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)
-            ])
-            self.rope = TimeRoPE(n_embd // n_head)
-            self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)
-
-        self.final_ln = nn.LayerNorm(n_embd)
-        self.risk_head = nn.Linear(n_embd, vocab_size, bias=False)
-        if target_mode == "next_token":
-            self.risk_head.weight = self.token_embedding.weight
-        self.query_token = nn.Parameter(torch.zeros(n_embd))
+        # Event and query time are encoded once before shared reasoning. In
+        # relative mode, cross-attention additionally uses TimeRoPE and RBF.
+        self.age_encoding = AgeSinusoidalEncoding(d_model)
+        self.event_projection = nn.Linear(d_model, d_model, bias=False)
+        self.event_norm = nn.LayerNorm(d_model)
+        self.query_projection = nn.Linear(d_model, d_model, bias=False)
+        self.trajectory_prototypes = nn.Parameter(
+            torch.empty(n_trajectory, self.trajectory_dim)
+        )
+        self.query_token = nn.Parameter(torch.empty(d_model))
+        nn.init.normal_(self.event_projection.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.query_projection.weight, mean=0.0, std=0.02)
+        nn.init.normal_(self.trajectory_prototypes, mean=0.0, std=0.02)
         nn.init.normal_(self.query_token, mean=0.0, std=0.02)
 
-    def _make_history_attn_mask(
+        use_relative_time = time_mode == "relative"
+        self.reasoning_core = SharedEventTrajectoryCore(
+            d_model=d_model,
+            n_trajectory=n_trajectory,
+            n_reasoning_rounds=n_reasoning_rounds,
+            dropout=dropout,
+            n_rbf_bases=16,
+            use_time_rope=use_relative_time,
+            use_rbf_bias=use_relative_time,
+        )
+        if use_relative_time:
+            self.rope = TimeRoPE(self.trajectory_dim)
+            self.rbf = GaussianRBFTimeBasis(
+                n_bases=16,
+                max_time_diff=40.0,
+            )
+        else:
+            self.rope = None
+            self.rbf = None
+
+        self.final_ln = nn.LayerNorm(d_model)
+        self.risk_head = nn.Linear(d_model, vocab_size, bias=False)
+        if target_mode == "next_token":
+            self.risk_head.weight = self.token_embedding.weight
+
+    def _make_event_invalid_mask(
         self,
-        padding_mask: torch.Tensor,
-        time_seq: torch.Tensor,
-        dtype: torch.dtype,
+        event_valid_mask: torch.Tensor,
+        event_time: torch.Tensor,
+        query_time: torch.Tensor,
+        query_position: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        valid_key = padding_mask[:, None, :]  # (B, 1, L)
-        visible_by_time = time_seq[:, None, :] <= time_seq[:, :, None]
-        valid = valid_key & visible_by_time
-        return torch.zeros(
-            valid.shape,
-            device=valid.device,
-            dtype=dtype,
-        ).masked_fill(~valid, -1e4)[:, None, :, :]
+        valid_key = event_valid_mask[:, None, :]
+        key_time = event_time[:, None, :]
+        query_time = query_time[:, :, None]
+        if query_position is None:
+            visible_by_time = key_time <= query_time
+        else:
+            key_position = torch.arange(
+                event_time.size(1),
+                device=event_time.device,
+            ).view(1, 1, -1)
+            visible_by_time = (key_time < query_time) | (
+                (key_time == query_time)
+                & (key_position <= query_position[:, :, None])
+            )
+        return ~(valid_key & visible_by_time)
 
     def _pool_other_by_time(
         self,
@@ -356,8 +582,8 @@ class DeepHealth(nn.Module):
             padding_mask = padding_mask.to(device=event_seq.device, dtype=torch.bool)
 
         event_len = event_seq.size(1)
-        h_disease = self.token_embedding(event_seq)
-        t_disease = time_seq
+        event_features = self.token_embedding(event_seq)
+        event_time = time_seq
 
         if other_time.shape != other_type.shape:
             raise ValueError(
@@ -365,64 +591,120 @@ class DeepHealth(nn.Module):
                 f"{tuple(other_time.shape)} vs {tuple(other_type.shape)}"
             )
         other_time = other_time.to(device=event_seq.device, dtype=time_seq.dtype)
-        h_other, other_mask = self.tokenizer(
+        other_features, other_mask = self.tokenizer(
             other_type=other_type,
             other_value=other_value,
             other_value_kind=other_value_kind,
         )
-        h_other = h_other.to(device=event_seq.device)
+        other_features = other_features.to(device=event_seq.device)
         other_mask = other_mask.to(device=event_seq.device, dtype=torch.bool)
 
-        h_disease = torch.cat([h_disease, h_other], dim=1)
-        t_disease = torch.cat([t_disease, other_time], dim=1)
-        padding_mask = torch.cat([padding_mask, other_mask], dim=1)
-        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+        event_features = torch.cat([event_features, other_features], dim=1)
+        event_time = torch.cat([event_time, other_time], dim=1)
+        event_valid_mask = torch.cat([padding_mask, other_mask], dim=1)
+
+        batch_size = event_seq.size(0)
+        sex_context = self.gender_embedding(sex)[:, None, :]
+        event_features = (
+            event_features
+            + sex_context
+            + self.age_encoding(event_time)
+        )
+        event_features = event_features * event_valid_mask.unsqueeze(-1).to(
+            event_features.dtype
+        )
+        event_memory = self.event_norm(
+            self.event_projection(event_features)
+        )
+        event_memory = event_memory * event_valid_mask.unsqueeze(-1).to(
+            event_memory.dtype
+        )
 
         if mode == "all_future":
-            batch_size = event_seq.size(0)
-            query = self.query_token.view(1, 1, -1).expand(batch_size, 1, -1)
-            h_disease = torch.cat([h_disease, query], dim=1)
-            t_disease = torch.cat([t_disease, t_query[:, None]], dim=1)
-            query_mask = torch.ones(
+            query_time = t_query[:, None]
+            query_position = None
+            query_features = (
+                self.query_token.view(1, 1, -1)
+                + sex_context
+                + self.age_encoding(query_time)
+            )
+            query_valid_mask = torch.ones(
                 batch_size,
                 1,
                 dtype=torch.bool,
                 device=event_seq.device,
             )
-            padding_mask = torch.cat([padding_mask, query_mask], dim=1)
+        else:
+            # Each event position is an independent parallel query. Including
+            # its event feature preserves token-level next-step semantics.
+            # Equal-time memory is additionally position-causal so a token
+            # cannot read a later token that may be its Delphi2M target.
+            query_time = event_time
+            query_position = torch.arange(
+                event_time.size(1),
+                device=event_time.device,
+            ).view(1, -1).expand(batch_size, -1)
+            query_features = event_features
+            query_valid_mask = event_valid_mask
 
-        sex_emb = self.gender_embedding(sex)[:, None, :]
-        h_disease = h_disease + sex_emb
-        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-
-        rope_cache = None
-        rbf_cache = None
-        if self.time_mode == "absolute":
-            h_disease = h_disease + self.age_encoding(t_disease)
-            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
-        elif self.time_mode == "relative":
-            rope_cache = self.rope.precompute_cache(t_disease)
-            rbf_cache = self.rbf.precompute_cache(t_disease)
-
-        attn_mask = self._make_history_attn_mask(
-            padding_mask=padding_mask,
-            time_seq=t_disease,
-            dtype=h_disease.dtype,
+        n_query = query_time.size(1)
+        query_context = self.query_projection(query_features).reshape(
+            batch_size,
+            n_query,
+            self.n_trajectory,
+            self.trajectory_dim,
         )
-        for block in self.blocks:
-            h_disease = block(
-                h_disease,
-                rope_cache=rope_cache,
-                rbf_cache=rbf_cache,
-                attn_mask=attn_mask,
+        trajectory_state = (
+            self.trajectory_prototypes.view(
+                1,
+                1,
+                self.n_trajectory,
+                self.trajectory_dim,
             )
-            h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+            + query_context
+        )
+        event_invalid_mask = self._make_event_invalid_mask(
+            event_valid_mask=event_valid_mask,
+            event_time=event_time,
+            query_time=query_time,
+            query_position=query_position,
+        )
 
-        h_disease = self.final_ln(h_disease)
-        h_disease = h_disease * padding_mask.unsqueeze(-1).to(h_disease.dtype)
+        event_rope_cache = None
+        query_rope_cache = None
+        rbf_cache = None
+        if self.time_mode == "relative":
+            if self.rope is None or self.rbf is None:
+                raise RuntimeError("Relative-time modules are not initialized")
+            event_rope_cache = self.rope.precompute_cache(event_time)
+            query_rope_cache = self.rope.precompute_cache(query_time)
+            rbf_cache = self.rbf.precompute_cross_cache(
+                query_time,
+                event_time,
+            )
+
+        event_key_value = self.reasoning_core.project_event_memory(
+            event_memory,
+            event_rope_cache=event_rope_cache,
+        )
+        for _ in range(self.n_reasoning_rounds):
+            trajectory_state = self.reasoning_core(
+                trajectory_state=trajectory_state,
+                event_key_value=event_key_value,
+                event_invalid_mask=event_invalid_mask,
+                query_rope_cache=query_rope_cache,
+                rbf_cache=rbf_cache,
+            )
+
+        hidden_sequence = self.final_ln(
+            trajectory_state.reshape(batch_size, n_query, self.d_model)
+        )
+        hidden_sequence = hidden_sequence * query_valid_mask.unsqueeze(-1).to(
+            hidden_sequence.dtype
+        )
 
         if mode == "all_future":
-            hidden = h_disease[:, -1, :]
+            hidden = hidden_sequence[:, 0, :]
             if return_output:
                 return DeepHealthOutput(
                     hidden=hidden,
@@ -437,13 +719,13 @@ class DeepHealth(nn.Module):
                 )
             return hidden
         if return_output:
-            h_event = h_disease[:, :event_len, :]
-            t_event = t_disease[:, :event_len]
-            event_mask = padding_mask[:, :event_len]
+            h_event = hidden_sequence[:, :event_len, :]
+            t_event = event_time[:, :event_len]
+            event_mask = event_valid_mask[:, :event_len]
             h_extra, t_extra, extra_mask = self._pool_other_by_time(
-                h_other=h_disease[:, event_len:, :],
-                other_time=t_disease[:, event_len:],
-                other_mask=padding_mask[:, event_len:],
+                h_other=hidden_sequence[:, event_len:, :],
+                other_time=event_time[:, event_len:],
+                other_mask=event_valid_mask[:, event_len:],
             )
             return DeepHealthOutput(
                 hidden=torch.cat([h_event, h_extra], dim=1),
@@ -451,7 +733,7 @@ class DeepHealth(nn.Module):
                 padding_mask=torch.cat([event_mask, extra_mask], dim=1),
                 event_len=event_len,
             )
-        return h_disease[:, :event_len, :]
+        return hidden_sequence[:, :event_len, :]
 
     def forward_next_token(self, **kwargs) -> torch.Tensor:
         return self._forward_shared(mode="next_token", **kwargs)

test_event_trajectory_backbone.py

@@ -0,0 +1,330 @@
+import math
+import unittest
+
+import torch
+
+from backbones import (
+    SharedEventTrajectoryCore,
+    SharedTrajectoryMixer,
+    TrajectoryCrossAttention,
+)
+from models import (
+    EVENT_TRAJECTORY_ARCHITECTURE,
+    MODEL_SIZE_PRESETS,
+    DeepHealth,
+    resolve_model_size,
+    validate_event_trajectory_config,
+    validate_event_trajectory_state_dict,
+)
+from train_util import get_model_parameter_counts
+
+
+def build_test_model(
+    *,
+    target_mode: str = "next_token",
+    time_mode: str = "absolute",
+    n_reasoning_rounds: int = 3,
+) -> DeepHealth:
+    return DeepHealth(
+        vocab_size=32,
+        model_size="nano",
+        n_reasoning_rounds=n_reasoning_rounds,
+        n_types=2,
+        n_cont_types=0,
+        n_categories=2,
+        cont_type_ids=[],
+        target_mode=target_mode,
+        time_mode=time_mode,
+    )
+
+
+def model_inputs() -> dict[str, torch.Tensor]:
+    return {
+        "event_seq": torch.tensor([[1, 2, 3, 4], [5, 6, 0, 0]]),
+        "time_seq": torch.tensor(
+            [[1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 0.0, 0.0]]
+        ),
+        "sex": torch.tensor([0, 1]),
+        "padding_mask": torch.tensor(
+            [[True, True, True, True], [True, True, False, False]]
+        ),
+        "other_type": torch.zeros(2, 1, dtype=torch.long),
+        "other_value": torch.zeros(2, 1),
+        "other_value_kind": torch.zeros(2, 1, dtype=torch.long),
+        "other_time": torch.zeros(2, 1),
+    }
+
+
+class EventTrajectoryBackboneTest(unittest.TestCase):
+    def test_model_size_presets(self) -> None:
+        expected = {
+            "nano": (120, 6, 20, 24),
+            "tiny": (256, 8, 32, 32),
+            "small": (512, 8, 64, 32),
+            "medium": (768, 12, 64, 48),
+            "huge": (1024, 16, 64, 64),
+        }
+        self.assertEqual(set(MODEL_SIZE_PRESETS), set(expected))
+        for name, values in expected.items():
+            preset = resolve_model_size(name)
+            self.assertEqual(
+                (
+                    preset.d_model,
+                    preset.n_trajectory,
+                    preset.trajectory_dim,
+                    preset.traj_hidden,
+                ),
+                values,
+            )
+
+    def test_default_mixer_shapes_and_parameter_count(self) -> None:
+        mixer = SharedTrajectoryMixer(
+            n_trajectory=8,
+            trajectory_dim=32,
+        )
+        state = torch.randn(2, 5, 8, 32)
+        self.assertEqual(mixer(state).shape, state.shape)
+        self.assertEqual(mixer.traj_hidden, 32)
+        self.assertEqual(tuple(mixer.gate_proj.shape), (32, 8, 32))
+        self.assertEqual(tuple(mixer.value_proj.shape), (32, 8, 32))
+        self.assertEqual(tuple(mixer.output_proj.shape), (32, 32, 8))
+        self.assertEqual(
+            get_model_parameter_counts(mixer),
+            {
+                "model_parameter_count": 24_576,
+                "trainable_parameter_count": 24_576,
+            },
+        )
+
+    def test_reasoning_rounds_share_one_core_parameter_set(self) -> None:
+        core_one = SharedEventTrajectoryCore(
+            d_model=256,
+            n_trajectory=8,
+            n_reasoning_rounds=1,
+        )
+        core_twelve = SharedEventTrajectoryCore(
+            d_model=256,
+            n_trajectory=8,
+            n_reasoning_rounds=12,
+        )
+        self.assertEqual(
+            sum(p.numel() for p in core_one.parameters()),
+            sum(p.numel() for p in core_twelve.parameters()),
+        )
+        self.assertAlmostEqual(core_one.attn_scale.item(), 1.0)
+        self.assertAlmostEqual(
+            core_twelve.attn_scale.item(),
+            1.0 / math.sqrt(12),
+            places=6,
+        )
+        self.assertAlmostEqual(
+            core_twelve.mixer_scale.item(),
+            1.0 / math.sqrt(12),
+            places=6,
+        )
+
+    def test_all_masked_attention_is_finite_and_zero(self) -> None:
+        attention = TrajectoryCrossAttention(
+            d_model=32,
+            n_trajectory=4,
+        )
+        memory = torch.randn(2, 3, 32)
+        key_value = attention.project_event_memory(memory)
+        state = torch.randn(2, 2, 4, 8)
+        invalid_mask = torch.ones(2, 2, 3, dtype=torch.bool)
+        output = attention(
+            trajectory_state=state,
+            event_key_value=key_value,
+            event_invalid_mask=invalid_mask,
+        )
+        self.assertTrue(torch.isfinite(output).all())
+        torch.testing.assert_close(output, torch.zeros_like(output))
+
+    def test_next_token_future_events_do_not_change_earlier_query(self) -> None:
+        torch.manual_seed(0)
+        model = build_test_model(n_reasoning_rounds=2)
+        model.eval()
+        inputs = model_inputs()
+        original = model(**inputs)
+        changed_inputs = dict(inputs)
+        changed_inputs["event_seq"] = inputs["event_seq"].clone()
+        changed_inputs["event_seq"][0, 3] = 9
+        changed = model(**changed_inputs)
+        torch.testing.assert_close(original[0, 1], changed[0, 1])
+
+    def test_next_token_later_equal_time_event_is_not_visible(self) -> None:
+        torch.manual_seed(0)
+        model = build_test_model(n_reasoning_rounds=2)
+        model.eval()
+        inputs = model_inputs()
+        inputs["time_seq"] = inputs["time_seq"].clone()
+        inputs["time_seq"][0] = torch.tensor([1.0, 1.0, 2.0, 3.0])
+        original = model(**inputs)
+        changed_inputs = dict(inputs)
+        changed_inputs["event_seq"] = inputs["event_seq"].clone()
+        changed_inputs["event_seq"][0, 1] = 9
+        changed = model(**changed_inputs)
+        torch.testing.assert_close(original[0, 0], changed[0, 0])
+
+    def test_padding_content_does_not_change_valid_queries(self) -> None:
+        torch.manual_seed(0)
+        model = build_test_model(
+            time_mode="relative",
+            n_reasoning_rounds=2,
+        )
+        model.eval()
+        inputs = model_inputs()
+        original = model(**inputs)
+        changed_inputs = dict(inputs)
+        changed_inputs["event_seq"] = inputs["event_seq"].clone()
+        changed_inputs["time_seq"] = inputs["time_seq"].clone()
+        changed_inputs["event_seq"][1, 2:] = torch.tensor([9, 10])
+        changed_inputs["time_seq"][1, 2:] = torch.tensor([30.0, 40.0])
+        changed = model(**changed_inputs)
+        torch.testing.assert_close(original[1, :2], changed[1, :2])
+
+    def test_next_token_and_all_future_output_contracts(self) -> None:
+        inputs = model_inputs()
+        next_model = build_test_model(target_mode="next_token")
+        next_hidden = next_model(**inputs)
+        self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
+        next_output = next_model(**inputs, return_output=True)
+        self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
+        self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
+
+        future_model = build_test_model(target_mode="all_future")
+        future_hidden = future_model(
+            **inputs,
+            t_query=torch.tensor([5.0, 3.0]),
+        )
+        self.assertEqual(tuple(future_hidden.shape), (2, 120))
+
+    def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
+        model = build_test_model(n_reasoning_rounds=12)
+        self.assertFalse(hasattr(model, "blocks"))
+        reasoning_keys = [
+            key
+            for key in model.state_dict()
+            if key.startswith("reasoning_core.")
+        ]
+        self.assertTrue(reasoning_keys)
+        self.assertFalse(any("blocks." in key for key in model.state_dict()))
+        self.assertFalse(any("out_proj" in key for key in reasoning_keys))
+        self.assertFalse(any("group_align" in key for key in reasoning_keys))
+
+    def test_event_key_and_value_are_projected_once_per_forward(self) -> None:
+        model = build_test_model(n_reasoning_rounds=12)
+        call_counts = {"key": 0, "value": 0}
+
+        def count_key(*_args) -> None:
+            call_counts["key"] += 1
+
+        def count_value(*_args) -> None:
+            call_counts["value"] += 1
+
+        key_handle = (
+            model.reasoning_core.cross_attention.k_proj
+            .register_forward_hook(count_key)
+        )
+        value_handle = (
+            model.reasoning_core.cross_attention.v_proj
+            .register_forward_hook(count_value)
+        )
+        try:
+            model(**model_inputs())
+        finally:
+            key_handle.remove()
+            value_handle.remove()
+        self.assertEqual(call_counts, {"key": 1, "value": 1})
+
+    def test_relative_time_forward_and_backward_are_finite(self) -> None:
+        torch.manual_seed(0)
+        model = build_test_model(
+            target_mode="all_future",
+            time_mode="relative",
+            n_reasoning_rounds=2,
+        )
+        hidden = model(
+            **model_inputs(),
+            t_query=torch.tensor([5.0, 3.0]),
+        )
+        (hidden * torch.randn_like(hidden)).sum().backward()
+        self.assertTrue(torch.isfinite(hidden).all())
+        self.assertIsNotNone(model.event_projection.weight.grad)
+        self.assertTrue(torch.isfinite(model.event_projection.weight.grad).all())
+        time_scale = model.reasoning_core.cross_attention.time_bias_scale
+        self.assertIsNotNone(time_scale)
+        self.assertIsNotNone(time_scale.grad)
+        self.assertGreater(abs(float(time_scale.grad)), 0.0)
+
+    def test_architecture_marker_and_checkpoint_are_required(self) -> None:
+        validate_event_trajectory_config(
+            {
+                "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+                "model_size": "nano",
+                "d_model": 120,
+                "n_trajectory": 6,
+                "trajectory_dim": 20,
+                "traj_hidden": 24,
+                "n_reasoning_rounds": 3,
+            }
+        )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_event_trajectory_config(
+                {"model_architecture": "traj_mixer_v2"}
+            )
+        with self.assertRaisesRegex(ValueError, "only accepts models trained"):
+            validate_event_trajectory_config(
+                {"model_architecture": "event_trajectory_shared_v1"}
+            )
+        with self.assertRaisesRegex(ValueError, "trajectory_dim"):
+            validate_event_trajectory_config(
+                {
+                    "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+                    "model_size": "nano",
+                    "d_model": 120,
+                    "n_trajectory": 6,
+                    "trajectory_dim": 10,
+                    "traj_hidden": 24,
+                    "n_reasoning_rounds": 3,
+                }
+            )
+
+        model = build_test_model()
+        state_dict = model.state_dict()
+        validate_event_trajectory_state_dict(
+            state_dict,
+            expected_d_model=120,
+            expected_n_trajectory=6,
+            expected_n_reasoning_rounds=3,
+        )
+        with self.assertRaisesRegex(
+            ValueError,
+            "Checkpoint architecture does not match",
+        ):
+            validate_event_trajectory_state_dict(
+                state_dict,
+                expected_n_reasoning_rounds=12,
+            )
+        state_dict.pop("reasoning_core.attn_scale")
+        with self.assertRaisesRegex(
+            ValueError,
+            "not a shared event-trajectory checkpoint",
+        ):
+            validate_event_trajectory_state_dict(state_dict)
+
+    def test_unknown_model_size_is_rejected(self) -> None:
+        with self.assertRaisesRegex(ValueError, "Unknown model_size"):
+            DeepHealth(
+                vocab_size=32,
+                model_size="giant",
+                n_reasoning_rounds=2,
+                n_types=2,
+                n_cont_types=0,
+                n_categories=2,
+                cont_type_ids=[],
+            )
+
+
+if __name__ == "__main__":
+    unittest.main()

train_all_future.py

@@ -27,12 +27,18 @@ from tqdm.auto import tqdm
 
 from dataset import AllFutureHealthDataset, all_future_collate_fn
 from losses import build_loss
-from models import DeepHealth
+from models import (
+    EVENT_TRAJECTORY_ARCHITECTURE,
+    MODEL_SIZE_NAMES,
+    DeepHealth,
+    resolve_model_size,
+)
 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,
@@ -77,10 +83,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--min_future_events", type=int, default=1)
     parser.add_argument("--validation_query_seed", type=int, default=None)
 
-    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_tab_layer", type=int, default=4)
+    parser.add_argument(
+        "--model_size",
+        type=str,
+        default="nano",
+        choices=MODEL_SIZE_NAMES,
+    )
+    parser.add_argument("--n_reasoning_rounds", type=int, default=12)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
@@ -145,10 +154,8 @@ def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -
 def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        n_embd=args.n_embd,
-        n_head=args.n_head,
-        n_hist_layer=args.n_hist_layer,
-        n_tab_layer=args.n_tab_layer,
+        model_size=args.model_size,
+        n_reasoning_rounds=args.n_reasoning_rounds,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -293,11 +300,17 @@ def build_metadata(
     val_subset,
     test_subset,
 ) -> Dict[str, Any]:
+    size_config = resolve_model_size(args.model_size)
     return {
         "run_name": run_name,
         "dataset_class": "AllFutureHealthDataset",
         "collate_fn": "all_future_collate_fn",
         "model_class": "DeepHealth",
+        "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+        "d_model": size_config.d_model,
+        "n_trajectory": size_config.n_trajectory,
+        "trajectory_dim": size_config.trajectory_dim,
+        "traj_hidden": size_config.traj_hidden,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
         "dist_mode": args.dist_mode,
@@ -335,12 +348,26 @@ 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.model_size}_r{args.n_reasoning_rounds}_"
+            f"{args.time_mode}_{args.dist_mode}_"
+            f"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}")
+    size_config = resolve_model_size(args.model_size)
+    logger.info(
+        "Model size: "
+        f"{args.model_size} "
+        f"(d_model={size_config.d_model}, "
+        f"n_trajectory={size_config.n_trajectory}, "
+        f"trajectory_dim={size_config.trajectory_dim}, "
+        f"traj_hidden={size_config.traj_hidden}); "
+        f"reasoning_rounds={args.n_reasoning_rounds}"
+    )
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
 
     logger.info("Loading all-future datasets...")
@@ -434,6 +461,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 +476,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,13 +24,20 @@ from tqdm.auto import tqdm
 
 from dataset import HealthDataset, collate_fn
 from losses import build_loss
-from models import DeepHealth, DeepHealthOutput
+from models import (
+    EVENT_TRAJECTORY_ARCHITECTURE,
+    MODEL_SIZE_NAMES,
+    DeepHealth,
+    DeepHealthOutput,
+    resolve_model_size,
+)
 from readouts import build_readout
 from targets import CHECKUP_IDX, NO_EVENT_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,
@@ -73,10 +80,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
     parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
 
-    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_tab_layer", type=int, default=4)
+    parser.add_argument(
+        "--model_size",
+        type=str,
+        default="nano",
+        choices=MODEL_SIZE_NAMES,
+    )
+    parser.add_argument("--n_reasoning_rounds", type=int, default=12)
     parser.add_argument("--n_bins", type=int, default=16)
     parser.add_argument("--extra_pool_reduce", type=str, default="mean",
                         choices=["mean", "sum"])
@@ -150,10 +160,8 @@ def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -
 def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
     return DeepHealth(
         vocab_size=dataset.vocab_size,
-        n_embd=args.n_embd,
-        n_head=args.n_head,
-        n_hist_layer=args.n_hist_layer,
-        n_tab_layer=args.n_tab_layer,
+        model_size=args.model_size,
+        n_reasoning_rounds=args.n_reasoning_rounds,
         n_types=dataset.n_types,
         n_cont_types=dataset.n_cont_types,
         n_categories=dataset.n_categories,
@@ -479,11 +487,17 @@ def build_metadata(
     val_subset,
     test_subset,
 ) -> Dict[str, Any]:
+    size_config = resolve_model_size(args.model_size)
     return {
         "run_name": run_name,
         "dataset_class": "NextStepHealthDataset",
         "collate_fn": "next_step_collate_fn",
         "model_class": "DeepHealth",
+        "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
+        "d_model": size_config.d_model,
+        "n_trajectory": size_config.n_trajectory,
+        "trajectory_dim": size_config.trajectory_dim,
+        "traj_hidden": size_config.traj_hidden,
         "model_target_mode": "next_token",
         "target_mode": args.target_mode,
         "dist_mode": "exponential",
@@ -519,7 +533,9 @@ def main() -> None:
 
     run_dir, run_name = create_unique_run_dir(
         lambda timestamp: (
-            f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
+            f"{args.model_size}_r{args.n_reasoning_rounds}_"
+            f"{args.time_mode}_exponential_"
+            f"next_token_{args.target_mode}_"
             f"gap_{args.no_event_interval_years:g}y_{timestamp}"
         )
     )
@@ -527,6 +543,16 @@ def main() -> None:
 
     logger.info(f"Starting next-step training run: {run_name}")
     logger.info(f"Device: {device}")
+    size_config = resolve_model_size(args.model_size)
+    logger.info(
+        "Model size: "
+        f"{args.model_size} "
+        f"(d_model={size_config.d_model}, "
+        f"n_trajectory={size_config.n_trajectory}, "
+        f"trajectory_dim={size_config.trajectory_dim}, "
+        f"traj_hidden={size_config.traj_hidden}); "
+        f"reasoning_rounds={args.n_reasoning_rounds}"
+    )
     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 +622,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 +638,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 parameter-count fields for logs and train_config.json."""
+    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