Add exposure cache and keep absolute time only
This commit is contained in:
130
backbones.py
130
backbones.py
@@ -5,114 +5,24 @@ import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class TimeRoPE(nn.Module):
|
||||
def __init__(self, dim: int, base: float = 10000.0):
|
||||
super().__init__()
|
||||
assert dim % 2 == 0, "RoPE dim must be even"
|
||||
self.dim = dim
|
||||
# inv_freq is not trainable, but should move with device.
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
|
||||
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
||||
|
||||
def precompute_cache(self, tau: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
t = tau.unsqueeze(-1) # (B, L, 1)
|
||||
angles = t * self.inv_freq # (B, L, dim//2)
|
||||
# Pre-expand for heads and interleave once (avoids N_layers repeats)
|
||||
cos = angles.cos().unsqueeze(1).repeat_interleave(2, dim=-1)
|
||||
sin = angles.sin().unsqueeze(1).repeat_interleave(2, dim=-1)
|
||||
return cos, sin # (B, 1, L, dim)
|
||||
|
||||
@staticmethod
|
||||
def _rotate_half(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Rotate pairs: ``[-x2, x1, -x4, x3, ...]``."""
|
||||
x1 = x[..., 0::2]
|
||||
x2 = x[..., 1::2]
|
||||
return torch.stack((-x2, x1), dim=-1).flatten(-2)
|
||||
|
||||
@staticmethod
|
||||
def apply_from_cache(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cos, sin = rope_cache # each (B, 1, L, dim)
|
||||
q_rot = q * cos + TimeRoPE._rotate_half(q) * sin
|
||||
k_rot = k * cos + TimeRoPE._rotate_half(k) * sin
|
||||
return q_rot, k_rot
|
||||
|
||||
def forward(
|
||||
self,
|
||||
tau: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cache = self.precompute_cache(tau)
|
||||
return self.apply_from_cache(q, k, cache)
|
||||
|
||||
|
||||
class GaussianRBFTimeBasis(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_bases: int = 16,
|
||||
max_time_diff: float = 40.0,
|
||||
):
|
||||
super().__init__()
|
||||
self.n_bases = n_bases
|
||||
|
||||
# Evenly spaced RBF centres for non-negative linear time differences.
|
||||
# Causal masking enforces query_time >= key_time, so diff is >= 0.
|
||||
centers = torch.linspace(0.0, max_time_diff, n_bases)
|
||||
self.register_buffer("centers", centers,
|
||||
persistent=False) # (n_bases,)
|
||||
|
||||
# Learnable log-widths (initialized to center spacing on linear scale).
|
||||
init_width = max(max_time_diff / max(n_bases - 1, 1), 1e-3)
|
||||
init_log_width = math.log(init_width)
|
||||
self.log_widths = nn.Parameter(torch.full((n_bases,), init_log_width))
|
||||
|
||||
def precompute_cache(self, tau: torch.Tensor) -> torch.Tensor:
|
||||
|
||||
time_coord = tau.float() # (B, L)
|
||||
# Pairwise signed difference: query_i - key_j.
|
||||
diff = time_coord.unsqueeze(
|
||||
2) - time_coord.unsqueeze(1) # (B, L_q, L_k)
|
||||
# Gaussian RBF: exp(-0.5 * ((diff - c) / w)^2)
|
||||
diff = diff.unsqueeze(-1) # (B, L, L, 1)
|
||||
widths = self.log_widths.exp() # (n_bases,)
|
||||
rbf_acts = torch.exp(
|
||||
-0.5 * ((diff - self.centers) / widths).square()
|
||||
# (B, L, L, n_bases)
|
||||
)
|
||||
return rbf_acts
|
||||
|
||||
|
||||
class TemporalAttention(nn.Module):
|
||||
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,
|
||||
):
|
||||
super().__init__()
|
||||
assert n_embd % n_head == 0, "n_embd must be divisible by n_head"
|
||||
self.n_head = n_head
|
||||
self.d_head = n_embd // n_head
|
||||
self.scale = 1.0 / math.sqrt(self.d_head)
|
||||
self.use_time_rope = use_time_rope
|
||||
self.use_rbf_bias = use_rbf_bias
|
||||
|
||||
# QKV projection (fused for efficiency)
|
||||
self.qkv = nn.Linear(n_embd, 3 * n_embd, bias=False)
|
||||
# Output projection
|
||||
self.out_proj = nn.Linear(n_embd, n_embd, bias=False)
|
||||
|
||||
# Layer-specific projection from shared RBF basis activations to per-head attention bias.
|
||||
self.rbf_proj = nn.Linear(n_rbf_bases, n_head, bias=False)
|
||||
self.time_bias_scale = nn.Parameter(torch.tensor(0.0))
|
||||
|
||||
self.resid_drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
@@ -120,20 +30,12 @@ class TemporalAttention(nn.Module):
|
||||
"""Match the previous version's GPT-style weight initialization."""
|
||||
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)
|
||||
|
||||
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"
|
||||
|
||||
B, L, _ = x.shape
|
||||
H, D = self.n_head, self.d_head
|
||||
|
||||
@@ -141,31 +43,11 @@ class TemporalAttention(nn.Module):
|
||||
qkv = self.qkv(x).reshape(B, L, 3, H, D).permute(2, 0, 3, 1, 4)
|
||||
q, k, v = qkv.unbind(0) # each (B, H, L, D)
|
||||
|
||||
# --- Apply RoPE (from shared cache) --------------------------------
|
||||
if self.use_time_rope:
|
||||
q, k = TimeRoPE.apply_from_cache(q, k, rope_cache)
|
||||
|
||||
# Build additive attention bias mask: time bias + causal/padding mask.
|
||||
time_bias = None
|
||||
if self.use_rbf_bias:
|
||||
time_bias = self.rbf_proj(rbf_cache).permute(
|
||||
0, 3, 1, 2) # (B, H, L, L)
|
||||
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
|
||||
|
||||
out = F.scaled_dot_product_attention(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
attn_mask=attn_bias,
|
||||
attn_mask=attn_mask,
|
||||
dropout_p=0.0,
|
||||
is_causal=False,
|
||||
scale=self.scale,
|
||||
@@ -218,18 +100,12 @@ class GPTBlock(nn.Module):
|
||||
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
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,
|
||||
)
|
||||
self.mlp = SwiGLU(n_embd=n_embd, dropout=mlp_dropout)
|
||||
self.ln1 = nn.LayerNorm(n_embd)
|
||||
@@ -238,11 +114,9 @@ class GPTBlock(nn.Module):
|
||||
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:
|
||||
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
|
||||
x = x + self.attn(self.ln1(x), attn_mask)
|
||||
x = x + self.mlp(self.ln2(x))
|
||||
return x
|
||||
|
||||
|
||||
Reference in New Issue
Block a user