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

Add nano and tiny model presets

Browse Source
This commit is contained in:
Jiarui Li
2026-07-23 15:01:33 +08:00
parent 06f29c0f0a
commit 6f7b5be405
3 changed files with 38 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md
View File

@@ -1,7 +1,7 @@
# Event–Trajectory Shared Reasoning Backbone
> 状态：**Frozen implementation baseline**
> 架构标识：`event_trajectory_shared_v1`
> 架构标识：`event_trajectory_shared_v2`
> 固化日期：**2026-07-23**
## 1. 核心定义
@@ -16,11 +16,12 @@ self.reasoning_core = SharedEventTrajectoryCore(...)
禁止为不同推理轮创建独立 Transformer blocks。参数只保存一套，计算上顺序运行多轮。
模型规模固定为四档：
模型规模固定为五档：
| model_size | d_model | n_trajectory | trajectory_dim | traj_hidden |
|---|---:|---:|---:|---:|
| nano | 256 | 8 | 32 | 32 |
| nano | 120 | 6 | 20 | 24 |
| tiny | 256 | 8 | 32 | 32 |
| small | 512 | 8 | 64 | 32 |
| medium | 768 | 12 | 64 | 48 |
| huge | 1024 | 16 | 64 | 64 |
@@ -70,10 +71,11 @@ event_key_value = reasoning_core.project_event_memory(E)
## 3. Trajectory State
每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用8个：
每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用6个：
\[
S\in\mathbb{R}^{B\times Q\times8\times32}.
S\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
d_{\mathrm{trajectory}}}.
\]
其中：
@@ -84,7 +86,7 @@ S\in\mathbb{R}^{B\times Q\times8\times32}.
定义可学习原型：
\[
P\in\mathbb{R}^{8\times32}.
P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}.
\]
查询上下文经过投影并 reshape：
@@ -92,7 +94,8 @@ P\in\mathbb{R}^{8\times32}.
\[
C_Q
=\operatorname{QueryProjection}(\text{query features})
\in\mathbb{R}^{B\times Q\times8\times32},
\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times
d_{\mathrm{trajectory}}},
\]
\[
@@ -340,12 +343,12 @@ next-token 的 query 位置全部并行，只有 reasoning rounds 顺序执行
训练配置必须写入：
```yaml
model_architecture: event_trajectory_shared_v1
model_architecture: event_trajectory_shared_v2
model_size: nano
d_model: 256
n_trajectory: 8
trajectory_dim: 32
traj_hidden: 32
d_model: 120
n_trajectory: 6
trajectory_dim: 20
traj_hidden: 24
n_reasoning_rounds: 12
model_parameter_count: <runtime count>
trainable_parameter_count: <runtime count>
@@ -354,7 +357,7 @@ trainable_parameter_count: <runtime count>
评估和导出入口必须同时验证：
1. `model_architecture` 完全匹配；
2. `model_size` 属于 `nano / small / medium / huge`；
2. `model_size` 属于 `nano / tiny / small / medium / huge`；
3. `d_model`、`n_trajectory`、`trajectory_dim` 和 `traj_hidden`
与对应规模预设完全匹配；
4. checkpoint 包含一套且仅一套 `reasoning_core` 关键参数；

5
models.py
View File

@@ -15,7 +15,7 @@ from backbones import (
from targets import PAD_IDX
EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v1"
EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v2"
@dataclass(frozen=True)
@@ -33,7 +33,8 @@ class EventTrajectoryModelSize:
MODEL_SIZE_PRESETS = {
"nano": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
"nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6),
"tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8),
"small": EventTrajectoryModelSize(d_model=512, n_trajectory=8),
"medium": EventTrajectoryModelSize(d_model=768, n_trajectory=12),
"huge": EventTrajectoryModelSize(d_model=1024, n_trajectory=16),

33
test_event_trajectory_backbone.py
View File

@@ -58,7 +58,8 @@ def model_inputs() -> dict[str, torch.Tensor]:
class EventTrajectoryBackboneTest(unittest.TestCase):
def test_model_size_presets(self) -> None:
expected = {
"nano": (256, 8, 32, 32),
"nano": (120, 6, 20, 24),
"tiny": (256, 8, 32, 32),
"small": (512, 8, 64, 32),
"medium": (768, 12, 64, 48),
"huge": (1024, 16, 64, 64),
@@ -186,9 +187,9 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
inputs = model_inputs()
next_model = build_test_model(target_mode="next_token")
next_hidden = next_model(**inputs)
self.assertEqual(tuple(next_hidden.shape), (2, 4, 256))
self.assertEqual(tuple(next_hidden.shape), (2, 4, 120))
next_output = next_model(**inputs, return_output=True)
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 256))
self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120))
self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4))
future_model = build_test_model(target_mode="all_future")
@@ -196,7 +197,7 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
**inputs,
t_query=torch.tensor([5.0, 3.0]),
)
self.assertEqual(tuple(future_hidden.shape), (2, 256))
self.assertEqual(tuple(future_hidden.shape), (2, 120))
def test_model_contains_one_shared_core_and_no_block_stack(self) -> None:
model = build_test_model(n_reasoning_rounds=12)
@@ -261,10 +262,10 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
{
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
"model_size": "nano",
"d_model": 256,
"n_trajectory": 8,
"trajectory_dim": 32,
"traj_hidden": 32,
"d_model": 120,
"n_trajectory": 6,
"trajectory_dim": 20,
"traj_hidden": 24,
"n_reasoning_rounds": 3,
}
)
@@ -272,15 +273,19 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
validate_event_trajectory_config(
{"model_architecture": "traj_mixer_v2"}
)
with self.assertRaisesRegex(ValueError, "only accepts models trained"):
validate_event_trajectory_config(
{"model_architecture": "event_trajectory_shared_v1"}
)
with self.assertRaisesRegex(ValueError, "trajectory_dim"):
validate_event_trajectory_config(
{
"model_architecture": EVENT_TRAJECTORY_ARCHITECTURE,
"model_size": "nano",
"d_model": 256,
"n_trajectory": 8,
"trajectory_dim": 16,
"traj_hidden": 32,
"d_model": 120,
"n_trajectory": 6,
"trajectory_dim": 10,
"traj_hidden": 24,
"n_reasoning_rounds": 3,
}
)
@@ -289,8 +294,8 @@ class EventTrajectoryBackboneTest(unittest.TestCase):
state_dict = model.state_dict()
validate_event_trajectory_state_dict(
state_dict,
expected_d_model=256,
expected_n_trajectory=8,
expected_d_model=120,
expected_n_trajectory=6,
expected_n_reasoning_rounds=3,
)
with self.assertRaisesRegex(