How Neural Networks Learn (Gradient-Based Optimization)
You’ve built MLPs and compiled them with model.compile(optimizer="sgd", ...)model.compile(optimizer="sgd", ...) — but
what does the optimizer actually do? This page opens that black box: derivative →
gradient → chain rule → backpropagation, ending with the same four lines of code
TensorFlow runs millions of times during every model.fit()model.fit() call.
The training loop, once more
Every neural network trains the same way, on repeat, until the loss is low enough:
- Draw a batch of training samples
xxand targetsy_truey_true. - Run the model on
xx— the forward pass — to get predictionsy_predy_pred. - Compute the loss: how far
y_predy_predis fromy_truey_true. - Update every weight a little, in the direction that makes the loss smaller.
Steps 1–3 are just tensor operations you already know. Step 4 — figuring out which direction reduces the loss, for potentially millions of weights at once — is the hard part, and it’s what the rest of this page is about.
flowchart LR X["Batch x"] --> F["Forward pass"] --> P["Prediction"] P --> L["Loss(prediction, target)"] Y["Target y_true"] --> L L --> G["Gradient (backward pass)"] G --> U["Optimizer updates weights"] U -.-> F
What’s a derivative?
For a smooth function f(x)f(x), a tiny change in xx (call it epsilon_xepsilon_x) produces a
tiny, predictable change in f(x)f(x). Near any point, ff can be approximated by a
straight line with some slope aa:
f(x + epsilon_x) ≈ f(x) + a * epsilon_xf(x + epsilon_x) ≈ f(x) + a * epsilon_x
That slope aa is the derivative of ff at that point. If aa is negative,
increasing xx decreases f(x)f(x) — so to minimize ff, you move xx in the opposite
direction from the derivative.
def f(x):
return x ** 2
x = 3.0
epsilon = 1e-4
slope = (f(x + epsilon) - f(x)) / epsilon
print(round(slope, 2)) # 6.0 -- matches the true derivative, 2*xdef f(x):
return x ** 2
x = 3.0
epsilon = 1e-4
slope = (f(x + epsilon) - f(x)) / epsilon
print(round(slope, 2)) # 6.0 -- matches the true derivative, 2*xFrom derivative to gradient
Now imagine ff doesn’t take a single number but a whole tensor WW (a model’s
weights). The generalization of “derivative” to tensor inputs is called the
gradient: grad(loss_value, W)grad(loss_value, W) is a tensor the same shape as WW, where each entry
tells you the direction and size of the loss’s change if you nudge that one weight.
Just like the 1D case, you reduce the loss by moving WW a small step in the
opposite direction from its gradient:
W_new = W - learning_rate * grad(loss, W)W_new = W - learning_rate * grad(loss, W)
learning_ratelearning_rate controls how big that step is. Too small, and training crawls (or
gets stuck); too large, and updates can bounce around randomly instead of settling
down.
Mini-batch stochastic gradient descent (SGD)
Solving for the exact minimum analytically is impossible for a network with millions of weights, so instead you take small, repeated steps downhill:
- Draw a batch of training samples
xxand targetsy_truey_true. - Run the forward pass to get
y_predy_pred. - Compute the loss between
y_predy_predandy_truey_true. - Compute the gradient of the loss with respect to every weight (the backward pass).
- Move every weight a little against its gradient:
w -= learning_rate * gradientw -= learning_rate * gradient.
Because each batch is a random sample of the data, this is called stochastic (random) gradient descent — mini-batch SGD when you use small batches instead of one sample or the whole dataset at once.
Visualize it
A single learnable parameter and its loss trace out a curve. Gradient descent is nothing more than repeatedly stepping downhill along that curve, using the local slope to decide which way “downhill” is:
The chain rule and backpropagation
A neural network is a chain of simple, differentiable operations: dotdot, ++, relurelu,
softmaxsoftmax. Calculus gives you a shortcut for differentiating a chain of functions —
the chain rule: if fg(x) = f(g(x))fg(x) = f(g(x)), then
grad(y, x) == grad(y, x1) * grad(x1, x)grad(y, x) == grad(y, x1) * grad(x1, x)
Applying this repeatedly, back through every layer of a network, from the loss all the way to the very first weight, is exactly what backpropagation does. It’s not a separate algorithm bolted onto neural networks — it’s just the chain rule, applied mechanically to a chain of tensor operations.
Consider a tiny computation graph: x1 = w * xx1 = w * x, x2 = x1 + bx2 = x1 + b, loss = abs(y_true - x2)loss = abs(y_true - x2). With x = 2, w = 3, b = 1, y_true = 4x = 2, w = 3, b = 1, y_true = 4, the forward pass gives x1 = 6x1 = 6,
x2 = 7x2 = 7, loss = 3loss = 3. Walking backward through the graph and multiplying local
derivatives along the way:
grad(loss, w) = grad(loss, x2) * grad(x2, x1) * grad(x1, w) = 1 * 1 * 2 = 2grad(loss, w) = grad(loss, x2) * grad(x2, x1) * grad(x1, w) = 1 * 1 * 2 = 2
grad(loss, b) = grad(loss, x2) * grad(x2, b) = 1 * 1 = 1grad(loss, b) = grad(loss, x2) * grad(x2, b) = 1 * 1 = 1
Modern frameworks never make you do this by hand — they implement automatic differentiation: record every operation as it runs, then walk the recording backward applying the chain rule for you.
The GradientTape in TensorFlow
TensorFlow’s automatic differentiation API is tf.GradientTapetf.GradientTape — a scope that
records every tensor operation run inside it, so you can later ask for the gradient
of any result with respect to any tf.Variabletf.Variable:
import tensorflow as tf
x = tf.Variable(0.)
with tf.GradientTape() as tape:
y = 2 * x + 3
grad_of_y_wrt_x = tape.gradient(y, x)
print(grad_of_y_wrt_x) # 2.0 -- the slope of y = 2x + 3import tensorflow as tf
x = tf.Variable(0.)
with tf.GradientTape() as tape:
y = 2 * x + 3
grad_of_y_wrt_x = tape.gradient(y, x)
print(grad_of_y_wrt_x) # 2.0 -- the slope of y = 2x + 3It works just as well with whole weight tensors:
import tensorflow as tf
W = tf.Variable(tf.random.uniform((2, 2)))
b = tf.Variable(tf.zeros((2,)))
x = tf.random.uniform((2, 2))
with tf.GradientTape() as tape:
y = tf.matmul(x, W) + b
grad_of_y_wrt_W_and_b = tape.gradient(y, [W, b])
print(grad_of_y_wrt_W_and_b[0].shape) # (2, 2) -- same shape as Wimport tensorflow as tf
W = tf.Variable(tf.random.uniform((2, 2)))
b = tf.Variable(tf.zeros((2,)))
x = tf.random.uniform((2, 2))
with tf.GradientTape() as tape:
y = tf.matmul(x, W) + b
grad_of_y_wrt_W_and_b = tape.gradient(y, [W, b])
print(grad_of_y_wrt_W_and_b[0].shape) # (2, 2) -- same shape as WIn practice you’d never write the weight update by hand either — a Keras optimizer does it for you:
import tensorflow as tf
from tensorflow.keras import optimizers
optimizer = optimizers.SGD(learning_rate=1e-3)
def one_training_step(model, x, y_true):
with tf.GradientTape() as tape:
y_pred = model(x)
loss = tf.keras.losses.mean_squared_error(y_true, y_pred)
loss = tf.reduce_mean(loss)
gradients = tape.gradient(loss, model.trainable_weights)
optimizer.apply_gradients(zip(gradients, model.trainable_weights))
return lossimport tensorflow as tf
from tensorflow.keras import optimizers
optimizer = optimizers.SGD(learning_rate=1e-3)
def one_training_step(model, x, y_true):
with tf.GradientTape() as tape:
y_pred = model(x)
loss = tf.keras.losses.mean_squared_error(y_true, y_pred)
loss = tf.reduce_mean(loss)
gradients = tape.gradient(loss, model.trainable_weights)
optimizer.apply_gradients(zip(gradients, model.trainable_weights))
return lossThis is precisely what happens inside model.fit()model.fit(), once per batch, for every epoch.
Mini-checkpoint
If the learning rate is far too large, what typically happens to training?
- Weight updates overshoot the minimum and the loss can bounce around or even increase instead of settling down.
Next
You now understand what model.fit()model.fit() is really doing under the hood. Time to put
the whole toolkit to work: First Example: Classifying Movie Reviews (IMDB,
Binary) — a complete, small Keras project from raw data to a trained classifier.
🧪 Try It Yourself
Exercise 1 – Approximate a Derivative
Exercise 2 – One Manual SGD Step
Exercise 3 – GradientTape Basics
If this helped you, consider buying me a coffee ☕
Buy me a coffeeWas this page helpful?
Let us know how we did
