Biki Kalita

Posted on Jul 23

Building Your Own Deep Learning Framework: A 3–5 Month Journey From Zero to Tiny Transformer

#python #machinelearning #deeplearning #tutorial

Prologue: The Machine That Learns to Learn

Somewhere in a research lab in the early 2010s, a small group of engineers sat staring at a whiteboard covered in arrows. They weren't drawing a neural network. They were drawing a graph of computation — a map of every multiplication, every addition, every function that data would pass through on its way to becoming a prediction. That whiteboard sketch, refined over years, eventually became the beating heart of PyTorch: a system that doesn't just compute forward, but remembers how it got there so it can compute backward — and learn.

Here's the secret almost nobody tells beginners: PyTorch is not magic. It's not some impenetrable fortress of genius that only PhDs can understand. At its core, it's a few ideas, executed carefully:

Track every operation performed on a number.
Remember how to undo that operation mathematically (its derivative).
Walk backward through that history, multiplying derivatives together (the chain rule).
Use the result to nudge numbers in a direction that reduces error.

That's it. That's the whole trick. Everything else — tensors, GPUs, layers, optimizers, transformers — is engineering built on top of that one idea.

Over the next 3 to 5 months, you are going to rebuild that idea yourself, from the ground up, in Python. Not because the world needs another deep learning framework. But because there is no faster way to actually understand deep learning than to build the thing that makes deep learning possible.

This is not a copy-paste tutorial. I will not hand you a finished library. What I will hand you is a map, a set of checkpoints, the theory you need at each stage, the mistakes you're likely to make, and the tests that will tell you whether you're right. The actual code — the thinking, the debugging, the "aha" moments — that has to be yours. That's where the learning lives.

Let's begin.

Part 1: Why Build This At All?

There are three ways to learn deep learning:

Use a framework. You import torch, call .backward(), and a neural network trains. You learn what deep learning does, but the internals stay a black box.
Read papers and textbooks. You learn the math in isolation, but it often stays abstract — symbols on a page that never quite connect to working code.
Build the framework yourself. You are forced to confront every hidden assumption, every edge case, every "wait, why does this even work?" moment, because nothing works until you make it work.

The third path is slower. It is also the only one that produces the kind of deep, load-bearing understanding that shows up in job interviews, research intuition, and your ability to debug a model that refuses to train. When you've implemented backpropagation by hand, you will never again see a loss.backward() call as a magic incantation. You'll see it as: "gradients flowing backward through a graph I understand, because I've built one."

This project is also a phenomenal way to practice a skill that's rarely taught explicitly: reading the world's collective knowledge and turning it into your own working system. You won't invent autodiff from nothing. You'll read blog posts, skim papers, study existing tiny frameworks like micrograd and tinygrad for inspiration (not for copying), and synthesize that knowledge into code that is unmistakably yours. That synthesis skill — reading, understanding, rebuilding — is exactly what separates engineers who can only follow tutorials from engineers who can solve novel problems.

By the end, you will have a working library that can:

Automatically differentiate arbitrary chains of operations (autodiff)
Represent and manipulate multi-dimensional arrays with broadcasting (a tensor engine)
Compose neural network layers (Linear, activations, normalization)
Compute losses and optimize weights (SGD, Adam)
Train a small transformer to generate text, character by character

And more importantly, you will understand why every piece exists.

Part 2: The Big Picture — What You're Actually Building

Before writing a single line of code, hold the whole system in your head, even blurrily. Deep learning frameworks are built in layers, like geological strata. Each layer depends on the one below it:

Layer 5:  Transformer blocks (attention, feedforward, positional encoding)
Layer 4:  Training loop (forward → loss → backward → optimizer step)
Layer 3:  Losses & Optimizers (MSE, cross-entropy, SGD, Adam)
Layer 2:  Neural network layers (Linear, ReLU, Softmax, LayerNorm)
Layer 1:  Tensor engine (N-dimensional arrays, broadcasting, matmul)
Layer 0:  Autodiff engine (the core: track operations, compute gradients)

Notice the direction of dependency: a transformer is just layers, arranged cleverly. Layers are just tensor operations, arranged into reusable objects. Tensor operations are just autodiff-tracked math. Everything above Layer 0 is organization. Layer 0 is the one true foundation. This is why you'll spend real, unhurried time there before moving on — everything else stands or falls on whether your autodiff engine is correct.

By the end of this project, you should be able to say, in plain language and without hand-waving:

"I built a system where every math operation remembers its inputs and knows its own derivative. When I ask for gradients, the system walks backward through the history of operations it recorded, applying the chain rule at each step, until every parameter knows exactly how much it contributed to the final error — and therefore how it should change."

If you can say that sentence and mean it, mission accomplished. Everything else is detail.

Part 3: The Roadmap — Nine Phases Over 3 to 5 Months

Here is the full journey. Don't worry about memorizing it — we'll walk through each phase in depth below. Treat this as the trail map you'll keep returning to.

Phase	Focus	Approx. Duration
1	Python + NumPy foundations, mathematical warm-up	1–2 weeks
2	Scalar autodiff engine (the heart of the project)	2–3 weeks
3	Tensor engine with broadcasting & matmul	3–4 weeks
4	Neural network layers (Linear, activations, init)	2 weeks
5	Losses and optimizers (SGD, momentum, Adam)	1–2 weeks
6	Training loop + first real experiments (XOR, MNIST-lite)	2 weeks
7	Transformer components (attention, embeddings, LayerNorm)	3–4 weeks
8	Mini language model demo (character-level generation)	2–3 weeks
9	Cleanup, tests, documentation, packaging	1–2 weeks

That totals roughly 17–24 weeks — comfortably inside your 3–5 month window, with room for the inevitable week where nothing works and you have to slow down. Budget for that week. It will happen, probably more than once, and it is not a sign of failure. It's a sign you've hit a real idea.

A sane weekly rhythm: 3–5 focused sessions per week, 1–3 hours each, mixing reading/study with hands-on building. Deep learning concepts reward spaced repetition — you'll understand backprop better the third time you re-derive it than the first.

Phase 1: Python + NumPy Foundations

What you're learning

Comfort with NumPy's array operations, broadcasting rules, and vectorized thinking — because your tensor engine will essentially be a thin, gradient-aware wrapper around NumPy.

Why it matters

If you don't deeply understand how NumPy broadcasts shapes like (3, 1) and (1, 4) into (3, 4), you will not understand why your tensor engine's gradients sometimes come out the wrong shape later. This bug bites everyone, and the fix is always the same: understand broadcasting now, thoroughly, so you recognize the symptom later.

What to study online

NumPy's official broadcasting documentation (read it twice — once now, once after Phase 3)
Any solid "vectorization vs. loops" tutorial that shows the same computation written both ways
A refresher on matrix multiplication, dot products, and the chain rule from calculus (Khan Academy or 3Blue1Brown's "Essence of Calculus" series are excellent for building visual intuition)

What to build yourself

A handful of small NumPy exercises: matrix multiply without np.dot, broadcast a (3,1) array against a (1,4) array and predict the output shape before running it, implement softmax and sigmoid by hand.
A "shape tracer" habit: for every array operation you write, say out loud what shape goes in and what shape comes out.

Tiny illustrative idea

a = np.array([[1], [2], [3]])   # shape (3, 1)
b = np.array([10, 20, 30, 40])  # shape (4,)
# predict the output shape BEFORE running a + b

How to test yourself

Before running any broadcasting expression, write down the predicted output shape on paper. Then check. If you're wrong more than once, stop and reread the broadcasting rules until you can predict shapes reliably.

Common mistakes

Assuming broadcasting "just works" without understanding the alignment rule (dimensions are compared from the right, and must either match or be 1).
Confusing element-wise multiplication (*) with matrix multiplication (@ or np.matmul).

You're ready to move on when…

You can look at two array shapes and confidently predict whether they'll broadcast together, and what shape the result will be — without running code to check.

Mini exercises

Implement softmax for a batch of vectors (shape (batch, classes)) using only NumPy, handling numerical stability (subtract the max before exponentiating).
Write a function that multiplies two matrices using nested Python loops, then verify it matches A @ B.

Phase 2: Scalar Autodiff — The Heart of Everything

This is the phase that matters most. Take it slow. Everything you build later is a generalization of what you build here.

What you're learning

How to build a system where every number "remembers" the operations that created it, so that gradients can be computed automatically via the chain rule. This is the single idea that makes deep learning trainable.

Why it matters

Backpropagation is often taught as an abstract algorithm on paper. Building it yourself turns it into something concrete: a graph of tiny objects, each holding a value, a gradient, and a memory of its parents. Once you've built this for scalars, tensors are "just" the same idea applied to arrays.

What to study online

Andrej Karpathy's "The spelled-out intro to neural networks and backpropagation" video and the micrograd repository. Watch it, understand the ideas, then close the tab and build your own version from memory and understanding — don't transcribe his code.
A clear explanation of computational graphs (search "computational graph backpropagation explained") — focus on the idea of a DAG (directed acyclic graph) and topological sort, because that's how you'll decide the correct order to walk backward through the graph.
The chain rule, revisited specifically in the context of composed functions: if z = f(y) and y = g(x), then dz/dx = dz/dy * dy/dx.

What to build yourself

Design a Scalar (or Value) class that wraps a single float and:

Stores its numeric value
Stores a gradient (grad), initialized to zero
Stores references to the "parent" scalars that created it (if any)
Stores a function that knows how to propagate gradient backward to those parents
Overloads Python operators (__add__, __mul__, __pow__, and later things like tanh or relu) so that writing normal-looking math (c = a * b + a) automatically builds a graph behind the scenes
Implements a .backward() method that performs a topological sort of the graph, then walks it in reverse, accumulating gradients using the chain rule

Tiny illustrative idea

Don't copy this — but here's the shape of the thinking, for a single operation:

def __mul__(self, other):
    out = Scalar(self.value * other.value, parents=(self, other))
    def _backward():
        self.grad += other.value * out.grad
        other.grad += self.value * out.grad
    out._backward_fn = _backward
    return out

Notice the pattern: every operation defines its own tiny local rule for how gradient flows backward to its inputs. That's the whole idea, repeated for every operation you support.

How to test it

This is where you build the habit that will save you for the rest of the project: compare your analytical gradient against a numerical gradient.

The numerical gradient of a function f at point x can be approximated as:

(f(x + h) - f(x - h)) / (2h)

for a small h (like 1e-5). If your hand-derived, backprop-computed gradient doesn't match this numerical approximation within a small tolerance, your backward pass has a bug — full stop, no exceptions. This technique is called gradient checking, and it is the single most important debugging tool in this entire project. Build it early. Use it constantly.

Common mistakes

Forgetting to zero out gradients between training steps (gradients accumulate by design, so you must explicitly reset them, or they'll silently corrupt every subsequent step).
Getting the local derivative wrong for a specific operation (e.g., writing the derivative of x**2 as 2 instead of 2*x).
Forgetting that when a value is used multiple times in a graph (e.g., y = x + x), gradients from both usages must be summed, not overwritten. This is the single most common bug in every autodiff implementation ever written by a beginner.
Building the backward traversal in the wrong order (not respecting topological order), which causes gradients to be computed using stale, not-yet-fully-accumulated values.

You're ready to move on when…

You can build an arbitrary expression using +, *, **, and tanh, call .backward() on the output, and have every intermediate scalar's .grad match a numerically-computed gradient within 1e-4 or so — for at least five different hand-built expressions, including ones where a variable is reused multiple times.

Mini exercises

Build the expression L = (a * b + c) ** 2 for scalar values a, b, c. Call backward and verify all three gradients against numerical gradient checking.
Deliberately introduce the "reused variable" bug (don't accumulate, just overwrite gradients) and observe how the numerical check catches it. This will burn the failure mode into your memory permanently.

Phase 3: The Tensor Engine — Scaling Up From Scalars

What you're learning

How to generalize your scalar autodiff engine to operate on N-dimensional arrays (tensors), including the trickiest and most bug-prone part of any framework: gradient broadcasting.

Why it matters

Real neural networks don't operate on individual numbers — they operate on batches of vectors and matrices, for speed and for the mathematical elegance of representing whole layers as matrix multiplications. This phase is where your framework becomes genuinely useful rather than a toy.

What to study online

Revisit NumPy broadcasting rules, this time thinking specifically about the reverse problem: if a tensor was broadcast during the forward pass (say, a (1,4) bias added to a (32,4) batch), how must the gradient be summed back down to the original smaller shape during the backward pass? (Hint: you sum gradients over exactly the dimensions that were broadcast.)
Matrix calculus basics — specifically the gradient of a matrix multiplication C = A @ B with respect to A and B. Search for "gradient of matrix multiplication backpropagation" and work through the derivation on paper before coding it.
The tinygrad or micrograd source code, again for inspiration and comparison after you've attempted your own design — not as a template to copy line by line.

What to build yourself

A Tensor class wrapping a NumPy array, with the same structure as your Scalar class: value, gradient, parent references, and a local backward function per operation.
Core operations: elementwise +, -, *, /, matrix multiplication (@), sum, mean, reshape, and broadcasting-aware arithmetic.
A .backward() method that performs topological sort and reverse traversal, just like your scalar engine — but now the "local backward rule" for each operation must handle array shapes correctly, including reducing broadcasted gradients back down.

Tiny illustrative idea

The conceptual shape of the fix for broadcast gradients (not literal code to paste, but the idea to implement yourself):

"After computing the gradient at the operation's output shape, if the input was broadcast to get there, sum the gradient over the broadcasted axes until it matches the input's original shape."

How to test it

Gradient checking again — but now applied to whole tensors. Perturb a single entry of a tensor by h, recompute the forward pass, measure the change in the loss, and compare to your backprop-computed gradient at that entry. Do this for a random sample of entries across several test cases: elementwise ops, broadcasting ops, and matrix multiplication.

Also write shape assertions liberally during development: after every backward operation, assert that the gradient's shape exactly matches the corresponding tensor's shape. Shape mismatches are the number one symptom of a broadcasting bug, and catching them immediately (rather than three layers later) will save you hours.

Common mistakes

Forgetting to sum gradients back down after broadcasting (the gradient tensor ends up the wrong shape, or numerically wrong even if the shape happens to match).
Getting the matrix multiplication gradient transposed incorrectly (a very common and very confusing bug — if C = A @ B, then the gradient w.r.t. A involves B transposed, and the gradient w.r.t. B involves A transposed; get this backwards and your loss will often still decrease, just slower or incorrectly, which makes the bug sneaky).
Silent shape bugs: NumPy broadcasting is forgiving, so a wrong-shape gradient often doesn't crash — it just silently gives you wrong numbers. This is why explicit shape assertions matter more here than almost anywhere else in the project.

You're ready to move on when…

You can build a small computation involving matrix multiplication, broadcasting (like adding a bias vector to a batch), and elementwise operations, call backward, and have every gradient match numerical gradient checking — and you understand, in your own words, why broadcasted gradients need to be summed back down.

Mini exercises

Implement y = X @ W + b where X is (batch, in_features), W is (in_features, out_features), and b is (out_features,). Verify all three gradients numerically.
Deliberately break the broadcasting-gradient-reduction step and observe what kind of numerical mismatch results — get familiar with what "wrong" looks like on this specific bug.

Phase 4: Neural Network Layers

What you're learning

How to organize tensor operations into reusable, composable objects — the Module abstraction that every framework uses (PyTorch calls it nn.Module).

Why it matters

This is where your project starts to feel like a real framework. You'll design an API elegant enough that stacking layers to build a network feels natural, not clunky.

What to study online

PyTorch's nn.Module documentation — not to copy the implementation, but to study the interface design: how parameters are registered, how forward() is called, how submodules nest inside each other.
Weight initialization strategies (Xavier/Glorot, Kaiming/He) — read the original reasoning for why initialization scale matters (poor initialization causes vanishing or exploding activations before training even starts).

What to build yourself

A base Module class with a parameters() method that can recursively collect every trainable tensor from itself and its submodules.
A Linear layer (y = xW + b) with properly initialized weights.
Activation functions as their own small modules or functions: ReLU, Tanh, Sigmoid, Softmax.
A way to compose modules sequentially (a Sequential container, or simply nesting modules inside a custom class with a forward method).

How to test it

Forward-pass shape tests: build a small network, push a batch through it, and confirm the output shape is exactly what you expect at every layer.
Initialization sanity checks: after initializing a Linear layer, check that the variance of its weights roughly matches what your chosen initialization scheme predicts (this catches typos in initialization formulas early, before they cause mysterious training failures three phases from now).
Gradient flow sanity check: run a forward and backward pass through a multi-layer network and confirm gradients exist (are non-zero and non-NaN) for every parameter, at every layer, not just the last one. This catches "dead" layers early.

Common mistakes

Initializing weights to all zeros (a classic mistake — with zero weights, every neuron in a layer computes the exact same output and gradient, so they never differentiate from each other, a failure mode called symmetry).
Forgetting to zero gradients before each new backward pass, causing gradients to silently accumulate across training steps.
Building forward() methods that accidentally break the autodiff graph — for example, converting a Tensor to a raw NumPy array partway through a computation, which severs the graph and makes .backward() unable to reach earlier layers.

You're ready to move on when…

You can build a 2–3 layer network with your Module system, push a batch of random data through it, call backward, and confirm every single parameter (from the first layer to the last) receives a sensible, non-zero, non-NaN gradient.

Mini exercises

Build a Sequential container that chains Linear → ReLU → Linear and verify the output shape for a batch input.
Deliberately zero-initialize a layer's weights and observe the symmetry problem directly — confirm that all neurons in that layer produce identical gradients.

Phase 5: Losses and Optimizers

What you're learning

How to measure "wrongness" (loss functions) and how to translate gradients into actual parameter updates (optimizers).

Why it matters

Loss functions define what the network is trying to learn. Optimizers define how it gets there. A subtle bug in either one produces the most frustrating kind of failure: a model that trains, technically, but never quite gets good — and you won't immediately know which piece is at fault.

What to study online

Mean Squared Error and Cross-Entropy loss — read a derivation of why cross-entropy is the natural loss for classification (its connection to maximum likelihood estimation is worth understanding, not just memorizing the formula).
The original SGD-with-momentum intuition (think of it as a ball rolling downhill, accumulating velocity) and the Adam paper's core idea (adaptive per-parameter learning rates using running estimates of the gradient's mean and variance).

What to build yourself

MSELoss and CrossEntropyLoss (careful with the numerically stable combination of softmax + cross-entropy — computing them separately can cause overflow; study how frameworks combine them into one numerically stable operation).
An Optimizer base structure, then SGD (with optional momentum) and Adam, each implementing a step() method that reads .grad off every parameter and updates .value accordingly, plus a zero_grad() method.

How to test it

Compare your loss function's output against a hand-computed value on a tiny example with numbers small enough to check by hand or calculator.
Overfit a tiny synthetic dataset (even just 4–8 data points) with your optimizer and confirm the loss goes essentially to zero. If it can't even overfit a handful of points, something is fundamentally broken before you try anything harder.
For Adam specifically, verify the bias-correction terms are implemented — a very common subtle bug is forgetting bias correction on the first few steps, which causes unstable early training.

Common mistakes

Not zeroing gradients before each step(), so gradients from previous batches contaminate the current update.
Using the wrong sign in the update rule (subtracting instead of adding, or vice versa — always double check: parameters move in the negative gradient direction to reduce loss).
Learning rate too high (loss explodes or oscillates) or too low (loss barely moves) — learn to recognize both failure signatures by deliberately inducing them once.

You're ready to move on when…

You can take a tiny synthetic dataset, run several hundred optimization steps with either SGD or Adam, and watch the loss curve monotonically (or nearly so) decrease toward zero — and if it doesn't, you know how to diagnose why.

Mini exercises

Implement plain SGD and use it to fit a single Linear layer to noisy linear data (y = 3x + 2 + noise) — check that the learned weight and bias converge close to 3 and 2.
Implement Adam and compare its convergence speed against plain SGD on the same tiny problem — observe (don't just read about) why adaptive methods often converge faster on ill-conditioned problems.

Phase 6: The Training Loop and First Real Experiments

What you're learning

How to assemble everything into the actual training loop — the rhythm of forward pass, loss computation, backward pass, and optimizer step that every deep learning system repeats millions of times.

Why it matters

This is the moment your framework stops being a pile of isolated pieces and becomes a system. It's also, often, the first moment something genuinely surprising happens: a model that actually learns a nontrivial function, using code you wrote entirely yourself.

What to study online

Classic explanations of the XOR problem and why it historically mattered (a single-layer perceptron cannot solve it — this is a famous, important lesson in why depth and nonlinearity matter).
General training loop patterns across frameworks — notice the near-universal shape: zero_grad() → forward() → compute loss → backward() → step().

What to build yourself

A full training loop for the XOR problem: 4 data points, a small 2-layer network with a nonlinear activation, and your loss + optimizer from Phase 5.
Basic logging: print or plot the loss every N steps so you can watch it decrease (or fail to).
If you're feeling ambitious, a tiny subset of MNIST (even just a few hundred examples, downsampled) as a slightly harder next test.

How to test it

Watch the loss curve. It should decrease, generally, though not necessarily perfectly monotonically.
Check final predictions against ground truth directly — for XOR, after training, does the network correctly predict 0 for (0,0) and (1,1), and 1 for (0,1) and (1,0)?
This is a good moment to also test overfitting a single batch as a diagnostic reflex you'll use for the rest of your career: if a model can't overfit a tiny amount of data, it has a bug, not a "needs more data" problem.

Common mistakes

Forgetting zero_grad() at the start of each loop iteration (the single most common training-loop bug in every framework, home-built or professional).
Using a network without a nonlinearity between layers, which collapses multiple linear layers into a mathematically equivalent single linear layer — and a single linear layer cannot solve XOR, no matter how long you train it. If your XOR loss refuses to go below a certain plateau, this is the first thing to check.
Learning rate mis-set for the specific optimizer (Adam and SGD often want quite different learning rate scales).

The documentary moment

There will be a specific instant — maybe late at night, maybe after your third attempt at fixing a shape bug — when you run your training loop and watch the loss actually fall, and the XOR network actually gets every example right, using nothing but code you wrote from first principles. That moment is worth the whole project. Let yourself notice it when it happens.

You're ready to move on when…

Your framework can solve XOR reliably, and you can explain, without notes, why a single linear layer fundamentally cannot.

Mini exercises

Train your XOR network and plot the decision boundary it learns (even a crude ASCII or matplotlib grid) — see the nonlinear boundary emerge visually.
Remove the nonlinear activation from your network and confirm it now fails to solve XOR, exactly as theory predicts.

Phase 7: Transformer Components

This is the most ambitious and mathematically rich phase. Budget extra time here, and don't be discouraged if it takes longer than planned — this is genuinely advanced material, and moving slowly here is the difference between "I copied a transformer" and "I understand a transformer."

What you're learning

The architecture that underlies essentially all modern large language models: token embeddings, positional information, self-attention, multi-head attention, layer normalization, and residual connections.

Why it matters

Understanding attention from the ground up — implementing the query/key/value mechanism yourself, watching attention weights sum to 1 via softmax, seeing residual connections stabilize deep networks — turns transformers from a buzzword into a mechanism you genuinely understand. This is the difference that shows in technical interviews and in your ability to read modern ML papers without getting lost.

What to study online

The original "Attention Is All You Need" paper — read it once quickly for the shape of the ideas, then again slowly, focusing specifically on the scaled dot-product attention formula and why the scaling factor (dividing by the square root of the key dimension) exists (it prevents softmax saturation for large dimensions).
A visual walkthrough of self-attention (there are several excellent illustrated blog posts and diagrams online — search for "illustrated transformer" or "illustrated self-attention") to build geometric intuition before diving into the formula.
Layer normalization and residual connections — understand these as stabilization techniques that make deep networks trainable, not as arbitrary architectural decoration.

What to build yourself

An Embedding layer (a lookup table mapping token indices to vectors — this is really just indexed access into a trainable matrix, with the indexing operation needing its own backward rule).
Positional encoding (either the fixed sinusoidal version from the original paper, or a simpler learned positional embedding — the learned version is easier to implement correctly first).
Scaled dot-product self-attention: compute queries, keys, and values via linear projections; compute attention scores as Q @ K^T / sqrt(d_k); apply softmax; multiply by V.
Multi-head attention: run several attention "heads" in parallel on different projected subspaces, then concatenate and project back.
LayerNorm and residual (skip) connections wrapping your attention and feedforward sublayers.
A causal mask for autoregressive generation (so a token at position i can only attend to positions ≤ i, not the future).

Tiny illustrative idea

The conceptual shape of attention, not literal code:

"For each token, compute how much it should 'attend to' every other token, using a similarity score between its query and every key. Normalize those scores with softmax so they form a valid weighting. Use those weights to compute a weighted average over the value vectors — that weighted average becomes the token's new, context-aware representation."

How to test it

Shape tests, relentlessly: attention involves several reshapes and transposes for multi-head splitting, and shape bugs here are extremely common. At every step, know exactly what shape you expect and assert it.
Attention weight sanity check: after softmax, confirm attention weights for each query sum to exactly 1 across the key dimension.
Causal mask sanity check: confirm that after masking, a token's attention weights to future positions are exactly zero, not just small.
Gradient checking one more time, on a small attention block in isolation, before trusting it inside a larger model — attention has enough moving parts that isolating it for testing is worth the extra effort.
Overfit a tiny synthetic sequence task (like copying a short sequence, or predicting the next character in a short repeating pattern) to confirm the whole transformer block can actually learn something before scaling up.

Common mistakes

Forgetting the causal mask, which lets the model "cheat" by looking at future tokens during training — a model trained this way often looks like it's learning beautifully, until you try to generate text with it and discover it never actually learned to predict anything, it just memorized answers using information it shouldn't have had access to.
Getting the transpose wrong when splitting into multiple attention heads, silently mixing information across heads that should stay separate.
Applying LayerNorm at the wrong point relative to the residual connection (pre-norm vs. post-norm architectures behave differently, and mixing conventions inconsistently causes subtle instability).
Numerical instability in softmax over attention scores when the scaling factor is missing or wrong.

You're ready to move on when…

You can build a single transformer block, push a small batch of token sequences through it, confirm shapes and attention-weight sums are correct at every step, confirm the causal mask genuinely blocks future information, and overfit a tiny toy sequence task.

Mini exercises

Implement attention without multi-head splitting first (single head), verify it fully, then generalize to multi-head — resist the urge to implement multi-head attention directly, since debugging the simpler version first will save you real time.
Visualize attention weights (as a heatmap) for a short sequence and manually inspect whether they look sensible before and after training.

Phase 8: The Mini Language Model Demo

What you're learning

How to assemble embeddings, positional information, stacked transformer blocks, and a final projection layer into a complete, trainable, generative language model — and how to actually generate new text from it.

Why it matters

This is the payoff phase. Everything from Phase 2 onward converges here: autodiff, tensors, layers, losses, optimizers, and attention, working together to do something that feels genuinely impressive — a small AI system, built by you, generating text one character at a time.

What to study online

Character-level language modeling as a task — read about why it's a good learning demo (small vocabulary, no tokenizer complexity, fast to train, easy to sanity-check output).
Autoregressive sampling strategies: greedy decoding, temperature sampling, top-k sampling — understand at least greedy and temperature sampling before implementing generation.
Cross-entropy loss specifically as applied to next-token prediction (the standard training objective for language models).

What to build yourself

A small character-level dataset (a public domain text file works well — something like a short story or a chunk of Shakespeare).
A vocabulary built from unique characters, with encode/decode functions between characters and integer indices.
A small GPT-style model: embedding layer + positional encoding + a few stacked transformer blocks + a final linear layer projecting to vocabulary size.
A training loop that samples random chunks of text, predicts the next character at every position, and computes cross-entropy loss.
A generate() function that starts from a seed string and autoregressively samples new characters one at a time, feeding each generated character back in as input for the next step.

How to test it

Training loss should decrease steadily over time — plot it.
Periodically, during training, generate a short sample of text and read it. Early on, it will be gibberish. Partway through, you'll start seeing real word-shapes and plausible letter combinations, even if the content is nonsensical. This progression — noise, to word-shapes, to something reading almost like language — is one of the most satisfying things you'll observe in this entire project.
Compare generation with different temperatures: very low temperature should produce repetitive, "safe" text; higher temperature should produce more varied (and eventually incoherent) text. Seeing this tradeoff directly builds real intuition about sampling.

Common mistakes

Off-by-one errors in constructing input/target pairs for next-character prediction (the target sequence should be the input sequence shifted by one position — get this wrong and the model trains on a nonsensical objective, usually still producing decreasing loss, which makes the bug sneaky).
Forgetting the causal mask here specifically, which lets the model see future characters during training and then produce oddly perfect-looking loss curves that mean nothing, since real generation won't have access to the future.
Training on a dataset too small to generalize, then being surprised that generated text is mostly memorized rather than genuinely learned patterns — a useful phase to observe and understand, not necessarily to "fix."

The documentary moment

This is the second big emotional payoff of the project. The first time you seed your model with a few characters and watch it generate a full paragraph of text that looks like language — even if it's utter nonsense semantically — you are watching a system you built, atom by atom, produce something that feels alive. Sit with that moment. It's earned.

You're ready to move on when…

You have a full training run showing decreasing loss, and generated samples that show clear qualitative improvement from early training to late training — moving from random noise toward recognizable word-shapes and plausible local structure.

Mini exercises

Train on two very different text sources (say, one is repetitive and structured, one is more free-form) and compare how quickly and how well the model picks up each one's patterns.
Implement and compare greedy decoding vs. temperature sampling on the same trained model, and describe in your own words the qualitative difference you observe.

Phase 9: Cleanup, Tests, Documentation, Packaging

What you're learning

How to turn a working research prototype into something that looks and behaves like real, professional software — because a framework nobody else (including future-you) can read or trust isn't really finished.

Why it matters

This phase is often skipped by beginners eager to move on to the next shiny project, but it's exactly the phase that turns "a thing I hacked together" into "a project I can show in an interview, link on my resume, and actually reuse." It also forces you to re-read your own code with fresh eyes, which reliably surfaces bugs you missed the first time.

What to study online

Basic pytest usage, if you haven't used a testing framework before — how to structure test files, write assertions, and run a full test suite with one command.
Good README structure for open-source projects — look at a few well-regarded small ML repos on GitHub for inspiration on structure, not content.
Basic Python packaging (a pyproject.toml or setup.py, so your library can be pip install -e .'d locally).

What to build yourself

A formal test suite covering: gradient checks for every operation, shape tests for every layer, an overfitting test for a tiny dataset, and a small end-to-end test that trains for a few steps and confirms loss decreases.
A README explaining what the project is, how to install it, a quickstart example, and a short explanation of the architecture (this is also a superb exercise in explaining your own system clearly — if you struggle to explain a part, that's a signal you don't fully understand it yet).
Docstrings and type hints across your public API.
A couple of example scripts (XOR training, character-level generation) that a stranger could run to see the library work.

How to test it

Run your full test suite from a clean checkout, ideally in a fresh virtual environment, to confirm there are no hidden dependencies on leftover state from your development process. Ask a friend, if possible, to follow your README's quickstart from scratch and tell you where they got confused — this is the single best way to find documentation gaps.

Common mistakes

Tests that only check "does it run" without checking "is the output correct" (a test that passes even when the math is subtly wrong is worse than no test at all, because it creates false confidence).
A README that describes what you intended to build rather than what actually exists in the code — keep them in sync.

You're ready to call it done when…

A stranger could clone your repository, follow your README, install the package, run the examples, and understand — at least at a high level — how the pieces fit together, without needing to ask you anything.

Part 4: The Debugging Mindset

You will spend a significant fraction of this project's time debugging, not writing new code. That is not a sign you're doing it wrong — it's the nature of building something this foundational. A few principles that will serve you well:

Numerical gradient checking is your best friend. Reach for it constantly, especially any time you add a new operation to your autodiff engine. It is objective, mechanical, and doesn't lie.
Overfit small before you scale up. If your model can't perfectly memorize four data points, it has a bug — it does not need more data, a bigger model, or more training time.
Assert shapes aggressively during development. Most bugs in tensor code are shape bugs. Catching them the instant they occur, rather than three layers downstream, saves enormous time.
Isolate before you integrate. Test each new component (an operation, a layer, an attention block) completely on its own, with its own small test cases, before wiring it into the larger system. Debugging a fully assembled transformer is much harder than debugging one attention block in isolation.
When something is "almost right," be suspicious, not relieved. A model that trains but plateaus at a mediocre loss, or a gradient that's close but not quite matching the numerical check, is very often hiding a real bug — broadcasting summed over the wrong axis, a transpose in the wrong place — rather than being "good enough." Chase these down; they compound.
Reproduce with the smallest possible example. When something breaks in your full training loop, don't debug it there. Strip the problem down to the smallest snippet that reproduces the bug — often just two or three lines — and debug that instead.

Part 5: What Not To Do

Don't copy an existing tiny framework's source code line by line. Reading micrograd, tinygrad, or similar projects for conceptual inspiration is not just fine but encouraged. Transcribing their code without building your own understanding defeats the entire purpose and will leave you unable to debug your own system when it inevitably breaks in a way the original project didn't.
Don't skip gradient checking "because it seems to be working." Loss curves that go down can mask real bugs (a common one: gradients computed with a sign error can still sometimes decrease loss, just less efficiently, hiding the bug for a long time).
Don't jump straight to the transformer phase before your autodiff and tensor engine are solid. Every bug you haven't found in Phase 2 or 3 will resurface, disguised and much harder to diagnose, somewhere deep inside your attention mechanism in Phase 7.
Don't optimize for speed prematurely. Your framework will be slow compared to PyTorch — that's fine and expected. This project is about correctness and understanding, not competing with a production framework backed by compiled CUDA kernels.
Don't work in total isolation. Read forum discussions, ask questions in ML-focused communities when you're stuck, and compare notes with other builders. Building alone doesn't mean learning alone.

Part 6: How to Actually Learn From the Internet While Building This

Since you won't be handed finished code, your primary skill throughout this project is turning scattered online knowledge into working understanding. Here's how to do that well:

Read broadly before you write code, narrowly while you're stuck. Before starting a new phase, skim two or three different explanations of the core concept (a blog post, a paper, a video) to build a rounded mental model. Once you're implementing and hit a specific wall, narrow your search to that exact problem.
Treat existing tiny frameworks as reference material, not templates. Open micrograd or tinygrad after you've made your own honest attempt, specifically to compare your design decisions against theirs — not before, and not to copy.
Papers are denser than blog posts but often clearer once you know what to look for. Read the "Attention Is All You Need" paper's method section only after you already have a rough intuition from a blog post or video — the paper will then click into place much faster.
GitHub issues and forum threads are goldmines for edge cases. If you hit a subtle bug (say, in broadcasting gradients or in causal masking), search for how others have described that exact symptom — someone has almost certainly hit it before you.
Write your own notes as you go. A short paragraph per phase, in your own words, explaining what you built and why, is one of the highest-leverage habits in this entire project. It forces synthesis, and it becomes both your documentation and your future interview prep material.

Epilogue: What You'll Have, and Where To Go Next

By the end of this project, you will have built — with your own hands, your own debugging sessions, your own late-night "oh, THAT'S why" moments — a working deep learning framework that can automatically differentiate arbitrary computations, represent and manipulate tensors with broadcasting, compose neural network layers, optimize them with SGD and Adam, and train a small transformer to generate text.

More importantly, you will have built a kind of understanding that's very hard to get any other way. You will look at a PyTorch stack trace and recognize exactly what kind of shape mismatch or graph-disconnection bug is happening, because you've caused and fixed that exact bug in your own code. You'll read a new architecture paper and be able to mentally sketch how you'd implement its core mechanism, because you've implemented the building blocks it's made from. In interviews, when someone asks "how does backpropagation actually work," you won't recite a definition — you'll explain a system you built.

This project also naturally sets up several strong next steps, whenever you're ready for them:

GPU acceleration: Extend your tensor engine to dispatch operations to a GPU backend, learning the basics of CUDA or a framework like CuPy along the way.
A larger, real transformer: Scale your mini language model up on a real dataset, and explore techniques like learning rate scheduling, gradient clipping, and mixed precision.
Convolutional networks: Apply the same "build it yourself" approach to a CNN, implementing convolution and pooling operations and their gradients from scratch.
Contributing to a real open-source ML framework: With this foundation, reading and contributing to real frameworks like PyTorch or JAX becomes dramatically more approachable — you'll recognize the patterns immediately.

Whatever you build after this, you'll build it with a kind of foundational confidence that's genuinely rare — the confidence of someone who has personally taken deep learning apart, piece by piece, and put it back together again with their own hands.

Now go write your first Scalar class. The journey starts there.