diff --git a/backbones.py b/backbones.py index ea61077..58d76cd 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_temporal_attention.py b/test_temporal_attention.py new file mode 100644 index 0000000..a1a3475 --- /dev/null +++ b/test_temporal_attention.py @@ -0,0 +1,46 @@ +import unittest + +import torch + +from backbones import TemporalAttention + + +class TemporalAttentionTest(unittest.TestCase): + def test_zero_rbf_bias_has_live_projection_gradient(self) -> None: + torch.manual_seed(0) + 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)) + + (initial_bias * target).sum().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) + + +if __name__ == "__main__": + unittest.main()