Add switchable DIFF V1 attention

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