Revert "Add switchable DIFF V1 attention"
This reverts commit bf5cae8758.
This commit is contained in:
203
backbones.py
203
backbones.py
@@ -4,11 +4,6 @@ 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,
|
||||
@@ -190,188 +185,6 @@ 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,
|
||||
@@ -543,19 +356,15 @@ 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 = build_attention(
|
||||
attention_type,
|
||||
self.attn = TemporalAttention(
|
||||
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)
|
||||
@@ -583,19 +392,15 @@ 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 = build_attention(
|
||||
attention_type,
|
||||
self.attn = TemporalAttention(
|
||||
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,
|
||||
@@ -625,8 +430,6 @@ 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)
|
||||
@@ -645,8 +448,6 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user