Bigram Lab

running

◷train_gpt.py▸get_bigram_hash◎learner＄no GPU

Blocks

Quick summary

Hash every (prev token, curr token) pair into a 251 520-bucket table. Position 0 gets reserved slot 251 519 (no prev…

Full explanation below the code →

fba-lab — train_gpt.py · get_bigram_hashexecuting

// block: get_bigram_hash · lines 1405–1420$ study train_gpt.py --block bigram-hash-fnHash every (prev token, curr token) pair into a 251 520-bucket table. Position 0 gets reserved slot … ✓

Explanation

Hash every (prev token, curr token) pair into a 251 520-bucket table. Position 0 gets reserved slot 251 519 (no prev token). Two odd integers XOR-combined → fast collision spread.

Think about

What problem does a 251K-bucket table cause when you have 2.5B possible bigrams?

// architecture

Live diagram

01 · train_gpt.py line 1405

get_bigram_hash

A cheap shortcut: inject (prev token, curr token) into every layer

A standard token embedding gives each token ID a learned vector, but it knows nothing about the previous token. The transformer layers learn context through attention — but for a tiny speedrun model with limited capacity, that's expensive. get_bigram_hash hard-injects a previous-token clue before the model even starts.

def get_bigram_hash(x):
"""For each position t: hash(token[t-1], token[t]) → bigram index."""
rand_int_1 = 36313
rand_int_2 = 27191
mod = args.bigram_vocab_size - 1 # 251519 = 50304×5 - 1
x = x.to(torch.int32).clone()
x[0] = mod # position 0: no prev token → reserved slot
x[1:] = torch.bitwise_xor(
rand_int_1 * x[1:], # scale curr token
rand_int_2 * x[:-1] # scale prev token
) % mod # hash into 251K-bucket table
return x

Why this helps small models

Many next-token predictions are highly local: "New→York", "torch→.nn". A normal transformer learns this through attention, but that takes many parameters. The bigram embedding gives the model a direct feature for (prev, curr) before it even enters the transformer stack.

251 520-bucket hash table (vs 50 304² ≈ 2.5B exact pairs)
Hash collisions are intentional — smaller table, acceptable overlap
Learned scalar bigram_lambdas[i] per layer controls weight
Initialized to zero — the model can ignore it if not useful

The collision problem

GPT-2's vocabulary has 50 304 tokens → 2.53B possible bigram pairs → 251K buckets → 10 059 bigrams per bucket on average. "New York", "New Delhi", "New Orleans" all got the same vector. The sign trick (@TrianX, WR #68, 79.7s) fixes this: each bigram gets the shared row × a unique ±1 mask → 2.06B reachable vectors.

Sign patterns act like random projections → near-orthogonal directions
bigram_dim: 768 → 192, sign table: nn.Parameter(8192, 192, fixed)
Entry #62 (bigram base): ~81.2s → Entry #68 (sign trick): 79.7s WR

A cheap shortcut: inject (prev token, curr token) into every layer

Why this helps small models

251 520-bucket hash table (vs 50 304² ≈ 2.5B exact pairs)
Hash collisions are intentional — smaller table, acceptable overlap
Learned scalar bigram_lambdas[i] per layer controls weight
Initialized to zero — the model can ignore it if not useful

The collision problem

Sign patterns act like random projections → near-orthogonal directions
bigram_dim: 768 → 192, sign table: nn.Parameter(8192, 192, fixed)
Entry #62 (bigram base): ~81.2s → Entry #68 (sign trick): 79.7s WR