diff --git a/attention_types.py b/attention_types.py new file mode 100644 index 0000000..f595c7c --- /dev/null +++ b/attention_types.py @@ -0,0 +1,26 @@ +"""Attention implementation identifiers and validation helpers.""" + + +STANDARD_ATTENTION = "standard" +DIFF_ATTENTION = "diff" +DEFAULT_ATTENTION_TYPE = STANDARD_ATTENTION +SUPPORTED_ATTENTION_TYPES = ( + STANDARD_ATTENTION, + DIFF_ATTENTION, +) + + +def resolve_attention_type(attention_type: str | None = None) -> str: + """Resolve an attention selector, defaulting old configs to standard.""" + resolved = DEFAULT_ATTENTION_TYPE if attention_type is None else attention_type + if not isinstance(resolved, str): + raise ValueError( + "attention_type must be one of " + f"{SUPPORTED_ATTENTION_TYPES}, got {resolved!r}" + ) + if resolved not in SUPPORTED_ATTENTION_TYPES: + raise ValueError( + f"Unsupported attention_type={resolved!r}; " + f"expected one of {SUPPORTED_ATTENTION_TYPES}." + ) + return resolved diff --git a/backbones.py b/backbones.py index 86e4789..62c0cd1 100644 --- a/backbones.py +++ b/backbones.py @@ -4,6 +4,11 @@ import torch import torch.nn as nn import torch.nn.functional as F +from attention_types import ( + DIFF_ATTENTION, + STANDARD_ATTENTION, + resolve_attention_type, +) from model_architectures import ( TRAJ_MIXER_ARCHITECTURE, TRANSFORMER_FFN_ARCHITECTURE, @@ -185,6 +190,188 @@ class TemporalAttention(nn.Module): return self.resid_drop(self.out_proj(out)) +def diff_lambda_init(depth: int) -> float: + """DIFF V1 layer-depth initialization (``depth`` is zero-based).""" + if depth < 0: + raise ValueError(f"depth must be >= 0, got {depth}") + return 0.8 - 0.6 * math.exp(-0.3 * depth) + + +class DifferentialAttention(nn.Module): + """DIFF V1 attention with paired Q/K heads and no per-head norm. + + ``n_head`` retains the baseline head count. Adjacent Q/K heads are paired, + so the module has ``n_head // 2`` differential heads. Each paired V head + has width ``2 * d_head`` and the merged output remains ``n_embd`` wide. + """ + + def __init__( + self, + n_embd: int, + n_head: int, + n_rbf_bases: int = 16, + dropout: float = 0.0, + use_time_rope: bool = True, + use_rbf_bias: bool = True, + depth: int = 0, + ): + super().__init__() + if n_head % 2 != 0: + raise ValueError( + "DifferentialAttention requires an even baseline n_head, " + f"got {n_head}" + ) + if n_embd % n_head != 0: + raise ValueError( + f"n_embd must be divisible by n_head, got {n_embd} and {n_head}" + ) + + self.n_head = n_head + self.n_diff_head = n_head // 2 + self.d_head = n_embd // n_head + self.d_value_head = 2 * self.d_head + self.scale = 1.0 / math.sqrt(self.d_head) + self.use_time_rope = use_time_rope + self.use_rbf_bias = use_rbf_bias + self.lambda_init = diff_lambda_init(depth) + + # These projection widths intentionally match TemporalAttention. + self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False) + self.out_proj = nn.Linear(n_embd, n_embd, bias=False) + + # Preserve one relative-time bias channel per baseline Q/K head. After + # pairing, each differential head therefore has one bias for each map. + self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False) + self.time_bias_scale = nn.Parameter(torch.tensor(1.0)) + + # DIFF V1 uses four layer-level vectors, shared across differential + # heads. There is deliberately no per-head RMSNorm/LayerNorm here. + self.lambda_q1 = nn.Parameter(torch.empty(self.d_head)) + self.lambda_k1 = nn.Parameter(torch.empty(self.d_head)) + self.lambda_q2 = nn.Parameter(torch.empty(self.d_head)) + self.lambda_k2 = nn.Parameter(torch.empty(self.d_head)) + + self.resid_drop = nn.Dropout(dropout) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.qkv.weight, mean=0.0, std=0.02) + nn.init.normal_(self.out_proj.weight, mean=0.0, std=0.02) + nn.init.zeros_(self.rbf_proj.weight) + nn.init.normal_(self.lambda_q1, mean=0.0, std=0.1) + nn.init.normal_(self.lambda_k1, mean=0.0, std=0.1) + nn.init.normal_(self.lambda_q2, mean=0.0, std=0.1) + nn.init.normal_(self.lambda_k2, mean=0.0, std=0.1) + + def _lambda(self, dtype: torch.dtype) -> torch.Tensor: + lambda_1 = torch.exp( + torch.sum(self.lambda_q1 * self.lambda_k1).float() + ) + lambda_2 = torch.exp( + torch.sum(self.lambda_q2 * self.lambda_k2).float() + ) + return (lambda_1 - lambda_2 + self.lambda_init).to(dtype=dtype) + + def forward( + self, + x: torch.Tensor, + rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + rbf_cache: torch.Tensor | None = None, + attn_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if self.use_time_rope: + assert rope_cache is not None, "rope_cache must be provided when use_time_rope is True" + if self.use_rbf_bias: + assert rbf_cache is not None, "rbf_cache must be provided when use_rbf_bias is True" + + batch_size, seq_len, n_embd = x.shape + + # q/k: (B, 2 * H_diff, L, d_head) + # v: (B, H_diff, L, 2 * d_head) + qkv = self.qkv(x).reshape( + batch_size, seq_len, 3, self.n_head, self.d_head + ).permute(2, 0, 3, 1, 4) + q, k, v_unpaired = qkv.unbind(0) + v = v_unpaired.transpose(1, 2).reshape( + batch_size, seq_len, self.n_diff_head, self.d_value_head + ).transpose(1, 2) + + # A single cache is broadcast over all heads, so both Q/K maps receive + # exactly the same TimeRoPE coordinates and rotation rule. + if self.use_time_rope: + q, k = TimeRoPE.apply_from_cache(q, k, rope_cache) + + time_bias = None + if self.use_rbf_bias: + time_bias = self.rbf_proj(rbf_cache).permute(0, 3, 1, 2) + time_bias = self.time_bias_scale.tanh() * time_bias + + if time_bias is not None and attn_mask is not None: + attn_bias = time_bias + attn_mask.to(time_bias.dtype) + elif time_bias is not None: + attn_bias = time_bias + elif attn_mask is not None: + attn_bias = attn_mask + else: + attn_bias = None + + attn_logits = torch.matmul(q, k.transpose(-1, -2)) * self.scale + if attn_bias is not None: + attn_logits = attn_logits + attn_bias + attn_weights = F.softmax( + attn_logits, dim=-1, dtype=torch.float32 + ).to(q.dtype) + attn_weights = attn_weights.reshape( + batch_size, + self.n_diff_head, + 2, + seq_len, + seq_len, + ) + diff_weights = ( + attn_weights[:, :, 0] + - self._lambda(q.dtype) * attn_weights[:, :, 1] + ) + + # No V1 per-head RMSNorm and no replacement normalization: the + # differential result goes directly to head merge and output projection. + out = torch.matmul(diff_weights, v) + out = out.transpose(1, 2).reshape(batch_size, seq_len, n_embd) + return self.resid_drop(self.out_proj(out)) + + +# Concise public alias for callers that prefer the shorter name. +DiffAttention = DifferentialAttention + + +def build_attention( + attention_type: str, + *, + n_embd: int, + n_head: int, + n_rbf_bases: int = 16, + dropout: float = 0.0, + use_time_rope: bool = True, + use_rbf_bias: bool = True, + depth: int = 0, +) -> nn.Module: + """Build a standard or DIFF V1 attention module.""" + resolved = resolve_attention_type(attention_type) + common_kwargs = { + "n_embd": n_embd, + "n_head": n_head, + "n_rbf_bases": n_rbf_bases, + "dropout": dropout, + "use_time_rope": use_time_rope, + "use_rbf_bias": use_rbf_bias, + } + if resolved == STANDARD_ATTENTION: + return TemporalAttention(**common_kwargs) + if resolved == DIFF_ATTENTION: + return DifferentialAttention(**common_kwargs, depth=depth) + raise ValueError(f"Unsupported attention_type: {resolved!r}") + + class SwiGLU(nn.Module): def __init__( self, @@ -356,15 +543,19 @@ class TransformerFFNBlock(nn.Module): use_time_rope: bool = False, use_rbf_bias: bool = False, n_rbf_bases: int = 16, + attention_type: str = STANDARD_ATTENTION, + depth: int = 0, ): super().__init__() - self.attn = TemporalAttention( + self.attn = build_attention( + attention_type, n_embd=n_embd, n_head=n_head, n_rbf_bases=n_rbf_bases, dropout=attn_dropout, use_time_rope=use_time_rope, use_rbf_bias=use_rbf_bias, + depth=depth, ) self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout) self.ln1 = nn.LayerNorm(n_embd) @@ -392,15 +583,19 @@ class TrajMixerBlock(nn.Module): use_time_rope: bool = False, use_rbf_bias: bool = False, n_rbf_bases: int = 16, + attention_type: str = STANDARD_ATTENTION, + depth: int = 0, ): super().__init__() - self.attn = TemporalAttention( + self.attn = build_attention( + attention_type, n_embd=n_embd, n_head=n_head, n_rbf_bases=n_rbf_bases, dropout=attn_dropout, use_time_rope=use_time_rope, use_rbf_bias=use_rbf_bias, + depth=depth, ) self.mlp = TrajMixer( n_embd=n_embd, @@ -430,6 +625,8 @@ def build_backbone_block( use_time_rope: bool = False, use_rbf_bias: bool = False, n_rbf_bases: int = 16, + attention_type: str = STANDARD_ATTENTION, + depth: int = 0, ) -> nn.Module: """Build one history block for a supported model architecture.""" architecture = resolve_model_architecture(model_architecture) @@ -448,6 +645,8 @@ def build_backbone_block( use_time_rope=use_time_rope, use_rbf_bias=use_rbf_bias, n_rbf_bases=n_rbf_bases, + attention_type=attention_type, + depth=depth, ) diff --git a/eval_data.py b/eval_data.py index 4c598e2..ab60670 100644 --- a/eval_data.py +++ b/eval_data.py @@ -231,6 +231,7 @@ def build_model_from_dataset( dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")), dropout=float(cfg_get(args, cfg, "dropout", 0.0)), model_architecture=model_architecture, + attention_type=str(cfg_get(args, cfg, "attention_type", "standard")), ) diff --git a/models.py b/models.py index c2c1b30..52aef66 100644 --- a/models.py +++ b/models.py @@ -11,6 +11,7 @@ from backbones import ( TokenAutoDiscretization, build_backbone_block, ) +from attention_types import resolve_attention_type from model_architectures import resolve_model_architecture from targets import PAD_IDX @@ -221,6 +222,7 @@ class DeepHealth(nn.Module): extra_pool_reduce: str = "mean", dropout: float = 0.0, model_architecture: str | None = None, + attention_type: str | None = None, ): super().__init__() if target_mode not in ["next_token", "all_future"]: @@ -242,6 +244,7 @@ class DeepHealth(nn.Module): if n_layer < 1: raise ValueError(f"n_layer must be >= 1, got {n_layer}") model_architecture = resolve_model_architecture(model_architecture) + attention_type = resolve_attention_type(attention_type) self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0) self.gender_embedding = nn.Embedding( 2, n_embd) # Assuming binary gender @@ -261,6 +264,7 @@ class DeepHealth(nn.Module): self.dist_mode = dist_mode self.extra_pool_reduce = extra_pool_reduce self.model_architecture = model_architecture + self.attention_type = attention_type self.n_layer = n_layer self.n_embd = n_embd self.vocab_size = vocab_size @@ -282,7 +286,9 @@ class DeepHealth(nn.Module): use_time_rope=False, use_rbf_bias=False, mlp_dropout=dropout, - ) for _ in range(n_layer) + attention_type=attention_type, + depth=layer_index, + ) for layer_index in range(n_layer) ]) self.rope = None self.rbf = None @@ -296,7 +302,9 @@ class DeepHealth(nn.Module): use_time_rope=True, use_rbf_bias=True, mlp_dropout=dropout, - ) for _ in range(n_layer) + attention_type=attention_type, + depth=layer_index, + ) for layer_index in range(n_layer) ]) self.rope = TimeRoPE(n_embd // n_head) self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0) diff --git a/tests/test_differential_attention.py b/tests/test_differential_attention.py new file mode 100644 index 0000000..f35db1d --- /dev/null +++ b/tests/test_differential_attention.py @@ -0,0 +1,195 @@ +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() diff --git a/train_all_future.py b/train_all_future.py index ab7ac3c..8fc4f01 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -25,6 +25,7 @@ from torch.optim import AdamW from torch.utils.data import DataLoader, RandomSampler from tqdm.auto import tqdm +from attention_types import DEFAULT_ATTENTION_TYPE, SUPPORTED_ATTENTION_TYPES from dataset import ( DISEASE_HISTORY_MODES, DISEASE_HISTORY_MODE_TIMED, @@ -119,6 +120,12 @@ def parse_args() -> argparse.Namespace: default=DEFAULT_MODEL_ARCHITECTURE, choices=SUPPORTED_MODEL_ARCHITECTURES, ) + parser.add_argument( + "--attention_type", + type=str, + default=DEFAULT_ATTENTION_TYPE, + choices=SUPPORTED_ATTENTION_TYPES, + ) parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--base_lr", type=float, default=3e-4) @@ -204,6 +211,7 @@ def build_model( dist_mode=args.dist_mode, dropout=args.dropout, model_architecture=args.model_architecture, + attention_type=args.attention_type, ) @@ -333,6 +341,7 @@ def build_metadata( "collate_fn": "all_future_collate_fn", "model_class": "DeepHealth", "model_architecture": args.model_architecture, + "attention_type": args.attention_type, "model_target_mode": "all_future", "target_mode": "all_future", "event_stream_version": "disease_death_only_v1", @@ -393,6 +402,7 @@ def main() -> None: logger.info(f"Starting all-future training run: {run_name}") logger.info(f"Device: {device}") logger.info(f"Model architecture: {args.model_architecture}") + logger.info(f"Attention type: {args.attention_type}") logger.info(f"Disease history mode: {args.disease_history_mode}") logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info("Continuous value scaling: RobustScale (required)") diff --git a/train_next_step.py b/train_next_step.py index 77b232d..5813be7 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -16,6 +16,7 @@ from torch.optim import AdamW from torch.utils.data import DataLoader, RandomSampler from tqdm.auto import tqdm +from attention_types import DEFAULT_ATTENTION_TYPE, SUPPORTED_ATTENTION_TYPES from dataset import HealthDataset, collate_fn from losses import build_loss from model_architectures import ( @@ -88,6 +89,12 @@ def parse_args() -> argparse.Namespace: default=DEFAULT_MODEL_ARCHITECTURE, choices=SUPPORTED_MODEL_ARCHITECTURES, ) + parser.add_argument( + "--attention_type", + type=str, + default=DEFAULT_ATTENTION_TYPE, + choices=SUPPORTED_ATTENTION_TYPES, + ) parser.add_argument("--t_min", type=float, default=0.0027378507871321013) parser.add_argument("--max_exp_input", type=float, default=60.0) @@ -151,6 +158,7 @@ def build_model( dist_mode="exponential", dropout=args.dropout, model_architecture=args.model_architecture, + attention_type=args.attention_type, ) @@ -373,6 +381,7 @@ def build_metadata( "collate_fn": "next_step_collate_fn", "model_class": "DeepHealth", "model_architecture": args.model_architecture, + "attention_type": args.attention_type, "model_target_mode": "next_token", "target_mode": "delphi2m", "event_stream_version": "disease_death_only_v1", @@ -425,6 +434,7 @@ def main() -> None: logger.info(f"Starting next-step training run: {run_name}") logger.info(f"Device: {device}") logger.info(f"Model architecture: {args.model_architecture}") + logger.info(f"Attention type: {args.attention_type}") logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}") logger.info("Continuous value scaling: RobustScale (required)") logger.info("time_mode=absolute, readout=token, target_mode=delphi2m")