DeepHealth/tests/test_differential_attention.py

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