JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add switchable DIFF V1 attention

Browse Source
This commit is contained in:
Jiarui Li
2026-08-21 11:48:16 +08:00
parent 75c9f06114
commit bf5cae8758
7 changed files with 453 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
attention_types.py Normal file
View File

@@ -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

203
backbones.py
View File

@@ -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,
)

1
eval_data.py
View File

@@ -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")),
)

12
models.py
View File

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

195
tests/test_differential_attention.py Normal file
View File

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

10
train_all_future.py
View File

@@ -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)")

10
train_next_step.py
View File

@@ -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")