← All work
Systems/Deep learning/2024/Built

Transformer from Scratch

Decoder-only language model with no high-level LLM libraries — BPE tokenizer, causal multi-head attention, scheduled training, then sampling with temperature, top-k and nucleus.

PyTorchBPEAttentionSampling
Personal research work
01

The problem

Using language models daily without being able to derive what happens between a token and a logit felt like a gap. Building a decoder-only transformer with nothing above raw tensor operations closed it.

02

Who uses it

2 groups

Me, as an applied-AI engineer

Reason about context length, attention cost and sampling behaviour from first principles.

The RAG and assistant work

Know why a temperature or top-p change moves output the way it does.

03

Workflow

6 stages
01Train BPE tokenizer
02Embed + positional encoding
03Causal multi-head attention
04Feed-forward + residual
05Train with schedule
06Sample
01 · TokenizeTriggerCorpus preparedSystemByte-pair encoding trained on the corpusOutputA vocabulary and merge rules; text becomes token IDs
02 · EmbedTriggerToken IDs batchedSystemToken embeddings plus positional encodingOutputA sequence of vectors carrying content and position
03 · AttendTriggerEach transformer blockSystemMulti-head self-attention with a causal maskOutputEach position mixed with earlier positions only
04 · TransformTriggerAfter attention in each blockSystemPosition-wise feed-forward, residual connections, layer normOutputA trainable deep stack that does not collapse
05 · TrainTriggerTraining loopSystemCross-entropy on next-token prediction, warmup and decay, gradient clippingOutputConverging loss with checkpoints
06 · GenerateTriggerInference from a promptSystemTemperature, top-k and nucleus sampling over logitsOutputGenerated text with controllable determinism and diversity
04

How it was built

5 parts

BPE tokenizer

Byte-pair encoding trained on the corpus — count adjacent pairs, merge the most frequent, repeat to a target vocabulary. Written directly, which makes vocabulary size a decision with visible consequences for sequence length rather than a config value.

Causal multi-head attention

Queries, keys and values projected per head, scaled dot-product scores, an upper-triangular mask so a position can never attend forward, softmax, then a weighted sum and output projection. The mask is what makes it a language model rather than an encoder.

Blocks that train

Attention and a position-wise feed-forward network per block, each wrapped in residual connections and layer norm. Without those the stack stops training past a handful of layers — the fastest way to learn why they are there is to leave them out once.

Training

Cross-entropy on next-token prediction, learning-rate warmup then decay, gradient clipping and checkpointing. Loss curves and held-out samples together — the loss can improve while the text does not.

Sampling

Greedy decoding, then temperature, top-k and nucleus sampling implemented over the logits. Seeing the same model swing between repetitive and incoherent as top-p moves makes those parameters stop being magic numbers.

05

What it is built with

Framework

  • PyTorch
  • Tensor ops only

Model

  • Decoder-only transformer
  • Multi-head attention
  • Causal mask
  • Layer norm

Training

  • Cross-entropy
  • LR warmup / decay
  • Gradient clipping
  • Checkpoints

Inference

  • Temperature
  • Top-k
  • Nucleus sampling
06

Outcome