Vision Transformer, Part 6 — How It Learns: Logits → Loss → Backprop → Update

Abstract. The forward pass turns pixels into a prediction; this final part runs the same path backwards — how a single number, the loss, reshapes all 1.2 million weights. It builds cross-entropy loss, backpropagation, the chain rule, and the AdamW weight update from first principles, showing with real numbers exactly how the loss reaches back through the network and decides which way every weight moves. By the end, training is no longer a black box.

In this 6-part series of vision transformer:

1. Part 1 — Embeddings: from pixels to a token sequence
2. Part 2 — Attention: how tokens relate to each other (single head)
3. Part 3 — The full block: multi-head, LayerNorm, residuals, MLP, and the prediction
4. Part 4 — Building the ViT in PyTorch, trained on CIFAR-10
5. Part 5 — Do overlapping patches help? A controlled experiment
6. Part 6 — How it learns: loss, backprop, the chain rule, the weight update — you are here

Part 6 of the series. Parts 1–3 explained the mechanics (embeddings, attention, the full block); Part 4 built and trained the model; Part 5 ran the overlap experiment. Here we answer the question underneath all of them: where do the trained weights come from?

Every “real number” shown across this series — the learned positional grid, the attention maps, γ ≈ 0.64, the horse logit of 3.80 — is the output of the loop in this article, run ~19,500 times.

The shape of learning

The forward pass ended at 10 logits. Learning is four steps that turn those logits into a better model:

logits → LOSS → BACKPROP → CHAIN RULE → UPDATE → (repeat)
  6.1     6.1      6.3         6.4         6.5      6.6

Everything hinges on one fact: every operation from pixels to loss is differentiable. That unbroken chain is what lets a single number at the end reach back and adjust a weight at the very beginning.

6.1 · From logits to the loss

Figure 1: Logits to Loss Significance

First, turn the logits into probabilities with softmax (the same softmax from Part 1, now over the 10 classes instead of over tokens). For our horse image: horse 88%, everything else tiny.

Then compare to the one-hot label with cross-entropy:

Because the target y is all zeros except a single 1 on the true class, every term drops out except one, and it collapses to:

For our horse: L = −log(0.88) = 0.13.

Why −log is the right shape — look at the curve:

Being right costs almost nothing; being confidently wrong is punished savagely. That asymmetry is exactly the training signal you want — it pushes hardest where the model is most wrong.

Footnote: our model trained with label smoothing 0.1, so the real target was slightly softened (0.9 on the true class, a sliver on the rest) to discourage over-confidence. The −log(p_true) form above is plain cross-entropy; label smoothing is a small modification on top.

6.2 · One loss per image, averaged over the batch

Two facts you already met, now doing their job:

• Only the CLS token is supervised. The loss is computed from the CLS logits alone. The 64 patch tokens never get their own target — they are trained only indirectly, through gradients that flow back into them because CLS attended to them.
• One image → one scalar loss. In a batch of 128 images, the 128 scalars are averaged into a single number, which produces one gradient and one weight update. So each step learns from 128 images at once.

6.3 · Backpropagation

Figure 3: Backpropagation Flow

The starting gradient is p − y

For softmax + cross-entropy, the gradient of the loss at each logit is exactly predicted minus target:

No messy derivative — it just falls out. On our horse:

• horse: 0.88 − 1 = −0.12 → negative → "raise this logit"
• every other class: p − 0 > 0 → positive → "lower these logits"

And the gradients sum to zero — the model can only redistribute certainty, so pushing horse up necessarily means pushing the others down.

Then δ flows backward, and grad = δ · x everywhere

That output error δ propagates right-to-left through the network. At every layer, two things happen:

• each weight’s gradient = δ · x — the error arriving from above times the input that weight saw during the forward pass
• δ is transformed and passed further back (as Wᵀδ) to become the error signal for the layer before it

One backward pass produces a gradient for all ~1.2M parameters. (This is why the forward pass had to stash every activation: backprop needs each layer’s x to form δ · x. That stored activation memory is the reason training uses far more memory than inference.)

6.4 · The chain rule — how the loss decides each weight’s direction

Figure 4.1: Chain rule — calculation of weight update direction

The gradient is a product of local slopes

For any weight, the chain rule says its gradient is the product of the local slope of every operation on the path from that weight to the loss:

The chain always starts at the loss (∂L/∂logit = p − y) and multiplies backward. For a weight in a linear layer, the last factor ∂(act)/∂w is just that weight's input x — which is exactly why the whole thing collapses to the tidy grad = δ · x you saw in 6.3.

The sign decides the direction

Here it is on a real head weight — the one connecting CLS feature 48 to the horse logit:

∂L/∂logit(horse) = p − y   =  0.88 − 1        = −0.12     (from 6.3)
∂logit/∂w        = CLS[48]  =            +1.48            (the forward input)
∂L/∂w = δ × input           = (−0.12)(+1.48) = −0.178

The gradient is negative. The update steps opposite the gradient (w ← w − lr·∂L/∂w), so the weight increases. And increasing this weight raises the horse logit next time (its input was positive) → the loss goes down.

Get Chandrashekhar Padole’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

The rule is simple and always loss-reducing:

So the loss at the output literally reaches back through the chain and tells every one of the 1.2M weights which way to move.

The full sign table

Three signs decide everything for a weight: the error δ (is the output pushing the loss up or down?), the input x, and the increment Δw we apply. Here are all 8 combinations:

Figure 4.2: Diffrent scanarios of Weight update direction

Read it and the whole thing clicks:

• The loss impact of an increment is ΔL ≈ δ · Δw · x — the product of the three signs. Flip any one sign and the loss impact flips.
• Gradient descent doesn’t search for the right direction. It sets Δw = −lr·(δ·x), which automatically lands on the four green-bordered rows — exactly the four where the loss goes down.
• So “increase or decrease?” is never a guess. It falls straight out of the signs of δ (from the loss) and x (the forward input).

6.5 · The weight update

Figure 5: Significance of Weight and gradient

Step downhill

Each weight moves opposite its gradient, down the loss surface:

The learning rate sets the step size — too large overshoots the valley, too small crawls.

What AdamW adds

Plain gradient descent is twitchy. AdamW — what we actually use — refines every step three ways:

1. Momentum — average recent gradients, so the step follows a steady direction instead of jittering on each noisy batch.
2. Per-parameter rate — divide each weight’s step by its own gradient magnitude, so consistently-large-gradient weights take smaller, safer steps.
3. Decoupled weight decay — gently shrink every weight toward 0 each step (regularization). This is the “W” in AdamW, and applying decay separately from the gradient step is what makes it generalize better than plain Adam.

Our settings: lr 3e-4 (cosine-decayed), weight decay 0.05, gradient clip 1.0. This update runs once per batch, on all ~1.2M parameters at once.

6.6 · The loop closes

for each batch:  forward → loss → backward → AdamW update      (5.1 → 5.3 → 5.4)
    × 390 batches                 = 1 epoch
    × 50 epochs                   ≈ 19,500 updates

One update barely moves anything. But over ~19,500 of them:

• the loss falls from ~2.3 (random guessing is −log(1/10) = 2.3) to ~0.8
• validation accuracy climbs to 82.68%
• the cosine schedule shrinks the learning rate along the way, so early steps explore boldly and late steps fine-tune

The gradient is not the difference between two iterations' losses — a common misconception. It is the analytic slope of the loss with respect to each weight, computed fresh each step by backprop from a single batch.

The full circle

FORWARD  (Parts 1–4):  pixels → … → logits → loss        "read the trained model"
BACKWARD (Part 5):     loss → grad = δ·x → AdamW update   "make the trained model"
                              └────────── repeat ──────────┘

That is the whole of training. Everything we dissected in the forward-pass parts — every learned weight, every attention pattern, every γ and β — is what this loop produces. We spent Parts 1–4 reading those numbers; Part 5 is where they come from.

You can now trace any value in the model in both directions: what it computes (forward) and how it was learned (backward).

Code and full experiment notes were built from scratch. The point was never to beat a benchmark; it was to understand every number — in both directions.