From 20c99484f33cdf306c3f0278ea6f28a4ff00d044 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 24 Jul 2026 10:51:41 +0800 Subject: [PATCH] Add two-stage TrajMixer mixing --- TrajMixer_设计方案.md | 444 ++++++++++++++++++++++++------------------ backbones.py | 53 ++++- models.py | 9 +- test_traj_mixer.py | 103 +++++++++- 4 files changed, 401 insertions(+), 208 deletions(-) diff --git a/TrajMixer_设计方案.md b/TrajMixer_设计方案.md index 961a893..c42625b 100644 --- a/TrajMixer_设计方案.md +++ b/TrajMixer_设计方案.md @@ -2,332 +2,390 @@ > 状态：**Frozen implementation baseline** > -> 版本：**v1.0** +> 版本：**v2.0 / traj_mixer_v3** > -> 固化日期：**2026-07-22** +> 固化日期：**2026-07-24** -本文档是 TrajMixer 后续实现与实验的唯一结构基线。除显式标记为消融项的配置外，所有实现均应遵循本文档；若结构发生变化，应先更新版本和实验记录。 +本文档是 TrajMixer 后续实现与实验的结构基线。本版本将原先仅含跨轨迹交互的 TrajMixer 扩展为“组内混合 + 跨组混合”的两阶段结构。 ## 1. 目标 -在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的轨迹交互模块替换 FFN。 +在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的两阶段 TrajMixer 替换 FFN。 保持不变的组件包括： - 原始 causal mask； - 原始 TimeRoPE / Relative Time Attention Bias； - 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)； -- 原始序列建模与训练目标。 +- 原始序列建模和训练目标。 -TrajMixer 不修改 Attention，只替换每个 Transformer block 中的 FFN residual branch。 +TrajMixer 不沿序列维度混合，也不引入时间递归。 ## 2. Block 总体结构 -概念结构： - ```text PreNorm Causal Multi-Head Attention -→ Residual -→ Standard Mixer PreNorm +→ Attention Residual +→ reshape [B, L, n_group, d_group] +→ Intra-Group PreNorm +→ Per-Group SwiGLU: d_group → 4d_group → d_group +→ Intra-Group Residual +→ Cross-Group PreNorm → Group-wise Feature Alignment -→ SwiGLU Cross-Group Mixer -→ Residual +→ Cross-Group SwiGLU: n_group → 4n_group → n_group +→ Cross-Group Residual +→ reshape [B, L, d] ``` -完整计算为： +Attention 阶段保持原样： \[ -U = X^{(l)} + \operatorname{Dropout}\!\left( -\operatorname{CausalMHA}\left( -\operatorname{LN}_{\mathrm{attn}}(X^{(l)}), -\text{time information} -\right)\right), +U=X^{(l)}+\operatorname{CausalMHA} +\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right). +\] + +随后： + +\[ +H^{(0)} +=\operatorname{reshape}(U) +\in\mathbb{R}^{B\times L\times G\times D}, +\] + +其中 \(G=n_{\mathrm{group}}\)，\(D=d_{\mathrm{group}}\)。 + +两阶段 TrajMixer 为： + +\[ +H^{(1)} +=H^{(0)} ++\operatorname{Dropout} +\left(\operatorname{IntraMixer} +\left(\operatorname{LN}_{D}(H^{(0)})\right)\right), \] \[ -N = \operatorname{LN}_{\mathrm{mixer}}(U), +H^{(2)} +=H^{(1)} ++\operatorname{Dropout} +\left(\operatorname{CrossMixer} +\left(\operatorname{LN}_{G}(H^{(1)})\right)\right), \] \[ -\Delta = \operatorname{TrajMixer}(N), +X^{(l+1)}=\operatorname{reshape}(H^{(2)}) +\in\mathbb{R}^{B\times L\times d}. \] -\[ -X^{(l+1)} = U + \operatorname{Dropout}(\Delta). -\] - -首版中的 \(\operatorname{LN}_{\mathrm{mixer}}\) 是作用于完整 \(d=120\) 维 residual representation 的标准 LayerNorm。 +`TrajMixer.forward()` 返回的是已经完成两次 residual update 的完整状态，而不是单个 residual delta。因此 `GPTBlock` 在 Attention residual 后直接返回 `TrajMixer(U)`，不得再写成 `U + TrajMixer(U)`。 ## 3. Latent Trajectory Group 定义 Attention 输出经过 \(W_O\) 后仍是标准 residual representation： \[ -N\in\mathbb{R}^{B\times L\times d},\qquad d=120. +U\in\mathbb{R}^{B\times L\times d}. \] -将 hidden dimension 划分为与 Attention head 数量相同的 group 数量： +固定： \[ -n_{\mathrm{group}}:=n_{\mathrm{head}}=10, -\qquad d_{\mathrm{group}}=\frac{d}{n_{\mathrm{head}}}=12, +G:=n_{\mathrm{group}}=n_{\mathrm{head}}, +\qquad +D:=d_{\mathrm{group}}=\frac{d}{G}, +\qquad +d=GD. \] -`n_group` 不再是独立超参数，代码统一使用 `n_head` 确定 residual group 数量。二者只共享数量；这些 residual groups 在语义和张量来源上仍不等同于原始 Attention heads。 - -并 reshape 为： +默认配置： \[ -N_{\mathrm{group}}in -\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}. +d=120,\qquad G=10,\qquad D=12. \] -这些 group 是 residual space 中的 **latent trajectory groups**，不等同于原始 Attention heads。本文中的 group、trajectory group 均指这一 residual-channel partition。 - -## 4. Group-wise Feature Alignment - -为缓解不同 group 内部坐标不对齐的问题，每个 group 使用独立的小矩阵： +reshape 后： \[ -B_i\in\mathbb{R}^{d_{\mathrm{group}}\times d_{\mathrm{group}}}, -\qquad i=1,\ldots,n_{\mathrm{group}}. +H^{(0)}\in\mathbb{R}^{B\times L\times G\times D}. \] -对每个 group 内的特征进行可学习对齐： +`n_group` 由 `n_head` 决定，但 residual groups 只是 residual space 的连续分区，不等同于 Attention heads。 + +## 4. 第一阶段：组内 SwiGLU Mixer + +第一阶段对每个 group 独立进行特征变换。不同 group 使用各自的投影参数，不发生 group 间信息交换。 + +先对每个 \((b,t,g)\) 的 \(D\) 维向量独立执行 LayerNorm： \[ -Z_{b,t,i,:}=N_{\mathrm{group},b,t,i,:}B_i. +\widetilde H^{(0)} +=\operatorname{LN}_{D}(H^{(0)}). \] -因此： +归一化统计量在每个 group 内独立计算；为保持轻量，LayerNorm 的 affine 参数在各 group 间共享。 + +对于 \(g=1,\ldots,G\)，定义： \[ -Z\in -\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}. +W_{g,\mathrm{intra}}^{(g)}, +W_{v,\mathrm{intra}}^{(g)} +\in\mathbb{R}^{D\times 4D}, \] -首版实现约定： - -- \(B_i\) 不带 bias； -- \(B_i\) 使用单位矩阵初始化； -- Alignment 只作用于 Mixer residual branch，不改变 Attention residual stream； -- 首版不增加逆变换或额外的 group 内输出投影。 - -Alignment 每层权重参数量为： - \[ -n_{\mathrm{group}}d_{\mathrm{group}}^2 -=10\times12^2 -=1{,}440. +W_{o,\mathrm{intra}}^{(g)} +\in\mathbb{R}^{4D\times D}. \] -## 5. SwiGLU Cross-Group Mixer - -Mixer 只沿 group 维度交互，不沿序列维度交互，因此不会引入时间递归或未来信息泄漏。 - -对于每个 group 内特征维度： +计算： \[ -r=1,\ldots,d_{\mathrm{group}}, +P_g +=\operatorname{SiLU} +\left(\widetilde H^{(0)}_g +W_{g,\mathrm{intra}}^{(g)}\right) +\odot +\left(\widetilde H^{(0)}_g +W_{v,\mathrm{intra}}^{(g)}\right), \] -定义： +\[ +\Delta_{\mathrm{intra},g} +=P_gW_{o,\mathrm{intra}}^{(g)}, +\] + +\[ +H^{(1)} +=H^{(0)} ++\operatorname{Dropout}(\Delta_{\mathrm{intra}}). +\] + +该阶段完成： + +\[ +D\rightarrow4D\rightarrow D, +\] + +用于增强每条潜在轨迹内部的非线性特征组合能力。 + +实现张量形状： + +```text +intra_norm: LayerNorm(d_group) +intra_gate_proj: [n_group, d_group, 4 * d_group] +intra_value_proj: [n_group, d_group, 4 * d_group] +intra_output_proj: [n_group, 4 * d_group, d_group] +``` + +三个 projection 均不带 bias。 + +## 5. 第二阶段：跨组 TrajMixer + +第二阶段沿 group 维度进行交互。对于每个内部坐标 \(r\)，独立执行： + +\[ +G\rightarrow4G\rightarrow G. +\] + +首先将 \(H^{(1)}\) 的最后两个维度交换，并在 group 维度执行 LayerNorm： + +\[ +\widetilde H^{(1)}_{b,t,:,r} +=\operatorname{LN}_{G} +\left(H^{(1)}_{b,t,:,r}\right). +\] + +归一化统计量对每个内部坐标 \(r\) 独立计算；LayerNorm 的 affine 参数在各内部坐标间共享。 + +### 5.1 Group-wise Feature Alignment + +沿用现有的可学习 group 特征对齐矩阵： + +\[ +B_g\in\mathbb{R}^{D\times D}, +\qquad g=1,\ldots,G, +\] + +\[ +Z_{b,t,g,:} +=\widetilde H^{(1)}_{b,t,g,:}B_g. +\] + +\(B_g\) 不带 bias，并使用单位矩阵初始化。 + +### 5.2 Cross-Group SwiGLU + +对每个内部坐标 \(r=1,\ldots,D\)，定义： \[ A_g^{(r)},A_v^{(r)} -\in\mathbb{R}^{n_{\mathrm{group}}\times h_{\mathrm{group}}}, -\] - -\[ +\in\mathbb{R}^{G\times4G}, +\qquad A_o^{(r)} -\in\mathbb{R}^{h_{\mathrm{group}}\times n_{\mathrm{group}}}. +\in\mathbb{R}^{4G\times G}. \] -隐藏宽度不再独立配置，固定为： +计算： \[ -h_{\mathrm{group}}=4n_{\mathrm{head}} -=4n_{\mathrm{group}}. -\] - -当前 \(n_{\mathrm{head}}=10\)，因此 \(h_{\mathrm{group}}=40\)。 - -对固定的 batch、时间位置和内部特征维度 \(r\)，将： - -\[ -Z_{b,t,:,r}\in\mathbb{R}^{n_{\mathrm{group}}} -\] - -视为 row vector，计算： - -\[ -G_{b,t,:,r}=Z_{b,t,:,r}A_g^{(r)}, +Q_{b,t,:,r} +=\operatorname{SiLU} +\left(Z_{b,t,:,r}A_g^{(r)}\right) +\odot +\left(Z_{b,t,:,r}A_v^{(r)}\right), \] \[ -V_{b,t,:,r}=Z_{b,t,:,r}A_v^{(r)}, +\Delta_{\mathrm{cross},b,t,:,r} +=Q_{b,t,:,r}A_o^{(r)}, \] \[ -M_{b,t,:,r}=\operatorname{SiLU}(G_{b,t,:,r})\odot V_{b,t,:,r}, +H^{(2)} +=H^{(1)} ++\operatorname{Dropout}(\Delta_{\mathrm{cross}}). \] -\[ -Y_{b,t,:,r}=M_{b,t,:,r}A_o^{(r)}. -\] - -其中： - -- gate 分支控制信息写入； -- value 分支提供交互内容； -- output matrix 将隐藏 group 表示投影回原始 group 数量； -- hidden group 表示固定扩展为 group 数量的 4 倍。 - -所有 \(r\) 的输出组合为： - -\[ -Y\in -\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}, -\] - -再 reshape 为： - -\[ -\Delta\in\mathbb{R}^{B\times L\times d}. -\] - -## 6. 参数张量与无歧义索引 - -建议的实现存储形状为： +实现张量形状： ```text +cross_norm: LayerNorm(n_group) group_align: [n_group, d_group, d_group] -gate_proj: [d_group, n_group, hidden_group] -value_proj: [d_group, n_group, hidden_group] -output_proj: [d_group, hidden_group, n_group] +gate_proj: [d_group, n_group, 4 * n_group] +value_proj: [d_group, n_group, 4 * n_group] +output_proj: [d_group, 4 * n_group, n_group] ``` -对应的索引公式为： +三个 projection 均不带 bias。 + +## 6. PreNorm 与残差约束 + +本版本固定使用两个独立的 PreNorm residual stage： + +1. `intra_norm` 只服务于组内 Mixer； +2. `cross_norm` 只服务于跨组 Mixer； +3. 第一阶段 residual 的输出是第二阶段的输入； +4. 两个 residual 都在 `TrajMixer` 内部完成； +5. 不再保留 block 外部的全维度 `ln2` 或额外 Mixer residual。 + +因此信息流必须是： + +```text +U +→ U + IntraMixer(IntraNorm(U)) +→ H1 + CrossMixer(CrossNorm(H1)) +→ output +``` + +## 7. 参数量 + +默认 \(d=120,G=10,D=12\)。 + +### 7.1 组内阶段 + +投影权重： \[ -G_{b,t,q,r} -=\sum_i Z_{b,t,i,r}\,A_{g,r,i,q}, +3G D(4D) +=12GD^2 +=17{,}280. \] +`LayerNorm(D)`： + \[ -V_{b,t,q,r} -=\sum_i Z_{b,t,i,r}\,A_{v,r,i,q}, +2D=24. \] -\[ -Y_{b,t,i,r} -=\sum_q -\left[\operatorname{SiLU}(G_{b,t,q,r})V_{b,t,q,r}\right] -A_{o,r,q,i}. -\] +### 7.2 跨组阶段 -首版的三个 Mixer projection 均不带 bias。 - -## 7. Mixer Hidden Width 与参数量 - -Mixer 的 group 维变换为： +跨组投影权重： \[ -n_{\mathrm{group}} -\rightarrow -h_{\mathrm{group}} -\rightarrow -n_{\mathrm{group}}. -\] - -固定 \(h_{\mathrm{group}}=4n_{\mathrm{group}}=40\) 时，Mixer 每层权重参数量为： - -\[ -3d_{\mathrm{group}}n_{\mathrm{group}}h_{\mathrm{group}} -=3\times12\times10\times40 +3D G(4G) +=12DG^2 =14{,}400. \] -加上 Group Feature Alignment 后，TrajMixer residual branch 每层共有： +Group Feature Alignment： \[ -14{,}400+1{,}440=15{,}840 +GD^2 +=1{,}440. \] -个主要权重参数。作为对照，原始 \(120\rightarrow480\rightarrow120\) FFN 每层约有 115,800 个参数。 +`LayerNorm(G)`： -参数对照口径说明：上面的 115,800 对应结构方案中的标准两层 FFN。当前代码库在 TrajMixer 替换前实际使用的是隐藏宽度 300 的全维度 SwiGLU（gate/value/output 三个线性层），每层共有 108,720 个参数（含 bias）。代码实验和 checkpoint 参数量比较必须以 108,720 作为历史实现基线，不能与概念方案中的标准 FFN 参数量混用。 +\[ +2G=20. +\] -## 8. LayerNorm 基线与消融 +### 7.3 每个 TrajMixer 合计 -为保持与原始 Transformer 的可比性，首版固定使用： +\[ +17{,}280+24+14{,}400+1{,}440+20 +=\boxed{33{,}164}. +\] -```text -原始 FFN baseline：FFN + 标准 LayerNorm -TrajMixer baseline：Mixer + 标准 LayerNorm -``` +相对于 `traj_mixer_v2` 的跨组单阶段结构 \(15{,}840\)，每层增加 \(17{,}324\) 个参数。作为历史实现对照，代码库原全维度 SwiGLU FFN 每层为 \(108{,}720\) 个参数。 -以下配置不属于首版主实验，只作为独立消融： +## 8. 初始化 -```text -Mixer + Group-wise LayerNorm -``` +固定初始化约定： -不得将 Group-wise LayerNorm 的结果直接作为“仅替换 FFN”的对照结果。 +- 组内 `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform； +- 组内 `intra_output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布； +- Group Alignment：单位矩阵； +- 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform； +- 跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布； +- 两个 LayerNorm：PyTorch 默认 affine 初始化； +- 两个 residual stage 的 Dropout 均沿用 `mlp_dropout`。 -## 9. 初始化 +两个 output projection 的小方差初始化使两阶段在训练初期都接近恒等 residual update。 -首版初始化约定： +## 9. 信息流与语义 -- Group Alignment \(B_i\)：单位矩阵初始化； -- \(A_g/A_v\)：Xavier uniform 初始化； -- \(A_o\)：均值为 0、标准差为 \(10^{-3}\) 的正态初始化； -- Dropout 概率沿用原始 FFN residual branch 的配置。 +**Attention**：从历史疾病事件中选择和整合相关信息。 -小初始化的 \(A_o\) 使新增分支在训练初期接近恒等残差更新，同时允许模型逐步学习轨迹交互。 +**Intra-Group Mixer**：学习每条潜在轨迹内部的非线性特征组合。 -## 10. 核心设计思想 +**Group Feature Alignment**：对齐不同潜在轨迹的内部坐标。 -**Attention**：负责从历史疾病序列中选择并整合相关信息。 - -**Group Feature Alignment**：负责学习不同 latent trajectory groups 的内部特征对齐。 - -**Cross-Group Mixer**：负责不同潜在疾病轨迹之间的非线性门控交互。 +**Cross-Group Mixer**：学习不同潜在轨迹在相同内部坐标上的门控交互。 整个模块保持： - 无时间递归； +- 不混合序列位置； - 序列维度完全并行； -- 参数量远低于原始 FFN； -- 保留 Transformer 的因果历史建模能力； -- 不把 residual groups 误解释为原始 Attention heads。 +- 保留原始因果 Attention； +- residual groups 不等同于 Attention heads。 -## 11. 首版固定配置 +## 10. 固定配置与 checkpoint 约束 ```yaml -model_architecture: traj_mixer_v2 +model_architecture: traj_mixer_v3 d_model: 120 -n_head: 10 # 同时决定 residual group 数量 -d_group: 12 -hidden_group_rule: 4 * n_head # 不单独配置 +n_head: 10 +n_group_rule: n_head +d_group_rule: d_model / n_group +intra_hidden_rule: 4 * d_group +cross_hidden_rule: 4 * n_group attention: unchanged -attention_output_projection: unchanged -mixer_norm: standard_layer_norm -group_alignment: per_group_12x12 -group_alignment_bias: false -group_alignment_init: identity -mixer_bias: false +intra_norm: layer_norm_over_d_group +cross_norm: layer_norm_over_n_group +group_alignment: per_group_d_group_x_d_group +projection_bias: false gate_value_init: xavier_uniform output_init_std: 0.001 -group_wise_layer_norm: false ``` -训练时必须将 `model_architecture: traj_mixer_v2`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印总参数量与可训练参数量。本分支的评估和导出入口只接受带有该架构标识、且 checkpoint 中包含 TrajMixer 参数张量的模型；其他版本或分支生成的模型应直接拒绝加载。 - 必须满足： \[ d=n_{\mathrm{group}}d_{\mathrm{group}}. \] -后续实现、单元测试、参数量核验和主实验均以以上配置为默认基线。 +训练时必须将 `model_architecture: traj_mixer_v3`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印参数量。 + +本分支的评估和导出入口只接受 `traj_mixer_v3` checkpoint，并检查两阶段 Norm、组内 projection、Group Alignment 和跨组 projection 参数是否齐全。`traj_mixer_v2` 及更早 checkpoint 不向后兼容，直接拒绝加载。 diff --git a/backbones.py b/backbones.py index ed11813..3e94262 100644 --- a/backbones.py +++ b/backbones.py @@ -177,7 +177,7 @@ class TemporalAttention(nn.Module): class TrajMixer(nn.Module): - """Lightweight gated interaction across latent residual-space groups. + """Two-stage 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. @@ -205,8 +205,24 @@ class TrajMixer(nn.Module): # are still residual-space partitions rather than attention heads. self.n_group = n_head self.d_group = n_embd // n_head + self.intra_hidden = 4 * self.d_group self.hidden_group = 4 * n_head + # Stage 1: each group independently mixes its internal features. + self.intra_norm = nn.LayerNorm(self.d_group) + 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) + ) + + # Stage 2: each internal coordinate independently mixes groups. + self.cross_norm = nn.LayerNorm(self.n_group) + # Per-group feature alignment: [group, input feature, output feature]. self.group_align = nn.Parameter( torch.empty(self.n_group, self.d_group, self.d_group) @@ -227,6 +243,11 @@ class TrajMixer(nn.Module): 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.normal_(self.intra_output_proj, mean=0.0, std=1e-3) + with torch.no_grad(): identity = torch.eye( self.d_group, @@ -243,7 +264,7 @@ class TrajMixer(nn.Module): nn.init.normal_(self.output_proj, mean=0.0, std=1e-3) def forward(self, x: torch.Tensor) -> torch.Tensor: - """Map ``(B, L, n_embd)`` to an equally shaped residual update.""" + """Apply two PreNorm residual stages without mixing sequence positions.""" if x.ndim != 3: raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}") if x.size(-1) != self.n_embd: @@ -255,8 +276,27 @@ class TrajMixer(nn.Module): grouped = x.reshape( batch_size, seq_len, self.n_group, self.d_group ) + + # Stage 1: d_group -> 4*d_group -> d_group, independently per group. + intra_input = self.intra_norm(grouped) + intra_gate = torch.einsum( + "blgd,gdh->blgh", intra_input, self.intra_gate_proj + ) + intra_value = torch.einsum( + "blgd,gdh->blgh", intra_input, self.intra_value_proj + ) + intra_hidden = F.silu(intra_gate) * intra_value + intra_update = torch.einsum( + "blgh,ghd->blgd", intra_hidden, self.intra_output_proj + ) + grouped = grouped + self.drop(intra_update) + + # Stage 2: n_group -> 4*n_group -> n_group for each coordinate. + cross_input = self.cross_norm( + grouped.transpose(-1, -2) + ).transpose(-1, -2) aligned = torch.einsum( - "blgd,gde->blge", grouped, self.group_align + "blgd,gde->blge", cross_input, self.group_align ) gate = torch.einsum( @@ -269,7 +309,8 @@ class TrajMixer(nn.Module): mixed = torch.einsum( "blhr,rhg->blgr", hidden, self.output_proj ) - return self.drop(mixed.reshape(batch_size, seq_len, self.n_embd)) + grouped = grouped + self.drop(mixed) + return grouped.reshape(batch_size, seq_len, self.n_embd) class GPTBlock(nn.Module): @@ -299,7 +340,6 @@ class GPTBlock(nn.Module): dropout=mlp_dropout, ) self.ln1 = nn.LayerNorm(n_embd) - self.ln2 = nn.LayerNorm(n_embd) def forward( self, @@ -309,8 +349,7 @@ class GPTBlock(nn.Module): 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 + return self.mlp(x) class TokenAutoDiscretization(nn.Module): diff --git a/models.py b/models.py index 1d2bbd4..204c490 100644 --- a/models.py +++ b/models.py @@ -15,7 +15,7 @@ from backbones import ( from targets import PAD_IDX -TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v2" +TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3" def validate_traj_mixer_config(config: Mapping[str, object]) -> None: @@ -29,6 +29,13 @@ def validate_traj_mixer_config(config: Mapping[str, object]) -> None: def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None: required_keys = { + "blocks.0.mlp.intra_norm.weight", + "blocks.0.mlp.intra_norm.bias", + "blocks.0.mlp.intra_gate_proj", + "blocks.0.mlp.intra_value_proj", + "blocks.0.mlp.intra_output_proj", + "blocks.0.mlp.cross_norm.weight", + "blocks.0.mlp.cross_norm.bias", "blocks.0.mlp.group_align", "blocks.0.mlp.gate_proj", "blocks.0.mlp.value_proj", diff --git a/test_traj_mixer.py b/test_traj_mixer.py index 015b4ac..c3957db 100644 --- a/test_traj_mixer.py +++ b/test_traj_mixer.py @@ -21,14 +21,98 @@ class TrajMixerTest(unittest.TestCase): x = torch.randn(2, 7, 120) self.assertEqual(mixer(x).shape, x.shape) - self.assertEqual(sum(p.numel() for p in mixer.parameters()), 15_840) + self.assertEqual(sum(p.numel() for p in mixer.parameters()), 33_164) expected = torch.eye(12).expand(10, 12, 12) torch.testing.assert_close(mixer.group_align.detach(), expected) + 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)) + self.assertEqual(tuple(mixer.intra_norm.normalized_shape), (12,)) + self.assertEqual(tuple(mixer.cross_norm.normalized_shape), (10,)) + + def test_zero_output_projections_make_both_stages_identity(self) -> None: + torch.manual_seed(0) + mixer = TrajMixer(120, n_head=10, dropout=0.0) + with torch.no_grad(): + mixer.intra_output_proj.zero_() + mixer.output_proj.zero_() + x = torch.randn(2, 5, 120) + torch.testing.assert_close(mixer(x), x) + + def test_intra_stage_is_independent_across_groups(self) -> None: + torch.manual_seed(0) + mixer = TrajMixer(120, n_head=10, dropout=0.0) + mixer.eval() + with torch.no_grad(): + mixer.output_proj.zero_() + + grouped = torch.randn(2, 4, 10, 12) + changed = grouped.clone() + changed[:, :, 3, :] += torch.randn_like(changed[:, :, 3, :]) + + original_out = mixer(grouped.reshape(2, 4, 120)).reshape( + 2, 4, 10, 12 + ) + changed_out = mixer(changed.reshape(2, 4, 120)).reshape( + 2, 4, 10, 12 + ) + unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9]) + torch.testing.assert_close( + original_out.index_select(2, unchanged_groups), + changed_out.index_select(2, unchanged_groups), + ) + + 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.intra_output_proj.zero_() + mixer.gate_proj.zero_() + mixer.value_proj.zero_() + mixer.output_proj.zero_() + + # For coordinate 0 only, read group 0 through hidden unit 0 and + # write the resulting gated value into group 1. + mixer.gate_proj[0, 0, 0] = 1.0 + mixer.value_proj[0, 0, 0] = 1.0 + mixer.output_proj[0, 0, 1] = 1.0 + + 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(grouped.reshape(1, 1, 6)).reshape(1, 1, 3, 2) + changed_out = mixer(changed.reshape(1, 1, 6)).reshape(1, 1, 3, 2) + + 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) @@ -58,11 +142,12 @@ class TrajMixerTest(unittest.TestCase): self.assertIsNotNone(parameter.grad, name) self.assertTrue(torch.isfinite(parameter.grad).all(), name) - def test_gpt_block_defaults_to_traj_mixer_and_standard_layer_norm(self) -> None: + def test_gpt_block_delegates_both_mixer_residuals_to_traj_mixer(self) -> None: block = GPTBlock(n_embd=120, n_head=10) self.assertIsInstance(block.mlp, TrajMixer) - self.assertIsInstance(block.ln2, torch.nn.LayerNorm) - self.assertEqual(tuple(block.ln2.normalized_shape), (120,)) + self.assertFalse(hasattr(block, "ln2")) + self.assertIsInstance(block.mlp.intra_norm, torch.nn.LayerNorm) + self.assertIsInstance(block.mlp.cross_norm, torch.nn.LayerNorm) x = torch.randn(2, 6, 120) self.assertEqual(block(x).shape, x.shape) @@ -75,6 +160,10 @@ class TrajMixerTest(unittest.TestCase): validate_traj_mixer_config({}) with self.assertRaisesRegex(ValueError, "only accepts models trained"): validate_traj_mixer_config({"model_architecture": "delphi_swiglu"}) + with self.assertRaisesRegex(ValueError, "only accepts models trained"): + validate_traj_mixer_config( + {"model_architecture": "traj_mixer_v2"} + ) def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None: block = GPTBlock(n_embd=120, n_head=10) @@ -84,7 +173,7 @@ class TrajMixerTest(unittest.TestCase): } validate_traj_mixer_state_dict(state_dict) - state_dict.pop("blocks.0.mlp.group_align") + state_dict.pop("blocks.0.mlp.intra_gate_proj") with self.assertRaisesRegex(ValueError, "not a TrajMixer checkpoint"): validate_traj_mixer_state_dict(state_dict) @@ -97,8 +186,8 @@ class TrajMixerTest(unittest.TestCase): self.assertEqual( get_model_parameter_counts(mixer), { - "model_parameter_count": 15_840, - "trainable_parameter_count": 15_840, + "model_parameter_count": 33_164, + "trainable_parameter_count": 33_164, }, )