File size: 7,600 Bytes
98c193b 2220ce4 | 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 | """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
# ---------------------------------------------------------------------------
# Addition: the canonical UNIQUE-answer task with chained inter-token
# dependence. Given the operands the answer is fully determined, so every
# position's true marginal is a point mass and parallel decoding faces no
# information-theoretic barrier - unlike `agree`/`parity`, where the value is
# genuinely free and the bound is unbeatable. What remains is a computation
# problem: digit i needs the carry out of digit i-1, and sequential decoding
# amortises that chain across passes while parallel decoding must fit it inside
# one fixed-depth forward pass.
# ---------------------------------------------------------------------------
DIGITS10 = "0123456789"
class AddTokenizer:
def __init__(self):
self.itos = ["[MASK]", "[PAD]", "+", "="] + list(DIGITS10)
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 build_addition(n, digits, seq_len, seed=0, exact=False):
"""Returns (X, answer_mask, tok).
Layout: operand a, '+', operand b, '=', then the answer written
units-first. `answer_mask` marks the positions the model must produce; the
prompt is always visible, so the task is conditional and the answer unique.
"""
tok = AddTokenizer()
rng = np.random.default_rng(seed)
X = np.full((n, seq_len), tok.pad, dtype=np.int64)
A = np.zeros((n, seq_len), dtype=bool)
for i in range(n):
d = digits if exact else int(rng.integers(1, digits + 1))
def draw(nd):
ds = rng.integers(0, 10, size=nd)
if nd > 1 and ds[0] == 0:
ds[0] = rng.integers(1, 10)
return "".join(str(int(v)) for v in ds)
pa, pb = draw(d), draw(d)
s = str(int(pa) + int(pb))[::-1] # units first
seq = list(pa) + ["+"] + list(pb) + ["="] + list(s)
if len(seq) > seq_len:
raise ValueError(f"seq_len {seq_len} too small for {d} digits")
ids = [tok.stoi[c] for c in seq]
X[i, : len(ids)] = ids
A[i, len(pa) + len(pb) + 2: len(ids)] = True # answer positions only
return X, A, tok
|