diff --git a/TrajMixer_设计方案.md b/TrajMixer_设计方案.md index 48ba7f2..961a893 100644 --- a/TrajMixer_设计方案.md +++ b/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 参数张量的模型；其他版本或分支生成的模型应直接拒绝加载。 必须满足： diff --git a/backbones.py b/backbones.py index 6d6e6ab..ed11813 100644 --- a/backbones.py +++ b/backbones.py @@ -189,7 +189,6 @@ class TrajMixer(nn.Module): self, n_embd: int, n_head: int = 10, - hidden_group: int = 20, dropout: float = 0.0, ): super().__init__() @@ -201,17 +200,12 @@ class TrajMixer(nn.Module): raise ValueError( f"n_embd must be divisible by n_head, got {n_embd} and {n_head}" ) - if hidden_group <= 0: - raise ValueError( - f"hidden_group must be > 0, got {hidden_group}" - ) - self.n_embd = n_embd # The residual-group count is tied to n_head, but the resulting groups # are still residual-space partitions rather than attention heads. self.n_group = n_head self.d_group = n_embd // n_head - self.hidden_group = hidden_group + self.hidden_group = 4 * n_head # Per-group feature alignment: [group, input feature, output feature]. self.group_align = nn.Parameter( @@ -221,13 +215,13 @@ class TrajMixer(nn.Module): # Per-feature cross-group projections. The feature index is kept # independent, exactly as specified by the TrajMixer baseline. self.gate_proj = nn.Parameter( - torch.empty(self.d_group, self.n_group, hidden_group) + torch.empty(self.d_group, self.n_group, self.hidden_group) ) self.value_proj = nn.Parameter( - torch.empty(self.d_group, self.n_group, hidden_group) + torch.empty(self.d_group, self.n_group, self.hidden_group) ) self.output_proj = nn.Parameter( - torch.empty(self.d_group, hidden_group, self.n_group) + torch.empty(self.d_group, self.hidden_group, self.n_group) ) self.drop = nn.Dropout(dropout) self.reset_parameters() @@ -286,7 +280,6 @@ class GPTBlock(nn.Module): attn_dropout: float = 0.0, mlp_dropout: float = 0.0, - hidden_group: int = 20, use_time_rope: bool = False, use_rbf_bias: bool = False, n_rbf_bases: int = 16, @@ -303,7 +296,6 @@ class GPTBlock(nn.Module): self.mlp = TrajMixer( n_embd=n_embd, n_head=n_head, - hidden_group=hidden_group, dropout=mlp_dropout, ) self.ln1 = nn.LayerNorm(n_embd) diff --git a/evaluate_auc.py b/evaluate_auc.py index dbc1de7..c28b5b0 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -336,7 +336,6 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), - hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)), ) diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index f7abae3..a45b331 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -205,7 +205,6 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data time_mode=str(cfg_get(args, cfg, "time_mode", "relative")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), - hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)), ) diff --git a/models.py b/models.py index 9359ee7..1d2bbd4 100644 --- a/models.py +++ b/models.py @@ -15,7 +15,7 @@ from backbones import ( from targets import PAD_IDX -TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v1" +TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v2" def validate_traj_mixer_config(config: Mapping[str, object]) -> None: @@ -188,7 +188,6 @@ class DeepHealth(nn.Module): dist_mode: str = "exponential", # "exponential", "weibull" or "mixed" extra_pool_reduce: str = "mean", dropout: float = 0.0, - hidden_group: int = 20, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -243,7 +242,6 @@ class DeepHealth(nn.Module): use_time_rope=False, use_rbf_bias=False, mlp_dropout=dropout, - hidden_group=hidden_group, ) for _ in range(n_hist_layer) ]) self.rope = None @@ -257,7 +255,6 @@ class DeepHealth(nn.Module): use_time_rope=True, use_rbf_bias=True, mlp_dropout=dropout, - hidden_group=hidden_group, ) for _ in range(n_hist_layer) ]) self.rope = TimeRoPE(n_embd // n_head) diff --git a/test_traj_mixer.py b/test_traj_mixer.py index 5dc3ae2..015b4ac 100644 --- a/test_traj_mixer.py +++ b/test_traj_mixer.py @@ -16,23 +16,23 @@ class TrajMixerTest(unittest.TestCase): mixer = TrajMixer( n_embd=120, n_head=10, - hidden_group=20, dropout=0.0, ) x = torch.randn(2, 7, 120) self.assertEqual(mixer(x).shape, x.shape) - self.assertEqual(sum(p.numel() for p in mixer.parameters()), 8_640) + self.assertEqual(sum(p.numel() for p in mixer.parameters()), 15_840) expected = torch.eye(12).expand(10, 12, 12) torch.testing.assert_close(mixer.group_align.detach(), expected) - self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 20)) - self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 20)) - self.assertEqual(tuple(mixer.output_proj.shape), (12, 20, 10)) + 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)) def test_mixer_does_not_mix_sequence_positions(self) -> None: torch.manual_seed(0) - mixer = TrajMixer(120, n_head=10, hidden_group=20, dropout=0.0) + mixer = TrajMixer(120, n_head=10, dropout=0.0) mixer.eval() x = torch.randn(2, 5, 120) changed = x.clone() @@ -48,7 +48,7 @@ class TrajMixerTest(unittest.TestCase): def test_gradients_reach_all_projection_families(self) -> None: torch.manual_seed(1) - mixer = TrajMixer(120, n_head=10, hidden_group=20, dropout=0.0) + mixer = TrajMixer(120, n_head=10, dropout=0.0) x = torch.randn(2, 4, 120, requires_grad=True) mixer(x).square().mean().backward() @@ -90,15 +90,15 @@ class TrajMixerTest(unittest.TestCase): def test_invalid_group_partition_is_rejected(self) -> None: with self.assertRaisesRegex(ValueError, "divisible"): - TrajMixer(n_embd=121, n_head=10, hidden_group=20) + TrajMixer(n_embd=121, n_head=10) def test_parameter_counts_match_traj_mixer_parameters(self) -> None: - mixer = TrajMixer(n_embd=120, n_head=10, hidden_group=20) + mixer = TrajMixer(n_embd=120, n_head=10) self.assertEqual( get_model_parameter_counts(mixer), { - "model_parameter_count": 8_640, - "trainable_parameter_count": 8_640, + "model_parameter_count": 15_840, + "trainable_parameter_count": 15_840, }, ) diff --git a/train_all_future.py b/train_all_future.py index 83879cc..fce1c05 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -90,7 +90,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--dist_mode", type=str, default="exponential", choices=["exponential", "weibull", "mixed"]) parser.add_argument("--dropout", type=float, default=0.0) - parser.add_argument("--hidden_group", type=int, default=20) parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) @@ -161,7 +160,6 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De time_mode=args.time_mode, dist_mode=args.dist_mode, dropout=args.dropout, - hidden_group=args.hidden_group, ) diff --git a/train_next_step.py b/train_next_step.py index 23ef9a9..58732c3 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -84,7 +84,6 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--time_mode", type=str, default="relative", choices=["relative", "absolute"]) parser.add_argument("--dropout", type=float, default=0.0) - parser.add_argument("--hidden_group", type=int, default=20) parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) @@ -166,7 +165,6 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth: time_mode=args.time_mode, dist_mode="exponential", dropout=args.dropout, - hidden_group=args.hidden_group, )