Skip to content
Skip to lesson
Back to Topics
LearnMath & StatisticsGradients and Backprop
⚡EasyFine-Tuning & Training

Gradients and Backprop

Train a one-weight latency predictor by reading a slope, taking a downhill step, tracing chain-rule paths, and checking the same gradient in PyTorch.

25 min read
Learning path
Step 10 of 199 in the full curriculum

Personalize this lesson

Adapt explanations and teaching visuals to your background and preferred voice.

A request takes 5 seconds, but a model predicts 6. The model's rule is simple: multiply the number of 100-token blocks by an adjustable number. For this 200-token request, that number is currently 3. Should we increase it or decrease it to make a better prediction?

The NumPy lesson taught how to organize and calculate with model inputs. Training adds another step: use a prediction's error to decide how the model's stored numbers should change. Start with one number so you can trace the whole update by hand.

A model with one adjustable number

A parameter is a stored number a model adjusts as it learns from data. Start with a single parameter, w: predicted response seconds per 100 input . Tokens are the text chunks processed by a language model. This latency predictor is intentionally stripped down so you can trace every moving piece without getting lost in tensor bookkeeping.

For one completed request:

QuantityMeaningValue
xprompt length in hundreds of tokens2.0
yobserved response seconds5.0
wcurrent seconds-per-100-token-block parameter3.0
prediction = w * xmodel estimate6.0
loss = (prediction - y)²penalty for being wrong1.0

The loss measures prediction error with one number. Squaring makes both a one-second overestimate and a one-second underestimate cost 1; an error of two seconds costs 4. Zero means a perfect prediction for this row. This squared-error loss has squared-seconds units. Lowering it improves our latency predictions; measuring service response time still requires the serving system.

Before running the cell, predict its two printed values. The table already gives you enough: the rule predicts 3 * 2 = 6, so the error is one second.

one-parameter-loss.py
def squared_loss(weight: float, token_blocks: float, actual_seconds: float) -> float:
    prediction = weight * token_blocks
    return (prediction - actual_seconds) ** 2

w = 3.0
x = 2.0
y = 5.0
prediction = w * x
loss = squared_loss(w, x, y)

print("prediction_seconds", prediction)
print("loss", loss)
Output
prediction_seconds 6.0
loss 1.0

That code scores a model, but it doesn't yet learn. The loss tells us how wrong its prediction is; the next question is directional: should w go up or down?

Read a slope by nudging one weight

Change only w by a tiny amount and recompute the loss. Before looking at the second row, predict what should happen: the current prediction is already too high, so an upward nudge should make the miss worse.

WeightPredictionLossWhat changed?
3.0006.0001.000000starting point
3.0016.0021.004004weight moved upward

Loss increased by 0.004004 when the weight increased by 0.001. Divide the change in loss by the change in weight:

slope estimate=1.004004−1.0000000.001=4.004\text{slope estimate} = \frac{1.004004 - 1.000000}{0.001} = 4.004

A positive slope means increasing w makes loss worse near this point. To lower the loss, move w down. The sign gives a direction; the size estimates the change in loss per unit change in weight. Check that ratio with the same nudge in Python:

finite-difference-nudge.py
def loss(weight: float) -> float:
    return (weight * 2.0 - 5.0) ** 2

w = 3.0
eps = 0.001
slope_estimate = (loss(w + eps) - loss(w)) / eps

print("loss_before", loss(w))
print("loss_after_nudge", loss(w + eps))
print("slope_estimate", round(slope_estimate, 3))
Output
loss_before 1.0
loss_after_nudge 1.0040039999999995
slope_estimate 4.004

This is a finite difference, a practical approximation to a derivative. A forward evaluation means running the prediction and loss calculation at chosen parameters. A one-sided check needs a new evaluation for each parameter you nudge; a centered check needs two. That becomes expensive for large models, but it's useful for checking a small formula.

The derivative gives the exact local slope

The nudge estimated 4.004. What would remain if we made the nudge smaller and smaller? A derivative is the slope that the difference ratio approaches, when that limit exists.[1] At our current w = 3, a nudge of size h gives prediction 6 + 2h, error 1 + 2h, and loss (1 + 2h)². Expand the square and divide the loss change by h:

(1+2h)2−1h=4h+4h2h=4+4h(h≠0)\frac{(1 + 2h)^2 - 1}{h} = \frac{4h + 4h^2}{h} = 4 + 4h \qquad (h \ne 0)

For h = 0.001, this is 4.004, exactly the estimate above. As h approaches zero, 4h vanishes and the slope approaches 4. We never divide by zero; we ask what the ratio approaches for nonzero nudges.

Now write a rule that works at other weights too. Let L name the loss:

L(w)=(2w−5)2L(w) = (2w - 5)^2

For any error e, the same expansion gives ((e + h)² - e²) / h = 2e + h, which approaches 2e. That's why the derivative of a square is twice its input. Our error is e = 2w - 5, which changes twice as fast as w. Multiply those two sensitivities:

dLdw=2(2w−5)×2\frac{dL}{dw} = 2(2w - 5) \times 2

Read dL/dw as "the derivative of loss with respect to weight." Differentiating the square gives the first factor. The final 2 comes from the prediction 2w: changing w by a small amount changes the prediction by twice that amount. This multiplication is the chain rule: combine sensitivities along a sequence of dependent calculations.[2]

At w = 3:

dLdw=2(6−5)×2=4\frac{dL}{dw} = 2(6 - 5) \times 2 = 4

The finite-difference estimate was 4.004; the derivative is exactly 4 at this point. Check it by comparing losses on both sides of w. The two weights are 2 * eps apart, so divide by that distance:

analytical-versus-numerical-slope.py
def loss(weight: float) -> float:
    return (weight * 2.0 - 5.0) ** 2

def analytical_gradient(weight: float) -> float:
    prediction = weight * 2.0
    return 2 * (prediction - 5.0) * 2.0

w = 3.0
eps = 1e-5
centered = (loss(w + eps) - loss(w - eps)) / (2 * eps)
exact = analytical_gradient(w)

print("analytical_gradient", exact)
print("finite_difference", round(centered, 6))
print("agree", abs(centered - exact) < 1e-6)
Output
analytical_gradient 4.0
finite_difference 4.0
agree True

A centered finite difference uses the losses at w + eps and w - eps. For this quadratic, the extra nudge terms cancel exactly in real arithmetic; the remaining discrepancy comes from floating-point rounding. For general smooth functions, it usually reduces approximation error compared with the same-sized one-sided nudge.

Gradient descent takes a downhill step

With one parameter, the derivative is the one-entry . Its sign tells whether increasing w raises or lowers loss locally. Gradient descent moves opposite that slope. A is the number we choose to scale the update.

Here the slope is positive. Adding the slope would move right, toward higher loss, so the update must subtract it.

wnew=w−learning rate×dLdww_{\text{new}} = w - \text{learning rate} \times \frac{dL}{dw}

With w = 3, derivative 4, and learning rate 0.1:

wnew=3−0.1×4=2.6w_{\text{new}} = 3 - 0.1 \times 4 = 2.6

Now the prediction is 2.6 * 2 = 5.2, and loss is (5.2 - 5)² = 0.04. One update reduced the loss from 1.0 to 0.04. Repeat the calculation, recomputing the slope at the new weight each time. Each printed row shows the state before its update:

repeat-gradient-descent.py
w = 3.0
x = 2.0
y = 5.0
learning_rate = 0.1

print("step weight prediction loss gradient")
for step in range(5):
    prediction = w * x
    loss = (prediction - y) ** 2
    gradient = 2 * (prediction - y) * x
    print(step, round(w, 4), round(prediction, 4), round(loss, 6), round(gradient, 4))
    w = w - learning_rate * gradient
Output
step weight prediction loss gradient
0 3.0 6.0 1.0 4.0
1 2.6 5.2 0.04 0.8
2 2.52 5.04 0.0016 0.16
3 2.504 5.008 6.4e-05 0.032
4 2.5008 5.0016 3e-06 0.0064

Read down the rows rather than only the final weight. Each update lowers the loss, and the gradient shrinks as w approaches 2.5. At that weight, the prediction is exactly 5, loss is zero, and the slope is zero. This is a minimum because squared loss can't be negative. In other functions, a zero derivative alone doesn't prove you've found a minimum.

The slope is local. It doesn't tell you how far to walk.

A correct gradient can still take a bad step

Calculus supplies direction. It doesn't choose a safe learning rate for you.

Keep the starting point and gradient fixed, then change only the learning rate. Predict the failure before calculating it: a rate ten times larger follows the same downhill direction, but travels much farther.

From the same starting point, a learning rate of 1.0 produces:

wnew=3−1.0×4=−1w_{\text{new}} = 3 - 1.0 \times 4 = -1

The new prediction is -2 seconds and loss becomes (-2 - 5)² = 49. The derivative was right; the step jumped straight past the good region.

learning-rate-overshoot.py
def one_step(learning_rate: float) -> tuple[float, float]:
    w, x, y = 3.0, 2.0, 5.0
    prediction = w * x
    gradient = 2 * (prediction - y) * x
    updated_w = w - learning_rate * gradient
    updated_loss = (updated_w * x - y) ** 2
    return updated_w, updated_loss

for lr in (0.1, 0.125, 1.0):
    updated_w, updated_loss = one_step(lr)
    print("lr", lr, "new_w", updated_w, "new_loss", round(updated_loss, 4))
Output
lr 0.1 new_w 2.6 new_loss 0.04
lr 0.125 new_w 2.5 new_loss 0.0
lr 1.0 new_w -1.0 new_loss 49.0
Two scales of the loss curve L(w)=(2w−5)² make the update visible. A close view at w=3 shows the tangent slope 4 and the short steps: rate 0.1 lands at w=2.6, loss 0.04; rate 0.125 reaches w=2.5, loss 0. A full view shows rate 1 taking the same gradient far past the minimum to w=−1, loss 49.
Each arrow starts at w = 3 and uses the same gradient, 4. Rate 0.1 takes a small downhill step; 0.125 reaches this quadratic's exact minimum; 1 overshoots and increases the loss to 49. Other models need learning rates chosen for their own loss geometry.

Why 0.125? The minimum is 0.5 weight units below the start, so the required step satisfies learning_rate * 4 = 0.5. Try 0.25 in the cell: it moves to w = 2, where loss is back to 1. The derivative gives the local downhill direction; this finite step overshoots the minimum and fails to lower the starting loss.

For this quadratic, we can say exactly which fixed rates converge. Since the derivative is 8(w - 2.5), subtracting a step gives

wnew−2.5=(1−8η)(w−2.5).w_{\text{new}} - 2.5 = (1 - 8\eta)(w - 2.5).

Here η\eta names the learning rate. Distance from the minimum shrinks when ∣1−8η∣<1|1-8\eta|<1, which means 0 < eta < 0.25. At 0.1, distance shrinks by a factor of 0.2 and loss by 0.2² = 0.04 each update. At 0.25, the weight oscillates between 3 and 2; a larger positive rate makes that distance grow. These statements use this exact quadratic, fixed inputs, and exact arithmetic. They aren't a convergence guarantee for a neural network.

Even changing units changes the safe rate. Counting raw tokens instead of 100-token blocks replaces x = 2, w = 3 with x = 200, w = 0.03, preserving the prediction 6. The derivative becomes 400, and the distance multiplier becomes 1 - 80000*eta. The stable interval is now 0 < eta < 0.000025, ten thousand times smaller. Reusing the old learning rate after a feature-scale change can destabilize the same prediction rule.

In a real training loop, treat a non-finite loss such as NaN or infinity as a diagnosis signal. It can come from an unstable step or invalid input. Adaptive and schedules change how a gradient is used across many steps. They still need a backward calculation you can trust.

One weight was enough to see a downhill step. A latency model still needs to separate prompt work from queue delay.

Two weights, two slopes

The same request has one second of recorded queue delay. Add a second parameter to scale that delay, and reset the two weights to [2, 2]. The prediction is 2 * 2 + 2 * 1 = 6, still one second too high. In notation, let x_p count prompt blocks, x_q record queue seconds, and put a hat on y to distinguish the prediction from the observed time:

y^=wpxp+wqxq\hat{y} = w_p x_p + w_q x_q

Here w_p has units of seconds per prompt block, while w_q is a unitless multiplier on queue seconds. The feature values are 2 and 1, so in these chosen units the same loss signal produces a prompt slope twice as large. A larger numerical gradient doesn't by itself mean a more important feature: changing the feature's units changes its gradient.

ValueNumber
prompt blocks x_p2
queue delay x_q1
current weights [w_p, w_q][2, 2]
prediction2*2 + 2*1 = 6
actual seconds5
error1
squared loss1

A partial derivative measures sensitivity to one parameter while holding the others fixed. The symbol ∂ replaces d to show that the loss has several adjustable inputs:

∂L∂wp=2(y^−y)xp=2×1×2=4\frac{\partial L}{\partial w_p} = 2(\hat{y} - y)x_p = 2 \times 1 \times 2 = 4 ∂L∂wq=2(y^−y)xq=2×1×1=2\frac{\partial L}{\partial w_q} = 2(\hat{y} - y)x_q = 2 \times 1 \times 1 = 2

Collect those slopes in parameter order and you have the gradient: [4, 2]. This ordered list is a ; no extra machinery is needed yet. Apply the same update rule separately to each weight:

two-parameter-gradient.py
x_prompt = 2.0
x_queue = 1.0
w_prompt = 2.0
w_queue = 2.0
actual_seconds = 5.0

prediction = w_prompt * x_prompt + w_queue * x_queue
error = prediction - actual_seconds
grad_prompt = 2 * error * x_prompt
grad_queue = 2 * error * x_queue
updated_prompt = w_prompt - 0.1 * grad_prompt
updated_queue = w_queue - 0.1 * grad_queue
new_loss = (updated_prompt * x_prompt + updated_queue * x_queue - actual_seconds) ** 2

print("prediction", prediction)
print("gradient", [grad_prompt, grad_queue])
print("updated_weights", [updated_prompt, updated_queue])
print("new_loss", new_loss)
Output
prediction 6.0
gradient [4.0, 2.0]
updated_weights [1.6, 1.8]
new_loss 0.0

For this row and learning rate, the update lands exactly on the target. But [2, 1] and [1.5, 2] would both predict 5 too. One row can't tell us how much prompt length and queue delay should each contribute across future requests. We'll need more rows after we've checked how these two slopes were calculated.

You just calculated two slopes separately. Backpropagation traces both from the loss, one step at a time.

Backpropagation is chain-rule bookkeeping

The two-parameter gradient came from a chain of small operations:

  1. Predict: prediction = w_p * x_p + w_q * x_q
  2. Subtract target: error = prediction - y
  3. Square error: loss = error ** 2

During the forward pass, these operations produce prediction = 6, error = 1, and loss = 1. Now predict the first signal sent backward: differentiating error² should give 2 * error = 2. Start at the scalar loss and move backward:

Local stepDerivativeValue
loss = error²d loss / d error = 2 * error2
error = prediction - yd error / d prediction1
prediction = w_p*x_p + w_q*x_qd prediction / d w_p = x_p2
same predictiond prediction / d w_q = x_q1

Multiply local factors along each path. Prompt length gets 2 * 1 * 2 = 4; queue delay gets 2 * 1 * 1 = 2. The route tells you which feature scales each parameter's correction. At the loss itself, start with dL/dL = 1: a unit change in loss changes loss by one unit.

Two aligned computation graphs trace one latency example. The forward pass carries feature and weight values through prediction 6, error 1, and squared loss 1. The backward pass starts with loss sensitivity 1, multiplies by 2e and by the subtraction derivative 1, then branches to separate gradients 4 for the prompt weight and 2 for the queue weight.
The forward pass keeps the values needed by the backward rules. Starting at dL/dL = 1, the backward pass multiplies local derivatives along each route: 1 × 2 × 1 × 2 = 4 for the prompt weight and 1 × 2 × 1 × 1 = 2 for the queue weight.

is that reverse walk through the computation. Operations retain the forward values their backward rules need, then multiply the arriving signal by a local derivative. Rumelhart, Hinton, and Williams used back-propagated error derivatives to learn internal representations in multilayer networks in their 1986 Nature paper.[3] The explicit calculation below keeps each factor visible:

trace-chain-rule-factors.py
x_prompt = 2.0
x_queue = 1.0
w_prompt = 2.0
w_queue = 2.0
actual = 5.0

prediction = w_prompt * x_prompt + w_queue * x_queue
error = prediction - actual
loss = error ** 2

d_loss_d_error = 2 * error
d_error_d_prediction = 1.0
d_prediction_d_prompt_weight = x_prompt
d_prediction_d_queue_weight = x_queue

d_loss_d_prompt_weight = d_loss_d_error * d_error_d_prediction * d_prediction_d_prompt_weight
d_loss_d_queue_weight = d_loss_d_error * d_error_d_prediction * d_prediction_d_queue_weight

print("forward", prediction, error, loss)
print("prompt_path", d_loss_d_error, "*", d_prediction_d_prompt_weight, "=", d_loss_d_prompt_weight)
print("queue_path", d_loss_d_error, "*", d_prediction_d_queue_weight, "=", d_loss_d_queue_weight)
Output
forward 6.0 1.0 1.0
prompt_path 2.0 * 2.0 = 4.0
queue_path 2.0 * 1.0 = 2.0

The same bookkeeping works when another operation sits between a weight and the loss. So far, each predicted term changes proportionally with its weight. A nonlinear operation behaves differently across parts of its input range.

For example, wrap the prompt term in ReLU (rectified linear unit), max(0, h). It passes a positive input and returns zero for a negative input. Name the value before ReLU h and the value after it z:

h=wpxp,z=ReLU(h),y^=z+wqxq,L=(y^−y)2h = w_p x_p,\quad z = \mathrm{ReLU}(h),\quad \hat{y} = z + w_q x_q,\quad L = (\hat{y} - y)^2

ReLU's local derivative is 1 when h > 0: a small increase passes through unchanged. It's 0 when h < 0: a small increase still leaves the output at zero. At h = 0, the left and right slopes disagree, so there's no ordinary derivative. This code uses a backward value of 0 at the corner, matching PyTorch's convention.[4]

Keep the same latency row, with w_p = 2 so h = 4 stays positive. The ReLU is a pass-through, and the prompt weight still receives 4.

Now flip the prompt weight to -0.5. h goes negative, ReLU outputs 0 seconds from the prompt path, and that weight gets no update. Predict that zero before running the cell.

relu-chain-backprop.py
x_prompt = 2.0
x_queue = 1.0
w_queue = 2.0
actual = 5.0

def relu(pre_activation: float) -> float:
    return max(0.0, pre_activation)

w_prompt = 2.0
h = w_prompt * x_prompt
z = relu(h)
prediction = z + w_queue * x_queue
error = prediction - actual
loss = error ** 2
d_loss_d_prediction = 2 * error
d_z_d_h = 1.0 if h > 0 else 0.0
d_loss_d_w_prompt = d_loss_d_prediction * d_z_d_h * x_prompt

print("forward", h, z, prediction, loss)
print("dL_dw_prompt", d_loss_d_w_prompt)

w_prompt_off = -0.5
h_off = w_prompt_off * x_prompt
z_off = relu(h_off)
prediction_off = z_off + w_queue * x_queue
d_z_d_h_off = 1.0 if h_off > 0 else 0.0
d_loss_d_w_prompt_off = (2 * (prediction_off - actual)) * d_z_d_h_off * x_prompt
print("relu_off", h_off, z_off, prediction_off, d_loss_d_w_prompt_off)
Output
forward 4.0 4.0 6.0 1.0
dL_dw_prompt 4.0
relu_off -1.0 0.0 2.0 -0.0

The printed -0.0 compares equal to zero; its sign comes from multiplying a negative loss signal by zero. Later deep-learning lessons stack many such local factors. The rule stays the same: store needed forward values and multiply local derivatives on the way back.

The corner needs a different check. Set w_prompt = 0 in this same row. Then L(w) = (ReLU(2w) - 3)²: on the negative side, loss stays 9; on the positive side, it is (2w - 3)². The left derivative is 0, while the right derivative is -12. A centered difference mixes the two sides and approaches -6, although PyTorch's backward rule returns 0 at the corner:

relu-corner-gradient-check.py
import torch

def loss(weight: float) -> float:
    return (max(0.0, 2 * weight) - 3) ** 2

eps = 1e-6
left = (loss(0.0) - loss(-eps)) / eps
right = (loss(eps) - loss(0.0)) / eps
centered = (loss(eps) - loss(-eps)) / (2 * eps)
w = torch.tensor(0.0, dtype=torch.float64, requires_grad=True)
torch_loss = (torch.relu(2 * w) - 3) ** 2
torch_loss.backward()

print("one_sided", round(left, 4), round(right, 4))
print("centered", round(centered, 4))
print("autograd", float(w.grad))
assert float(w.grad) == 0.0
Output
one_sided 0.0 -12.0
centered -6.0
autograd 0.0

This disagreement doesn't identify a backward bug: no ordinary derivative exists at this point. Nor does the zero returned by autograd prove a minimum. A small positive weight lowers loss. Check smooth points on each side to test the local rules, and treat the corner's backward value as the operation's convention.[4][5]

When paths meet, add their signals

The example above sends the loss signal to two different parameters. A computation graph can also use the same parameter more than once. Changing a shared parameter changes both paths, so their contributions to the loss must add.

For this example only, tie the two numerical coefficients to one stored value w. Treat the inputs as normalized numbers: x_p is prompt length divided by 100 tokens, x_q is queue time divided by one second, and the output is measured in seconds. Both inputs are then unitless, so a shared seconds-valued coefficient is consistent:

y^=wxp+wxq\hat{y} = w x_p + w x_q

With w = 2, x_p = 2, x_q = 1, and y = 5, prediction is 6 and error is 1. The prompt path contributes 2 * 1 * 2 = 4; the queue path contributes 2 * 1 * 1 = 2. Both paths reach w, so its gradient is 4 + 2 = 6.

You can also combine the forward terms first: w * 2 + w * 1 = w * 3. Its local derivative is 3, so the complete derivative is 2 * error * 3 = 6. The two routes agree. Verify the separate contributions in code:

sum-shared-parameter-paths.py
w = 2.0
x_prompt = 2.0
x_queue = 1.0
actual = 5.0

prediction = w * x_prompt + w * x_queue
error = prediction - actual
prompt_path = 2 * error * x_prompt
queue_path = 2 * error * x_queue
gradient = prompt_path + queue_path

print("prediction", prediction)
print("path_contributions", [prompt_path, queue_path])
print("shared_weight_gradient", gradient)
assert gradient == 6.0
Output
prediction 6.0
path_contributions [4.0, 2.0]
shared_weight_gradient 6.0
A backward computation graph runs from squared loss L=1 through error e=1 and prediction ŷ=6, then splits into prompt and queue paths before returning to shared weight w=2. The prompt contribution is 2×1×2=4, the queue contribution is 2×1×1=2, and the shared weight gradient adds them to 6.
When a single parameter feeds multiple calculation paths, the total derivative is the sum of path contributions. Both branches receive the upstream adjoint 2.0 from ŷ. The prompt path contributes 2.0 × 2.0 = 4.0, the queue path contributes 2.0 × 1.0 = 2.0, and the accumulator at w adds them to yield 6.0.

Real models repeat this pattern across rows and token positions, and their computation graphs can have many branches. We'll keep the arithmetic visible with four requests next, adding the row contributions and dividing by four because we chose mean squared error.

Four completed requests

Return to two independent weights and use four completed requests. Each row has prompt length and queue delay, and the target is observed response seconds. These are constructed data: targets follow the rule 2 * prompt_blocks + queue_seconds exactly.

RequestPrompt blocksQueue secondsObserved seconds
010.52.5
1215
230.56.5
34210

Use the mean squared error: compute each row's squared error, add them, then divide by four. Its gradient follows the same pattern. Compute 2 * error * feature for each row, add the contributions for each weight, and divide by four.

At initial weights [0, 0], all predictions are zero. The prompt-weight contributions are [-5, -20, -39, -80], whose mean is -36. The queue-weight contributions are [-2.5, -10, -6.5, -40], whose mean is -14.75. With learning rate 0.03, the first updated weights must be [1.08, 0.4425].

In , store the four rows as X with shape (4, 2). From the earlier array lesson, X @ weights multiplies each row's two features by the corresponding weights and adds them, giving four predictions. The gradient below uses explicit feature columns, so you can match it to the arithmetic above:

train-latency-predictor.py
import numpy as np

X = np.array([
    [1.0, 0.5],
    [2.0, 1.0],
    [3.0, 0.5],
    [4.0, 2.0],
])
y = np.array([2.5, 5.0, 6.5, 10.0])
weights = np.array([0.0, 0.0])
learning_rate = 0.03

print("step loss weights")
for step in range(101):
    predictions = X @ weights
    errors = predictions - y
    loss = float(np.mean(errors ** 2))
    gradient = np.array([
        np.mean(2 * errors * X[:, 0]),
        np.mean(2 * errors * X[:, 1]),
    ])
    if not np.isfinite(loss) or not np.all(np.isfinite(gradient)):
        raise FloatingPointError("non-finite loss or gradient")
    if step in (0, 1, 10, 100):
        print(step, round(loss, 6), np.round(weights, 4).tolist())
    weights = weights - learning_rate * gradient

print("learned_weights", np.round(weights, 4).tolist())
print("final_predictions", np.round(X @ weights, 3).tolist())
Output
step loss weights
0 43.375 [0.0, 0.0]
1 9.852759 [1.08, 0.4425]
10 0.003643 [2.0574, 0.8557]
100 0.000709 [2.0259, 0.9364]
learned_weights [2.0257, 0.937]
final_predictions [2.494, 4.988, 6.546, 9.977]

The known weights are [2, 1], and 101 updates stop near that solution. The loss falls quickly at first, then more slowly. A fixed iteration count sets the stopping point; inspect the loss and gradient to judge whether optimization has settled. Evaluate unseen requests to assess how well the rule predicts beyond these four training rows.

Why does the last part take longer? Prompt length and queue delay often increase together in this table, so changes to the two weights can partly cancel. Compare two offsets from the exact solution. Both change the weights by the same absolute amounts, but one produces much smaller errors:

Offset from [2, 1]Prediction errors on rows 0–3Mean squared loss
[+0.1, -0.25][-0.025, -0.050, +0.175, -0.100]0.0109375
[+0.1, +0.25][+0.225, +0.450, +0.425, +0.900]0.3109375

The loss is flatter along some directions than others. This is the optimization issue called conditioning: a step size that avoids overshooting a steep direction can make slow progress along a flatter one. In this quadratic, the late weight error lies mostly along a flatter direction. Feature scaling changes this geometry, but it doesn't remove every correlation between columns. Later linear algebra will quantify it; here the two error lists show the mechanism without another optimizer.[6]

You'll often see the two column means written compactly as (2 / len(X)) * (X.T @ errors). Here X.T has shape (2, 4): each row now holds one feature's values across the four requests. Multiplying by the four errors adds feature-weighted errors, then 2 / len(X) supplies the square's factor and the mean. Both versions compute the same two slopes; the next linear-algebra lessons develop that notation further.

A finite-value guard doesn't make training stable by itself. It stops the loop before a NaN or infinity update silently corrupts the weights. If it fires, inspect the learning rate and input scale first.

This tiny model keeps every prediction, loss, derivative, and update visible in code you can explain. A production latency service requires the additional data, serving, and reliability work described below.

A missing factor can hide in a manual backward formula. Before trusting one on a larger batch, check each parameter with a small nudge.

Gradient checking catches backward bugs

For each parameter, nudge only that value by eps, compute the two losses, and compare the numerical slope to the formula from your backward pass. Keep the data and random choices fixed between those evaluations; otherwise their difference includes changes beyond the parameter nudge. Use weights [1.3, 1.1], away from the minimum so the check tests nonzero slopes. The two gradients should agree component by component, not just in overall direction.

gradient-check-two-parameters.py
import numpy as np

X = np.array([[1.0, 0.5], [2.0, 1.0], [3.0, 0.5], [4.0, 2.0]])
y = np.array([2.5, 5.0, 6.5, 10.0])
weights = np.array([1.3, 1.1])

def mean_squared_loss(current: np.ndarray) -> float:
    errors = X @ current - y
    return float(np.mean(errors ** 2))

errors = X @ weights - y
manual_gradient = (2 / len(X)) * (X.T @ errors)

eps = 1e-5
numerical_gradient = np.zeros_like(weights)
for i in range(len(weights)):
    direction = np.zeros_like(weights)
    direction[i] = eps
    numerical_gradient[i] = (
        mean_squared_loss(weights + direction) - mean_squared_loss(weights - direction)
    ) / (2 * eps)

print("manual", np.round(manual_gradient, 6).tolist())
print("numerical", np.round(numerical_gradient, 6).tolist())
print("check_passed", np.allclose(manual_gradient, numerical_gradient, atol=1e-6))
assert np.allclose(manual_gradient, numerical_gradient, atol=1e-6)
Output
manual [-9.9, -3.925]
numerical [-9.9, -3.925]
check_passed True

Try removing the factor 2 from manual_gradient and rerun the cell. The assertion should fail: numerical differences evaluate the squared loss and expose the missing factor in the proposed derivative. Restore the factor before continuing. Passing at one point supports the formula for that case. Check more inputs to exercise other values and edge cases.

That comparison is only trustworthy in enough precision. Single-precision floats can make a correct backward rule look wrong.

Catastrophic cancellation in single precision

Use double precision (float64) for these small gradient checks. The check subtracts nearly equal loss values, and those values have already been rounded by floating-point arithmetic. After the matching leading digits cancel, the tiny difference can contain a large relative error. This is catastrophic cancellation.

Making eps smaller reduces approximation error only until rounding starts to dominate. Keep the one-weight latency loss and use the same eps = 1e-6 in two dtypes. Both analytical derivatives are 4, but the numerical checks needn't be equally accurate:

gradcheck-precision.py
import numpy as np

for dtype in (np.float64, np.float32):
    w = np.array([3.0], dtype=dtype)
    eps = np.array([1e-6], dtype=dtype)
    loss_plus = (2 * (w + eps) - 5) ** 2
    loss_minus = (2 * (w - eps) - 5) ** 2
    numerical = float(((loss_plus - loss_minus) / (2 * eps))[0])
    print(dtype.__name__, "numerical", round(numerical, 6))
    print("close to 4?", abs(numerical - 4) < 1e-5)
Output
float64 numerical 4.0
close to 4? True
float32 numerical 3.814697
close to 4? False

The formula hasn't changed, only its numerical evaluation. Framework gradient checkers face the same limit: PyTorch's gradcheck defaults are intended for double-precision inputs.[5] Also avoid testing exactly at ReLU's corner, where no ordinary derivative exists.

Let PyTorch repeat the backward pass

The manual check agrees. Now let a framework repeat the same calculation. Reverse-mode automatic differentiation propagates chain-rule products backward from outputs to inputs. It's especially useful when one scalar loss depends on many parameters, which is the training shape we have here: one reverse pass computes all its parameter derivatives.[7] Autograd uses the operations' local derivative rules, rather than nudging each parameter or manipulating an entire symbolic formula. Its results still use floating-point arithmetic, so automatic differentiation doesn't eliminate rounding error.

arrays are called . Its autograd system records differentiable operations on tensors that request gradients, then computes the reverse pass when you call backward().[4] Setting requires_grad=True on the weights asks for their derivatives. This cell repeats the four-row latency table and checks PyTorch against the manual NumPy derivative.

The weights are a leaf tensor: we created them directly, rather than calculating them from another tracked tensor. PyTorch stores their derivatives in .grad. Calling backward again adds to that buffer; it doesn't replace it. Predict the second value before running the cell. The detach().clone() calls below save independent snapshots of the gradient without tracking the copies themselves.

pytorch-confirms-the-gradient.py
import numpy as np
import torch

X_np = np.array([[1.0, 0.5], [2.0, 1.0], [3.0, 0.5], [4.0, 2.0]], dtype=np.float32)
y_np = np.array([2.5, 5.0, 6.5, 10.0], dtype=np.float32)
w_np = np.array([1.3, 1.1], dtype=np.float32)

manual = (2 / len(X_np)) * (X_np.T @ (X_np @ w_np - y_np))

X = torch.tensor(X_np)
y = torch.tensor(y_np)
w = torch.tensor(w_np, requires_grad=True)
loss = torch.mean((X @ w - y) ** 2)
loss.backward()
first_backward = w.grad.detach().clone()

same_loss_again = torch.mean((X @ w - y) ** 2)
same_loss_again.backward()
after_second_backward = w.grad.detach().clone()

w.grad = None
fresh_loss = torch.mean((X @ w - y) ** 2)
fresh_loss.backward()
after_reset_backward = w.grad.detach().clone()

def rounded(values: torch.Tensor | np.ndarray) -> list[float]:
    return [round(float(value), 3) for value in values]

print("loss", round(float(loss.item()), 6))
print("manual", rounded(manual))
print("first_backward", rounded(first_backward))
print("agree", np.allclose(manual, first_backward.numpy(), atol=1e-6))
print("after_second_backward", rounded(after_second_backward))
print("after_reset_backward", rounded(after_reset_backward))
Output
loss 3.268751
manual [-9.9, -3.925]
first_backward [-9.9, -3.925]
agree True
after_second_backward [-19.8, -7.85]
after_reset_backward [-9.9, -3.925]

The second backward pass doubles the stored numbers because it adds the same gradient again. A training loop normally calls optimizer.zero_grad() before the next batch, or sets buffers to None, unless accumulating across several batches is intentional.[8]

Adding contributions within one backward pass follows the chain rule. Keeping contributions across separate backward calls is a framework accumulation policy. Both use addition, but clearing .grad controls only the second behavior. Also notice that backward() hasn't updated w: an optimizer or explicit subtraction must apply the gradient afterward.

What if the output isn't one number?

Our loss gave backward one starting signal, dL/dL = 1. The four predictions by themselves don't specify a scalar objective. To differentiate several outputs, tell backward how much sensitivity arrives at each one. That list is called the backward seed.

Use requests 0 and 2 from the table: both have queue delay 0.5, but prompt blocks are 1 and 3. Suppose we want the derivative of their prediction gap, prediction[1] - prediction[0]. Its seed is [-1, +1]. The queue-weight contributions cancel, and the prompt-weight contributions add to -1*1 + 1*3 = 2:

seed-a-vector-backward-pass.py
import torch

X = torch.tensor([[1.0, 0.5], [3.0, 0.5]], dtype=torch.float64)
weights = torch.tensor([1.3, 1.1], dtype=torch.float64, requires_grad=True)
predictions = X @ weights
seed = torch.tensor([-1.0, 1.0], dtype=torch.float64)
predictions.backward(seed)

print("predictions", [round(float(p), 3) for p in predictions.detach()])
print("prediction_gap", round(float((predictions[1] - predictions[0]).detach()), 3))
print("seeded_gradient", weights.grad.tolist())
print("manual", (X.T @ seed).tolist())
assert torch.allclose(weights.grad, X.T @ seed)
Output
predictions [1.85, 4.45]
prediction_gap 2.6
seeded_gradient [2.0, 0.0]
manual [2.0, 0.0]

The Jacobian is the table of each output's partial derivatives with respect to each input parameter. Here its two rows are [1, 0.5] and [3, 0.5], so its shape is (2 outputs, 2 weights). Backward combines those rows using the seed, yielding one derivative per weight. This is a vector-Jacobian product, written as vTJv^T J with row vectors, or JTvJ^T v with column vectors. Autograd computes that product without needing to materialize a full Jacobian.[7][9]

Changing the seed changes the quantity differentiated. [1, 1] would differentiate the sum of predictions and produce [4, 1], not the gap's [2, 0]. For our four-row mean squared loss, the arriving prediction seed is 2 * errors / 4, which explains the same X.T @ seed formula used earlier. Calling backward() on a tensor with several elements requires this explicit seed or a prior reduction to a scalar.

In-place writes can break the backward pass

PyTorch tracks operations on tensors to construct the computation graph for backpropagation. In-place writes can invalidate saved values or mutate aliased tensors.[4]

To compute some derivatives, autograd saves forward-pass values. Overwriting a saved value before .backward() can make that derivative impossible to compute correctly, so PyTorch checks for changes and raises when the backward pass needs an invalidated value. Not every in-place operation is forbidden; the dependency on the overwritten value determines the problem.

Views are a second trap, similar to the shared buffers in NumPy. Expanding a size-1 axis with .expand() uses an of 0, so several output indices refer to the same memory location. Updating all those entries at once with add_() would write repeatedly to that location. PyTorch rejects this overlapping write; use .clone() first when you need independent entries.[10]

The two cells below isolate storage behavior rather than fit latency. The first uses sigmoid, a smooth function that squashes a number into the range (0, 1). If its output is b, its derivative is b * (1 - b), so backward needs the original b. Changing that output breaks the dependency. The second cell writes through a view whose entries share storage. Read each error as a diagnosis of a different problem.

pytorch-autograd-traps.py
import torch

# Trap 1: Mutating a variable needed for backward in-place
a = torch.tensor([2.0], requires_grad=True)
b = torch.sigmoid(a)
c = b * 3
b.add_(1.0)  # Mutates b in-place
try:
    c.backward()
except RuntimeError as error:
    print("Trap 1 error:", str(error).split("\n")[0])
else:
    raise AssertionError("Expected backward to reject the overwritten sigmoid output")

# Trap 2: Mutating an expanded view sharing storage in-place
x = torch.tensor([1.0], requires_grad=True)
y = x * 2
z = y.expand(3)  # Expanded view (stride 0)
try:
    z.add_(1.0)  # Try to mutate expanded view in-place
except RuntimeError as error:
    print("Trap 2 error:", str(error).split("\n")[0])
else:
    raise AssertionError("Expected the overlapping in-place write to fail")
Output
Trap 1 error: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1]], which is output 0 of Sigmoid, is at version 1; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True, check_nan=False).
Trap 2 error: unsupported operation: more than one element of the written-to tensor refers to a single memory location. Please clone() the tensor before performing the operation.

Test the loop on another request

Try one more latency row before opening the solution sketches: x_prompt = 3, x_queue = 1, actual seconds y = 7, and weights [w_prompt, w_queue] = [2, 2].

  1. Compute prediction, error, and squared loss.
  2. Compute both partial derivatives by multiplying 2 * error by the relevant feature.
  3. Apply one gradient-descent update with learning rate 0.05. What is the new prediction and loss?
  4. If a centered finite-difference check for the prompt parameter returns 5.99999 while your manual derivative says 6, is that evidence of a bug?
  5. If loss becomes NaN after a very large update, what two things should you inspect first?

Solution sketches

  1. Prediction is 2*3 + 2*1 = 8, error is 8 - 7 = 1, and loss is 1.
  2. The prompt partial is 2*1*3 = 6; the queue partial is 2*1*1 = 2. The gradient is [6, 2].
  3. New weights are [2 - 0.05*6, 2 - 0.05*2] = [1.7, 1.9]. Prediction is 1.7*3 + 1.9 = 7.0, so loss is 0.
  4. Not by itself. The absolute difference is about 0.00001, or roughly 0.00017% of the derivative. Compare it with your chosen absolute and relative tolerances; matching signs alone would not be enough. The NumPy check above accepts this difference with its default relative tolerance.
  5. Inspect the learning rate and the scale or validity of input values. A huge step or non-finite data can make the forward loss invalid before a later optimizer can help.

Keep evidence from each update

When a training step misbehaves, leave a trail you can inspect:

StepQuestion you can answer
Forward passWhat prediction did current parameters make?
LossHow wrong was that prediction?
DerivativeHow would loss respond to a small change in one parameter?
GradientWhat is the correction signal for every parameter?
BackpropagationHow do local derivatives carry that signal backward?
Gradient checkHow do I test a handwritten backward formula?

For this small deterministic model, verify four things: the forward loss is finite, the gradient has the expected sign and shape, a centered finite difference agrees within tolerance at smooth points, and one small step lowers loss on the checked example. Numerical checks belong on small debugging cases; they aren't a requirement to perturb every parameter during ordinary training.

If those checks pass but training diverges, inspect the learning rate and input scale before rewriting backpropagation. Practice the sequence on the latency table: compute one finite-difference check, then confirm that the step lowers loss.

Momentum, Adam, and learning-rate schedules control how to use gradients over many steps. Those tools wait until vectors, , and deeper linear algebra are in place. Next you need names for the lists of corrections you computed.

Complete the lesson

Mastery Check

Answer every question, then check your score. Score 75% or higher to mark this lesson complete.

1.A scalar predictor uses x = 3, w = 2 and target 8 with squared loss. After one gradient-descent step at rate 0.05, which gradient, new weight and new loss follow?
2.For a nonzero scalar feature x = 4 and squared error, a fixed rate 0.08 is used repeatedly from a nonoptimal weight. Which exact-arithmetic distance multiplier and behavior follow?
3.Two prediction rows are [2,1] and [5,1]. Backward receives seed [−2,+1]. Which selected scalar combination and parameter gradient follow?
4.Three graph routes reach one stored parameter with loss-gradient contributions +7, −4 and +1. No mean reduction appears in the loss. Which gradient belongs in that parameter’s buffer?
5.A mean-squared-error batch has X rows [1,2], [2,0], [0,3], targets [3,4,6] and zero weights. Which two-entry gradient follows?
6.For L(w) = (ReLU(4w) − 3)² at w = 0, which left derivative, right derivative and PyTorch backward value are correct?
7.A leaf gradient starts empty. Two freshly recomputed identical losses each contribute [−3,5], then the buffer is set to None and one fresh identical backward runs. What are the buffer values after the second and third passes?
8.Case A computes b = sigmoid(a), c = 5b, overwrites b with add_(1), then differentiates c. Case B clones six expanded copies of 3x, adds 1 to the clone and differentiates their sum. a and x are scalar leaves requiring gradients. Which result pair follows?
9.A hidden-unit model is h = w1×x+b1, z = ReLU(h), prediction = w2×z+b2, with squared loss. Use x=2, target5, w1=1, b1=0.5, w2=1.5 and b2=0.25. Which gradient in order [w1,b1,w2,b2] follows?
10.A linear model has rows [1,1] and [2,2] with targets 3 and 6. Compare weights [1,2] with [2,1]. Which training losses and identification conclusion follow?

10 questions remaining.

Discussion

Questions and insights from fellow learners.

Discussion loads when you reach this section.