Skip to content

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:

  1. Draw a batch of training samples xx and targets y_truey_true.
  2. Run the model on xx — the forward pass — to get predictions y_predy_pred.
  3. Compute the loss: how far y_predy_pred is from y_truey_true.
  4. 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.

diagram One training step mermaid
Forward pass produces a prediction, the loss compares it to the target, and the gradient tells the optimizer how to nudge every weight.

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.

Approximating a derivative with a tiny epsilon
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*x
Approximating a derivative with a tiny epsilon
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*x

From 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:

  1. Draw a batch of training samples xx and targets y_truey_true.
  2. Run the forward pass to get y_predy_pred.
  3. Compute the loss between y_predy_pred and y_truey_true.
  4. Compute the gradient of the loss with respect to every weight (the backward pass).
  5. 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:

sketch A ball descending a loss surface p5.js
The white ball starts high on the loss curve and steps downhill each frame, using the local slope (gradient) to decide which way to move -- then eases back to the start and does it again.

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:

A tiny GradientTape example
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 + 3
A tiny GradientTape example
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 + 3

It works just as well with whole weight tensors:

GradientTape over a weight matrix and bias
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 W
GradientTape over a weight matrix and bias
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 W

In practice you’d never write the weight update by hand either — a Keras optimizer does it for you:

One manual training step, the way an optimizer does it
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 loss
One manual training step, the way an optimizer does it
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 loss

This 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 coffee

Was this page helpful?

Let us know how we did