File size: 7,553 Bytes
2f06744 | 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 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | #!/usr/bin/env python3
"""Characterize addressing limits: adversarial keys, packing, collision scaling.
Tests three properties:
1. Random-key collision rate at increasing pair counts (how does the lattice
handle natural collisions?)
2. Adversarial keys: find key pairs whose addresses collide, plant them in
the same sequence, measure interference
3. Memory scaling: measure bytes allocated vs unique facts stored
"""
from __future__ import annotations
import json
import os
import sys
import time
import torch
HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.join(HERE, "..")
sys.path.insert(0, os.path.join(ROOT, "src"))
from mulattice.protocol import SynthCfg, build_model, force_fast_scan, logits_at, synth_batch
from mulattice.kernel import phi_2d
DEVICE = os.environ.get("DEVICE", "cuda" if torch.cuda.is_available() else "cpu")
CKPT = os.path.join(ROOT, "checkpoints", "capacity8k.pt")
def key_addresses(model, tokens, key_pos):
"""Write addresses of tokens at key_pos."""
mx = model.blocks[0].mixer
with torch.no_grad():
xc = mx.conv(model.blocks[0].n1(model.emb(tokens)))
kaddr = mx.cfg.addr_scale * torch.tanh(mx.Wa(xc))
ar = torch.arange(tokens.shape[0], device=tokens.device).unsqueeze(1)
return kaddr[ar, key_pos]
def test_random_collisions(cfg, model):
"""Natural collision rate at increasing pair counts."""
results = []
for n_pairs in [64, 256, 512, 1024, 2048, 4096, 8192]:
g = torch.Generator().manual_seed(42)
b = synth_batch(1, n_pairs, 4, cfg, filler=0, generator=g)
tokens = b["tokens"].to(DEVICE)
key_pos = b["key_pos"].to(DEVICE)
addr = key_addresses(model, tokens, key_pos)[0].cpu()
cells = torch.floor(addr).long()
codes = cells[:, 0] * 10000 + cells[:, 1]
unique_cells = len(torch.unique(codes))
collisions = n_pairs - unique_cells
density = n_pairs / (2 * cfg.addr_scale + 3) ** 2
results.append({
"n_pairs": n_pairs, "unique_cells": unique_cells,
"collisions": collisions, "collision_rate": collisions / n_pairs,
"keys_per_cell": density,
})
print(f" {n_pairs:5d} keys: {collisions:4d} collisions "
f"({collisions/n_pairs:.1%}), {unique_cells:5d} cells, "
f"density={density:.3f} keys/cell", flush=True)
return results
def test_adversarial(cfg, model):
"""Find key pairs whose addresses collide, measure interference."""
g = torch.Generator().manual_seed(123)
# Generate many keys, compute their addresses, find colliding pairs
n_candidates = 2000
key_ids = torch.randperm(cfg.NK, generator=g)[:n_candidates]
tokens = key_ids.unsqueeze(0).to(DEVICE)
pos = torch.arange(n_candidates).unsqueeze(0).to(DEVICE)
addr = key_addresses(model, tokens, pos)[0].cpu()
cells = torch.floor(addr).long()
codes = cells[:, 0] * 10000 + cells[:, 1]
# Find pairs that map to the same cell
from collections import defaultdict
cell_map = defaultdict(list)
for i, c in enumerate(codes.tolist()):
cell_map[c].append(i)
colliding_pairs = []
for c, indices in cell_map.items():
if len(indices) >= 2:
for j in range(min(3, len(indices) - 1)):
colliding_pairs.append((indices[j], indices[j + 1], c))
print(f"\n Found {len(colliding_pairs)} colliding key pairs "
f"from {n_candidates} candidates", flush=True)
# Measure interference: store both keys, query both
results = []
n_test = min(20, len(colliding_pairs))
for idx in range(n_test):
i, j, cell = colliding_pairs[idx]
k1, k2 = int(key_ids[i]), int(key_ids[j])
v1, v2 = 100, 200 # distinct values
# Build a small sequence: k1→v1, k2→v2, query k1, query k2
seq = [k1, v1 + cfg.NK, k2, v2 + cfg.NK, k1, cfg.SEP, k2, cfg.SEP]
tokens = torch.tensor([seq]).to(DEVICE)
ans = torch.tensor([[5, 7]]).to(DEVICE) # SEP positions
gold = torch.tensor([[v1 + cfg.NK, v2 + cfg.NK]])
with torch.no_grad():
pred = logits_at(model, tokens, ans).argmax(-1).cpu()
ok1 = int(pred[0, 0]) == gold[0, 0]
ok2 = int(pred[0, 1]) == gold[0, 1]
results.append({"k1": k1, "k2": k2, "cell": cell,
"k1_correct": bool(ok1), "k2_correct": bool(ok2)})
if results:
acc1 = sum(r["k1_correct"] for r in results) / len(results)
acc2 = sum(r["k2_correct"] for r in results) / len(results)
print(f" Adversarial collision test ({n_test} pairs): "
f"key1 recall={acc1:.3f}, key2 recall={acc2:.3f}", flush=True)
return results
def test_memory_scaling(cfg, model):
"""Measure bytes allocated vs unique facts stored."""
from mulattice.lattice import LatticeMemory
results = []
for n_pairs in [64, 256, 1024, 4096, 8192, 16384]:
g = torch.Generator().manual_seed(99)
keys = torch.randperm(cfg.NK, generator=g)[:n_pairs]
vals = torch.randperm(cfg.NV, generator=g)[:n_pairs]
# Compute addresses for these keys
tokens = keys.unsqueeze(0).to(DEVICE)
pos = torch.arange(n_pairs).unsqueeze(0).to(DEVICE)
addr = key_addresses(model, tokens, pos)[0]
# Count unique cells
cells = torch.floor(addr).long()
codes = cells[:, 0] * 10000 + cells[:, 1]
unique_cells = len(torch.unique(codes.cpu()))
bytes_per_cell = 2 * 32 * 32 * 4 # 2 heads × 32×32 × 4 bytes
total_bytes = unique_cells * bytes_per_cell
results.append({
"n_pairs": n_pairs, "unique_cells": unique_cells,
"bytes": total_bytes, "mb": total_bytes / 1e6,
"bytes_per_pair": total_bytes / n_pairs,
})
print(f" {n_pairs:5d} pairs → {unique_cells:5d} cells → "
f"{total_bytes/1e6:.1f} MB ({total_bytes/n_pairs:.0f} B/pair)",
flush=True)
if n_pairs >= 16384:
# Extrapolate to 1M and 131K pairs
cells_per_pair = unique_cells / n_pairs
for target in [131072, 1000000]:
est_cells = int(target * cells_per_pair)
est_bytes = est_cells * bytes_per_cell
print(f" extrapolated {target:,} pairs → "
f"{est_cells:,} cells → {est_bytes/1e9:.1f} GB", flush=True)
return results
def main():
cfg = SynthCfg()
ck = torch.load(CKPT, map_location="cpu", weights_only=False)
cfg.addr_scale = float(ck["proto_cfg"]["addr_scale"])
model = build_model(cfg)
model.load_state_dict(ck["model"])
model = model.to(DEVICE).eval()
force_fast_scan(model)
print("=" * 60)
print("1. NATURAL COLLISION RATE (random keys)")
print("=" * 60)
collision_results = test_random_collisions(cfg, model)
print("\n" + "=" * 60)
print("2. ADVERSARIAL COLLISIONS (keys mapped to same cell)")
print("=" * 60)
adversarial_results = test_adversarial(cfg, model)
print("\n" + "=" * 60)
print("3. MEMORY SCALING")
print("=" * 60)
memory_results = test_memory_scaling(cfg, model)
out = os.path.join(ROOT, "results", "addressing_characterization.json")
with open(out, "w") as f:
json.dump({
"collisions": collision_results,
"adversarial": adversarial_results,
"memory": memory_results,
}, f, indent=2)
print(f"\nsaved {out}")
if __name__ == "__main__":
main()
|