diff --git a/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md b/Event_Trajectory_Shared_Reasoning_Backbone_设计方案.md index 79962b4..45ad831 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_v2` +> 架构标识：`event_trajectory_shared_v1` > 固化日期：**2026-07-23** ## 1. 核心定义 @@ -16,12 +16,11 @@ self.reasoning_core = SharedEventTrajectoryCore(...) 禁止为不同推理轮创建独立 Transformer blocks。参数只保存一套，计算上顺序运行多轮。 -模型规模固定为五档： +模型规模固定为四档： | model_size | d_model | n_trajectory | trajectory_dim | traj_hidden | |---|---:|---:|---:|---:| -| nano | 120 | 6 | 20 | 24 | -| tiny | 256 | 8 | 32 | 32 | +| nano | 256 | 8 | 32 | 32 | | small | 512 | 8 | 64 | 32 | | medium | 768 | 12 | 64 | 48 | | huge | 1024 | 16 | 64 | 64 | @@ -71,11 +70,10 @@ event_key_value = reasoning_core.project_event_memory(E) ## 3. Trajectory State -每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用6个： +每个查询维护 `n_trajectory` 个显式 trajectory slots；nano 默认使用8个： \[ -S\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times -d_{\mathrm{trajectory}}}. +S\in\mathbb{R}^{B\times Q\times8\times32}. \] 其中： @@ -86,7 +84,7 @@ d_{\mathrm{trajectory}}}. 定义可学习原型： \[ -P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}. +P\in\mathbb{R}^{8\times32}. \] 查询上下文经过投影并 reshape： @@ -94,8 +92,7 @@ P\in\mathbb{R}^{n_{\mathrm{trajectory}}\times d_{\mathrm{trajectory}}}. \[ C_Q =\operatorname{QueryProjection}(\text{query features}) -\in\mathbb{R}^{B\times Q\times n_{\mathrm{trajectory}}\times -d_{\mathrm{trajectory}}}, +\in\mathbb{R}^{B\times Q\times8\times32}, \] \[ @@ -343,12 +340,12 @@ next-token 的 query 位置全部并行，只有 reasoning rounds 顺序执行 训练配置必须写入： ```yaml -model_architecture: event_trajectory_shared_v2 +model_architecture: event_trajectory_shared_v1 model_size: nano -d_model: 120 -n_trajectory: 6 -trajectory_dim: 20 -traj_hidden: 24 +d_model: 256 +n_trajectory: 8 +trajectory_dim: 32 +traj_hidden: 32 n_reasoning_rounds: 12 model_parameter_count: trainable_parameter_count: @@ -357,7 +354,7 @@ trainable_parameter_count: 评估和导出入口必须同时验证： 1. `model_architecture` 完全匹配； -2. `model_size` 属于 `nano / tiny / small / medium / huge`； +2. `model_size` 属于 `nano / 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 43eb50b..c72d188 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_v2" +EVENT_TRAJECTORY_ARCHITECTURE = "event_trajectory_shared_v1" @dataclass(frozen=True) @@ -33,8 +33,7 @@ class EventTrajectoryModelSize: MODEL_SIZE_PRESETS = { - "nano": EventTrajectoryModelSize(d_model=120, n_trajectory=6), - "tiny": EventTrajectoryModelSize(d_model=256, n_trajectory=8), + "nano": 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 952bc8e..d173e0d 100644 --- a/test_event_trajectory_backbone.py +++ b/test_event_trajectory_backbone.py @@ -58,8 +58,7 @@ def model_inputs() -> dict[str, torch.Tensor]: class EventTrajectoryBackboneTest(unittest.TestCase): def test_model_size_presets(self) -> None: expected = { - "nano": (120, 6, 20, 24), - "tiny": (256, 8, 32, 32), + "nano": (256, 8, 32, 32), "small": (512, 8, 64, 32), "medium": (768, 12, 64, 48), "huge": (1024, 16, 64, 64), @@ -187,9 +186,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, 120)) + self.assertEqual(tuple(next_hidden.shape), (2, 4, 256)) next_output = next_model(**inputs, return_output=True) - self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 120)) + self.assertEqual(tuple(next_output.hidden.shape), (2, 4, 256)) self.assertEqual(tuple(next_output.padding_mask.shape), (2, 4)) future_model = build_test_model(target_mode="all_future") @@ -197,7 +196,7 @@ class EventTrajectoryBackboneTest(unittest.TestCase): **inputs, t_query=torch.tensor([5.0, 3.0]), ) - self.assertEqual(tuple(future_hidden.shape), (2, 120)) + self.assertEqual(tuple(future_hidden.shape), (2, 256)) def test_model_contains_one_shared_core_and_no_block_stack(self) -> None: model = build_test_model(n_reasoning_rounds=12) @@ -262,10 +261,10 @@ class EventTrajectoryBackboneTest(unittest.TestCase): { "model_architecture": EVENT_TRAJECTORY_ARCHITECTURE, "model_size": "nano", - "d_model": 120, - "n_trajectory": 6, - "trajectory_dim": 20, - "traj_hidden": 24, + "d_model": 256, + "n_trajectory": 8, + "trajectory_dim": 32, + "traj_hidden": 32, "n_reasoning_rounds": 3, } ) @@ -273,19 +272,15 @@ 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": 120, - "n_trajectory": 6, - "trajectory_dim": 10, - "traj_hidden": 24, + "d_model": 256, + "n_trajectory": 8, + "trajectory_dim": 16, + "traj_hidden": 32, "n_reasoning_rounds": 3, } ) @@ -294,8 +289,8 @@ class EventTrajectoryBackboneTest(unittest.TestCase): state_dict = model.state_dict() validate_event_trajectory_state_dict( state_dict, - expected_d_model=120, - expected_n_trajectory=6, + expected_d_model=256, + expected_n_trajectory=8, expected_n_reasoning_rounds=3, ) with self.assertRaisesRegex(