File size: 4,576 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 | """Address-diversity regularisation for μLattice (experimental).
A mixture-of-experts-style load-balancing term applied to addresses: maximise the entropy of
the soft cell-usage distribution so addresses stay spread instead of collapsing onto a few
cells.
Status: the *collapse* hypothesis that motivated this was later contradicted -- on CUDA a run
with every key in one cell still generalised (0.875 at 256 keys), and the term did not
reproducibly help. It is kept as a documented experiment, not a recommended part of the mixer.
"""
from __future__ import annotations
import torch
def usage_entropy(addr: torch.Tensor, scale: float, grid: int = 24) -> torch.Tensor:
"""Entropy (nats) of the soft cell-usage distribution of `addr` (..., 2).
Addresses are normalised into a `grid` x `grid` histogram by differentiable bilinear
assignment, so the entropy carries gradient back to the address projection. Maximum is
log(grid**2); a collapsed map gives near zero.
"""
x = (addr[..., 0].clamp(-scale, scale) + scale) / (2 * scale) # -> [0, 1]
y = (addr[..., 1].clamp(-scale, scale) + scale) / (2 * scale)
gx = (x * grid).clamp(0.0, float(grid) - 1e-6) # [0, grid)
gy = (y * grid).clamp(0.0, float(grid) - 1e-6)
ix = gx.floor().long(); iy = gy.floor().long()
fx = (gx - ix.float()).reshape(-1); fy = (gy - iy.float()).reshape(-1)
ix1 = (ix + 1).clamp(max=grid - 1); iy1 = (iy + 1).clamp(max=grid - 1)
p = torch.zeros(grid * grid, device=addr.device, dtype=torch.float32)
p.index_add_(0, (ix * grid + iy).reshape(-1), ((1 - fx) * (1 - fy)).float())
p.index_add_(0, (ix1 * grid + iy).reshape(-1), (fx * (1 - fy)).float())
p.index_add_(0, (ix * grid + iy1).reshape(-1), ((1 - fx) * fy).float())
p.index_add_(0, (ix1 * grid + iy1).reshape(-1), (fx * fy).float())
p = p / p.sum().clamp_min(1e-9)
return -(p * (p + 1e-12).log()).sum()
def address_diversity_loss(addr: torch.Tensor, scale: float, grid: int = 24) -> torch.Tensor:
"""Minimise this to keep addresses spread: negative entropy of cell usage."""
return -usage_entropy(addr, scale, grid)
def write_overlap(addr: torch.Tensor) -> torch.Tensor:
"""Mean pairwise published-kernel overlap of write addresses, (B, N, 2) -> (B,).
D(k_i, k_j) = D1(x) * D1(y) from kernel.py measures exactly what the scan interferes
on: shared 3x3 cell neighbourhoods. This is the quantity the collision gate uses --
NOT a grid histogram of the address box.
"""
from .kernel import D_1
B, N, _ = addr.shape
x = addr[..., 0]
y = addr[..., 1]
dx = D_1(x[:, :, None], x[:, None, :]) # (B, N, N)
dy = D_1(y[:, :, None], y[:, None, :])
D = dx * dy
diag = N # D(i,i) == 1 each
off = (D.sum(dim=(1, 2)) - diag) / max(N * (N - 1), 1)
return off
def gated_spread_loss(addr: torch.Tensor, scale: float, gate_thresh: float = 0.05) -> torch.Tensor:
"""Collision-gated spread: lambda_t = clamp(overlap - 0.05, 0, 1) times the NEGATIVE
entropy of centre-cell usage, where usage is the real phi weights on the real integer
cells (not a histogram grid). Zero when the keys are already spread; positive only
while they actually collide."""
from .kernel import phi_2d
B, N, _ = addr.shape
# Floor 0.3 + gated escalation to 1.0. A pure gate self-disables at TRAIN-time key
# counts (64 keys spread well enough that overlap < thresh), leaving eval-time packing
# (256 keys) unaddressed -- measured: gated-only plateaued @256 at 0.66 vs 0.91 baseline.
gate = (write_overlap(addr) - gate_thresh).clamp(0.0, 1.0).detach()
lam = 0.3 + 0.7 * gate
w, z = phi_2d(addr) # (B,N,3,3), (B,N,3,3,2) int coords
size = int(2 * scale + 3)
off = int(scale) + 1 # keep integer arithmetic
gi = (z + off).clamp(0, size - 1)
idx = (gi[..., 0] * size + gi[..., 1]).reshape(B, -1).to(torch.int64)
mass = (w / w.sum(dim=(-2, -1), keepdim=True)).reshape(B, -1) # normalised per key
hist = torch.zeros(B, size * size, device=addr.device, dtype=addr.dtype)
idx_f = idx.reshape(B, -1)
mass_f = mass.reshape(B, -1)
for b in range(B): # index_add_ needs 1-D indices
hist[b].index_add_(0, idx_f[b], mass_f[b])
p = hist / hist.sum(dim=1, keepdim=True).clamp_min(1e-9)
entropy = -(p * (p + 1e-12).log()).sum(dim=1)
return (lam * (-entropy)).mean()
|