Spaces:

Ex0bit
/

mLattice

Running

mLattice / experiments /eval_adversarial.py

Docker Space: actual demo server, 3 checkpoints, readable dark text

2f06744 verified 28 days ago

7.55 kB

#!/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()

	#!/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()