DeepHealth/TrajMixer_设计方案.md

# TrajMixer Block 最终设计方案

> 状态：**Frozen implementation baseline**
>
> 版本：**v3.0 / traj_mixer_v5**
>
> 固化日期：**2026-07-24**

本文档是当前 TrajMixer 的实现与实验基线。本版本采用单 PreNorm、单外层 residual、静态门控组内融合和跨 group SwiGLU。

## 1. 目标

在不改变 Delphi Transformer Attention 的前提下，用轻量、完全并行的 TrajMixer 替换 FFN。

保持不变：

- causal mask；
- TimeRoPE；
- Relative Time Attention Bias；
- Multi-Head Attention，包括 \(W_Q/W_K/W_V/W_O\)；
- 序列建模和训练目标。

TrajMixer 不沿序列维度混合，也不引入时间递归。

## 2. Block 结构

```text
PreNorm Causal Multi-Head Attention
→ Attention Residual
→ Full-width TrajMixer PreNorm
→ reshape [B, L, n_group, d_group]
→ Per-Group SwiGLU: d_group → 4d_group → d_group
→ Static Gated Fusion
→ Cross-Group SwiGLU: n_group → 4n_group → n_group
→ reshape [B, L, n_embd]
→ Dropout
→ One TrajMixer Residual
```

Attention 阶段：

\[
X
=X^{(l)}
+\operatorname{CausalMHA}
\left(\operatorname{LN}_{\mathrm{attn}}(X^{(l)})\right).
\]

TrajMixer 阶段：

\[
N=\operatorname{LN}_{d}(X),
\]

\[
G=\operatorname{reshape}(N)
\in\mathbb{R}^{B\times L\times n_{\mathrm{group}}\times d_{\mathrm{group}}},
\]

\[
P=\operatorname{IntraMixer}(G),
\]

\[
U=G+\sigma(\Theta)\odot P,
\]

\[
\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：

\[
X\in\mathbb{R}^{B\times L\times d}.
\]

定义：

\[
n_{\mathrm{group}}:=n_{\mathrm{head}},
\qquad
d_{\mathrm{group}}=\frac{d}{n_{\mathrm{group}}},
\qquad
d=n_{\mathrm{group}}d_{\mathrm{group}}.
\]

默认：

\[
d=120,\qquad
n_{\mathrm{group}}=10,\qquad
d_{\mathrm{group}}=12.
\]

这些 group 是 residual space 的连续分区，不等同于 Attention heads；二者只共享数量。

## 4. 唯一的 Full-Width PreNorm

TrajMixer 只使用一个：

```text
norm: LayerNorm(n_embd)
```

LayerNorm 作用于完整 \(d\) 维 residual representation，然后才 reshape：

\[
G=\operatorname{reshape}
\left(\operatorname{LN}_{d}(X)\right).
\]

本版本明确删除：

```text
intra_norm
cross_norm
group_align
```

不得在组内或跨组阶段再增加额外 LayerNorm。

## 5. 组内 SwiGLU

每个 group 使用独立参数，对其 \(d_{\mathrm{group}}\) 维内部特征执行：

\[
d_{\mathrm{group}}
\rightarrow
4d_{\mathrm{group}}
\rightarrow
d_{\mathrm{group}}.
\]

对 group \(g\)：

\[
W_{g,\mathrm{intra}}^{(g)},
W_{v,\mathrm{intra}}^{(g)}
\in
\mathbb{R}^{d_{\mathrm{group}}\times4d_{\mathrm{group}}},
\]

\[
W_{o,\mathrm{intra}}^{(g)}
\in
\mathbb{R}^{4d_{\mathrm{group}}\times d_{\mathrm{group}}}.
\]

计算：

\[
H_g
=
\operatorname{SiLU}
\left(G_gW_{g,\mathrm{intra}}^{(g)}\right)
\odot
\left(G_gW_{v,\mathrm{intra}}^{(g)}\right),
\]

\[
P_g=H_gW_{o,\mathrm{intra}}^{(g)}.
\]

实现形状：

```text
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。

## 6. 静态门控融合

定义可学习 gate logits：

\[
\Theta\in
\mathbb{R}^{n_{\mathrm{group}}\times d_{\mathrm{group}}}.
\]

实际门值为：

\[
\Gamma=\sigma(\Theta).
\]

初始化：

\[
\Theta_{g,r}
=\operatorname{logit}(0.1)
=\log\frac{0.1}{0.9}
\approx-2.1972,
\]

因此：

\[
\Gamma_{g,r}\approx0.1.
\]

融合：

\[
U=G+\Gamma\odot P.
\]

\(\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}^{n_{\mathrm{group}}\times4n_{\mathrm{group}}},
\]

\[
A_o^{(r)}
\in
\mathbb{R}^{4n_{\mathrm{group}}\times n_{\mathrm{group}}}.
\]

计算：

\[
Q_{:,r}
=
\operatorname{SiLU}\left(U_{:,r}A_g^{(r)}\right)
\odot
\left(U_{:,r}A_v^{(r)}\right),
\]

\[
\Delta_{:,r}=Q_{:,r}A_o^{(r)}.
\]

实现形状：

```text
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。不同内部坐标拥有独立的跨 group 参数，且不沿序列维度交互。

## 8. 唯一的外层 Residual

跨 group 输出 reshape 回：

\[
\Delta\in\mathbb{R}^{B\times L\times d}.
\]

最终：

\[
\operatorname{TrajMixer}(X)
=X+\operatorname{Dropout}(\Delta).
\]

固定约束：

- 组内阶段后不执行独立 residual；
- 跨组阶段后不执行独立 residual；
- `GPTBlock` 不再额外执行 `X + TrajMixer(X)`；
- 整个 TrajMixer 只有一次主 residual。

## 9. 初始化

固定初始化：

- `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}\) 的正态分布；
- Full-width LayerNorm：PyTorch 默认 affine 初始化；
- Dropout：沿用 `mlp_dropout`。

组内输出使用正常 Xavier 初始化以保证其具有完整表达能力；静态门控将其初始贡献限制在约 0.1。最终跨 group 输出投影保持小值初始化，使整个 TrajMixer residual update 在训练初期接近零。

Relative Time Attention Bias 初始化固定为：

- `rbf_proj.weight`：零初始化；
- `time_bias_scale`：初始化为 \(1.0\)；
- 初始 RBF attention bias 严格为零；
- `rbf_proj.weight` 从第一个优化步骤即可获得梯度。

## 10. 参数量

默认 \(d=120\)、\(n_{\mathrm{group}}=10\)、\(d_{\mathrm{group}}=12\)。

Full-width LayerNorm：

\[
2d=240.
\]

组内 projections：

\[
3n_{\mathrm{group}}d_{\mathrm{group}}
\left(4d_{\mathrm{group}}\right)
=17{,}280.
\]

静态门控：

\[
n_{\mathrm{group}}d_{\mathrm{group}}
=120.
\]

跨 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_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
group_alignment: false
intra_residual: false
cross_residual: false
traj_mixer_outer_residual: true
projection_bias: false
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`，并在日志中打印参数量。

评估和导出入口只接受 `traj_mixer_v5` checkpoint，并检查 Full-width LayerNorm、组内 projections、静态门控和跨 group projections 是否齐全。`traj_mixer_v4` 及更早 checkpoint 不向后兼容，直接拒绝加载。