| """Replace the dense V x C word->class matrix with a hand-written binary code. |
| |
| word -> d-bit code of its class, emitted as +-1 Linear(V, d, bias=False) |
| class rows = +-1 codes; score_c = d - 2*hamming Linear(d, C) (+0.5 bias on default) |
| |
| Params drop from V*C to V*d + d*C. Price: multi-word prompts sum two codes |
| and can land on a third class. Both are measured here. |
| """ |
|
|
| import itertools |
| import sys |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| sys.path.insert(0, ".") |
| from experiments.scale_classes import build, REAL_CLASSES, REAL_SYN |
| import pixelgen.model as M |
|
|
|
|
| def code_layer(m): |
| C, V = len(m.classes), len(m.vocab) |
| d = max(1, int(np.ceil(np.log2(C)))) |
| codes = torch.tensor([[(ci >> b) & 1 for b in range(d)] for ci in range(C)], dtype=torch.float32) * 2 - 1 |
| enc = nn.Linear(V, d, bias=False); dec = nn.Linear(d, C) |
| with torch.no_grad(): |
| enc.weight.zero_(); dec.bias.zero_() |
| for ci, cls in enumerate(m.classes): |
| for w in M.CLASS_SYNONYMS[cls]: |
| enc.weight[:, m.vocab[w]] = codes[ci] |
| dec.weight.copy_(codes) |
| dec.bias[m.classes.index(M.DEFAULT_CLASS)] = 0.5 |
| seq = nn.Sequential(enc, dec) |
| for p in seq.parameters(): p.requires_grad_(False) |
| return seq, d |
|
|
|
|
| def interference(m): |
| """Fraction of 2-class-word prompts whose argmax is neither word's class.""" |
| words = [(c, M.CLASS_SYNONYMS[c][0]) for c in REAL_CLASSES] |
| bad = 0; n = 0 |
| for (ca, wa), (cb, wb) in itertools.combinations(words, 2): |
| toks = m.tokenize(f"{wa} {wb}") |
| bag = torch.zeros(1, len(m.vocab)); bag[0, toks] += 1 |
| cls = m.classes[m.word_to_class(bag).argmax(-1).item()] |
| bad += cls not in (ca, cb); n += 1 |
| return bad / n |
|
|
|
|
| Ns = [int(a) for a in sys.argv[1:]] or [0, 100, 1000, 3000, 10000] |
| print(f"{'classes':>8} {'d':>3} {'dense params':>13} {'coded params':>13} {'ratio':>7} {'isolated':>8} {'interf':>7}") |
| for n in Ns: |
| m, _ = build(n) |
| ref = {p: m.generate([p] * 2, [0, 1], scale=1) for p in REAL_CLASSES + ["red potion", "", "dragon", "bone"]} |
| dense_p = sum(p.numel() for p in m.word_to_class.parameters()) |
| seq, d = code_layer(m) |
| m.word_to_class = seq |
| coded_p = sum(p.numel() for p in seq.parameters()) |
| ok = all(np.array_equal(m.generate([p] * 2, [0, 1], scale=1), ref[p]) for p in ref) |
| if n: |
| nm = f"syn{n-1:05d}" |
| ok = ok and np.array_equal(m.generate([nm + "b"], [0], scale=1)[0], m.generate([nm], [0], scale=1)[0]) |
| print(f"{len(m.classes):>8} {d:>3} {dense_p:>13,} {coded_p:>13,} {dense_p/coded_p:>6.1f}x {str(ok):>8} {interference(m):>6.0%}") |
|
|
| print("\n(interf = share of 'wordA wordB' prompts over the 21 real classes that resolve to a THIRD class;") |
| print(" dense layer: 0% by construction, ties go to lower index)") |
|
|
