Scale TrajMixer hidden width with head count

TrajMixer_设计方案.md

@@ -143,12 +143,15 @@ A_o^{(r)}
 \in\mathbb{R}^{h_{\mathrm{group}}\times n_{\mathrm{group}}}.
 \]
 
-其中默认：
+隐藏宽度不再独立配置，固定为：
 
 \[
-h_{\mathrm{group}}=20.
+h_{\mathrm{group}}=4n_{\mathrm{head}}
+=4n_{\mathrm{group}}.
 \]
 
+当前 \(n_{\mathrm{head}}=10\)，因此 \(h_{\mathrm{group}}=40\)。
+
 对固定的 batch、时间位置和内部特征维度 \(r\)，将：
 
 \[
@@ -178,7 +181,7 @@ Y_{b,t,:,r}=M_{b,t,:,r}A_o^{(r)}.
 - gate 分支控制信息写入；
 - value 分支提供交互内容；
 - output matrix 将隐藏 group 表示投影回原始 group 数量；
-- 当 \(h_{\mathrm{group}}=n_{\mathrm{group}}=10\) 时，三类矩阵退化为原始的 \(10\times10\) 方阵形式。
+- hidden group 表示固定扩展为 group 数量的 4 倍。
 
 所有 \(r\) 的输出组合为：
 
@@ -237,18 +240,18 @@ h_{\mathrm{group}}
 n_{\mathrm{group}}.
 \]
 
-默认 \(h_{\mathrm{group}}=20\) 时，Mixer 每层权重参数量为：
+固定 \(h_{\mathrm{group}}=4n_{\mathrm{group}}=40\) 时，Mixer 每层权重参数量为：
 
 \[
 3d_{\mathrm{group}}n_{\mathrm{group}}h_{\mathrm{group}}
-=3\times12\times10\times20
-=7{,}200.
+=3\times12\times10\times40
+=14{,}400.
 \]
 
 加上 Group Feature Alignment 后，TrajMixer residual branch 每层共有：
 
 \[
-7{,}200+1{,}440=8{,}640
+14{,}400+1{,}440=15{,}840
 \]
 
 个主要权重参数。作为对照，原始 \(120\rightarrow480\rightarrow120\) FFN 每层约有 115,800 个参数。
@@ -302,11 +305,11 @@ Mixer + Group-wise LayerNorm
 ## 11. 首版固定配置
 
 ```yaml
-model_architecture: traj_mixer_v1
+model_architecture: traj_mixer_v2
 d_model: 120
 n_head: 10  # 同时决定 residual group 数量
 d_group: 12
-hidden_group: 20
+hidden_group_rule: 4 * n_head  # 不单独配置
 attention: unchanged
 attention_output_projection: unchanged
 mixer_norm: standard_layer_norm
@@ -319,7 +322,7 @@ output_init_std: 0.001
 group_wise_layer_norm: false
 ```
 
-训练时必须将 `model_architecture: traj_mixer_v1`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印总参数量与可训练参数量。本分支的评估和导出入口只接受带有该架构标识、且 checkpoint 中包含 TrajMixer 参数张量的模型；其他分支生成的模型应直接拒绝加载。
+训练时必须将 `model_architecture: traj_mixer_v2`、`model_parameter_count` 和 `trainable_parameter_count` 写入 `train_config.json`，并在训练日志中显式打印总参数量与可训练参数量。本分支的评估和导出入口只接受带有该架构标识、且 checkpoint 中包含 TrajMixer 参数张量的模型；其他版本或分支生成的模型应直接拒绝加载。
 
 必须满足：