Attention as a Differentiable Map

A weighted combination of values

$v_j$

stored value row $j$

$p_{ij}$

how much query $i$ uses item $j$

$o_i$

one output row for query $i$

$T_k$

how many items are stored

Each $p_{ij}$ is computed from the inputs, not stored, and $o_i$ lands in the convex hull of the value rows: between the smallest and the largest when the values are scalars.

Attention as a Differentiable Map

Scores from inner products

$Q$

query rows, $T_q\times d_k$

$K$

key rows, $T_k\times d_k$

$S$

scores, $T_q\times T_k$

$\sqrt{d_k}$

a fixed constant, $2$ at $d_k=4$

The constant is part of the operation; its derivative contribution is required here, its justification waits.

Attention as a Differentiable Map

Positive weights that sum to one

$P_{ij}>0$

positive, and row $i$ sums to $1$

$m_i$

the row maximum, so every exponential is at most $1$

row $i$

handled on its own

no mask

every key open to every query

A later mask must leave one valid key in each row, and needs its own derivative policy.

Vote

Attention as a Differentiable Map

Which score change leaves weights unchanged?

A. The shared factor cancels above and below, so with $s=(0,\ln 3)$ the row is $(1/4,3/4)$ for every $c$. B would make it add to $e^{c}$.

Attention as a Differentiable Map

The attention output

$P$

weights, $T_q\times T_k$

$V$

value rows, $T_k\times d_v$

$O$

one output row per query

$d_v$

how wide a value row is

Attention as a Differentiable Map

Backward through the value mixture

$G_O$

arrives from the loss, $T_q\times d_v$

$G_V$

$V$'s shape, $T_k\times d_v$

$G_P$

$P$'s shape, $T_q\times T_k$

$P,V$

saved by the forward pass

One multiplication, two inputs, two gradients. Each keeps its own input's shape, and that is what catches a dropped transpose.

Vote

Attention as a Differentiable Map

Credit to values and weights

A. Derivative in $v_j$ is its weight; in $p_j$ its value. Centred differences in each argument reproduce all four.

Attention as a Differentiable Map

A softmax output depends on all scores

$Z$

$\sum_l e^{s_l}$

$\delta_{ij}$

$1$ if $i=j$, else $0$

$p_i$

$e^{s_i}/Z$

$\partial p_i/\partial s_j$

$T_k\ge2$: nonzero for every $(i,j)$

Quotient rule: own numerator only at $i=j$, shared denominator at every $j$. No weight has a private score.

Attention as a Differentiable Map

Assemble the softmax Jacobian

$p$

one row's weights, as a column

$J_{ii}$

$p_i(1-p_i)>0$ if $T_k\ge2$

$J_{ij}$

$-p_ip_j<0$ off the diagonal

$J\mathbf1=\mathbf0$

rows: a shift changes nothing

$\mathbf1\T J=\mathbf0\T$

columns: the total stays $1$

Finite, unmasked scores, and one $p$ and one $J$ per score row. The two sums agree only because $J$ is symmetric.

Attention as a Differentiable Map

Apply the VJP without building a Jacobian

$g_p$

one row of $G_P$, as a column

$p\T g_p$

one number: the weighted mean

$\odot$

entry by entry

$g_s$

one row of $G_S$, same length

This is $J\T g_p$ rewritten, not approximated. Do not materialise a full attention Jacobian in code: the short form needs $T_k$ numbers, never $T_k^2$.

Vote

Attention as a Differentiable Map

A complete softmax backward example

B. $\tfrac14(-6)$ and $\tfrac34(2)$, which add to $0$: a shift common to all scores moves no weight, so $g_s$ has no part along $\mathbf1$.

Attention as a Differentiable Map

Backward through query-key scores

$G_S$

the score gradient, $T_q\times T_k$

$G_Q$

$Q$'s shape, $T_q\times d_k$

$G_K$

$K$'s shape, $T_k\times d_k$

$1/\sqrt{d_k}$

a constant, so it stays in both

$S$ is linear in $Q$ and in $K$, so each gradient is the other array contracted with $G_S$ and divided by the same constant.

Attention as a Differentiable Map

Shared inputs receive several contributions

$X$

one shared array of $T$ token rows

$W_\bullet$

output-by-input, as always here

$G_X$

$X$'s shape; a sum of three terms

one head

single-head self-attention, $T_q=T_k=T$

A variable used three times accumulates three gradients; drop two and a centred difference in $X$ disagrees.

Vote

Attention as a Differentiable Map

Reconstruct the complete reverse chain

A. $(2\times3)(3\times4)$ over $2$. C is off by exactly $2$. Verify by centred differences in all $35$ entries.