Skip to Content

The training loop of every neural network: compute the loss at the output, push the error backward through the layers using the chain rule, and step the weights in the negative-gradient direction.

Data Flow

   forward pass
      x --> L1 --> L2 --> ... --> Lk --> loss
                                          |
                                          v
                                   backward pass
                                          |
              g = d(loss)/d(W)            (chain rule)
                                          |
                                          v
                          W = W - lr * g

What it is

Variants: SGD, Momentum, Adam (most common), AdamW (with weight decay), Lion.

Code

import torch
 
x  = torch.randn(4, 3)
y  = torch.tensor([1., 0., 0., 1.])
m  = torch.nn.Linear(3, 2)
opt = torch.optim.AdamW(m.parameters(), lr=1e-3)
 
for step in range(1000):
    pred = m(x)
    loss = torch.nn.functional.cross_entropy(pred, y)
    opt.zero_grad()
    loss.backward()                 # backprop
    opt.step()                      # gradient descent

Pitfalls

Analogy

Hiking down a foggy hill: every step feel the slope, step the opposite way, repeat until the valley.

Interview tip: Always state both directions: forward (compute output), backward (chain rule to compute gradients).

Advertisement