From 6f7b5be4050373699e454365d5289a799daf4570 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Thu, 23 Jul 2026 15:01:33 +0800 Subject: [PATCH] Add nano and tiny model presets --- ...tory_Shared_Reasoning_Backbone_设计方案.md | 29 ++++++++-------- models.py | 5 +-- test_event_trajectory_backbone.py | 33 +++++++++++-------- 3 files changed, 38 insertions(+), 29 deletions(-) diff --git a/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md b/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md index 45ad831..79962b4 100644 --- a/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md +++ b/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md @@ -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: trainable_parameter_count: @@ -354,7 +357,7 @@ trainable_parameter_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` 关键参数； diff --git a/models.py b/models.py index c72d188..43eb50b 100644 --- a/models.py +++ b/models.py @@ -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), diff --git a/test_event_trajectory_backbone.py b/test_event_trajectory_backbone.py index d173e0d..952bc8e 100644 --- a/test_event_trajectory_backbone.py +++ b/test_event_trajectory_backbone.py @@ -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(