From Karpathy's micrograd to smoltorch: Understanding Autograd from First Principles
Why I Built My Own Deep Learning Framework?
After watching Andrej Karpathy’s micrograd lecture, I had a realization: I’d been using PyTorch for months, but I didn’t really understand how autograd worked. Sure, I could call .backward() and get gradients, but what was actually happening under the hood?
So, I decided to build my own - from scratch, with just NumPy.
This is the story of building smoltorch, a tiny but functional deep learning library that achieves 96.5% accuracy on real datasets
What did micrograd teach?
micrograd laid the groundwork for learning the inner workings of an autograd library. How the gradients flow and how does nudging the weights in a direction affect the flow of gradients. Karpathy sensei’s implementation of micrograd was focused on scalar values and the autograd implemented reflect the same.
This was a nice building ground to learn how a library like PyTorch does differentiation. The biggest revelation moment turned out to be the accumulation of gradients, which I never really put much though into, turned out to be the multiplication of gradients in chain rule.
Let’s consider a function $z$, which depends on $y$, which in turn depends on $x$
$$\frac{dz}{dx} = \frac{dz}{dy} \cdot \frac{dy}{dx}$$
Simplified, but this is accumulation of gradients.
The challenge: scaling from scalars to tensors
When tensors came into picture, things change. Local gradients involve matrix operations. Broadcasting compilcates gradient accumulation. I had to write an algorithm which “undo” the broadcasting, which could revert the direction of broadcast. If a matrix has broadcasted it’s values along the row, and un-row it and similarly do for column.
micrograd worked with individual numbers:
x = Value(2.0)
y = Value(3.0)
z = x * y # Just two numbers
but real neural nets work with matrices
x = Tensor([
[1, 2, 3],
[4, 5, 6]
]) # 2D array
y = Tensor([
[7, 8],
[9, 10],
[11, 12]
]) # Another 2D array
z = x @ y # Matrix multiplication!
This changes the way gradient’s are computed! Let’s understand this is some more code examples
x = np.array([
[1, 2, 3]
]) # shape (1, 3)
y = np.array([
[10],
[20]
]) # shape (2, 1)
z = x + y # shape (2, 3) due to broadcasting
'''
Result:
[
[11, 12, 13],
[21, 22, 23]
]
'''
How does the gradient flow with broadcasting? When z.grad flows back, shape (2, 3):
z.grad = [
[1, 1, 1],
[1, 1, 1]
] # shape (2, 3)
'''
x was broadcast from (1, 3) → (2, 3)
so, we sum across axis=0
'''
x.grad = np.sum(z.grad, axis=0, keepdims=True)
= [[2, 2, 2]] # shape (1, 3)
'''
y was broadcast from (2, 1) → (2, 3)
so, we sum across axis=1
'''
y.grad = np.sum(z.grad, axis=1, keepdims=True)
= [[3], [3]] # shape (2, 1)
The algorithm to unbroadcast goes something like:
def unbroadcast(grad, orig):
while len(grad.shape) > len(orig):
grad = grad.sum(axis=0)
for i, (g_dim, o_dim) in enumerate(zip(grad.shape, orig)):
if o_dim == 1 and g_dim > 1:
grad = grad.sum(axis=i, keepdims=True)
return grad
The Building Blocks
After understanding how gradients flow through matrix operations, I needed to implement the fundamental operations that make neural networks work. Each operation had its own challenges, but they all followed the same pattern: compute the forward pass, then define how gradients flow backward.
Addition: The Broadcasting puzzle
In micrograd, it was trivial:
# Scalar addition
z = x + y
# Gradient: both inputs get the upstream gradient
x.grad += out.grad # ∂z/∂x = 1
y.grad += out.grad # ∂z/∂y = 1
But with tensors, NumPy's broadcasting created a hidden complexity:
x = np.array([[1, 2, 3]]) # shape (1, 3)
y = np.array([[10], [20]]) # shape (2, 1)
z = x + y # shape (2, 3) - broadcast!
# Result:
# [[11, 12, 13], # x was copied down
# [21, 22, 23]] # y was copied across
Here's the insight that clicked for me: when a tensor gets broadcast during forward pass, its gradient needs to be "unbroadcast" by summing.
Think about it: x with shape (1, 3) affected all 6 positions in the output. So when z.grad comes back with shape (2, 3), we need to sum along the dimension that was broadcast:
z.grad = [[1, 1, 1],
[1, 1, 1]]
# x was broadcast from (1, 3) to (2, 3)
# Sum over axis 0 to get back to (1, 3)
x.grad = z.grad.sum(axis=0, keepdims=True) # [[2, 2, 2]]
# y was broadcast from (2, 1) to (2, 3)
# Sum over axis 1 to get back to (2, 1)
y.grad = z.grad.sum(axis=1, keepdims=True) # [[3], [3]]
I wrote an _unbroadcast helper that became essential for every operation:
def _unbroadcast(self, grad, original_shape):
"""Reverse broadcasting by summing over expanded dimensions"""
# Sum over leading dimensions that were added
while len(grad.shape) > len(original_shape):
grad = grad.sum(axis=0)
# Sum over dimensions that had size 1 but got broadcast
for i, (grad_dim, orig_dim) in enumerate(zip(grad.shape, original_shape)):
if orig_dim == 1 and grad_dim > 1:
grad = grad.sum(axis=i, keepdims=True)
return grad
Matrix Multiplication: The Big One
This was the operation I was most nervous about. Matrix multiplication is the heart of neural networks - every linear layer uses it. Getting the gradients wrong here would break everything.
The forward pass is just NumPy:
z = x @ y # That's it!
But the backward pass requires understanding the relationship between shapes. Here's the derivation that finally made sense to me:
# Given:
x.shape = (2, 3)
y.shape = (3, 5)
z = x @ y # z.shape = (2, 5)
# When z.grad flows back (shape 2, 5), we need:
# x.grad with shape (2, 3)
# y.grad with shape (3, 5)
# For x.grad: we need (2, 3) from (2, 5) and (3, 5)
# Pattern: (2, 5) @ (5, 3) = (2, 3)
# So: x.grad = z.grad @ y.T
# For y.grad: we need (3, 5) from (2, 3) and (2, 5)
# Pattern: (3, 2) @ (2, 5) = (3, 5)
# So: y.grad = x.T @ z.grad
The implementation:
def __matmul__(self, other):
out = Tensor(self.data @ other.data, (self, other), '@')
def _backward():
# The gradient formulas
grad_self = out.grad @ other.data.swapaxes(-2, -1)
grad_other = self.data.swapaxes(-2, -1) @ out.grad
# Handle broadcasting (yes, even matmul broadcasts!)
grad_self = self._unbroadcast(grad_self, self.data.shape)
grad_other = self._unbroadcast(grad_other, other.data.shape)
self.grad += grad_self
other.grad += grad_other
out._backward = _backward
return out
I used swapaxes(-2, -1) instead of .T to handle batched matrix multiplication correctly - this was a subtle bug that took me an hour to find!
Key Insights and “aha” moments
What I actually learnt?
Broadcasting is everywhere
- Every time NumPy broadcasts, your gradient code needs to “unbroadcast” by summing. Miss this step? Your gradients will have wrong shapes
The computational graph is everything
- Modern frameworks build a graph dynamically. Understanding topological sorting was the key to understanding backpropagation
Numerical stability matters
- My first BCE implementation broke the gradient flow! The fix taught me why frameworks have so many epsilon values and clipping operations
Weight initialization is critical
random * 0.1 gave me dead networks
Xavier initialization fixed it
The Results
Regression (Make Regression Dataset)
Started with MSE loss of over $8000$
After 100 epochs, dropped down to less than $100$
Predictions within 10-15% of true values
Binary Classification (Breast Cancer Dataset)
First attempt? 37.8% accuracy as the model was stuck and not learning
After fixing BCE? 96.5% accuracy
This is quite close to some production frameworks out there in under 400 lines of Python
What is missing and why it’s okay?
No GPU support, but Apple M4 + NumPy is blazingly fast
No convolution, coming soon
No Adam optimizer, SGD works for learning
No mini-batches, full-batch training as of now
None of that matters for understanding the fundamentals and you do learn a lot when building from scratch.
What’s next?
Now that the core works, the roadmap is as follows (cannot guarantree any timelines or order of implementation):
Adam optimizer (momentum + adaptive learning rates)
Softmax for multi-class classification
Convolutional layers
Metal performance shaders for M4 optimization
Conclusion
If you’re serious about deep learning, build your own framework. Not because you will use it in production, but because you will understand PyTorch 10x better afterward.
When you call loss.backward() in pyTorch now, you won’t just see magic, you will see:
Topological sort building the graph
Chain rule being applied in reverse
Gradients flowing through operations
Broadcasting being handled correctly
GitHub repo: https://github.com/kashifulhaque/smoltorch
