Spaces:

Ex0bit
/

mLattice

Running

μLattice: live demo — your own facts, store at scale

20a3a74 verified 28 days ago

5.94 kB

"""
μ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))

	"""
	μ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))