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}}}.
\]
其中默认：
隐藏宽度不再独立配置，固定为：
\[
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 参数张量的模型；其他版本或分支生成的模型应直接拒绝加载。
必须满足：

16
backbones.py
View File

@@ -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)

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")),
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)),
)

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")),
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)),
)

5
models.py
View File

@@ -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)

22
test_traj_mixer.py
View File

@@ -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,
},
)

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",
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,
)

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",
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,
)