Inside PyTorch Autograd: Building Gradient Flow from Scratch
Introduction
If you have ever trained a neural network in PyTorch, you have probably called .backward() without thinking twice. The loss goes down, gradients appear on tensors, and optimization just works. But what actually happens under the hood?
Modern deep learning frameworks not just manually compute gradients, they build computation graphs and let auto differentiation do the work.
Instead of relying on PyTorch, we’ll build a tiny automatic differentiation engine ourselves and watch backpropagation emerge naturally.
Why Gradients Matter?
A neural network learns by adjusting its parameters(weights and biases) so that its output gets closer to the correct answer.
To know how to adjust each weight, we ask:
If I change this weight a little bit how much does the final output change?
The question is answered by a derivative (gradient).
Also, its sign (positive or negative) tells us in which direction to move.
What is Automatic Differentiation?
Automatic Differentiation is a method where computer computes exact derivatives by breaking the big function into small operations.
Any neural network, no matter how complex, is made up of simple operations:
- Addition
- Multiplication
- Exponentials
- ReLU, sigmoid, tanh etc.
e.g) There is a function
which can be written down as
This can be represented as a sequence of steps, not one big formula.
Computation Graph (the backbone of AutoDiff)
Computation Graph is a directed graph which is made from these operations where
- Nodes → Operations or Variables
- Edges → Flow of Direction
Let us understand it by an example where we compute forward and backward pass:
The function we have is
where
Break the function into tiny operations
Forward Pass (Compute Values)
While doing this, the system stores:
- Operations (+, square)
- Input value
- Output value
This stored structure is Computation Graph.
Backward Pass (Applying chain rule backwards)
Now we want
Instead of differentiating the whole formula at once, we go backwards applying the chain rule locally.
Get Anway Kapoor’s stories in your inbox
Join Medium for free to get updates from this writer.
Start from output:
Derivative:
Value of a = 5:
Move one step backward:
Derivate:
Chain Rule:
Final Gradient:
It differentiated each tiny operation, multiplied gradient as it went backwards and used the stored values from forward pass.
Code Walkthrough
A one neuron model with simple equation:
where,
w → weight
x → input data
b → bias
Every number in our computation will be wrapped in an object that knows three things:
- its numerical value
- how final output depends on it
- how it was created
A Value That Knows Its Gradient
class Value:
def __init__(self, data, prev=(), op=""):
self.data = data
self.grad = 0.0
self.prev = prev
self.op = op
self._backward = lambda: None- data → forward pass value
- grad → gradient accumulated during backprop
- parents → which values produced this value
- backward_fn → local chain rule for this operation
Forward Computation — Addition
The forward pass computes a number and backward pass distributes gradients equally to both inputs
def __add__(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data + other.data, (self, other), "+")
def _backward():
self.grad += out.grad
other.grad += out.grad
out._backward = _backward
return out- out is a new Value
- it stores where it came from
Each operation defines its own local derivative. Backpropagation just chains them together.
Forward Computation — Multiplication
Here the gradients depend on the actual values computed during forward pass
def __mul__(self, other):
other = other if isinstance(other, Value) else Value(other)
out = Value(self.data * other.data, (self, other), "*")
def _backward():
self.grad += other.data * out.grad
other.grad += self.data * out.grad
out._backward = _backward
return out- The derivative of a*b depends on b (and vice versa)
- This is why forward values must be stored
Non Linearity — Square
def __pow__(self, power):
out = Value(self.data ** power, (self,), f"**{power}")
def _backward():
self.grad += power * (self.data ** (power - 1)) * out.grad
out._backward = _backward
return out- Local derivate: 2x for square
- This pattern generalizes for tanh, ReLU, sigmoid
Convenience Operators
def __neg__(self):
return self * -1
def __sub__(self, other):
return self + (-other)- __neg__ defines how -x works by reusing multiplication
- __sub__ defines x-y as x + (-y)
- No new gradient logic is required
Backward pass function
def backward(self):
topo = []
visited = set()
def build(v):
if v not in visited:
visited.add(v)
for parent in v.prev:
build(parent)
topo.append(v)
build(self)
self.grad = 1.0 # dy/dy = 1
for v in reversed(topo):
v._backward()- We compute each node in a topological order
- Reverse order = backward pass
- Each node applies its local rule
Computing the Gradients
x = Value(2)
w = Value(3)
b = Value(1)
y= (w * x + b) ** 2
y.backward()
print("Forward output:")
print("y =", y.data)
print("\nGradients:")
print("dy/dw =", w.grad)
print("dy/db =", b.grad)
print("dy/dx =", x.grad)Output
Forward output:
y = 49
Gradients:
dy/dw = 28.0
dy/db = 14.0
dy/dx = 42.0Practice Problem
Solution:
import torch
from typing import List, Tuple, Union
def train_neuron(
features: Union[List[List[float]], torch.Tensor],
labels: Union[List[float], torch.Tensor],
initial_weights: Union[List[float], torch.Tensor],
initial_bias: float,
learning_rate: float,
epochs: int
) -> Tuple[List[float], float, List[float]]:
"""
Train a single neuron (sigmoid activation) with mean-squared-error loss.
Returns (updated_weights, updated_bias, mse_per_epoch)
â weights & bias are rounded to 4 decimals; each MSE value is rounded too.
"""
X = features.float()
y = labels.float().view(-1, 1)
W = initial_weights.float().view(-1, 1)
b = torch.tensor(initial_bias, dtype=torch.float32)
n = X.shape[0]
mse_values = []
for _ in range(epochs):
z = X @ W + b
y_hat = torch.sigmoid(z)
mse = torch.mean((y_hat - y) ** 2)
mse_values.append(round(mse.item(), 4))
dL_dyhat = 2 * (y_hat - y) / n
dyhat_dz = y_hat * (1 - y_hat)
dL_dz = dL_dyhat * dyhat_dz
dz_dW = X.T
dL_dW = dz_dW @ dL_dz
dL_db = torch.sum(dL_dz)
W = W - learning_rate * dL_dW
b = b - learning_rate * dL_db
updated_weights = [round(w.item(), 4) for w in W.view(-1)]
updated_bias = round(b.item(), 4)
return updated_weights, updated_bias, mse_valuespFinal Remarks
In this article, we went through the core ideas behind PyTorch’s autograd by building a minimal automatic differentiation engine from scratch.
