File size: 4,608 Bytes
2db32a1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # coding=utf-8
"""Spiral Rotary Positional Encoding (SRPE).
Implements Section VI of the Wiola paper. For position ``p`` and dimension-pair
index ``j in [d_h/2]``:
omega_j = theta_0 ** (-2j / d_h)
Theta_j(p)= p * omega_j * (1 + 1/k_s)
r_j(p) = 1 + a_s * sin(p * f_s * omega_j)
c_j(p) = r_j(p) * cos(Theta_j(p))
s_j(p) = r_j(p) * sin(Theta_j(p))
and the rotation is applied to q/k with the standard rotate-half trick, which
is algebraically identical to Eqs. (12)-(13):
out_lo = q_lo * c - q_hi * s
out_hi = q_hi * c + q_lo * s
Unlike RoPE, ``Theta`` carries a second winding angle (via ``1 + 1/k_s``) and
the magnitude is modulated by a radial term ``r_j(p)``, embedding positions on a
3D helical manifold with no learned parameters.
"""
import torch
import torch.nn as nn
def rotate_half(x: torch.Tensor) -> torch.Tensor:
"""Rotate the two halves of the last dimension: [x1, x2] -> [-x2, x1]."""
half = x.shape[-1] // 2
x1 = x[..., :half]
x2 = x[..., half:]
return torch.cat((-x2, x1), dim=-1)
class SpiralRotaryEmbedding(nn.Module):
def __init__(
self,
head_dim: int,
max_position_embeddings: int = 2048,
theta: float = 10000.0,
spiral_divisor: int = 8,
radial_amplitude: float = 0.1,
radial_frequency: float = 0.01,
):
super().__init__()
if head_dim % 2 != 0:
raise ValueError(f"head_dim must be even, got {head_dim}.")
self.head_dim = head_dim
self.max_position_embeddings = max_position_embeddings
self.theta = theta
self.k_s = spiral_divisor
self.a_s = radial_amplitude
self.f_s = radial_frequency
# omega_j = theta ** (-2j / d_h), j in [0, d_h/2).
j = torch.arange(0, head_dim, 2, dtype=torch.float32)
inv_freq = theta ** (-(j / head_dim)) # shape [d_h/2]
self.register_buffer("inv_freq", inv_freq, persistent=False)
self._cached_len = 0
self._build_cache(max_position_embeddings)
def _build_cache(self, seq_len: int, device=None, dtype=torch.float32):
positions = torch.arange(
seq_len,
dtype=torch.float32,
device=device if device is not None else self.inv_freq.device,
)
inv_freq = self.inv_freq.to(positions.device)
# Combined winding angle Theta_j(p) = p * omega_j * (1 + 1/k_s).
combined = 1.0 + 1.0 / self.k_s
angle = torch.outer(positions, inv_freq) * combined # [T, d_h/2]
# Radial modulation r_j(p) = 1 + a_s * sin(p * f_s * omega_j).
radial = 1.0 + self.a_s * torch.sin(torch.outer(positions, inv_freq) * self.f_s)
cos = radial * torch.cos(angle)
sin = radial * torch.sin(angle)
# Duplicate to full head_dim so it lines up with rotate_half.
cos = torch.cat((cos, cos), dim=-1) # [T, d_h]
sin = torch.cat((sin, sin), dim=-1)
self.register_buffer("cos_cached", cos.to(dtype), persistent=False)
self.register_buffer("sin_cached", sin.to(dtype), persistent=False)
self._cached_len = seq_len
@torch.no_grad()
def _maybe_extend(self, seq_len: int, device, dtype):
if seq_len > self._cached_len or self.cos_cached.device != device:
self._build_cache(max(seq_len, self._cached_len), device=device, dtype=dtype)
def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
"""Return (cos, sin) gathered for ``position_ids``.
Args:
x: any tensor only used to read dtype/device.
position_ids: LongTensor of shape [batch, seq_len].
Returns:
cos, sin each of shape [batch, seq_len, head_dim].
"""
max_pos = int(position_ids.max().item()) + 1
self._maybe_extend(max_pos, x.device, torch.float32)
cos = self.cos_cached[position_ids] # [B, S, d_h]
sin = self.sin_cached[position_ids]
return cos.to(x.dtype), sin.to(x.dtype)
def apply_srpe(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
"""Apply SRPE rotation to query and key tensors.
Args:
q, k: shape [batch, num_heads, seq_len, head_dim].
cos, sin: shape [batch, seq_len, head_dim].
"""
cos = cos.unsqueeze(1) # [B, 1, S, d_h] -> broadcast over heads
sin = sin.unsqueeze(1)
q_rot = (q * cos) + (rotate_half(q) * sin)
k_rot = (k * cos) + (rotate_half(k) * sin)
return q_rot, k_rot
|