Training a Small GPT on Indian Folk Tales: Tokenizer, Mask and Loss

A small GPT trained on Indian folk tales

Anyone who has built a Vision Transformer already knows roughly 70 percent of a language model. Embeddings, attention, the residual stack and the training loop are the same machinery in both. What's left is the delta, and this walk-through covers it by building a small GPT from first principles and training it from scratch on one laptop GPU. The corpus is Indian folk tales: Jataka stories, Panchatantra-style fables and tales from Punjab and Bengal.

A genuine snag appears early. GPT-2's stock tokenizer chops words such as Brahman and Rajah into fragments because it never encountered them, so a tokenizer is trained from the corpus instead. At the end sits a 9-million-parameter model that has moved from random noise to trained weights, and anyone can reproduce each number by running the roughly ten-minute training job.

Why vision came first

Looking at pictures before text may feel like the wrong order, as chatbots work with words. Yet a transformer is indifferent to whether its input is photo patches or sentence pieces. The attention heads, LayerNorm, skip connections, MLP and optimization loop don't change. The better question is where that common engine is simplest to learn, and images win for four reasons.

• A token is simply an image patch, a fixed grid cut from the picture, so there's no vocabulary, no tokenizer and no unknown-word problem to solve first.
• There's no causal mask. A ViT sees the entire image at once, so plain attention comes before language's extra rule against peeking at the future.
• The behavior is visible. Patches, attention maps and the regions the model focuses on can all be drawn, and pixels feel tangible compared with text.
• It trains from scratch on one GPU in minutes (on CIFAR-10), so each figure is a number you produced yourself and didn't take on faith.

The upshot is that the full transformer already exists once a ViT is built, just dressed for vision. A language model only adds how text becomes tokens (a lookup and not a computation), the causal mask and next-token prediction. The hard, abstract part, language itself, is much easier to reason about once the machine is familiar.

Training is one game: guess the next token

At training time, everything a language model does boils down to reading the tokens so far and predicting the one that follows.

The figure illustrates predicting a next token from the tokens before it. Feed it the opening "In a certain village there lived a poor Brahman", and out comes a probability for each candidate that could follow. A period, a comma, or a word like who or named would all be grammatical. Nobody has to annotate anything: the correct answer is whatever token the text contained next. That is why raw text alone can teach a language model, and it makes up the whole learning signal. Every other component exists to support this single guess, and the ViT already supplied most of them.

In code, the idea takes two lines: run the model, then normalize its scores into a distribution over the next token.

logits, _ = model(ids)                     # ids = the tokens so far, shape (1, T)
probs = F.softmax(logits[0, -1], dim=-1)   # one probability per vocabulary token

The rest covers how model is built and trained so that those probabilities turn out good.

Why the tokenizer matters

A ViT slices an image into patches, while a language model slices text into tokens. A patch is computed from pixel values, but a text token is fetched from a fixed vocabulary. Somebody has to build that vocabulary beforehand, and the builder determines which things the model finds easy.

The first attempt reused r50k_base, GPT-2's ready-made tokenizer, whose vocabulary was frozen from roughly 2019 Western web text. The next figure shows what it does to words from the tales.

Brahman splits into Brah and man, Rajah into Raj and ah, and Panchatantra shatters into five pieces. None of these is obscure. Rajah occurs 709 times in the corpus, yet the tokenizer never saw such words often enough in its own training data to give them a token each. Fragmentation lengthens sequences and forces the model to reassemble Brahman from Brah plus man, where man is also shared with woman and human. With only about a million tokens of training data, that's a real handicap.

The remedy is a tokenizer trained on the tales, using byte-pair encoding (BPE). GPT uses the same algorithm, only fitted to a better-matched corpus. The source material is twelve public-domain Indian folk tale collections from Project Gutenberg, including two Jataka books, two Indian fairy tale volumes, Old Deccan Days, Folk-Tales of Bengal, The Talking Thrush, Tales of the Punjab, Twenty-Two Goblins, two Santal collections and Tales of the Sun from Southern India, adding up to about 809,000 words.

The figure steps through the byte-pair merges used to train the tokenizer. BPE begins with the 256 raw byte values, so nothing can be unknown, then repeatedly fuses the commonest neighboring pair into a new token. After 7,935 merges over the tales, the vocabulary has 8,192 entries and includes whole units like Rajah, Brahman and jackal.

One design choice deserves a clear statement: GPT-2's tokenizer is replaced completely, not extended.

The figure displays the custom vocabulary of 8,192 tokens. Adding tokens onto r50k's existing 50,257 is a sensible route when a big pretrained model's English skills should be kept, as projects like OpenHathi do. Training from scratch on a small, specialized corpus is another situation, where r50k's Western-web vocabulary only takes space, and the tales touched just 37 percent of it. The fresh 8k vocabulary is smaller, 96.6 percent in use and tailored to the tales, so swapping is the clean option.

Training the tokenizer takes a handful of lines with the HuggingFace tokenizers library. The script reads every raw text book, strips the Project Gutenberg header and footer, and trains a BPE model.

import glob
from tokenizers import Tokenizer, models, trainers, pre_tokenizers
from llm_transformer.prep_tales import strip_gutenberg, RAW_DIR

# read every raw .txt book and strip its Project Gutenberg header/footer,
# leaving just the story text - one big string per book
tales_texts = [strip_gutenberg(open(p, encoding="utf-8", errors="ignore").read())
               for p in sorted(glob.glob(f"{RAW_DIR}/*.txt"))]   # 12 books, ~809k words
tok = Tokenizer(models.BPE())
tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)   # start from bytes
trainer = trainers.BpeTrainer(vocab_size=8192, special_tokens=["<|endoftext|>"],
                              initial_alphabet=pre_tokenizers.ByteLevel.alphabet())
tok.train_from_iterator(tales_texts, trainer)   # learn the 7,935 merges
tok.save("tales_bpe.json")

The output .json file is the finished tokenizer, a fixed list of merges. It isn't trained again. It's loaded, and used to convert text to ids and ids back to text. A single .encode() call links raw text to the neural network.

tok = Tokenizer.from_file("tales_bpe.json")   # reload the trained tokenizer
ids = tok.encode("Once upon a time").ids      # -> [412, 987, 15, 733]   (integers)
idx = torch.tensor(ids).unsqueeze(0)          # (1, 4) — a batch of one sequence

vocab_size = tok.get_vocab_size()             # 8192 - sets the NUMBER OF ROWS below

Nothing but those integers reaches the network, and the original words no longer exist for it. Each id is simply a position in a table, and vocab_size fixes how many entries the embedding layer reserves.

From ids to vectors

Once text is a list of ids, each id turns into a vector. This is the one place where an LLM is simpler than a ViT.

The figure connects a tokenizer index to an embedding vector. In a ViT, a small matrix multiplication computes each patch's token from pixels. In an LLM the token is a discrete id, so the vector is just row number id of a learned table, a plain lookup. Token 789 retrieves row 789, which is a 256-dimensional vector, and adding a positional embedding, as in the ViT, finishes the input stage.

In code, an nn.Embedding layer is the table, and calling it performs the lookup by indexing, not multiplication. Its .weight is a learnable 8192 by 256 tensor holding 2.1 million weights that backprop trains, and a given row only changes in steps where its token occurs.

self.token_embed = nn.Embedding(vocab_size, n_embed)   # the 8,192 x 256 table
# .weight IS a learnable (8192, 256) parameter tensor — 2.1M weights, trained by backprop;
# each id's row is nudged only on steps where that token appears in the batch
...
tok_emb = self.token_embed(idx)     # idx (B, T)  ->  vectors (B, T, 256), one row per id

Everything after this point, meaning attention, the residual stream and the MLP blocks, is identical to the Vision Transformer, with one exception.

The causal mask

A ViT allows each patch to attend to all the others. A language model can't do that, because its job is to guess what comes next. If position t saw position t+1, it would just copy the answer. The remedy is masking the attention so every token sees itself and earlier tokens only.

The mask in the figure is a lower triangle. Before softmax, all future positions get negative infinity, which softmax converts into a weight of exactly zero, so blocked tokens add nothing. This triangle is the only architectural gap between ViT attention and LLM attention. Dot products, scaling by the square root of the head size, softmax and the weighted sum of values all remain the same.

In code, it's one extra line on top of ordinary attention. A lower-triangular matrix of ones is made once, scores are computed as usual, and the zeroed positions are filled with negative infinity before the softmax.

# tril = torch.tril(torch.ones(T, T))  — a lower-triangular matrix of 1s, made once
attn = q @ k.transpose(-2, -1) / math.sqrt(head_size)          # scores, (B, T, T)
attn = attn.masked_fill(self.tril[:T, :T] == 0, float('-inf')) # blank out the future
attn = F.softmax(attn, dim=-1)                                  # -inf -> weight 0
out  = attn @ v

One pass trains every context length

The one-line mask conceals something that merits a closer look. People tend to imagine a language model working like it generates text: adding a token to a growing context and redoing attention each step. That's true at inference but not during training. In training, the sentence is fully known, so Q, K and V for all T tokens are computed together, forming the whole T by T attention matrix in a single step.

That single pass is only valid because of the causal mask. Consider row t of the masked matrix: all entries right of the diagonal are negative infinity, so they become zero after softmax. Computing out = attn @ V then gives row t a weighted sum over only the first t value vectors. Every context length is therefore handled at the same time.

The figure shows different runs and the shape of the mask in training. Read it row by row:

• Row 1 attends to token 1 only, so its output predicts token 2 from a one-token context.
• Row 3 attends to tokens 1 through 3 and predicts token 4 from three tokens of context.
• Row 6 attends to all six tokens and predicts token 7 from the full context.

Here's another view. Each row's visible context forms a lower-triangular matrix that grows by one each time, 1 by 1 for the first token, then 2 by 2 and 3 by 3. These are the nested top-left blocks of the large matrix, drawn in orange, so a single 6 by 6 pass holds all smaller runs, and the strip at the bottom unrolls them. They're all produced inside one attn @ V, and nothing loops over positions in Python. This explains transformers' training efficiency: a sequence with T tokens amounts to T training examples, with contexts running from length 1 to T, scored in parallel. The same triangular mask that supplies those examples also stops any position from looking ahead during generation, which proceeds row by row and keeps each step cheap by caching past keys and values (the KV-cache).

Real numbers through attention

So far each cell has only meant attend or don't attend. Opening up the multiply shows actual numbers for the same three runs, taken from the trained model's block 0, head 0 on the sentence "Once there lived a clever jackal".

The figure lays out the Q, K, V computations across the runs. Within a row, the attention weights are positive and add up to 1, so the output is a convex blend of the visible value vectors, which is a weighted average.

• Run 1 has one token in view with weight 1.00, so out1 = V1. The first position copies its own value, since there's nothing yet to mix.
• In run 3, the output mixes three values using weights of about 0.20 for the first token, 0.06 for the second and 0.74 for the third. It relies most on its own token, lived, but still takes in one fifth of Once, so it's a true mixture.
• In run 6, weight is spread over all six values, with 0.26 on each of there and jackal. This blended 32-dimensional vector summarizes the entire prefix for the head, and later layers convert it into a prediction.

That covers everything attention does: scores from queries and keys set the weights, and the weights mix the values. Masking guarantees that the mix never extends past the current token, so each row is a fair summary of everything read so far. Real heads carry 32 value dimensions, and the model has eight heads in each of six layers. The strips display the first 8 dimensions of a single head, though the same mechanism applies throughout.

T predictions, T losses

Next-token prediction brings a benefit the ViT lacked. The ViT made a single prediction per image, a class label. The causal mask lets a language model guess the next token at all positions at once, so a sequence gives as many learning signals as it has tokens.

The figure averages T predictions and loss values for a T-token input over every run. Each position receives a cross-entropy loss, which measures how surprised the model was by the real next token. Some are cheap, such as a after lived at only 0.58. Others are expensive, such as the adjective clever at 7.17, because plenty of words could have fit. Averaging the six gives 3.59, and backpropagation drives that one number down, using the same chain-rule machinery as the ViT.

In code, all T losses come from a single cross_entropy call. The (B, T) grid of positions is flattened into one long list, and every entry is scored against its true next token.

logits = self.lm_head(x)                       # (B, T, vocab)
B, T, C = logits.shape
loss = F.cross_entropy(logits.reshape(B*T, C), # every position ...
                       targets.reshape(B*T))   # ... vs the token that actually followed

Each loss is a vocabulary classification

It helps to be exact about what cross_entropy compares, because both prediction and target are vectors, not single words. At every position, the lm_head produces a logits vector with 8,192 entries, one score per vocabulary token, and softmax converts it into a probability distribution over all of them. The target is a one-hot vector of the same size, containing 1 at the word that really followed and 0 everywhere else. Predicting the next token is then a classification problem where the vocabulary words are the classes.

The figure compares logits in an LLM and in a ViT. It spells out one subtle point: the bars are ordered by probability just to be readable, while the actual output is a vector with fixed layout, where slot i always belongs to token id i. The word "son" stays at index 498 in every forward pass, whether training or inference, and softmax alters only the heights on the slots, never their order. This explains why the one-hot target's lone 1 is at index 498 rather than in a first position: the tokenizer assigned each token its identity once, as its slot.

That's exactly the classifier head of the ViT, which applies softmax to class scores and scores them with cross-entropy against a one-hot label. Two things differ: the class count (8,192 vocabulary tokens versus around 10 image labels) and the frequency (at all T positions versus once per image). Against a one-hot target, cross-entropy also reduces to minus the log of the probability assigned to the right word, which is why each position's loss reads like a surprise score. "son" at 9.8 percent costs 2.32, a word the model strongly expected would cost much less, and a real shock would cost much more.

How a batch is built

The reshape(B*T) quietly says a batch has B sequences, each contributing T position losses. So where do the B sequences come from, and are they all the same length?

People often stumble on one detail: pretraining sequences aren't padded and all have exactly the context length. The corpus is joined into one long token stream, with an <|endoftext|> token between documents, and windows of fixed length are cut at random positions. All windows come from one continuous stream, so none can be short, and no padding token is needed.

The figure follows the route from the training corpus to batches and on to the loss. The one special token in the stream is <|endoftext|> (EOT), which serves as a separator rather than padding and shows where one folk tale stops and the next starts. A window may cross such a border and contain an EOT mid-way, and that's how the model learns that a story has ended and a fresh one begins. Padding does show up later, during fine-tuning, when examples vary in length and padded positions are masked out of the loss, which is a different setting from packed pretraining.

The loss is the second area where the LLM departs from the ViT, and the difference is a welcome one. The ViT produced one prediction and one loss per image. A 128-token sequence here produces 128 losses, which average into that sample's loss (3.78 for sample 1 in the figure), and the 32 sample losses in a batch average into the batch loss of 3.92 that backprop minimizes. With these two rounds of averaging, a single step lowers the mean of 32 times 128, or 4,096, next-token losses, far more signal per step than a single loss per image.

Watching the loss fall

Wire that loss to an optimizer and repeat. This is the real training run on the tales corpus.

The plot charts training loss and held-out validation loss over epochs. Training begins at 9.2, which equals ln(8192), the loss of random guessing across an 8,192-token vocabulary, and the best held-out loss reaches 4.73. The curves remain close, about 1.0 apart, which means that with a narrow domain the model generalizes instead of just memorizing. The learner itself is small.

The figure ties the training flow to the count of learnable weights. With 9.0 million parameters, the model trains in roughly ten minutes on a 6 GB laptop GPU. A little more than half live in the six transformer blocks, and the remainder belong to the embedding table and output head. Training consists solely of adjusting these 9 million numbers repeatedly, so the next-token probabilities put more weight on the token that really followed.

In code, the loop is what you'd expect, and the same shape as a ViT's: sample random context and next-token windows, run the forward pass to get one mean loss, zero the gradients, backpropagate and step.

for step in range(train_steps):
    xb, yb = get_batch(train_data)     # random (context, next-token) windows
    _, loss = model(xb, yb)            # forward: T predictions -> one mean loss
    opt.zero_grad(); loss.backward()   # backprop the loss into all 9M parameters
    opt.step()                         # nudge them downhill

The full script adds a cosine learning-rate schedule and evaluation.

What's new and what's unchanged

Compared with the Vision Transformer, only a little is new in the LLM: a token embedding that works as a lookup table, a causal mask, and a next-token loss at each position that classifies over the full vocabulary. One practical lesson also stands out, which is that the tokenizer belongs to the model, so a domain-specific one repays the effort. Attention, positional embeddings, the residual stream, LayerNorm, the MLP blocks and the complete backprop-and-optimizer loop carry over without change.

Real numbers opened up the inner workings: how the causal mask lets a single forward pass score all context lengths together, what attn @ V produces (each output row blends the value vectors visible to that token), how training batches come from packing instead of padding, and how one position's loss, a classification across the vocabulary, joins the others in a single average.

The weights are done, but a trained model does nothing until it runs. Inference comes next: the generation loop, the KV-cache that keeps it quick, and what the small model writes when given the start of a tale.