File size: 6,259 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 | """
Batched, differentiable μLattice lattice scan.
Same published math as mulattice.lattice.LatticeMemory, but the sequential
recurrence is executed with tensor ops over the whole batch instead of a Python
loop over (token x 9 cells), so the synthetic MQAR mixers can be trained.
Sparsity is preserved: cells are keyed by (sequence, z_x, z_y) via
torch.unique, and only the keys that actually occur are allocated. No dense
grid over the address bounding box is ever materialised. Cell count is
unbounded -- it is the number of distinct addresses written.
Semantics identical to the reference implementation:
phi from the published bump kernel, exactly 9 cells per token,
v_bar = alpha * W_c kappa ; r = v - v_bar ; W_c <- alpha W_c + beta r kappa^T
read = sum over the 3x3 neighbourhood of phi_c * (W_c q2), empty cells 0.
"""
from __future__ import annotations
import torch
from .kernel import MAX_CELLS_PER_TOKEN, phi_2d
def compact_ids(keys: torch.Tensor) -> tuple[torch.Tensor, int]:
"""Map arbitrary int64 keys to dense contiguous ids, via sort.
Uses sort + run detection instead of torch.unique: on MPS, torch.unique corrupts int64
values outside the int32 range (measured: 4959 distinct keys reported as 2789,
with returned values like -30933022 instead of 549724880885). That scrambled
the cell ids, collapsed the lattice onto a few cells, and made training
diverge to inf gradients on MPS while CPU was fine. Sort + run detection is
exact on every backend.
"""
if keys.numel() == 0:
return keys.clone(), 0
order = torch.argsort(keys)
ks = keys[order]
first = torch.ones_like(ks, dtype=torch.bool)
first[1:] = ks[1:] != ks[:-1]
ids_sorted = torch.cumsum(first.to(torch.int64), 0) - 1
ids = torch.empty_like(ids_sorted)
ids[order] = ids_sorted
return ids, int(first.sum())
def _encode_keys(coords: torch.Tensor, n_seq: int, size: int = 1 << 20) -> torch.Tensor:
"""coords (B,T,9,2) int64 -> flat int64 keys, collision-free for |z| < size/2."""
z = coords + (size // 2)
b = torch.arange(n_seq, device=coords.device, dtype=torch.int64).view(n_seq, 1, 1)
return (b * (size * size)) + z[..., 0] * size + z[..., 1]
class SparseLatticeScan(torch.nn.Module):
def __init__(self, d_k: int, d_v: int, n_heads: int, size: int = 1 << 20,
weighted_write: bool = True):
super().__init__()
self.d_k, self.d_v, self.n_heads, self.size = d_k, d_v, n_heads, size
self.weighted_write = weighted_write
self.last_n_cells = 0
def forward(
self,
q_addr: torch.Tensor, # (B,T,2) real address of the query
k_addr: torch.Tensor, # (B,T,2) real address of the key
q2: torch.Tensor, # (B,T,H,d_k) content query
kappa: torch.Tensor, # (B,T,H,d_k) content key
v: torch.Tensor, # (B,T,H,d_v) value
alpha: torch.Tensor, # (B,T,H,1) decay in (0,1)
beta: torch.Tensor, # (B,T,H,1) write gate in (0,1)
read_mode: str = "last", # "last" | "all" | "none"
causal_read: bool = True,
):
B, T, H = kappa.shape[0], kappa.shape[1], self.n_heads
dev, dt = kappa.device, kappa.dtype
wk, zk = phi_2d(k_addr) # (B,T,3,3), (B,T,3,3,2)
wq, zq = phi_2d(q_addr)
wk = wk.reshape(B, T, MAX_CELLS_PER_TOKEN)
wq = wq.reshape(B, T, MAX_CELLS_PER_TOKEN)
zk = zk.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)
zq = zq.reshape(B, T, MAX_CELLS_PER_TOKEN, 2)
need_q = read_mode != "none"
keys_w = _encode_keys(zk, B, self.size).reshape(-1)
keys_q = _encode_keys(zq, B, self.size).reshape(-1) if need_q else keys_w[:0]
all_keys = torch.cat([keys_w, keys_q]) if need_q else keys_w
inv, n_cells = compact_ids(all_keys)
self.last_n_cells = int(n_cells)
inv_w = inv[: keys_w.numel()].reshape(B, T, MAX_CELLS_PER_TOKEN)
inv_q = (
inv[keys_w.numel():].reshape(B, T, MAX_CELLS_PER_TOKEN) if need_q else None
)
# cells allocated lazily: only distinct keys that occur, zero-initialised
cells = torch.zeros(n_cells, H, self.d_v, self.d_k, device=dev, dtype=dt)
outs = []
for t in range(T):
idx = inv_w[:, t] # (B,9)
W = cells.index_select(0, idx.reshape(-1)).reshape(B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
kap = kappa[:, t].unsqueeze(1) # (B,1,H,dk)
a = alpha[:, t].unsqueeze(1) # (B,1,H,1)
b = beta[:, t].unsqueeze(1)
v_bar = a * (W @ kap.unsqueeze(-1)).squeeze(-1) # (B,9,H,dv)
r = v[:, t].unsqueeze(1) - v_bar
gate = b * wk[:, t].unsqueeze(-1).unsqueeze(-1) if self.weighted_write else b
W = a.unsqueeze(-1) * W + gate.unsqueeze(-1) * (r.unsqueeze(-1) * kap.unsqueeze(-2))
cells.index_put_((idx.reshape(-1),), W.reshape(-1, H, self.d_v, self.d_k))
if read_mode == "all":
qi = inv_q[:, t] if causal_read else None
if qi is None: # non-causal: after full seq
outs.append(None)
else:
Wq = cells.index_select(0, qi.reshape(-1)).reshape(
B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
rq = (Wq @ q2[:, t].unsqueeze(1).unsqueeze(-1)).squeeze(-1) # (B,9,H,dv)
outs.append((wq[:, t].unsqueeze(-1).unsqueeze(-1) * rq).sum(1))
if read_mode == "last":
qi = inv_q[:, T - 1] # (B,9)
Wq = cells.index_select(0, qi.reshape(-1)).reshape(
B, MAX_CELLS_PER_TOKEN, H, self.d_v, self.d_k)
rq = (Wq @ q2[:, T - 1].unsqueeze(1).unsqueeze(-1)).squeeze(-1)
out = (wq[:, T - 1].unsqueeze(-1).unsqueeze(-1) * rq).sum(1) # (B,H,dv)
return out, n_cells
if read_mode == "all":
return torch.stack(outs, dim=1), n_cells # (B,T,H,dv)
return None, n_cells
|