diff --git a/TrajMixer_设计方案.md b/TrajMixer_设计方案.md index c42625b..006bec1 100644 --- a/TrajMixer_设计方案.md +++ b/TrajMixer_设计方案.md @@ -343,6 +343,15 @@ GD^2 两个 output projection 的小方差初始化使两阶段在训练初期都接近恒等 residual update。 +Relative Time Attention Bias 的初始化固定为: + +- `rbf_proj.weight`:零初始化; +- `time_bias_scale`:初始化为 \(1.0\); +- 初始 RBF attention bias 仍严格为零; +- `rbf_proj.weight` 从第一个优化步骤即可获得梯度。 + +不得同时将 `rbf_proj.weight` 和 `time_bias_scale` 初始化为零,否则两个相乘分支的梯度都会为零,RBF 时间偏置将无法开始学习。 + ## 9. 信息流与语义 **Attention**:从历史疾病事件中选择和整合相关信息。 diff --git a/backbones.py b/backbones.py index 3e94262..f4651d1 100644 --- a/backbones.py +++ b/backbones.py @@ -111,7 +111,10 @@ class TemporalAttention(nn.Module): # Layer-specific projection from shared RBF basis activations to per-head attention bias. self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False) - self.time_bias_scale = nn.Parameter(torch.tensor(0.0)) + # Keep the initial RBF attention bias exactly zero through the + # zero-initialized projection, while leaving that projection with a + # live gradient from the first optimization step. + self.time_bias_scale = nn.Parameter(torch.tensor(1.0)) self.resid_drop = nn.Dropout(dropout) self.reset_parameters() diff --git a/test_traj_mixer.py b/test_traj_mixer.py index c3957db..b20e991 100644 --- a/test_traj_mixer.py +++ b/test_traj_mixer.py @@ -2,7 +2,7 @@ import unittest import torch -from backbones import GPTBlock, TrajMixer +from backbones import GPTBlock, TemporalAttention, TrajMixer from models import ( TRAJ_MIXER_ARCHITECTURE, validate_traj_mixer_config, @@ -12,6 +12,42 @@ from train_util import get_model_parameter_counts class TrajMixerTest(unittest.TestCase): + def test_zero_rbf_bias_has_live_projection_gradient(self) -> None: + attention = TemporalAttention( + n_embd=12, + n_head=3, + use_time_rope=False, + use_rbf_bias=True, + ) + features = torch.randn(2, 4, 4, 16) + target = torch.randn(2, 4, 4, 3) + + initial_bias = ( + attention.time_bias_scale.tanh() + * attention.rbf_proj(features) + ) + torch.testing.assert_close(initial_bias, torch.zeros_like(initial_bias)) + + loss = (initial_bias * target).sum() + loss.backward() + + projection_grad = attention.rbf_proj.weight.grad + self.assertIsNotNone(projection_grad) + self.assertGreater(projection_grad.abs().sum().item(), 0.0) + + with torch.no_grad(): + attention.rbf_proj.weight.add_(projection_grad, alpha=-1e-3) + attention.zero_grad(set_to_none=True) + updated_bias = ( + attention.time_bias_scale.tanh() + * attention.rbf_proj(features) + ) + (updated_bias * target).sum().backward() + + scale_grad = attention.time_bias_scale.grad + self.assertIsNotNone(scale_grad) + self.assertGreater(scale_grad.abs().item(), 0.0) + def test_default_shape_parameters_and_initialization(self) -> None: mixer = TrajMixer( n_embd=120,