Gradient descent from scratch, with warmup and cosine learning-rate schedules

In this article

1. Before you start
2. Why interviewers ask this
3. Gradient descent from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the PyTorch and machine learning (ML) coding round guide, where coding questions on training loops are collected.

Why interviewers ask this

Implementing gradient descent is a short coding task that exposes whether a candidate understands training or only calls a library. Candidates report being asked to write a naive gradient descent step, to code an optimiser for a simple function, and to write a learning-rate decay function, in AI and big-tech loops. These are reported by candidates, not published rubrics.

A weak answer writes w -= lr * grad and cannot say what happens when lr is wrong. A strong answer states the update rule, shows that the learning rate controls convergence, and names two or three schedules with the reason each exists. A Staff-level answer adds the order of calls in a training loop, why large batches need a larger rate with a warmup, and which parts of the state a scheduler touches and which it does not: a scheduler changes the learning rate only, not the weights or the optimiser's momentum or moment buffers.

Gradient descent from first principles

Learning rate

The learning rate is a positive number that multiplies the gradient to set how far each update moves the parameters; too small makes training slow, too large makes it overshoot or diverge.

Learning-rate schedule

A learning-rate schedule is a rule that sets the learning rate as a function of the training step or epoch, so the rate can start small, hold, and then shrink as training proceeds.

In plain words

Picture standing on a hillside in thick fog. You cannot see the valley, but you can feel which way the ground slopes under your feet. You take a step downhill, feel again, and repeat. The learning rate is the length of your stride. The analogy breaks in one place. A real hill has two directions. A model can have millions of parameters, so "downhill" is a direction in a space with millions of dimensions.

In one precise sentence, gradient descent updates the parameters as θ ← θ − η·∇L(θ), where η is the learning rate and ∇L is the gradient, the list of slopes of the loss with respect to each parameter.

Step 1: Measure the slope with the gradient

The gradient is the list of partial derivatives of the loss, one per parameter, and it points in the direction in which the loss rises fastest. For the bowl $f(x)=x^2$ the derivative is 2x, so at $x=5$ the slope is 10, and the loss rises as you move right.¹

Input: the current parameters. Output: one slope per parameter. In a real model the framework computes this by backpropagation (worked by hand in the forward and backward passes of a small network), but the update rule does not care how the slope was found, which is why a from-scratch exercise often hands you a function and its derivative. You now know what to subtract; the next step is how much.

Step 2: Take a step against the gradient, scaled by the learning rate

The update subtracts the gradient times the learning rate, because the gradient points uphill and you want to go down. At $x=5$ with a learning rate of 0.1, the step is 0.1 times 10, which is 1, and the new position is 4.¹

Input: the parameters, the gradient and the rate. Output: new parameters. Written as code the step is one line per parameter, and repeating it is the whole algorithm. Batch, minibatch and stochastic variants differ only in how many examples are used to estimate the gradient. Stochastic gradient descent (SGD) is the version that estimates the gradient from a random sample of the data, so "stochastic" describes the gradient estimate, not the weights.

Step 3: The learning rate decides whether the loop converges

On $f(x)=x^2$ the update simplifies to $x\leftarrow (1-2\eta )x$, so after $t$ steps the position is $(1-2\eta {)}^t$ times the start. The factor (1 − 2η) decides everything. A rate between 0 and 0.5 shrinks $x$ smoothly; at exactly 0.5 the factor is 0 and the loop lands on the minimum in one step; between 0.5 and 1 the factor is negative, so $x$ jumps across the minimum each step but still shrinks; at 1 it bounces between 5 and −5 forever; above 1 it grows.¹

This is the reason "a smaller rate is always safer" is wrong in practice: it is safe from divergence but slow. With a rate of 0.1 the loop needs 28 steps to get within 0.01 of the minimum from a start of 5.¹ Real losses are not this simple bowl, so the safe rate depends on the curvature of the loss: on a quadratic, gradient descent converges only if , where $L$ is the largest curvature. The bowl $x^2$ has $L=2$, which is why its threshold is 1. The same three regimes appear.

Step 4: A schedule changes the rate during training

A schedule helps because the best rate changes over a run: a larger rate makes fast early progress, and a smaller one settles near a minimum without bouncing around it. Three common shapes are step decay, which multiplies the rate by a factor every few epochs; exponential decay, which multiplies it by a factor every step; and cosine annealing, which follows half a cosine wave from a maximum down to a minimum. The cosine form is ${\eta }_t={\eta }_{\min }+\frac{1}{2}({\eta }_{\max }-{\eta }_{\min })(1+\cos (\pi t/T))$.

Loshchilov and Hutter (2016) propose a warm restart technique for SGD, in which the rate is periodically raised again to improve its anytime performance. A restart resets the learning rate only; the weights are kept. Warmup is a short phase at the start of training in which the rate rises from a small value to its target. Goyal et al. (2017) adopt a linear scaling rule, which adjusts the learning rate as a function of the minibatch size, and report no loss of accuracy with minibatches up to 8192 images on ImageNet. Under that rule, a minibatch 32 times larger than a base of 256 means a rate 32 times larger: from an illustrative base of 0.1 to 3.2.¹

1 · Slope2 · Step3 · Rate check4 · ScheduleReplay

Step-by-step description

1. Measure the gradient of the loss.
2. Subtract rate times gradient.
3. The rate decides convergence.
4. Change the rate over training.

• The gradient is the slope of the loss for each parameter.
• The update subtracts the learning rate times the gradient.
• A schedule changes the learning rate during training.

Worked example

The loop below minimises $f(x)=x^2$ from a start of 5. It is the whole of gradient descent from scratch, and every number after it comes from running it.¹

def gradient_descent(grad, x0, lr, steps):
    x = x0
    path = [x]
    for _ in range(steps):
        x = x - lr * grad(x)   # step against the gradient
        path.append(x)
    return path

path = gradient_descent(lambda x: 2 * x, 5.0, 0.1, 5)
# rounded: [5.0, 4.0, 3.2, 2.56, 2.048, 1.6384]

With a rate of 0.1 the path is 5, 4, 3.2, 2.56, 2.048 and 1.6384, shrinking by a factor of 0.8 each step. With 0.5 it is at 0 after one step. With 0.9 the factor is −0.8, so the path is 5, −4, 3.2, −2.56, 2.048 and −1.6384: the loop oscillates but still converges. With 1.1 the factor is −1.2 and the path is 5, −6, 7.2, −8.64, 10.368 and −12.4416; after 20 steps $x$ has reached about 191.69 in size.¹

Text description of the diagram

• At rate 0.1 the position after 5 steps is 1.6384.
• At rate 1.1 the position after 5 steps is minus 12.4416.
• A rate of exactly 1 bounces between 5 and minus 5 forever.

The schedules can be written the same way. For an illustrative maximum rate of 0.1 over ten steps, step decay that halves every three steps gives 0.1, 0.1, 0.1, 0.05, 0.05, 0.05, 0.025, 0.025, 0.025 and 0.0125. Exponential decay with a factor of 0.9 gives 0.1, 0.09, 0.081 and then 0.0729, falling to 0.0387 at the tenth step. Cosine annealing from 0.1 to 0 with $T=10$ gives 0.1, 0.0976, 0.0905, 0.0794, 0.0655, 0.05, 0.0345, 0.0206, 0.0095, 0.0024 and 0 at step 10.¹ A warmup of three steps followed by cosine gives 0.0333, 0.0667, 0.1, then a decay to 0.005 at the tenth step.¹

A second run fits the line $y=2x+1$ to the four illustrative points (0, 1), (1, 3), (2, 5) and (3, 7) with mean squared error (MSE) as the loss. At a rate of 0.05, 200 steps give a slope of 2.0001 and an intercept of 0.9999; at 0.3 the same loop diverges, with a loss above 410 million after 20 steps.¹

Text description of the diagram

• Step decay gives 0.1 for three steps, then 0.05, then 0.025.
• Exponential decay with factor 0.9 reaches 0.0387 at the tenth step.
• Cosine annealing from 0.1 to 0 over T equal to 10 reaches 0.0024 at step 9.

Trade-offs and edge cases

A schedule trades early speed for late precision, and the right shape depends on the run. Step decay is easy to explain and tune but drops abruptly. Exponential decay is smooth but can shrink the rate too far if its factor is aggressive. Cosine annealing needs the total number of steps in advance. Restarts, in the spirit of stochastic gradient descent with warm restarts, trade a temporary loss increase for a chance to leave a poor region.

Edge cases that interviewers probe:

• Convexity. With a small enough rate, gradient descent reaches the global minimum for convex losses such as the bowl above; for a neural network it finds a stationary point, and the result depends on the start.
• Adam and schedules. Adaptive methods such as Adam (Kingma and Ba, 2014) still take a base learning rate. The original Transformer trained with Adam, a warmup of 4,000 steps and then inverse-square-root decay (Vaswani et al., 2017, §5.3), so schedules and warmup are still used with adaptive methods.
• Order of calls. In PyTorch the usual loop calls the optimiser's step first and the scheduler's step after it, as the "How to adjust learning rate" section of the torch.optim documentation describes (read October 2026; check your version).
• Large batches. A larger batch gives a less noisy gradient, so a larger rate can work, which is the logic behind the linear scaling rule of Goyal et al. (2017) and behind distributed setups such as fully sharded data parallel training.
• Numerical problems. A loss that turns into nan usually signals a rate that is too high or an unstable operation; the numerically stable softmax page covers one common cause.

How to answer it in an interview

The spoken structure: state the update rule, write the loop, explain what the rate controls, then name a schedule and the reason for it. If you are building a model from the same parts, the generative pre-trained transformer (GPT) from scratch tutorial shows where this loop sits in a full training script.

Common mistakes

• "A smaller learning rate is always safer." (wrong) It avoids divergence but can make training very slow or stall on a plateau. (right) "A smaller rate is stable but slow, so I tune it on a log scale."
• "Gradient descent finds the global minimum." (wrong) It reaches the global minimum for convex losses with a small enough rate, not in general. (right) "It converges to a minimum or stationary point, and for deep networks the start matters."
• "Stochastic means random weights." (wrong) The weights are updated deterministically from the gradient. (right) "Stochastic means the gradient is estimated from a random sample of the data."
• "Adam needs no schedule." (wrong) Adam still has a base rate, and warmup and decay are still common with it. (right) "Adam adapts per-parameter step sizes, but I still schedule the base rate."
• "A warm restart resets the model." (wrong) A restart raises the learning rate again and keeps the weights. (right) "Only the rate jumps back up; training continues from the same weights."

Follow-up questions

How do you implement a single gradient descent step?

Compute the gradient of the loss with respect to each parameter, multiply it by the learning rate, and subtract the result from the parameter. In a vector setting do it for the whole array at once, for example w = w - lr * grad.

How do you write a learning-rate decay function?

Write a function that takes the step number and returns the rate. For exponential decay it is base * gamma ** step; for step decay it is base * gamma ** (step // period); for cosine annealing it is the formula in Step 4. Test it by printing the first ten values, as in the worked example.

What happens if the learning rate is too large?

The update overshoots the minimum. On the bowl in this page the position changes sign and grows by a factor of 1.2 per step at a rate of 1.1, and for the line fit a rate of 0.3 sends the loss above 410 million within 20 steps.¹ In a real run the loss climbs or turns into nan.

How would you add momentum to gradient descent?

Keep a velocity for each parameter: $v\leftarrow \beta v+\mathrm{\nabla }L$, then $\theta \leftarrow \theta -\eta v$, with $\beta$ commonly around 0.9. The velocity damps oscillation across steep directions and builds speed along shallow ones; Sutskever et al. (2013) show its importance for training deep networks.

Does Adam need a learning-rate schedule?

Adam adapts the step size per parameter, but it still has a base learning rate that you choose, and the original Transformer paper pairs it with warmup and decay (Vaswani et al., 2017). Treat the schedule as an extra knob worth trying, not as something Adam makes unnecessary.

What does stochastic mean in stochastic gradient descent?

It means the gradient is estimated from a random sample of the data instead of the full dataset. The update rule is unchanged; only the quality and cost of the gradient estimate change. See the ML coding round guide for neighbouring from-scratch questions, such as multi-head attention in PyTorch and 2D convolution.

Sources

1. SGDR: Stochastic Gradient Descent with Warm Restarts (arXiv 1608.03983v5) — arXiv (Loshchilov and Hutter), 13 Aug 2016
2. Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour (arXiv 1706.02677v2) — arXiv (Goyal et al.), 8 Jun 2017
3. Adam: A Method for Stochastic Optimization (arXiv 1412.6980) — arXiv (Kingma and Ba), 22 Dec 2014
4. Attention Is All You Need (arXiv 1706.03762v7) — arXiv (Vaswani et al.), 12 Jun 2017
5. On the importance of initialization and momentum in deep learning — ICML 2013, PMLR (Sutskever et al.), 16 Jun 2013
6. How to adjust learning rate (torch.optim documentation) — PyTorch documentation, 26 Oct 2026