Gradients have shapes — Linear algebra for LLMs

One slope per parameter

\displaystyle \frac{\partial L}{\partial W} = \left[\frac{\partial L}{\partial W_{ij}}\right] \quad \text{same shape as } W

The loss $L$ is one number. Nudge one entry $W_{ij}$ and $L$ changes at some rate: that rate is entry $(i, j)$ of the gradient. So the gradient of a $4096 \times 4096$ weight is another $4096 \times 4096$ matrix.

Write

\partial L/\partial W

has the shape of

W

: one slope per entry.

The gradient of a dot product

\displaystyle L = w \cdot x = w_1 x_1 + w_2 x_2 + w_3 x_3

\displaystyle \frac{\partial L}{\partial w} = [x_1, x_2, x_3] = x

Each $w_i$ appears once, multiplied by $x_i$ , so its slope is $x_i$ . With $x = [1, 2, 3]$ , the gradient is $[1, 2, 3]$ : the weight paired with the biggest input matters most.

Write

\nabla_w (w \cdot x) = x

.

And of a squared length

\displaystyle L = \tfrac12 |w|^2 = \tfrac12 \textstyle\sum_i w_i^2 \;\Rightarrow\; \nabla_w L = w

Each term $\tfrac12 w_i^2$ has slope $w_i$ . This is weight decay’s pull: it shrinks every weight in proportion to its size.

Gradient descent

\displaystyle w \leftarrow w - \eta\, \nabla_w L

−

0.1

=

The gradient points uphill, so step the other way, scaled by the learning rate $\eta$ . With $L = w \cdot x$ : before the step $L = 6$ , after it $L = 0.9 + 1.6 + 2.1 = 4.6$ .

Write

w \leftarrow w - \eta\nabla L

. Here

L

:

6 \to 4.6

.

✏️ Your turn

One step

\displaystyle L = w \cdot x, \quad x = [2, -1], \quad w = [0, 0], \quad \eta = 0.5

Write the gradient and $w$ after one step.

Work it out on paper, then go to the next card to check.

Answer

\displaystyle \nabla_w L = [2, -1]

\displaystyle w = [0, 0] - 0.5\,[2, -1] = [-1, 0.5]

$L$ goes from 0 to $-2 - 0.5 = -2.5$ : down, as it should.

In PyTorch

loss.backward()
w.grad.shape == w.shape      # True: one slope per parameter
with torch.no_grad():
    w -= lr * w.grad

📌 Recap

Gradients have shapes in one card

$\partial L / \partial W$ has $W$ ’s shape.
$\nabla_w(w \cdot x) = x$ ; $\nabla_w \tfrac12|w|^2 = w$ .
Descent: $w \leftarrow w - \eta\nabla L$ .

Write

The three lines above.