Custom Transformer Decoder

This project is a decoder-only transformer implemented without a machine learning framework. Every component is written directly in NumPy: the query, key, and value projections, scaled dot-product attention with causal masking, layer normalization, residual connections, and the final projection to vocabulary logits. The objective was not to produce a competitive model but to work through the mechanics at a level where no detail is hidden behind a library.

The implementation is presented as a Jupyter notebook with Matplotlib for visualization. The tokenizer operates at the character level, which keeps the vocabulary small and the generated output legible, and an interactive demo allows the model to be prompted and observed token by token.

The KV cache

The component that warrants particular attention is the KV cache, because it is the point at which the implementation stops being a textbook exercise. Naive autoregressive generation recomputes attention over all previous tokens for each new one, so the cost per token grows linearly and the total cost grows quadratically. The cache stores the key and value tensors produced at each step: the first token runs a full forward pass to populate it, and every subsequent token computes only its own key and value, appends them to the cache, and attends over the accumulated entries.

The demo prints the time per generated token, which turns the speedup into a measurable quantity rather than a theoretical one. Observing per-token latency remain flat as the context grows is a direct confirmation that the cache behaves as the analysis predicts.

Discussion

Implementing attention by hand changes one's understanding of it. Tensor shapes and the order of operations, which are easy to treat as abstractions inside a framework, constitute the entire problem here. The notebook also provides optional attention heatmaps and a plot of cache growth per step; these were added because visualizing intermediate state proved to be the most effective debugging instrument available.