Build a Bigram Model in 10 Minutes

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.

Build a Bigram Model in 10 Minutes A bigram model looks at one token and predicts the next token. In Why KV Cache and Not QKV Cache? why-kv-cache.html , I used a 0-layer Transformer before adding attention. Here I make that model smaller: one weight table, \ \W\ . We will generate data from known bigram probabilities. Then we will see if nine learned numbers can recover them. Make the Data Our vocabulary has three tokens: a , b , and . The true transition matrix is \ \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, \ P x {t+1}=\texttt{b}\mid x t=\texttt{a} =0.20. \ python import torch from torch.utils.data import DataLoader, Dataset import torch.nn as nn import torch.nn.functional as F from torch.optim import AdamW import matplotlib.pyplot as plt def generate true data : true bigrams = { "a": torch.tensor 0.50, 0.20, 0.30 , "b": torch.tensor 0.10, 0.80, 0.10 , " ": torch.tensor 0.90, 0.05, 0.05 , } vocab = list true bigrams.keys n seq = 100 000 full seq = "a" all tokens = vocab.index full seq -1 for in range n seq : last token = full seq -1 sample arg = torch.multinomial input=true bigrams last token , num samples=1, sample id = sample arg.item all tokens.append sample id full seq += vocab sample id return full seq, torch.tensor all tokens This makes one long string and 100,000 transitions. Make the Pairs The input is one token. The target is the token after it: \ \X = x 0,x 1,\ldots,x {T-1} , \qquad \Y = x 1,x 2,\ldots,x T . \ python class BigramSequence Dataset : def init self : full seq, all tokens = generate true data self.X = all tokens :-1 self.Y = all tokens 1: def getitem self, index : return self.X index , self.Y index def len self : return len self.X Zero Layers The complete model is one embedding table: python class ZeroLayer nn.Module : def init self : super . init self.W = nn.Embedding num embeddings=3, embedding dim=3, def forward self, x : return self.W x nn.Embedding is a row lookup. If the current token has id \ i\ , the model returns These three numbers are logits. Softmax makes them probabilities: \ \hat{\ps} t = \textsf{softmax} \W {i,:} . \ There is no hidden layer. There is no attention. Train One pass over the data is enough for this example. model = ZeroLayer data = BigramSequence loader = DataLoader data, batch size=1 500 optimizer = AdamW model.parameters , lr=0.1 losses = for x, y in loader: optimizer.zero grad logits = model x batch, 3 loss = F.cross entropy logits, y loss.backward optimizer.step losses.append loss.item Cross entropy applies softmax internally. So the model returns raw logits during training. \ \mathcal{L} = -\frac{1}{T}\sum {t=0}^{T-1} \log \hat{P} x {t+1}\mid x t . \ Read the Weights Apply softmax to each row of \ \W\ : with torch.no grad : learned bigrams = model.W.weight.softmax dim=1 print learned bigrams plt.plot losses plt.xlabel "update" plt.ylabel "cross entropy" plt.show One run gives: tensor 0.5012, 0.1975, 0.3013 , 0.0959, 0.8052, 0.0988 , 0.9098, 0.0416, 0.0486