File size: 5,943 Bytes
20a3a74 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """
μLattice -- bump kernel and 2D lattice weighting.
The bump kernel and 2D lattice weighting:
b(t) = cos^2(pi t / 3) for |t| < 3/2 , 0 otherwise [bump, support 1.5]
phi_z(x) = b(x - z) / sqrt( sum_j b(x - j)^2 ) [normalized weights]
D_1(q,k) = sum_z phi_z(q) phi_z(k) [1D kernel]
D(q,k) = D_1(q_x,k_x) * D_1(q_y,k_y) [2D product kernel]
Consequences that the recipe elevates to invariants:
* |t| < 3/2 means an axis coordinate touches AT MOST THREE integers,
so a token touches at most 3 x 3 = 9 lattice cells. Never more.
* sum_z phi_z(x)^2 == 1 exactly (L2 partition of unity), and the same
holds for the 2D product weights, so D(q,q) == 1.
* Cell fan-out is 3^d for lattice dimension d. d = 2 -> 9.
"""
from __future__ import annotations
import math
import torch
# Support per bump variant. The BLOG publishes cos^2(pi t / 3) with |t| < 3/2.
# The reference implementation uses
# cos^2(pi t / 2) with |t| < 1 instead: a narrower bump, at most TWO nonzero cells
# per axis (<=4 in 2D), and exactly ONE when the coordinate is integral.
BUMP_SUPPORT_BY_KIND = {"cos2": 1.5, "shelf": 1.5, "cos2_half": 1.0}
BUMP_SUPPORT = 1.5
LATTICE_DIM = 2
MAX_CELLS_PER_TOKEN = 3 ** LATTICE_DIM # published: 9
def bump(t: torch.Tensor, kind: str = "cos2") -> torch.Tensor:
"""Bump function on |t| < 3/2, else 0.
Two variants are published (the second one is the one in the released
reference figure code):
"cos2" b(t) = cos^2(pi t / 3) [blog equation]
"cos2_half" b(t) = cos^2(pi t / 2), |t| < 1 [what the released reference C
actually uses; note the different support]
"shelf" b(t) = exp(-t^2/(2 s^2)) - exp(-2.25/(2 s^2)), s^2 = 0.72
(the subtractive term makes it exactly 0 at |t| = 3/2)
"""
if kind == "cos2":
inside = t.abs() < BUMP_SUPPORT_BY_KIND["cos2"]
return torch.where(inside, torch.cos(math.pi * t / 3.0) ** 2, torch.zeros_like(t))
if kind == "cos2_half":
inside = t.abs() < BUMP_SUPPORT_BY_KIND["cos2_half"]
return torch.where(inside, torch.cos(math.pi * t / 2.0) ** 2, torch.zeros_like(t))
if kind == "shelf":
s2 = 0.72
b = torch.exp(-(t * t) / (2 * s2)) - math.exp(-2.25 / (2 * s2))
return torch.where(t.abs() < BUMP_SUPPORT_BY_KIND["shelf"], b, torch.zeros_like(t))
raise ValueError(f"unknown bump kind {kind!r}")
def touched_cells(x: torch.Tensor, kind: str = "cos2") -> torch.Tensor:
"""Integer cells an axis-coordinate x can touch, shape (..., 3).
|x - z| < 3/2 <=> z in (x - 3/2, x + 3/2), an open interval of length 3,
which contains at most 3 integers: ceil(x - 3/2), ..., floor(x + 3/2).
Both candidate ends are returned; the kernel value is 0 (not
merely small) at |t| == 3/2, so no tolerance is required.
"""
lo = torch.ceil(x - BUMP_SUPPORT_BY_KIND.get(kind, BUMP_SUPPORT)).to(torch.int64)
return lo.unsqueeze(-1) + torch.arange(3, device=x.device, dtype=torch.int64)
def phi_1d(x: torch.Tensor, kind: str = "cos2") -> tuple[torch.Tensor, torch.Tensor]:
"""Normalized 1D weights. Returns (weights (...,3), cell coords (...,3)).
phi_z(x) = b(x - z) / sqrt(sum_j b(x - j)^2)
"""
z = touched_cells(x, kind) # (...,3)
b = bump(x.unsqueeze(-1) - z.to(x.dtype), kind) # (...,3)
denom = torch.sqrt((b * b).sum(dim=-1, keepdim=True)).clamp_min(1e-30)
return b / denom, z
def phi_2d(xy: torch.Tensor, kind: str = "cos2") -> tuple[torch.Tensor, torch.Tensor]:
"""Normalized 2D weights over the 3x3 neighborhood.
xy : (..., 2) real-valued address (q or k)
returns weights (..., 3, 3) and integer coords (..., 3, 3, 2)
Product form: phi_(zx,zy)(x,y) = phi_zx(x) * phi_zy(y), so
sum over the 9 cells of phi^2 == 1.
"""
wx, zx = phi_1d(xy[..., 0], kind) # (...,3)
wy, zy = phi_1d(xy[..., 1], kind)
w = wx.unsqueeze(-1) * wy.unsqueeze(-2) # (...,3,3)
cz = torch.stack(
torch.broadcast_tensors(zx.unsqueeze(-1), zy.unsqueeze(-2)), dim=-1
) # (...,3,3,2)
return w, cz
def D_1(q: torch.Tensor, k: torch.Tensor, kind: str = "cos2") -> torch.Tensor:
"""Published 1D kernel D_1(q,k) = sum_z phi_z(q) phi_z(k).
Evaluated in closed form (Zhuoran Yang / random-feature identity for the
raised-cosine bump) by summing over the union of the two touch sets; the
form below is the definition, kept literal for auditability.
"""
wq, zq = phi_1d(q, kind) # (...,3), (...,3)
# Sum runs only over cells where both weights are nonzero. Every such
# cell is in touch(q) (3 candidates), and phi_z(k) for an arbitrary z is
# b(k - z) / norm_k with norm_k taken over touch(k) -- the definition,
# since b(k - z) = 0 for every z outside touch(k).
zk = touched_cells(k, kind) # (...,3)
norm_k = (bump(k.unsqueeze(-1) - zk.to(k.dtype), kind) ** 2).sum(-1, keepdim=True).sqrt().clamp_min(1e-30)
phi_k_at_zq = bump(k.unsqueeze(-1) - zq.to(k.dtype), kind) / norm_k
return (wq * phi_k_at_zq).sum(dim=-1)
def D_2(q: torch.Tensor, k: torch.Tensor, kind: str = "cos2") -> torch.Tensor:
"""Published 2D kernel: D(q,k) = D_1(qx,kx) * D_1(qy,ky)."""
return D_1(q[..., 0], k[..., 0], kind) * D_1(q[..., 1], k[..., 1], kind)
def mix_read(cell_values: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
"""Mix per-cell reads by the bump weights.
cell_values : (..., 3, 3, d_v) read from each touched cell
weights : (..., 3, 3) normalized bump weights phi
returns (..., d_v)
"""
return (weights.unsqueeze(-1) * cell_values).sum(dim=(-3, -2))
|