JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Scale TrajMixer hidden width with head count

Browse Source
This commit is contained in:
Jiarui Li
2026-07-22 16:14:52 +08:00
parent 978c88a4ed
commit 68a6a3df88
8 changed files with 29 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
TrajMixer_设计方案.md
View File

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

16
backbones.py
View File

@@ -189,7 +189,6 @@ class TrajMixer(nn.Module):
self, self,
n_embd: int, n_embd: int,
n_head: int = 10, n_head: int = 10,
hidden_group: int = 20,
dropout: float = 0.0, dropout: float = 0.0,
): ):
super().__init__() super().__init__()
@@ -201,17 +200,12 @@ class TrajMixer(nn.Module):
raise ValueError( raise ValueError(
f"n_embd must be divisible by n_head, got {n_embd} and {n_head}" 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 self.n_embd = n_embd
# The residual-group count is tied to n_head, but the resulting groups # The residual-group count is tied to n_head, but the resulting groups
# are still residual-space partitions rather than attention heads. # are still residual-space partitions rather than attention heads.
self.n_group = n_head self.n_group = n_head
self.d_group = n_embd // 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]. # Per-group feature alignment: [group, input feature, output feature].
self.group_align = nn.Parameter( self.group_align = nn.Parameter(
@@ -221,13 +215,13 @@ class TrajMixer(nn.Module):
# Per-feature cross-group projections. The feature index is kept # Per-feature cross-group projections. The feature index is kept
# independent, exactly as specified by the TrajMixer baseline. # independent, exactly as specified by the TrajMixer baseline.
self.gate_proj = nn.Parameter( 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( 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( 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.drop = nn.Dropout(dropout)
self.reset_parameters() self.reset_parameters()
@@ -286,7 +280,6 @@ class GPTBlock(nn.Module):
attn_dropout: float = 0.0, attn_dropout: float = 0.0,
mlp_dropout: float = 0.0, mlp_dropout: float = 0.0,
hidden_group: int = 20,
use_time_rope: bool = False, use_time_rope: bool = False,
use_rbf_bias: bool = False, use_rbf_bias: bool = False,
n_rbf_bases: int = 16, n_rbf_bases: int = 16,
@@ -303,7 +296,6 @@ class GPTBlock(nn.Module):
self.mlp = TrajMixer( self.mlp = TrajMixer(
n_embd=n_embd, n_embd=n_embd,
n_head=n_head, n_head=n_head,
hidden_group=hidden_group,
dropout=mlp_dropout, dropout=mlp_dropout,
) )
self.ln1 = nn.LayerNorm(n_embd) self.ln1 = nn.LayerNorm(n_embd)

1
evaluate_auc.py
View File

@@ -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")), time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)), dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)),
) )

1
evaluate_auc_v2.py
View File

@@ -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")), time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)), dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
hidden_group=int(cfg_get(args, cfg, "hidden_group", 20)),
) )

5
models.py
View File

@@ -15,7 +15,7 @@ from backbones import (
from targets import PAD_IDX 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: 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" dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
extra_pool_reduce: str = "mean", extra_pool_reduce: str = "mean",
dropout: float = 0.0, dropout: float = 0.0,
hidden_group: int = 20,
): ):
super().__init__() super().__init__()
if target_mode not in ["next_token", "all_future"]: if target_mode not in ["next_token", "all_future"]:
@@ -243,7 +242,6 @@ class DeepHealth(nn.Module):
use_time_rope=False, use_time_rope=False,
use_rbf_bias=False, use_rbf_bias=False,
mlp_dropout=dropout, mlp_dropout=dropout,
hidden_group=hidden_group,
) for _ in range(n_hist_layer) ) for _ in range(n_hist_layer)
]) ])
self.rope = None self.rope = None
@@ -257,7 +255,6 @@ class DeepHealth(nn.Module):
use_time_rope=True, use_time_rope=True,
use_rbf_bias=True, use_rbf_bias=True,
mlp_dropout=dropout, mlp_dropout=dropout,
hidden_group=hidden_group,
) for _ in range(n_hist_layer) ) for _ in range(n_hist_layer)
]) ])
self.rope = TimeRoPE(n_embd // n_head) self.rope = TimeRoPE(n_embd // n_head)

22
test_traj_mixer.py
View File

@@ -16,23 +16,23 @@ class TrajMixerTest(unittest.TestCase):
mixer = TrajMixer( mixer = TrajMixer(
n_embd=120, n_embd=120,
n_head=10, n_head=10,
hidden_group=20,
dropout=0.0, dropout=0.0,
) )
x = torch.randn(2, 7, 120) x = torch.randn(2, 7, 120)
self.assertEqual(mixer(x).shape, x.shape) 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) expected = torch.eye(12).expand(10, 12, 12)
torch.testing.assert_close(mixer.group_align.detach(), expected) torch.testing.assert_close(mixer.group_align.detach(), expected)
self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 20)) self.assertEqual(mixer.hidden_group, 40)
self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 20)) self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
self.assertEqual(tuple(mixer.output_proj.shape), (12, 20, 10)) 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: def test_mixer_does_not_mix_sequence_positions(self) -> None:
torch.manual_seed(0) 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() mixer.eval()
x = torch.randn(2, 5, 120) x = torch.randn(2, 5, 120)
changed = x.clone() changed = x.clone()
@@ -48,7 +48,7 @@ class TrajMixerTest(unittest.TestCase):
def test_gradients_reach_all_projection_families(self) -> None: def test_gradients_reach_all_projection_families(self) -> None:
torch.manual_seed(1) 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) x = torch.randn(2, 4, 120, requires_grad=True)
mixer(x).square().mean().backward() mixer(x).square().mean().backward()
@@ -90,15 +90,15 @@ class TrajMixerTest(unittest.TestCase):
def test_invalid_group_partition_is_rejected(self) -> None: def test_invalid_group_partition_is_rejected(self) -> None:
with self.assertRaisesRegex(ValueError, "divisible"): 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: 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( self.assertEqual(
get_model_parameter_counts(mixer), get_model_parameter_counts(mixer),
{ {
"model_parameter_count": 8_640, "model_parameter_count": 15_840,
"trainable_parameter_count": 8_640, "trainable_parameter_count": 15_840,
}, },
) )

2
train_all_future.py
View File

@@ -90,7 +90,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--dist_mode", type=str, default="exponential", parser.add_argument("--dist_mode", type=str, default="exponential",
choices=["exponential", "weibull", "mixed"]) choices=["exponential", "weibull", "mixed"])
parser.add_argument("--dropout", type=float, default=0.0) 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("--batch_size", type=int, default=128)
parser.add_argument("--base_lr", type=float, default=3e-4) 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, time_mode=args.time_mode,
dist_mode=args.dist_mode, dist_mode=args.dist_mode,
dropout=args.dropout, dropout=args.dropout,
hidden_group=args.hidden_group,
) )

2
train_next_step.py
View File

@@ -84,7 +84,6 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--time_mode", type=str, default="relative", parser.add_argument("--time_mode", type=str, default="relative",
choices=["relative", "absolute"]) choices=["relative", "absolute"])
parser.add_argument("--dropout", type=float, default=0.0) 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", parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"]) choices=["delphi2m", "uts"])
@@ -166,7 +165,6 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
time_mode=args.time_mode, time_mode=args.time_mode,
dist_mode="exponential", dist_mode="exponential",
dropout=args.dropout, dropout=args.dropout,
hidden_group=args.hidden_group,
) )