{"slug": "build-a-bigram-model-in-10-minutes", "title": "Build a Bigram Model in 10 Minutes", "summary": "A tutorial by an unnamed author demonstrates that a bigram model with a single embedding table can recover true transition probabilities from 100,000 generated samples, achieving learned probabilities close to the true matrix (e.g., 0.5012 vs. 0.50 for a->a). The model, implemented in PyTorch with a zero-layer architecture, uses one weight table of nine parameters and trains in one pass with cross-entropy loss.", "body_md": "# Build a Bigram Model in 10 Minutes\n\nA bigram model looks at one token and predicts the next token.\n\nIn [Why KV Cache and Not QKV Cache?](why-kv-cache.html), I\nused a 0-layer Transformer before adding attention. Here I make that\nmodel smaller: one weight table, \\(\\W\\).\n\nWe will generate data from known bigram probabilities. Then we will see if nine learned numbers can recover them.\n\n## Make the Data\n\nOur vocabulary has three tokens: `a`\n\n, `b`\n\n, and\n`_`\n\n.\n\nThe true transition matrix is\n\n\\[ \\PB_{\\text{true}} = \\begin{bmatrix} 0.50 & 0.20 & 0.30 \\\\ 0.10 & 0.80 & 0.10 \\\\ 0.90 & 0.05 & 0.05 \\end{bmatrix}. \\]A row is the current token. A column is the next token. For example,\n\n\\[ P(x_{t+1}=\\texttt{b}\\mid x_t=\\texttt{a})=0.20. \\]\n\n``` python\nimport torch\nfrom torch.utils.data import DataLoader, Dataset\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.optim import AdamW\nimport matplotlib.pyplot as plt\n\ndef generate_true_data():\n    true_bigrams = {\n        \"a\": torch.tensor([0.50, 0.20, 0.30]),\n        \"b\": torch.tensor([0.10, 0.80, 0.10]),\n        \"_\": torch.tensor([0.90, 0.05, 0.05]),\n    }\n    vocab = list(true_bigrams.keys())\n    n_seq = 100_000\n\n    full_seq = \"a\"\n    all_tokens = [vocab.index(full_seq[-1])]\n\n    for _ in range(n_seq):\n        last_token = full_seq[-1]\n        sample_arg = torch.multinomial(\n            input=true_bigrams[last_token],\n            num_samples=1,\n        )\n        sample_id = sample_arg.item()\n        all_tokens.append(sample_id)\n        full_seq += vocab[sample_id]\n\n    return full_seq, torch.tensor(all_tokens)\n```\n\nThis makes one long string and 100,000 transitions.\n\n## Make the Pairs\n\nThe input is one token. The target is the token after it:\n\n\\[ \\X = [x_0,x_1,\\ldots,x_{T-1}], \\qquad \\Y = [x_1,x_2,\\ldots,x_T]. \\]\n\n``` python\nclass BigramSequence(Dataset):\n    def __init__(self):\n        full_seq, all_tokens = generate_true_data()\n        self.X = all_tokens[:-1]\n        self.Y = all_tokens[1:]\n\n    def __getitem__(self, index):\n        return self.X[index], self.Y[index]\n\n    def __len__(self):\n        return len(self.X)\n```\n\n## Zero Layers\n\nThe complete model is one embedding table:\n\n``` python\nclass ZeroLayer(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.W = nn.Embedding(\n            num_embeddings=3,\n            embedding_dim=3,\n        )\n\n    def forward(self, x):\n        return self.W(x)\n```\n\n`nn.Embedding`\n\nis a row lookup. If the current token has id\n\\(i\\), the model returns\n\nThese three numbers are logits. Softmax makes them probabilities:\n\n\\[ \\hat{\\ps}_t = \\textsf{softmax}(\\W_{i,:}). \\]There is no hidden layer. There is no attention.\n\n## Train\n\nOne pass over the data is enough for this example.\n\n```\nmodel = ZeroLayer()\ndata = BigramSequence()\nloader = DataLoader(data, batch_size=1_500)\noptimizer = AdamW(model.parameters(), lr=0.1)\n\nlosses = []\n\nfor x, y in loader:\n    optimizer.zero_grad()\n\n    logits = model(x)                 # [batch, 3]\n    loss = F.cross_entropy(logits, y)\n\n    loss.backward()\n    optimizer.step()\n    losses.append(loss.item())\n```\n\nCross entropy applies softmax internally. So the model returns raw logits during training.\n\n\\[ \\mathcal{L} = -\\frac{1}{T}\\sum_{t=0}^{T-1} \\log \\hat{P}(x_{t+1}\\mid x_t). \\]## Read the Weights\n\nApply softmax to each row of \\(\\W\\):\n\n```\nwith torch.no_grad():\n    learned_bigrams = model.W.weight.softmax(dim=1)\n    print(learned_bigrams)\n\nplt.plot(losses)\nplt.xlabel(\"update\")\nplt.ylabel(\"cross entropy\")\nplt.show()\n```\n\nOne run gives:\n\n```\ntensor([[0.5012, 0.1975, 0.3013],\n        [0.0959, 0.8052, 0.0988],\n        [0.9098, 0.0416, 0.0486]])\n```\n\n", "url": "https://wpnews.pro/news/build-a-bigram-model-in-10-minutes", "canonical_source": "https://www.sithankanna.com/posts/learning-bigram-zero-layers.html", "published_at": "2026-08-30 20:01:39+00:00", "updated_at": "2026-08-30 20:21:53.391780+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-research", "ai-tools"], "entities": ["PyTorch"], "alternates": {"html": "https://wpnews.pro/news/build-a-bigram-model-in-10-minutes", "markdown": "https://wpnews.pro/news/build-a-bigram-model-in-10-minutes.md", "text": "https://wpnews.pro/news/build-a-bigram-model-in-10-minutes.txt", "jsonld": "https://wpnews.pro/news/build-a-bigram-model-in-10-minutes.jsonld"}}