157 lines
6.6 KiB
Python
157 lines
6.6 KiB
Python
"""Positional encodings shared by every MuVAP+ transformer stack.
|
|
|
|
VAP is a causal model: a frame may only attend to frames at or before it.
|
|
Every encoding here is defined for that regime, so switching between them
|
|
changes how position is represented and never what is visible.
|
|
"""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
#: A learned absolute table is deliberately absent. It is defined only up to
|
|
#: `max_frames` and cannot extrapolate past it, and these models are trained on
|
|
#: fixed-length clips but deployed on unbounded streams -- the position it would
|
|
#: need at inference is one it never saw a gradient for. ALiBi's distance bias
|
|
#: and RoPE's rotation are both defined at any offset; sinusoidal is at least
|
|
#: closed-form rather than a lookup.
|
|
POS_ENCODINGS = ("alibi", "rope", "sinusoidal", "none")
|
|
|
|
|
|
def alibi_slopes(num_heads):
|
|
"""Per-head decay rates from the ALiBi paper (Press et al., 2022)."""
|
|
|
|
def power_of_two_slopes(n):
|
|
start = 2 ** (-(2 ** -(math.log2(n) - 3)))
|
|
return [start * start**i for i in range(n)]
|
|
|
|
if math.log2(num_heads).is_integer():
|
|
return power_of_two_slopes(num_heads)
|
|
closest = 2 ** math.floor(math.log2(num_heads))
|
|
extra = alibi_slopes(2 * closest)[0::2][: num_heads - closest]
|
|
return power_of_two_slopes(closest) + extra
|
|
|
|
|
|
def causal_alibi_bias(num_frames, slopes, device, dtype):
|
|
"""Additive `[1, heads, T, T]` mask holding both ALiBi decay and causality."""
|
|
positions = torch.arange(num_frames, device=device, dtype=dtype)
|
|
distance = positions.view(1, -1) - positions.view(-1, 1)
|
|
bias = distance.view(1, 1, num_frames, num_frames) * slopes.to(
|
|
device=device, dtype=dtype
|
|
).view(1, -1, 1, 1)
|
|
causal = torch.ones(num_frames, num_frames, device=device, dtype=torch.bool).tril()
|
|
return bias.masked_fill(~causal.view(1, 1, num_frames, num_frames), float("-inf"))
|
|
|
|
|
|
def causal_float_mask(num_queries, num_keys, device, dtype):
|
|
"""Additive `[1, 1, Tq, Tk]` mask: 0 where attendable, `-inf` ahead of the query.
|
|
|
|
Query `t` sits at absolute frame `Tk - Tq + t` -- a query stream shorter than
|
|
its key stream is streaming inference over a longer history -- so it may
|
|
attend keys `<= Tk - Tq + t`. Additive rather than boolean on purpose:
|
|
`F.scaled_dot_product_attention` reads a bool mask as "True = attend" and a
|
|
float one as a pre-softmax bias, and mixing the two is silent.
|
|
"""
|
|
offset = num_keys - num_queries
|
|
queries = torch.arange(num_queries, device=device).view(-1, 1) + offset
|
|
keys = torch.arange(num_keys, device=device).view(1, -1)
|
|
mask = torch.zeros(num_queries, num_keys, device=device, dtype=dtype)
|
|
return mask.masked_fill(keys > queries, float("-inf")).view(
|
|
1, 1, num_queries, num_keys
|
|
)
|
|
|
|
|
|
def as_additive_mask(mask, dtype):
|
|
"""A caller's mask as an additive float bias, whichever convention it used."""
|
|
if mask.dtype == torch.bool:
|
|
return torch.zeros_like(mask, dtype=dtype).masked_fill_(~mask, float("-inf"))
|
|
return mask.to(dtype)
|
|
|
|
|
|
class RotaryEmbedding(nn.Module):
|
|
"""Causal rotary position embedding (Su et al., 2024).
|
|
|
|
Rotates the query and key of every head by an angle proportional to the
|
|
absolute frame index, which makes the attention logit depend only on the
|
|
distance between two frames. Under a causal mask that distance is always
|
|
non-negative, so the model sees a strictly backward-looking clock and no
|
|
future frame can leak through the encoding.
|
|
"""
|
|
|
|
def __init__(self, head_dim, base=10_000.0, max_frames=4096):
|
|
super().__init__()
|
|
if head_dim % 2 != 0:
|
|
raise ValueError(f"rotary head_dim must be even, got {head_dim}")
|
|
self.head_dim = head_dim
|
|
self.base = float(base)
|
|
inv_freq = 1.0 / (
|
|
self.base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)
|
|
)
|
|
self.register_buffer("inv_freq", inv_freq, persistent=False)
|
|
self.register_buffer("cos", torch.empty(0), persistent=False)
|
|
self.register_buffer("sin", torch.empty(0), persistent=False)
|
|
self._build_cache(max_frames, inv_freq.device, torch.float32)
|
|
|
|
def _build_cache(self, num_frames, device, dtype):
|
|
positions = torch.arange(num_frames, device=device, dtype=torch.float32)
|
|
angles = torch.outer(positions, self.inv_freq.to(device=device))
|
|
angles = torch.cat([angles, angles], dim=-1)
|
|
self.cos = angles.cos().to(dtype)
|
|
self.sin = angles.sin().to(dtype)
|
|
|
|
def _cached(self, num_frames, device, dtype):
|
|
stale = (
|
|
self.cos.numel() == 0
|
|
or self.cos.shape[0] < num_frames
|
|
or self.cos.device != device
|
|
or self.cos.dtype != dtype
|
|
)
|
|
if stale:
|
|
self._build_cache(max(num_frames, 1), device, dtype)
|
|
return self.cos[:num_frames], self.sin[:num_frames]
|
|
|
|
@staticmethod
|
|
def _rotate_half(x):
|
|
first, second = x.chunk(2, dim=-1)
|
|
return torch.cat([-second, first], dim=-1)
|
|
|
|
def forward(self, query, key, offset=0):
|
|
"""Rotate `[batch, heads, frames, head_dim]` query and key in place of no-op."""
|
|
frames = query.shape[-2] + offset
|
|
cos, sin = self._cached(frames, query.device, query.dtype)
|
|
q_cos, q_sin = cos[offset:], sin[offset:]
|
|
k_cos, k_sin = cos[: key.shape[-2]], sin[: key.shape[-2]]
|
|
query = query * q_cos + self._rotate_half(query) * q_sin
|
|
key = key * k_cos + self._rotate_half(key) * k_sin
|
|
return query, key
|
|
|
|
|
|
class SinusoidalPositionEmbedding(nn.Module):
|
|
"""Fixed additive sinusoidal embedding, kept for ablation against RoPE."""
|
|
|
|
def __init__(self, dim, max_frames=4096):
|
|
super().__init__()
|
|
position = torch.arange(max_frames, dtype=torch.float32).unsqueeze(1)
|
|
div = torch.exp(
|
|
torch.arange(0, dim, 2, dtype=torch.float32) * (-math.log(10_000.0) / dim)
|
|
)
|
|
table = torch.zeros(max_frames, dim)
|
|
table[:, 0::2] = torch.sin(position * div)
|
|
table[:, 1::2] = torch.cos(position * div)
|
|
self.register_buffer("table", table, persistent=False)
|
|
|
|
def forward(self, x):
|
|
if x.shape[1] > self.table.shape[0]:
|
|
raise ValueError(
|
|
f"sequence of {x.shape[1]} frames exceeds the {self.table.shape[0]}-frame table"
|
|
)
|
|
return x + self.table[: x.shape[1]].to(dtype=x.dtype)
|
|
|
|
|
|
def build_input_position_embedding(pos_encoding, dim, max_frames=4096):
|
|
"""Additive input-side encoding, or None when position lives in attention."""
|
|
if pos_encoding == "sinusoidal":
|
|
return SinusoidalPositionEmbedding(dim, max_frames)
|
|
return None
|