File size: 5,322 Bytes
98c193b | 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 | """Datasets: text8 for language modelling, plus dependency-controlled probes.
text8 is the standard small-scale benchmark for discrete diffusion language
models (D3PM, SEDD, MDLM all report bits-per-character on it), which is what
makes a from-scratch comparison meaningful.
The probes exist because text8 alone cannot isolate the phenomenon we study.
Parallel decoding fails when simultaneously-revealed tokens depend on each
other, so we need tasks where that dependence is a knob rather than a property
of English.
"""
import os
import numpy as np
TEXT8_URL = "http://mattmahoney.net/dc/text8.zip"
# text8 is lowercase letters plus space: a 27-symbol alphabet, no tokenizer.
ALPHABET = " abcdefghijklmnopqrstuvwxyz"
SPECIALS = ["[MASK]", "[PAD]"]
class CharTokenizer:
def __init__(self):
self.itos = SPECIALS + list(ALPHABET)
self.stoi = {c: i for i, c in enumerate(self.itos)}
self.mask = self.stoi["[MASK]"]
self.pad = self.stoi["[PAD]"]
def __len__(self):
return len(self.itos)
def encode(self, s):
return [self.stoi[c] for c in s if c in self.stoi]
def decode(self, ids):
return "".join(self.itos[i] for i in ids)
def load_text8(root, seq_len=256):
"""Standard 90M/5M/5M character split used throughout the literature."""
raw = os.path.join(root, "text8")
if not os.path.exists(raw):
os.makedirs(root, exist_ok=True)
import zipfile, io, requests # only needed on first download
r = requests.get(TEXT8_URL, timeout=300)
r.raise_for_status()
with zipfile.ZipFile(io.BytesIO(r.content)) as z:
z.extract("text8", root)
text = open(raw).read()
assert len(text) == 100_000_000, f"unexpected text8 length {len(text)}"
tok = CharTokenizer()
cache = os.path.join(root, "text8.npy")
if os.path.exists(cache):
ids = np.load(cache)
else:
ids = np.array(tok.encode(text), dtype=np.uint8)
np.save(cache, ids)
tr, va = ids[:90_000_000], ids[90_000_000:95_000_000]
te = ids[95_000_000:]
def chunk(a):
n = len(a) // seq_len
return a[: n * seq_len].reshape(n, seq_len).astype(np.int64)
return chunk(tr), chunk(va), chunk(te), tok
# ---------------------------------------------------------------------------
# Dependency-controlled probes.
#
# Each item is a sequence whose tokens must agree with each other. `coupling`
# sets how many positions are bound together, which is the exact quantity that
# breaks parallel decoding: revealing two mutually-dependent positions in one
# step forces the model to guess them independently.
# ---------------------------------------------------------------------------
def build_probe(n, seq_len, coupling, seed=0, task="copy"):
"""
copy : second half must equal the first half. Each answer position is
fixed by one context position, so answer positions are mutually
INDEPENDENT given the context - parallel decoding should be safe.
agree : within each group of `coupling` positions all tokens must be equal
AND adjacent groups must differ. The value of a group is free, so
its members are mutually DEPENDENT: revealing two of them in one
step is a 1/V coin flip. This is the failure mode, isolated.
parity : groups are a run of bits whose last element is the XOR of the
rest, with adjacent groups forced to differ, so the group is
jointly constrained without any member being determined alone.
The adjacent-groups-differ rule matters: without it a model that emits one
constant token everywhere satisfies "all members agree" perfectly, and an
untrained model scores 100% at maximum parallelism. That degenerate solution
made the first version of this metric meaningless.
"""
tok = CharTokenizer()
rng = np.random.default_rng(seed)
V = 2 + 26 # specials + letters
lo, hi = 2, V # sample letters only
X = np.zeros((n, seq_len), dtype=np.int64)
for i in range(n):
if task == "copy":
half = seq_len // 2
src = rng.integers(lo, hi, size=half)
X[i, :half] = src
X[i, half:half * 2] = src
elif task == "agree":
g, prev = coupling, -1
for s in range(0, seq_len, g):
v = int(rng.integers(lo, hi))
while v == prev: # adjacent groups must differ
v = int(rng.integers(lo, hi))
X[i, s:s + g] = v
prev = v
elif task == "parity":
g, prev = coupling, None
for s in range(0, seq_len, g):
end = min(s + g, seq_len)
if end - s < 2:
X[i, s:end] = lo
continue
while True:
bits = list(rng.integers(0, 2, size=end - s - 1))
vals = bits + [int(sum(bits) % 2)]
if vals != prev: # adjacent groups must differ
break
X[i, s:end] = [lo + v for v in vals]
prev = vals
else:
raise ValueError(task)
return X, tok
|