Training a neural bigram model on Manas
Abstract
I am studying language models from the bottom up. After building a Kyrgyz byte-level BPE tokenizer, I wanted to follow the next stage myself: take real token IDs, turn them into predictions, calculate a loss, run backpropagation, update the weights, and generate new text.
I built a neural bigram language model and trained it on the Kyrgyz epic Manas. The model is one square table. It receives one current token and stores a score for every possible next token. It has no attention, hidden layers, or memory beyond the last token. This limitation makes the full learning process visible.
The final experiment used the first 10,000 BPE tokens of the epic itself. The model contained 6,375,625 trainable parameters and ran through PyTorch on the Apple mps device. Ten thousand training steps took 15.84 seconds. Validation cross-entropy fell from the uniform baseline of 7.834 to 6.380, while top-1 next-token accuracy reached 17.1%. The output learned names, endings, punctuation, and short patterns from Manas, but it could not preserve a subject or event across a sentence.
The most useful finding came from the data. My first runs learned scholarly prose because the digital source included a long introduction before the epic. The model exposed that mistake in its generated vocabulary. I corrected the corpus boundary and repeated the experiment on the epic text alone.
1. Why I built this
My previous project focused on tokenization. I collected a Kyrgyz corpus, implemented byte-level BPE, trained several tokenizers, and studied the tradeoff between vocabulary size and sequence length. By the end, I could trace text into UTF-8 bytes, follow BPE merges, inspect token IDs, and decode them back into the original text.
The tokenizer stops at a list of numbers. A language model starts there.
I had watched Andrej Karpathy build a small bigram model in his neural-network lecture series. The code is short, but several ideas inside it deserve careful attention. Why does an embedding table have the shape vocabulary × vocabulary? How can one loss number describe hundreds of separate mistakes? How does backward() know which row to change? Why does the input have two dimensions while the logits have three?
I wanted to answer those questions with a model I had prepared myself. I used the Kyrgyz tokenizer from the previous experiment and a passage from Manas. This made the project small enough to inspect and personal enough to keep interesting.
2. What a bigram language model knows
A bigram model predicts the next item from the current item. A word-level version might learn that New often precedes York. A character-level version might learn that q often precedes u. My model works with BPE tokens, so an item can be a word, a word fragment, punctuation, or a space joined to nearby text.
Suppose the tokenizer produces this simplified sequence:
["Түп", " атасы", " Түгөл", " кан", ","]
The training pairs are:
"Түп" -> " атасы"
" атасы" -> " Түгөл"
" Түгөл" -> " кан"
" кан" -> ","
When the model sees " Түгөл", it reads one row from its table. That row contains one score for " кан", another score for ",", another for " Манас", and scores for every other token in the pilot vocabulary.
The model sees no earlier text. It cannot use "Түп" or " атасы" while predicting after " Түгөл". During generation, each sampled token replaces the entire useful context.
This model therefore learns adjacency rather than meaning. It can discover that a name often receives a certain ending, that a comma follows a phrase, or that one poetic fragment often follows another. It cannot remember who entered a battle or why an event happened.
3. One table contains the whole network
The complete PyTorch model fits in one important line:
self.token_embedding_table = nn.Embedding(C, C)
C is the vocabulary size used by the model. nn.Embedding usually maps a token ID to a shorter vector that another layer processes. Here each vector already has length C. It directly represents the scores for the next token.
You can picture the weights as a matrix:
possible next token
0 1 2 ... C-1
current token 0 score score score ... score
current token 1 score score score ... score
current token 2 score score score ... score
...
current token C-1 score score score ... score
Selecting current token 42 returns row 42. The value at column 137 is the model's raw score for next token 137 after current token 42.
These numbers are logits. A logit does not need to stay between zero and one. Softmax turns the full row into probabilities whose sum equals one.
4. The full tokenizer vocabulary was the wrong size
The Kyrgyz tokenizer from my previous experiment contains 32,768 tokens. Using all of them in an exact bigram table would require:
32,768 × 32,768 = 1,073,741,824 parameters
Float32 weights alone would occupy more than 4 GiB. Training also needs gradients and framework memory. That scale would make a simple learning experiment wasteful, and most rows would never appear in a short passage.
I kept the original BPE segmentation and removed unused IDs from the local model. The process was mechanical:
- tokenize the selected passage with the 32,768-token Kyrgyz BPE;
- collect the original IDs that appear in that passage;
- sort them and assign local IDs from
0upward; - save the reversible local-to-original map.
The final 10,000-token passage used 2,525 distinct tokenizer IDs. The model table became:
2,525 × 2,525 = 6,375,625 parameters
The text was not tokenized again by a smaller algorithm. Each piece remained the same BPE token. The model simply stopped reserving rows and columns for tokens that could not appear in this closed pilot.
This choice has a clear limitation. The local vocabulary comes from both the selected training and validation sections. It reveals which token identities exist in validation, though it does not expose their order or frequency. The original tokenizer has no unknown token, so a strictly train-only vocabulary would make some validation tokens impossible to represent. I kept the closed vocabulary and documented the boundary.
5. Preparing Manas
I used Manas01 from the Manas-UdS corpus. Its metadata identifies the text as the 2010 edition of Manas in Sayakbai Karalaev's version. The complete extracted document contains 1,825,071 characters and 471,600 tokens under my BPE tokenizer.
The first plan was simple: take the first 10,000 tokens, reserve the final 10% for validation, and train on the first 90%.
That plan contained a hidden mistake. Manas01 represents a printed book, not a clean stream of the recited epic. It starts with a title page, publication data, forewords, and scholarly essays. The epic begins 30,877 characters later at this heading:
Манастын туула элегиндеги бабалары
I found the problem after training. The generated sample contained words about archaeology, scholarship, plots, and editorial work. Those words came from the source. The model had learned exactly what I gave it.
The final preparation command searches for the unique first epic heading and fails if it finds zero or multiple matches. It excludes everything before the heading, then performs a full BPE round-trip check. Decoding every encoded token must reproduce the selected text exactly.
After removing the leading material, the epic portion contains 465,069 BPE tokens. The pilot takes the first 10,000 of them. The first 9,000 form the training sequence; the last 1,000 form a chronological validation sequence.
I chose a chronological split because random overlapping windows can place almost identical neighboring text in both sets. That leakage would give the model an easier validation task and a misleading score.
6. Understanding B, T, and C
Three letters appear throughout small language-model code:
Bis batch size;Tis the number of token positions in each sampled sequence;Cis the number of possible token classes.
The final experiment used B = 16, T = 32, and C = 2525.
The input x has shape (B, T), which becomes (16, 32). It contains integer token IDs. It does not contain a C dimension.
The target y also has shape (16, 32). It contains the same sampled text shifted by one position. Every item in x points at the item in y that followed it in Manas.
The model creates the third dimension. Looking up every current ID returns one row of C logits:
x current IDs (16, 32)
y correct next IDs (16, 32)
logits scores for all IDs (16, 32, 2525)
T = 32 does not give this model 32-token memory. The 32 positions only let one optimization step process more independent pairs. Each prediction still uses one current token.
7. One complete forward pass
The forward pass moves from integer IDs to one scalar loss.
16 × 326.38M weights16 × 32 × 2525one numberThe loss is one number, but the computational graph keeps the contribution from all 512 positions.
First, the embedding lookup replaces each ID with its row of next-token scores. The code produces (16, 32, 2525) logits.
Cross-entropy expects a list of predictions and one correct class for each prediction. The code reshapes the tensors for that calculation:
logits: (16, 32, 2525) -> (512, 2525)
targets: (16, 32) -> (512)
No information disappears. The operation only places all batch positions into one long list.
8. How cross-entropy becomes one number
Consider one position where the correct next token has probability 0.9. Its negative log probability is small:
-log(0.9) ≈ 0.105
If the model assigns the correct token probability 0.01, the penalty grows:
-log(0.01) ≈ 4.605
Cross-entropy performs this calculation for all 512 positions and averages the results. PyTorch accepts raw logits and combines log-softmax with negative log likelihood in a stable implementation. I do not call softmax manually during training.
The whole batch loss can be written as:
loss = mean(-log(probability assigned to each correct next token))
At initialization, every table value equals zero. Every next token therefore starts with the same probability, 1 / C. The expected loss is log(C):
log(2525) ≈ 7.834
The measured initial validation loss was 7.833996, matching that prediction.
9. How backward works from one scalar
The scalar loss does not replace the 512 separate mistakes. It only summarizes them at the end of the forward graph.
Each loss contribution still points backward to the logit that produced it. Each logit points to a table entry. Each table entry belongs to the row selected by a current token ID.
For one prediction, the gradient of cross-entropy with respect to a logit has a simple form:
gradient = predicted_probability - correct_answer_indicator
For the correct next token, the indicator equals one. Its gradient is usually negative, so subtracting the gradient raises the correct score. For every incorrect token, the indicator equals zero. A predicted alternative receives a positive gradient proportional to its probability, so the update lowers its score.
PyTorch adds contributions when the same row appears several times in a batch. Rows that never appear in the batch receive no evidence from that step.
The optimizer then applies plain stochastic gradient descent:
optimizer.zero_grad(set_to_none=True)
loss.backward()
optimizer.step()
Conceptually, the last line performs:
weight = weight - learning_rate × gradient
I used SGD instead of Adam because the direct update matches the idea I wanted to inspect and avoids additional optimizer state. Zero initialization also works here. Hidden neurons would suffer from symmetry if all their weights started equal, but each table row represents a different current token and receives different training pairs.
10. Training and stopping
I used a batch size of 16 and sequence length of 32, so each step sampled 512 token transitions. The learning rate was 5.0. The run used seed 1337 for reproducibility.
Every 250 steps, the program stopped sampling and evaluated every available transition. Training evaluation covered 8,999 pairs, and validation evaluation covered 999 pairs.
The program saved a checkpoint whenever validation cross-entropy improved by at least 0.0001. It would stop after eight evaluations without improvement. A hard limit of 10,000 steps bounded the experiment.
The best validation result occurred at the final step, and the loss was still falling slowly. This means the model did not reach a demonstrated optimum. The 10,000-step limit served as a compute boundary. The experiment had already answered its main question, so I did not keep extending the run for a slightly lower number.
All tensor computation ran through PyTorch on Apple's mps device. The code checks that PyTorch includes MPS support and that the current machine can use it. If either check fails, training stops. It does not silently continue on the CPU.
11. The iterations that changed the experiment
The final run makes more sense when placed after the earlier attempts.
11.1 First run: prove the path works
The first version used the first 10,000 tokens of the full printed document. Its local vocabulary contained 3,755 IDs, so the table had 14.1 million parameters.
I set a 2,000-step limit. The run finished in 9.26 seconds. Validation cross-entropy fell from 8.231 to 7.558, which proved that MPS training, backpropagation, checkpointing, and evaluation worked.
The best result occurred at step 2,000, and the curve was still falling at every evaluation. I had set the ceiling too early.
11.2 Second run: extend the horizon
I raised the limit to 10,000 steps and reduced evaluation frequency. I kept the same seed, data, batch shape, initialization, optimizer, and learning rate.
The longer run reached validation cross-entropy 7.100. It still improved at the final step. The generated sample also revealed the corpus issue: it mixed epic names with academic terms.
11.3 Third attempt: correct the source and fail early
I removed the leading paratext and prepared the first 10,000 tokens of the epic itself. The new excerpt used only 2,525 distinct token IDs, cutting the table from 14.1 million to 6.38 million parameters.
Training completed, but generation failed before the program wrote its final report. The configured prompt Манас used an original BPE token ID that did not occur in the closed pilot vocabulary. The visible word appears inside longer forms, but BPE can assign a different token when capitalization, spacing, or an ending changes.
I kept the failure explicit. The program now validates a prompt before spending training compute. I changed the prompt to Манастын, whose tokenization occurs in the selected passage, and repeated the seeded run once.
These attempts were not hyperparameter search. Each one corrected a concrete boundary exposed by the previous run: insufficient steps, incorrect source region, and an invalid closed-vocabulary prompt.
12. Final results
The epic-only model had 6,375,625 float32 parameters, about 24.3 MiB of weights. It completed 10,000 steps in 15.84 seconds on an Apple M5. The run processed roughly 323,000 sampled transitions per second.
Validation cross-entropy improved by 1.454 nats from the uniform baseline. Perplexity fell from 2,525 to 590. In practical language, the trained distribution became about 4.28 times less confused than a distribution that treated every local token equally.
Top-1 validation accuracy reached 17.1%. The correct next token appeared among the five highest-scoring choices 32.5% of the time.
Training cross-entropy fell much further, from 7.834 to 2.646. The large gap between training and validation has a direct cause. Only 36.34% of validation token pairs appeared anywhere in training. The model had seen the current token at 84.28% of validation positions, but it often had not seen that exact next-token combination.
An exact table cannot share knowledge between similar rows. A larger neural model can learn that related words or endings deserve related representations. This bigram treats every row independently.
13. Generation after training
Generation starts from a prompt and repeats four operations:
- take the last token ID;
- read its row of logits;
- apply softmax and sample one next token;
- append the sampled ID and use it for the next step.
The saved run used prompt Манастын, temperature 0.8, and 160 new tokens. It began:
Манастын туй Түркүгүн жүрүп өлсөм - ал да жок, Кытайга салгануттан Атасы сазкандай кылганы бил Катыныныдеп жүгүрүп Алөөкөнүн алдына Адам билбес жообу бар...
The output contains clear traces of its source. It produces names such as Алөөкө, repeated references to Кытай, action fragments, dialogue punctuation, and formulaic endings. It also joins incompatible fragments, breaks words, and changes direction without warning.
The result matches the model's information boundary. A sentence can begin with a locally plausible pair, then lose its subject as soon as the next sampled token changes the row. More training cannot create memory that the architecture does not contain.
14. What the loss proves
The lower validation loss supports a narrow claim: the model learned useful statistics about which BPE token tends to follow another in the selected passage.
It does not prove that the model understands Kyrgyz or Manas. It does not measure factual recall, story consistency, grammar across a sentence, or performance on another domain. The validation tail contains only 999 transitions from one nearby section of the same epic.
The experiment also does not compare the custom tokenizer against another tokenizer inside the same model. That would require training matched models with controlled compute. This run used the Kyrgyz tokenizer as a known input so I could focus on the training mechanism.
15. From bigram to a transformer
The bigram table already performs next-token prediction, the same basic task used to pretrain autoregressive LLMs. The difference lies in the information available to each prediction and the way parameters share patterns.
A transformer replaces the isolated row lookup with several learned stages. Token embeddings map IDs into a smaller semantic space. Position embeddings tell the model where tokens occur. Attention lets each position read selected information from earlier positions. Feed-forward layers transform that combined representation, and an output layer produces logits over the vocabulary.
The loss can remain cross-entropy. Autograd still follows the graph backward, and an optimizer still updates parameters against their gradients. The computational graph becomes deeper, while the central loop stays recognizable:
tokens -> logits -> cross-entropy -> backward -> update
This is why the bigram experiment helped. It isolated the loop before attention and hidden representations added more moving parts.
16. What I learned
I started with a fuzzy picture of a batch as a three-dimensional input. Building the model made the boundary precise. The input is a matrix of integer IDs. The vocabulary dimension appears only after the model turns each ID into a vector of logits.
I also understood why one average loss remains useful for backward. The scalar sits at the end of a graph that still contains every prediction. Averaging changes the scale of the gradients; it does not erase their paths.
The data correction mattered as much as the network code. The first model sounded scholarly because I had trained it on a scholarly introduction. Looking at generated output helped me inspect the corpus rather than blame the optimizer.
Finally, the train-validation gap showed the weakness of an exact lookup table. The model memorized seen transitions well and had no machinery for relating an unseen pair to similar ones. That limitation creates the need for learned embeddings and context.
17. Conclusion
This project completed one small but real language-model experiment. A pinned source produced an epic-only passage. My Kyrgyz BPE tokenizer converted it into IDs. A PyTorch Embedding(C, C) table produced next-token logits. Cross-entropy measured 512 predictions at a time, autograd computed the gradients, and SGD changed 6.38 million weights on Apple MPS.
The validation result improved from 7.834 to 6.380, and the generator acquired local patterns from Manas. Its broken narrative showed exactly where one-token context ends.
The implementation, decision record, failed attempts, and machine-readable report are available in the project repository. The Kyrgyz tokenizer repository contains the earlier corpus and BPE experiment that supplied the tokenization.