Revert "Add switchable DIFF V1 attention"
This reverts commit bf5cae8758.
This commit is contained in:
@@ -1,195 +0,0 @@
|
||||
import math
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from attention_types import resolve_attention_type
|
||||
from backbones import (
|
||||
DifferentialAttention,
|
||||
TemporalAttention,
|
||||
TimeRoPE,
|
||||
TrajMixer,
|
||||
TrajMixerBlock,
|
||||
build_attention,
|
||||
diff_lambda_init,
|
||||
)
|
||||
from models import DeepHealth
|
||||
|
||||
|
||||
class AttentionSelectorTests(unittest.TestCase):
|
||||
def test_selector_builds_independent_implementations(self):
|
||||
standard = build_attention("standard", n_embd=120, n_head=10)
|
||||
diff = build_attention("diff", n_embd=120, n_head=10)
|
||||
self.assertIsInstance(standard, TemporalAttention)
|
||||
self.assertIsInstance(diff, DifferentialAttention)
|
||||
|
||||
def test_legacy_default_is_standard(self):
|
||||
self.assertEqual(resolve_attention_type(None), "standard")
|
||||
|
||||
def test_invalid_selector_is_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
resolve_attention_type("dynamic_diff_v2")
|
||||
|
||||
def test_standard_default_and_explicit_selector_are_exactly_compatible(self):
|
||||
torch.manual_seed(17)
|
||||
default = TrajMixerBlock(n_embd=24, n_head=4)
|
||||
torch.manual_seed(17)
|
||||
explicit = TrajMixerBlock(
|
||||
n_embd=24,
|
||||
n_head=4,
|
||||
attention_type="standard",
|
||||
)
|
||||
self.assertEqual(default.state_dict().keys(), explicit.state_dict().keys())
|
||||
for key, value in default.state_dict().items():
|
||||
torch.testing.assert_close(value, explicit.state_dict()[key])
|
||||
x = torch.randn(2, 3, 24)
|
||||
torch.testing.assert_close(default(x), explicit(x))
|
||||
|
||||
|
||||
class DifferentialAttentionTests(unittest.TestCase):
|
||||
def test_shapes_pairing_and_parameter_counts(self):
|
||||
standard = TemporalAttention(
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
use_time_rope=False,
|
||||
use_rbf_bias=False,
|
||||
)
|
||||
diff = DifferentialAttention(
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
use_time_rope=False,
|
||||
use_rbf_bias=False,
|
||||
)
|
||||
|
||||
self.assertEqual(diff.n_diff_head, 5)
|
||||
self.assertEqual(diff.d_head, 12)
|
||||
self.assertEqual(diff.d_value_head, 24)
|
||||
self.assertEqual(tuple(diff.qkv.weight.shape), (360, 120))
|
||||
|
||||
x = torch.randn(2, 7, 120)
|
||||
projected = diff.qkv(x).reshape(2, 7, 3, 10, 12)
|
||||
q, k, v = projected.unbind(2)
|
||||
self.assertEqual(tuple(q.shape), (2, 7, 10, 12))
|
||||
self.assertEqual(tuple(k.shape), (2, 7, 10, 12))
|
||||
self.assertEqual(tuple(v.reshape(2, 7, 5, 24).shape), (2, 7, 5, 24))
|
||||
self.assertEqual(tuple(diff(x).shape), (2, 7, 120))
|
||||
|
||||
standard_count = sum(p.numel() for p in standard.parameters())
|
||||
diff_count = sum(p.numel() for p in diff.parameters())
|
||||
self.assertEqual(standard_count, 57_761)
|
||||
self.assertEqual(diff_count, 57_809)
|
||||
self.assertEqual(diff_count - standard_count, 4 * 12)
|
||||
|
||||
def test_relative_time_mask_and_backward_are_supported(self):
|
||||
torch.manual_seed(3)
|
||||
diff = DifferentialAttention(
|
||||
n_embd=24,
|
||||
n_head=4,
|
||||
use_time_rope=True,
|
||||
use_rbf_bias=True,
|
||||
depth=2,
|
||||
)
|
||||
x = torch.randn(2, 4, 24, requires_grad=True)
|
||||
tau = torch.tensor([[0.0, 1.0, 3.0, 6.0], [0.0, 2.0, 2.0, 5.0]])
|
||||
rope_cache = TimeRoPE(6).precompute_cache(tau)
|
||||
diff_time = tau.unsqueeze(2) - tau.unsqueeze(1)
|
||||
centers = torch.linspace(0.0, 40.0, 16)
|
||||
widths = torch.full((16,), 40.0 / 15.0)
|
||||
rbf_cache = torch.exp(
|
||||
-0.5 * ((diff_time.unsqueeze(-1) - centers) / widths).square()
|
||||
)
|
||||
causal = torch.triu(torch.full((4, 4), -1e4), diagonal=1)
|
||||
attn_mask = causal.view(1, 1, 4, 4)
|
||||
|
||||
out = diff(x, rope_cache, rbf_cache, attn_mask)
|
||||
self.assertEqual(tuple(out.shape), (2, 4, 24))
|
||||
self.assertTrue(torch.isfinite(out).all())
|
||||
out.square().mean().backward()
|
||||
self.assertIsNotNone(diff.lambda_q1.grad)
|
||||
self.assertTrue(torch.isfinite(diff.lambda_q1.grad).all())
|
||||
|
||||
def test_lambda_parameterization_and_depth_initialization(self):
|
||||
diff = DifferentialAttention(n_embd=120, n_head=10, depth=3)
|
||||
for name in ("lambda_q1", "lambda_k1", "lambda_q2", "lambda_k2"):
|
||||
self.assertEqual(tuple(getattr(diff, name).shape), (12,))
|
||||
self.assertTrue(
|
||||
math.isclose(
|
||||
diff.lambda_init,
|
||||
0.8 - 0.6 * math.exp(-0.3 * 3),
|
||||
rel_tol=0.0,
|
||||
abs_tol=1e-12,
|
||||
)
|
||||
)
|
||||
self.assertTrue(math.isclose(diff_lambda_init(0), 0.2, abs_tol=1e-12))
|
||||
|
||||
def test_no_per_head_normalization_exists(self):
|
||||
diff = DifferentialAttention(n_embd=120, n_head=10)
|
||||
normalization_types = (nn.LayerNorm,)
|
||||
if hasattr(nn, "RMSNorm"):
|
||||
normalization_types = normalization_types + (nn.RMSNorm,)
|
||||
normalizations = [
|
||||
module
|
||||
for module in diff.modules()
|
||||
if isinstance(module, normalization_types)
|
||||
]
|
||||
self.assertEqual(normalizations, [])
|
||||
|
||||
def test_diff_attention_does_not_change_traj_mixer(self):
|
||||
standard = TrajMixerBlock(
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
attention_type="standard",
|
||||
)
|
||||
diff = TrajMixerBlock(
|
||||
n_embd=120,
|
||||
n_head=10,
|
||||
attention_type="diff",
|
||||
)
|
||||
self.assertIsInstance(standard.mlp, TrajMixer)
|
||||
self.assertIsInstance(diff.mlp, TrajMixer)
|
||||
self.assertEqual(standard.mlp.state_dict().keys(), diff.mlp.state_dict().keys())
|
||||
self.assertEqual(
|
||||
sum(p.numel() for p in standard.mlp.parameters()),
|
||||
sum(p.numel() for p in diff.mlp.parameters()),
|
||||
)
|
||||
|
||||
def test_odd_baseline_head_count_is_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
DifferentialAttention(n_embd=120, n_head=5)
|
||||
|
||||
def test_deephealth_selects_diff_per_layer_and_only_adds_lambda_vectors(self):
|
||||
common = {
|
||||
"vocab_size": 16,
|
||||
"n_embd": 24,
|
||||
"n_head": 4,
|
||||
"n_layer": 3,
|
||||
"n_types": 1,
|
||||
"n_cont_types": 0,
|
||||
"n_categories": 1,
|
||||
"cont_type_ids": [],
|
||||
"target_mode": "all_future",
|
||||
"time_mode": "relative",
|
||||
"dist_mode": "weibull",
|
||||
"model_architecture": "traj_mixer_v5",
|
||||
}
|
||||
standard = DeepHealth(**common, attention_type="standard")
|
||||
diff = DeepHealth(**common, attention_type="diff")
|
||||
|
||||
self.assertTrue(
|
||||
all(isinstance(block.attn, TemporalAttention) for block in standard.blocks)
|
||||
)
|
||||
self.assertTrue(
|
||||
all(isinstance(block.attn, DifferentialAttention) for block in diff.blocks)
|
||||
)
|
||||
self.assertEqual(
|
||||
[block.attn.lambda_init for block in diff.blocks],
|
||||
[diff_lambda_init(depth) for depth in range(3)],
|
||||
)
|
||||
standard_count = sum(p.numel() for p in standard.parameters())
|
||||
diff_count = sum(p.numel() for p in diff.parameters())
|
||||
self.assertEqual(diff_count - standard_count, 3 * 4 * (24 // 4))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user