Mini-GPT

Goal

End-to-end GPT-style language model: embeddings -> transformer blocks -> logits -> generation

Architecture

Token IDs
  -> Token Embedding + Position Embedding
  -> TransformerBlock x N
  -> Final LayerNorm
  -> Linear -> Vocab logits

Shapes: input (T), logits (T, V)

Next-token objective

Targets are input shifted by one. Loss is mean cross-entropy over positions.

Training loop

logits = model(x)
loss = cross_entropy_loss(logits, y)
zero_grads()
loss.backward()
update()

Generation

• Crop to max_seq_len
• Softmax last logits
• Sample with temperature

Note

Single-head attention in this pack for clarity.