Step 9 of 28Skill: AI Systems Optimization
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
Token sequence: "New"tok=1570 "York"tok=6849 "City"tok=2254 "torch"tok=22618 ··· get_bigram_hash(tokens) hash[t] = (36313 × tok[t] ⊕ 27191 × tok[t-1]) mod 251519 Bigram indices: 251519reserved 84312(hash) 127041(hash) 63205(hash) bigram_embed nn.Embedding(251520, 768) GPT.forward — 11 layers token_embed(tokens) → x x = λ·x + bigram_λ[i] · x0_bigram CausalSelfAttention(x) MLP(x) repeated × 11 layers
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