Fix RBF time-bias initialization

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,