Add two-stage TrajMixer mixing
This commit is contained in:
@@ -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 不向后兼容，直接拒绝加载。
|
||||
|
||||
Reference in New Issue
Block a user