File size: 2,868 Bytes
6fbb543 | 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 | """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 # noqa: E402
import pixelgen.model as M # noqa: E402
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 # noqa (tokenize returns ids)
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)")
|