Fix RBF time-bias initialization

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**：从历史疾病事件中选择和整合相关信息。

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

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,