From 8d0d71292efcbc3bf84a14f349ea813092027b5d Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Fri, 24 Jul 2026 15:25:02 +0800 Subject: [PATCH] Refactor TrajMixer to single residual --- TrajMixer_设计方案.md | 431 ++++++++++++++++++++---------------------- backbones.py | 97 +++++----- models.py | 10 +- test_traj_mixer.py | 69 ++++--- 4 files changed, 304 insertions(+), 303 deletions(-) diff --git a/TrajMixer_设计方案.md b/TrajMixer_设计方案.md index 006bec1..320d3b0 100644 --- a/TrajMixer_设计方案.md +++ b/TrajMixer_设计方案.md @@ -2,176 +2,178 @@ > 状态：**Frozen implementation baseline** > -> 版本：**v2.0 / traj_mixer_v3** +> 版本：**v3.0 / traj_mixer_v5** > > 固化日期：**2026-07-24** -本文档是 TrajMixer 后续实现与实验的结构基线。本版本将原先仅含跨轨迹交互的 TrajMixer 扩展为“组内混合 + 跨组混合”的两阶段结构。 +本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。 ## 1. 目标 -在保持原始 Delphi Transformer Attention 结构不变的前提下，用轻量、可并行的两阶段 TrajMixer 替换 FFN。 +在不改变 Delphi Transformer Attention 的前提下，用轻量、完全并行的 TrajMixer 替换 FFN。 -保持不变的组件包括： +保持不变： -- 原始 causal mask； -- 原始 TimeRoPE / Relative Time Attention Bias； -- 原始 Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)； -- 原始序列建模和训练目标。 +- causal mask； +- TimeRoPE； +- Relative Time Attention Bias； +- Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)； +- 序列建模和训练目标。 TrajMixer 不沿序列维度混合，也不引入时间递归。 -## 2. Block 总体结构 +## 2. Block 结构 ```text PreNorm Causal Multi-Head Attention → Attention Residual +→ Full-width TrajMixer PreNorm → 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 +→ Static Gated Fusion → Cross-Group SwiGLU: n_group → 4n_group → n_group -→ Cross-Group Residual -→ reshape [B, L, d] +→ reshape [B, L, n_embd] +→ Dropout +→ One TrajMixer Residual ``` -Attention 阶段保持原样： +Attention 阶段： \[ -U=X^{(l)}+\operatorname{CausalMHA} +X +=X^{(l)} ++\operatorname{CausalMHA} \left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right). \] -随后： +TrajMixer 阶段： \[ -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}_{d}(X), \] \[ -H^{(2)} -=H^{(1)} -+\operatorname{Dropout} -\left(\operatorname{CrossMixer} -\left(\operatorname{LN}_{G}(H^{(1)})\right)\right), +G=\operatorname{reshape}(N) +\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}}, \] \[ -X^{(l+1)}=\operatorname{reshape}(H^{(2)}) -\in\mathbb{R}^{B\times L\times d}. +P=\operatorname{IntraMixer}(G), \] -`TrajMixer.forward()` 返回的是已经完成两次 residual update 的完整状态，而不是单个 residual delta。因此 `GPTBlock` 在 Attention residual 后直接返回 `TrajMixer(U)`，不得再写成 `U + TrajMixer(U)`。 +\[ +U=G+\sigma(\Theta)\odot P, +\] -## 3. Latent Trajectory Group 定义 +\[ +\Delta=\operatorname{reshape} +\left(\operatorname{CrossMixer}(U)\right), +\] + +\[ +X^{(l+1)}=X+\operatorname{Dropout}(\Delta). +\] + +整个 TrajMixer 只有最后一次 `X + update` 是 residual。`U=G+\sigma(\Theta)\odot P` 是 update 分支内部的静态门控特征融合，不是相对于主 residual stream 的独立 residual stage。 + +## 3. Group 定义 Attention 输出经过 \(W_O\) 后仍是标准 residual representation： \[ -U\in\mathbb{R}^{B\times L\times d}. +X\in\mathbb{R}^{B\times L\times d}. \] -固定： +定义： \[ -G:=n_{\mathrm{group}}=n_{\mathrm{head}}, +n_{\mathrm{group}}:=n_{\mathrm{head}}, \qquad -D:=d_{\mathrm{group}}=\frac{d}{G}, +d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}}, \qquad -d=GD. +d=n_{\mathrm{group}}d_{\mathrm{group}}. \] -默认配置： +默认： \[ -d=120,\qquad G=10,\qquad D=12. +d=120,\qquad +n_{\mathrm{group}}=10,\qquad +d_{\mathrm{group}}=12. \] -reshape 后： +这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。 + +## 4. 唯一的 Full-Width PreNorm + +TrajMixer 只使用一个： + +```text +norm: LayerNorm(n_embd) +``` + +LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape： \[ -H^{(0)}\in\mathbb{R}^{B\times L\times G\times D}. +G=\operatorname{reshape} +\left(\operatorname{LN}_{d}(X)\right). \] -`n_group` 由 `n_head` 决定，但 residual groups 只是 residual space 的连续分区，不等同于 Attention heads。 +本版本明确删除： -## 4. 第一阶段：组内 SwiGLU Mixer +```text +intra_norm +cross_norm +group_align +``` -第一阶段对每个 group 独立进行特征变换。不同 group 使用各自的投影参数，不发生 group 间信息交换。 +不得在组内或跨组阶段再增加额外 LayerNorm。 -先对每个 \((b,t,g)\) 的 \(D\) 维向量独立执行 LayerNorm： +## 5. 组内 SwiGLU + +每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行： \[ -\widetilde H^{(0)} -=\operatorname{LN}_{D}(H^{(0)}). +d_{\mathrm{group}} +\rightarrow +4d_{\mathrm{group}} +\rightarrow +d_{\mathrm{group}}. \] -归一化统计量在每个 group 内独立计算；为保持轻量，LayerNorm 的 affine 参数在各 group 间共享。 - -对于 \(g=1,\ldots,G\)，定义： +对 group \(g\)： \[ W_{g,\mathrm{intra}}^{(g)}, W_{v,\mathrm{intra}}^{(g)} -\in\mathbb{R}^{D\times 4D}, +\in +\mathbb{R}^{d_{\mathrm{group}}\times4d_{\mathrm{group}}}, \] \[ W_{o,\mathrm{intra}}^{(g)} -\in\mathbb{R}^{4D\times D}. +\in +\mathbb{R}^{4d_{\mathrm{group}}\times d_{\mathrm{group}}}. \] 计算： \[ -P_g -=\operatorname{SiLU} -\left(\widetilde H^{(0)}_g -W_{g,\mathrm{intra}}^{(g)}\right) +H_g += +\operatorname{SiLU} +\left(G_gW_{g,\mathrm{intra}}^{(g)}\right) \odot -\left(\widetilde H^{(0)}_g -W_{v,\mathrm{intra}}^{(g)}\right), +\left(G_gW_{v,\mathrm{intra}}^{(g)}\right), \] \[ -\Delta_{\mathrm{intra},g} -=P_gW_{o,\mathrm{intra}}^{(g)}, +P_g=H_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] @@ -179,222 +181,205 @@ intra_output_proj: [n_group, 4 * d_group, d_group] 三个 projection 均不带 bias。 -## 5. 第二阶段：跨组 TrajMixer +## 6. 静态门控融合 -第二阶段沿 group 维度进行交互。对于每个内部坐标 \(r\)，独立执行： +定义可学习 gate logits： \[ -G\rightarrow4G\rightarrow G. +\Theta\in +\mathbb{R}^{n_{\mathrm{group}}\times d_{\mathrm{group}}}. \] -首先将 \(H^{(1)}\) 的最后两个维度交换，并在 group 维度执行 LayerNorm： +实际门值为： \[ -\widetilde H^{(1)}_{b,t,:,r} -=\operatorname{LN}_{G} -\left(H^{(1)}_{b,t,:,r}\right). +\Gamma=\sigma(\Theta). \] -归一化统计量对每个内部坐标 \(r\) 独立计算；LayerNorm 的 affine 参数在各内部坐标间共享。 - -### 5.1 Group-wise Feature Alignment - -沿用现有的可学习 group 特征对齐矩阵： +初始化： \[ -B_g\in\mathbb{R}^{D\times D}, -\qquad g=1,\ldots,G, +\Theta_{g,r} +=\operatorname{logit}(0.1) +=\log\frac{0.1}{0.9} +\approx-2.1972, \] +因此： + \[ -Z_{b,t,g,:} -=\widetilde H^{(1)}_{b,t,g,:}B_g. +\Gamma_{g,r}\approx0.1. \] -\(B_g\) 不带 bias，并使用单位矩阵初始化。 +融合： -### 5.2 Cross-Group SwiGLU +\[ +U=G+\Gamma\odot P. +\] -对每个内部坐标 \(r=1,\ldots,D\)，定义： +\(\Gamma\) 对 batch 和序列位置共享，但每个 group、每个内部坐标拥有独立可学习值。 + +## 7. 跨 Group SwiGLU + +对于每个内部坐标 \(r\)，独立沿 group 维度执行： + +\[ +n_{\mathrm{group}} +\rightarrow +4n_{\mathrm{group}} +\rightarrow +n_{\mathrm{group}}. +\] + +定义： \[ A_g^{(r)},A_v^{(r)} -\in\mathbb{R}^{G\times4G}, -\qquad +\in +\mathbb{R}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}}, +\] + +\[ A_o^{(r)} -\in\mathbb{R}^{4G\times G}. +\in +\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}. \] 计算： \[ -Q_{b,t,:,r} -=\operatorname{SiLU} -\left(Z_{b,t,:,r}A_g^{(r)}\right) +Q_{:,r} += +\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right) \odot -\left(Z_{b,t,:,r}A_v^{(r)}\right), +\left(U_{:,r}A_v^{(r)}\right), \] \[ -\Delta_{\mathrm{cross},b,t,:,r} -=Q_{b,t,:,r}A_o^{(r)}, +\Delta_{:,r}=Q_{:,r}A_o^{(r)}. \] -\[ -H^{(2)} -=H^{(1)} -+\operatorname{Dropout}(\Delta_{\mathrm{cross}}). -\] - -实现张量形状： +实现形状： ```text -cross_norm: LayerNorm(n_group) -group_align: [n_group, d_group, d_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。 +三个 projection 均不带 bias。不同内部坐标拥有独立的跨 group 参数，且不沿序列维度交互。 -## 6. PreNorm 与残差约束 +## 8. 唯一的外层 Residual -本版本固定使用两个独立的 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 组内阶段 - -投影权重： +跨 group 输出 reshape 回： \[ -3G D(4D) -=12GD^2 -=17{,}280. +\Delta\in\mathbb{R}^{B\times L\times d}. \] -`LayerNorm(D)`： +最终： \[ -2D=24. +\operatorname{TrajMixer}(X) +=X+\operatorname{Dropout}(\Delta). \] -### 7.2 跨组阶段 +固定约束： -跨组投影权重： +- 组内阶段后不执行独立 residual； +- 跨组阶段后不执行独立 residual； +- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`； +- 整个 TrajMixer 只有一次主 residual。 -\[ -3D G(4G) -=12DG^2 -=14{,}400. -\] +## 9. 初始化 -Group Feature Alignment： +固定初始化： -\[ -GD^2 -=1{,}440. -\] - -`LayerNorm(G)`： - -\[ -2G=20. -\] - -### 7.3 每个 TrajMixer 合计 - -\[ -17{,}280+24+14{,}400+1{,}440+20 -=\boxed{33{,}164}. -\] - -相对于 `traj_mixer_v2` 的跨组单阶段结构 \(15{,}840\)，每层增加 \(17{,}324\) 个参数。作为历史实现对照，代码库原全维度 SwiGLU FFN 每层为 \(108{,}720\) 个参数。 - -## 8. 初始化 - -固定初始化约定： - -- 组内 `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform； -- 组内 `intra_output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布； -- Group Alignment：单位矩阵； +- `intra_gate_proj/intra_value_proj`：每个 group 独立 Xavier uniform； +- `intra_output_proj`：每个 group 独立 Xavier uniform； +- `intra_gate_logits`：初始化为 \(\operatorname{logit}(0.1)\)； - 跨组 `gate_proj/value_proj`：每个内部坐标独立 Xavier uniform； -- 跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布； -- 两个 LayerNorm：PyTorch 默认 affine 初始化； -- 两个 residual stage 的 Dropout 均沿用 `mlp_dropout`。 +- 最终跨组 `output_proj`：均值 0、标准差 \(10^{-3}\) 的正态分布； +- Full-width LayerNorm：PyTorch 默认 affine 初始化； +- Dropout：沿用 `mlp_dropout`。 -两个 output projection 的小方差初始化使两阶段在训练初期都接近恒等 residual update。 +组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。 -Relative Time Attention Bias 的初始化固定为： +Relative Time Attention Bias 初始化固定为： - `rbf_proj.weight`：零初始化； - `time_bias_scale`：初始化为 \(1.0\)； -- 初始 RBF attention bias 仍严格为零； +- 初始 RBF attention bias 严格为零； - `rbf_proj.weight` 从第一个优化步骤即可获得梯度。 -不得同时将 `rbf_proj.weight` 和 `time_bias_scale` 初始化为零，否则两个相乘分支的梯度都会为零，RBF 时间偏置将无法开始学习。 +## 10. 参数量 -## 9. 信息流与语义 +默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。 -**Attention**：从历史疾病事件中选择和整合相关信息。 +Full-width LayerNorm： -**Intra-Group Mixer**：学习每条潜在轨迹内部的非线性特征组合。 +\[ +2d=240. +\] -**Group Feature Alignment**：对齐不同潜在轨迹的内部坐标。 +组内 projections： -**Cross-Group Mixer**：学习不同潜在轨迹在相同内部坐标上的门控交互。 +\[ +3n_{\mathrm{group}}d_{\mathrm{group}} +\left(4d_{\mathrm{group}}\right) +=17{,}280. +\] -整个模块保持： +静态门控： -- 无时间递归； -- 不混合序列位置； -- 序列维度完全并行； -- 保留原始因果 Attention； -- residual groups 不等同于 Attention heads。 +\[ +n_{\mathrm{group}}d_{\mathrm{group}} +=120. +\] -## 10. 固定配置与 checkpoint 约束 +跨 group projections： + +\[ +3d_{\mathrm{group}}n_{\mathrm{group}} +\left(4n_{\mathrm{group}}\right) +=14{,}400. +\] + +每层 TrajMixer 合计： + +\[ +240+17{,}280+120+14{,}400 +=\boxed{32{,}040}. +\] + +默认 relative-time、12 层、`vocab_size=1256`、无额外信息类型时，完整模型参数量为： + +\[ +\boxed{1{,}232{,}428}. +\] + +## 11. 固定配置与 checkpoint 约束 ```yaml -model_architecture: traj_mixer_v3 +model_architecture: traj_mixer_v5 d_model: 120 n_head: 10 n_group_rule: n_head d_group_rule: d_model / n_group +traj_mixer_norm: layer_norm_over_n_embd intra_hidden_rule: 4 * d_group +intra_gate_shape: [n_group, d_group] +intra_gate_initial_sigmoid: 0.1 cross_hidden_rule: 4 * n_group -attention: unchanged -intra_norm: layer_norm_over_d_group -cross_norm: layer_norm_over_n_group -group_alignment: per_group_d_group_x_d_group +group_alignment: false +intra_residual: false +cross_residual: false +traj_mixer_outer_residual: true projection_bias: false -gate_value_init: xavier_uniform -output_init_std: 0.001 +intra_output_init: xavier_uniform +cross_output_init_std: 0.001 ``` -必须满足： +训练时必须将 `model_architecture: traj_mixer_v5`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在日志中打印参数量。 -\[ -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 不向后兼容，直接拒绝加载。 +评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。 diff --git a/backbones.py b/backbones.py index f4651d1..f60b74c 100644 --- a/backbones.py +++ b/backbones.py @@ -180,7 +180,7 @@ class TemporalAttention(nn.Module): class TrajMixer(nn.Module): - """Two-stage gated mixing within and across latent trajectory groups. + """PreNorm gated mixing within and across latent trajectory groups. The groups are contiguous partitions of the post-``W_O`` residual representation. They are deliberately not treated as attention heads. @@ -211,8 +211,10 @@ class TrajMixer(nn.Module): self.intra_hidden = 4 * self.d_group self.hidden_group = 4 * n_head + # A single full-width PreNorm serves the entire TrajMixer branch. + self.norm = nn.LayerNorm(self.n_embd) + # Stage 1: each group independently mixes its internal features. - self.intra_norm = nn.LayerNorm(self.d_group) self.intra_gate_proj = nn.Parameter( torch.empty(self.n_group, self.d_group, self.intra_hidden) ) @@ -222,13 +224,8 @@ class TrajMixer(nn.Module): 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) + self.intra_gate_logits = nn.Parameter( + torch.empty(self.n_group, self.d_group) ) # Per-feature cross-group projections. The feature index is kept @@ -249,15 +246,11 @@ class TrajMixer(nn.Module): 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, - dtype=self.group_align.dtype, - device=self.group_align.device, - ) - self.group_align.copy_(identity.unsqueeze(0).expand_as(self.group_align)) + nn.init.xavier_uniform_(self.intra_output_proj[group_idx]) + nn.init.constant_( + self.intra_gate_logits, + math.log(0.1 / 0.9), + ) # Initialise each feature-specific matrix independently so Xavier's # fan-in/fan-out calculation sees a two-dimensional matrix. @@ -266,8 +259,34 @@ class TrajMixer(nn.Module): nn.init.xavier_uniform_(self.value_proj[feature_idx]) nn.init.normal_(self.output_proj, mean=0.0, std=1e-3) + def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor: + """Mix features independently inside each residual-space group.""" + intra_gate = torch.einsum( + "blgd,gdh->blgh", grouped, self.intra_gate_proj + ) + intra_value = torch.einsum( + "blgd,gdh->blgh", grouped, self.intra_value_proj + ) + intra_hidden = F.silu(intra_gate) * intra_value + return torch.einsum( + "blgh,ghd->blgd", intra_hidden, self.intra_output_proj + ) + + def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor: + """Mix groups independently for each within-group coordinate.""" + gate = torch.einsum( + "blgr,rgh->blhr", grouped, self.gate_proj + ) + value = torch.einsum( + "blgr,rgh->blhr", grouped, self.value_proj + ) + hidden = F.silu(gate) * value + return torch.einsum( + "blhr,rhg->blgr", hidden, self.output_proj + ) + def forward(self, x: torch.Tensor) -> torch.Tensor: - """Apply two PreNorm residual stages without mixing sequence positions.""" + """Apply one full-width PreNorm and one outer residual update.""" if x.ndim != 3: raise ValueError(f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}") if x.size(-1) != self.n_embd: @@ -276,44 +295,22 @@ class TrajMixer(nn.Module): ) batch_size, seq_len, _ = x.shape - grouped = x.reshape( + grouped = self.norm(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 + # The static per-channel gate starts at sigmoid(logit) ~= 0.1. + intra_output = self._intra_mix(grouped) + intra_gate = torch.sigmoid(self.intra_gate_logits).view( + 1, 1, self.n_group, self.d_group ) - 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) + mixed_input = grouped + intra_gate * intra_output # 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", cross_input, self.group_align + update = self._cross_mix(mixed_input).reshape( + batch_size, seq_len, self.n_embd ) - - gate = torch.einsum( - "blgr,rgh->blhr", aligned, self.gate_proj - ) - value = torch.einsum( - "blgr,rgh->blhr", aligned, self.value_proj - ) - hidden = F.silu(gate) * value - mixed = torch.einsum( - "blhr,rhg->blgr", hidden, self.output_proj - ) - grouped = grouped + self.drop(mixed) - return grouped.reshape(batch_size, seq_len, self.n_embd) + return x + self.drop(update) class GPTBlock(nn.Module): diff --git a/models.py b/models.py index 204c490..f859c9a 100644 --- a/models.py +++ b/models.py @@ -15,7 +15,7 @@ from backbones import ( from targets import PAD_IDX -TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3" +TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5" def validate_traj_mixer_config(config: Mapping[str, object]) -> None: @@ -29,14 +29,12 @@ 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.norm.weight", + "blocks.0.mlp.norm.bias", "blocks.0.mlp.intra_gate_proj", "blocks.0.mlp.intra_value_proj", "blocks.0.mlp.intra_output_proj", - "blocks.0.mlp.cross_norm.weight", - "blocks.0.mlp.cross_norm.bias", - "blocks.0.mlp.group_align", + "blocks.0.mlp.intra_gate_logits", "blocks.0.mlp.gate_proj", "blocks.0.mlp.value_proj", "blocks.0.mlp.output_proj", diff --git a/test_traj_mixer.py b/test_traj_mixer.py index b20e991..60feb53 100644 --- a/test_traj_mixer.py +++ b/test_traj_mixer.py @@ -57,10 +57,16 @@ 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()), 33_164) - - expected = torch.eye(12).expand(10, 12, 12) - torch.testing.assert_close(mixer.group_align.detach(), expected) + self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040) + self.assertFalse(hasattr(mixer, "group_align")) + self.assertFalse(hasattr(mixer, "intra_norm")) + self.assertFalse(hasattr(mixer, "cross_norm")) + self.assertEqual(tuple(mixer.norm.normalized_shape), (120,)) + self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12)) + torch.testing.assert_close( + torch.sigmoid(mixer.intra_gate_logits.detach()), + torch.full((10, 12), 0.1), + ) self.assertEqual(mixer.intra_hidden, 48) self.assertEqual( tuple(mixer.intra_gate_proj.shape), @@ -78,35 +84,42 @@ class TrajMixerTest(unittest.TestCase): 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: + def test_zero_final_output_projection_makes_mixer_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_forward_matches_single_outer_residual_formula(self) -> None: + torch.manual_seed(0) + mixer = TrajMixer(120, n_head=10, dropout=0.0) + mixer.eval() + x = torch.randn(2, 5, 120) + + grouped = mixer.norm(x).reshape(2, 5, 10, 12) + intra_output = mixer._intra_mix(grouped) + static_gate = torch.sigmoid(mixer.intra_gate_logits).view( + 1, 1, 10, 12 + ) + mixed_input = grouped + static_gate * intra_output + update = mixer._cross_mix(mixed_input).reshape(2, 5, 120) + + torch.testing.assert_close(mixer(x), x + update) + def test_intra_stage_is_independent_across_groups(self) -> None: torch.manual_seed(0) mixer = TrajMixer(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 - ) + original_out = mixer._intra_mix(grouped) + changed_out = mixer._intra_mix(changed) unchanged_groups = torch.tensor([0, 1, 2, 4, 5, 6, 7, 8, 9]) torch.testing.assert_close( original_out.index_select(2, unchanged_groups), @@ -117,7 +130,6 @@ class TrajMixerTest(unittest.TestCase): 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_() @@ -138,8 +150,8 @@ class TrajMixerTest(unittest.TestCase): 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) + original_out = mixer._cross_mix(grouped) + changed_out = mixer._cross_mix(changed) self.assertNotEqual( original_out[0, 0, 1, 0].item(), @@ -178,12 +190,13 @@ class TrajMixerTest(unittest.TestCase): self.assertIsNotNone(parameter.grad, name) self.assertTrue(torch.isfinite(parameter.grad).all(), name) - def test_gpt_block_delegates_both_mixer_residuals_to_traj_mixer(self) -> None: + def test_gpt_block_delegates_single_mixer_residual_to_traj_mixer(self) -> None: block = GPTBlock(n_embd=120, n_head=10) self.assertIsInstance(block.mlp, TrajMixer) self.assertFalse(hasattr(block, "ln2")) - self.assertIsInstance(block.mlp.intra_norm, torch.nn.LayerNorm) - self.assertIsInstance(block.mlp.cross_norm, torch.nn.LayerNorm) + self.assertIsInstance(block.mlp.norm, torch.nn.LayerNorm) + self.assertFalse(hasattr(block.mlp, "intra_norm")) + self.assertFalse(hasattr(block.mlp, "cross_norm")) x = torch.randn(2, 6, 120) self.assertEqual(block(x).shape, x.shape) @@ -200,6 +213,14 @@ class TrajMixerTest(unittest.TestCase): validate_traj_mixer_config( {"model_architecture": "traj_mixer_v2"} ) + with self.assertRaisesRegex(ValueError, "only accepts models trained"): + validate_traj_mixer_config( + {"model_architecture": "traj_mixer_v3"} + ) + with self.assertRaisesRegex(ValueError, "only accepts models trained"): + validate_traj_mixer_config( + {"model_architecture": "traj_mixer_v4"} + ) def test_checkpoint_must_contain_traj_mixer_parameters(self) -> None: block = GPTBlock(n_embd=120, n_head=10) @@ -222,8 +243,8 @@ class TrajMixerTest(unittest.TestCase): self.assertEqual( get_model_parameter_counts(mixer), { - "model_parameter_count": 33_164, - "trainable_parameter_count": 33_164, + "model_parameter_count": 32_040, + "trainable_parameter_count": 32_040, }, )