A language model from scratch: tokens, attention, and training

A language model can seem mysterious when its first introduction is a chat interface. Strip that interface away and the training problem is concrete: given some tokens, predict the next one. Repeating this prediction turns a learned probability distribution into a sequence generator.

This guide builds a small decoder-only Transformer from basic PyTorch modules. “From scratch” means that the model’s parameters begin untrained and the architecture is written explicitly. PyTorch still provides tensors, automatic differentiation, and an optimizer. Training a useful large language model requires vastly more data, compute, and evaluation than this educational example.

Start with a probability problem

For a token sequence $x_1,\ldots,x_T$, the chain rule gives:

We train a neural network to approximate each conditional distribution. If the prefix is “the capital of,” the output is not a single word. It is a score for every possible next token. A softmax converts those scores, or logits, into probabilities:

$V$ is the vocabulary size. If the target token receives probability $0.5$, its negative log-likelihood is about $0.693$ nats. If it receives $0.01$, the loss is about $4.605$. The training signal penalizes confident mistakes much more heavily than uncertain correct predictions.

The average next-token loss is:

It is not a direct score for truthfulness or conversational helpfulness. A model can learn to imitate text patterns while remaining unreliable at factual questions.

Tokens, data, and a fair validation split

A tokenizer maps text to integer IDs. Character vocabularies are intuitive but must handle unseen characters. Subword tokenizers trade vocabulary size against sequence length. For a minimal example, UTF-8 bytes provide a fixed vocabulary of 256 values and can represent any UTF-8 text without an unknown-token rule.

The trade-off is efficiency. A Nepali character may occupy several UTF-8 bytes, so the same visible passage can consume a different number of tokens than an English passage. Token-level loss is therefore not directly comparable across unrelated tokenizers or corpora. For multilingual work, inspect bytes per token, sequence lengths, and language-specific evaluation rather than assuming one context length means equal linguistic coverage.

Use text that you have permission to train on. Remove accidental secrets and obvious duplicates. For a real experiment, split by document, source, and time where relevant before constructing overlapping windows. Otherwise near-duplicate material can appear in training and validation and produce an optimistic estimate.

Our script uses a contiguous split of one local file to keep the example small. It is an educational validation split, not a complete contamination-control procedure. Supply a sufficiently long input.txt and do not interpret results from this file as a general benchmark.

Embeddings give tokens a working representation

An embedding table $E\in\mathbb R^{V\times d}$ turns token ID $i$ into row $E_i$. These vectors are learned through prediction; they are not fixed dictionary definitions. Tokens that help with similar prediction tasks may develop related representations.

Attention alone does not encode the order of a set of input vectors. Our small model therefore adds a learned position vector:

With batch size $B$, context length $T$, and embedding width $d$, the resulting tensor has shape $(B,T,d)$. Keep these dimensions visible while implementing the model. Many attention bugs are simply an unnoticed exchange of the head, time, and feature axes.

Attention is a learned weighted lookup

For one attention head, project the current representations into queries, keys, and values:

A query describes what a position is looking for. Keys describe what other positions offer. Values hold the information to combine. A dot product scores compatibility, and a softmax turns scores into weights:

Each row of $A$ sums to one. If one query gives two available keys scores of $2$ and $0$, their weights are approximately $0.881$ and $0.119$. The output is mostly the first value, with a smaller contribution from the second. This is a differentiable retrieval operation inside the network.

The scaling factor keeps dot products from growing simply because the key dimension is larger. The mask $M$ is zero for permitted positions and negative infinity for future positions. After softmax, those future positions receive zero weight. Without this causal mask, the model can see the answer during training.

The Transformer paper introduced scaled dot-product attention and multi-head attention. Our decoder-only, pre-normalized implementation is a deliberately simplified teaching architecture, not a line-by-line reproduction of the original encoder-decoder model.

Build the block around attention

Multiple heads let the model compute several weighted lookups in parallel. They split the representation into smaller subspaces, concatenate the results, and project them back to width $d$. More heads do not guarantee distinct human-interpretable roles.

Attention mixes information across positions. A feed-forward network then transforms the representation at each position independently. Residual connections carry an existing representation around each transformation, while layer normalization stabilizes its scale:

flowchart TD
  H["H: shape B x T x d"] --> LN1["LayerNorm"]
  LN1 --> Attn["Masked multi-head attention"]
  Attn --> Add1(("+"))
  H --> Add1
  Add1 --> LN2["LayerNorm"]
  LN2 --> MLP["Position-wise MLP"]
  MLP --> Add2(("+"))
  Add1 --> Add2
  Add2 --> Out["H'': same shape, passed to the next block"]

A stack of these blocks repeatedly exchanges and transforms information. A final linear projection maps each position to $V$ logits. At no point does the model need an explicit grammar rule or a table of facts. Whether it learns useful structure depends on the data, architecture, optimization, and scale.

A complete small training script

The script below is intentionally modest: two blocks, four heads, learned positions, and a byte vocabulary. Install PyTorch for your platform separately, save this as tiny_lm.py, and put a UTF-8 text file named input.txt beside it. It needs no downloaded model weights.

from pathlib import Path
import math
import torch
from torch import nn
from torch.nn import functional as F

torch.manual_seed(7)
device = "cuda" if torch.cuda.is_available() else "cpu"
context, batch_size = 64, 16
width, heads, layers = 128, 4, 2

raw = Path("input.txt").read_bytes()
data = torch.tensor(list(raw), dtype=torch.long)
split = int(0.9 * len(data))
train, valid = data[:split], data[split:]
if min(len(train), len(valid)) <= context:
    raise ValueError("Use a longer text file for both splits")


def get_batch(stream):
    starts = torch.randint(len(stream) - context, (batch_size,))
    x = torch.stack([stream[i:i + context] for i in starts])
    y = torch.stack([stream[i + 1:i + context + 1] for i in starts])
    return x.to(device), y.to(device)


class CausalAttention(nn.Module):
    def __init__(self):
        super().__init__()
        assert width % heads == 0
        self.qkv = nn.Linear(width, 3 * width, bias=False)
        self.proj = nn.Linear(width, width, bias=False)
        self.register_buffer(
            "mask", torch.tril(torch.ones(context, context, dtype=torch.bool))
        )

    def forward(self, x):
        b, t, c = x.shape
        q, k, v = self.qkv(x).chunk(3, dim=-1)
        def split_heads(a):
            return a.reshape(b, t, heads, c // heads).transpose(1, 2)
        q, k, v = map(split_heads, (q, k, v))
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(c // heads)
        scores = scores.masked_fill(~self.mask[:t, :t], float("-inf"))
        out = scores.softmax(dim=-1) @ v
        out = out.transpose(1, 2).contiguous().reshape(b, t, c)
        return self.proj(out)


class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.ln1, self.ln2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.attn = CausalAttention()
        self.mlp = nn.Sequential(
            nn.Linear(width, 4 * width), nn.GELU(),
            nn.Linear(4 * width, width),
        )

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))


class TinyLM(nn.Module):
    def __init__(self):
        super().__init__()
        self.tokens = nn.Embedding(256, width)
        self.positions = nn.Embedding(context, width)
        self.blocks = nn.Sequential(*[Block() for _ in range(layers)])
        self.norm = nn.LayerNorm(width)
        self.head = nn.Linear(width, 256, bias=False)

    def forward(self, ids):
        t = ids.shape[1]
        if t > context:
            raise ValueError("Sequence exceeds the context window")
        pos = torch.arange(t, device=ids.device)
        x = self.tokens(ids) + self.positions(pos)
        return self.head(self.norm(self.blocks(x)))


model = TinyLM().to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)


@torch.no_grad()
def estimate_loss(stream, batches=10):
    was_training = model.training
    model.eval()
    losses = []
    for _ in range(batches):
        x, y = get_batch(stream)
        logits = model(x)
        losses.append(F.cross_entropy(
            logits.reshape(-1, 256), y.reshape(-1)
        ).item())
    model.train(was_training)
    return sum(losses) / len(losses)


for step in range(500):
    x, y = get_batch(train)
    loss = F.cross_entropy(model(x).reshape(-1, 256), y.reshape(-1))
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    if step % 100 == 0:
        print(step, "train", estimate_loss(train), "valid", estimate_loss(valid))


@torch.no_grad()
def generate(prefix="\n", count=200, temperature=0.8):
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    model.eval()
    encoded = list(prefix.encode("utf-8")) or [10]
    ids = torch.tensor([encoded], dtype=torch.long, device=device)
    for _ in range(count):
        logits = model(ids[:, -context:])[:, -1, :] / temperature
        next_id = torch.multinomial(logits.softmax(dim=-1), num_samples=1)
        ids = torch.cat((ids, next_id), dim=1)
    return bytes(ids[0].tolist()).decode("utf-8", errors="replace")

print(generate())
torch.save(model.state_dict(), "tiny_lm.pt")

Inputs and targets differ by one token. A window containing IDs [10, 20, 30] has targets [20, 30, 40]. This shifting lets the model learn at every position simultaneously. Flattening the batch and time axes gives cross-entropy one row of vocabulary logits per target token. Pass logits directly; do not apply softmax before PyTorch’s cross-entropy loss.

Sampling is a separate operation from training. Temperature divides the logits before softmax. Lower positive temperatures concentrate probability on the highest-scoring choices; higher values spread it out. Neither setting makes a weak model knowledgeable. Byte sampling may generate invalid UTF-8 combinations, so decoding uses replacement characters rather than crashing.

Debug the learning process before making it larger

First, verify tensor shapes and target alignment. The output should be $(B,T,256)$ and the targets $(B,T)$. A causal test should change future input tokens and confirm that earlier output logits remain unchanged, up to numerical tolerance. Disable stochastic layers while running that test.

Next, try to overfit a very small batch. If the network cannot substantially reduce loss on a few examples, investigate gradients, masking, target construction, and the optimizer before adding data. Then restore a proper validation split and look for a gap between training and held-out performance.

A roughly uniform byte predictor has loss $\log(256)\approx5.545$ nats. This is a reference calculation, not an exact expected initialization loss: random logits are not perfectly uniform. Perplexity is $\exp(\mathcal L)$ when the loss uses natural logarithms, but it should be compared only under compatible tokenization and evaluation conditions.

Always read generated samples alongside losses. A falling loss can coexist with memorization, repetition, or failure to preserve a language’s spelling. Evaluate held-out documents, repetitive loops, sensitive-data memorization, and performance by language and source.

What changes when the model becomes large?

Naive attention materializes an attention matrix with $T^2$ entries per head. Doubling the context length quadruples that matrix’s size. This does not mean total training cost always grows by exactly four: projections, feed-forward layers, memory movement, and kernels also matter. PyTorch’s scaled dot-product attention provides optimized implementations suitable for a later version of this exercise.

Larger training runs add distributed optimization, mixed precision, checkpoint recovery, data-mixture management, throughput measurement, and careful compute budgeting. Empirical scaling-law research investigates relationships among model size, data, compute, and loss; it does not promise that making any one dimension larger will fix a poorly designed experiment.

Pretraining is also not the same as building an assistant. Instruction tuning, alignment methods, retrieval, tool integration, evaluation, and abuse prevention address different requirements. A small model that continues a text file is a successful demonstration of autoregressive learning. It is not a reliable general-purpose assistant.

Practice: make one change and explain its effect

Measure held-out loss, tokens processed per second, and parameter count for the small baseline. Change only the context length, then only the number of blocks. Keep the training-token budget comparable when comparing losses. Run several seeds before attributing a small difference to architecture.

For Nepali text, compare the byte baseline with a tokenizer trained on an authorized Nepali corpus. Report both token-level loss and a tokenizer-independent measure such as negative log-likelihood per byte on the same held-out text. Inspect generated text with a fluent reader. The exercise becomes much more informative when “the loss improved” is followed by an explanation of exactly what was measured.