| """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) |
| y = (addr[..., 1].clamp(-scale, scale) + scale) / (2 * scale) |
| gx = (x * grid).clamp(0.0, float(grid) - 1e-6) |
| 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, :]) |
| dy = D_1(y[:, :, None], y[:, None, :]) |
| D = dx * dy |
| diag = N |
| 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 |
| |
| |
| |
| gate = (write_overlap(addr) - gate_thresh).clamp(0.0, 1.0).detach() |
| lam = 0.3 + 0.7 * gate |
| w, z = phi_2d(addr) |
| size = int(2 * scale + 3) |
| off = int(scale) + 1 |
| 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) |
| 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): |
| 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() |
|
|
