G14: Attention, embeddings, and transformer reasoning

Prerequisites and observable goals#

The central question: which values should this position combine?#

Attention constructs a data-dependent weighted combination of value vectors. It does not look up a guaranteed fact or produce a proof that a token is important. The original transformer paper defines scaled dot-product attention and multi-head composition. This chapter develops the operation on small original examples; the lab implements forward attention, not a pretrained model or a complete transformer training system. [G-ATTENTION]

Let $Q\in\R^{n\times d_k}$ contain query rows, $K\in\R^{m\times d_k}$ key rows, and $V\in\R^{m\times d_v}$ value rows. The number of queries $n$ need not equal the number of keys $m$. Key and value row counts must agree because every score column weights the corresponding value row. Compute

$S$ contains scores, $P$ contains nonnegative row-normalized weights, and $O$ contains output vectors. A row of $O$ need not sum to one. Its coordinates are values, not probabilities. The square root is a scale on dot-product scores, not a normalization of the final output.

A numerical example before the general architecture#

Take one query and two keys with $d_k=d_v=1$, $Q=[0]$, $K=(0,0)^T$, and $V=(2,8)^T$. The scores are $(0,0)$, so softmax gives $(1/2,1/2)$. The output is $(1/2)2+(1/2)8=5$. The computation did not choose one value; it mixed both.

If scores are $(\log3,0)$, exponentials are $(3,1)$, weights are $(3/4,1/4)$, and the same values produce $3.5$. Do not substitute the unnormalized scores directly into the weighted sum. Adding any constant to all scores in one row leaves softmax unchanged because the common exponential factor cancels. Subtract the largest allowed score before exponentiating to reduce overflow:

Here $A$ is the set of allowed columns. Set disallowed weights to zero. If $A$ is empty, there is no probability distribution to normalize; the reference rejects the row rather than silently returning NaNs or arbitrary zeros.

The scale $1/\sqrt{d_k}$ has an intuitive variance motivation: under independent zero-mean unit-variance query/key coordinates with suitable independence between vectors, the dot product has variance $d_k$. Dividing by $\sqrt{d_k}$ keeps that variance near one. Learned vectors need not satisfy those assumptions; this is a motivating model, not a theorem that all network scores have unit variance.

Mask before normalization#

In causal self-attention, query position $i$ may use key positions $j\le i$, not future positions. For two zero-score positions with scalar values $(2,8)$, the masked weight matrix and output are

The first position cannot see value 8. If you instead normalize the first row over both columns and then zero the future entry, the row becomes $(1/2,0)$ and outputs 1. The remaining weights no longer sum to one. Renormalization after masking could repair this simple case, but masking before softmax expresses the intended contract directly and avoids accidentally allowing future scores to influence the denominator.

The diagram uses zero-based positions increasing downward and rightward. A useful test changes only a future value or key and checks that an earlier causal output is unchanged. Such tests check noninterference for selected inputs; the indexing argument explains why the implementation should preserve it generally.

Differentiate the mixing operation#

For a single softmax row, direct differentiation of the quotient gives

The Jacobian is $\operatorname{diag}(p)-pp^T$. You rarely need to construct that full matrix. Given upstream sensitivities $u_j=\partial L/\partial p_j$,

For $p=(1/4,3/4)$ and $u=(1,3)$, the weighted mean is $5/2$, so the score gradient is $(-3/8,3/8)$. Its sum is zero. This matches softmax invariance to a common score shift: that direction should have zero first-order effect.

For matrices, let $D_O$ be the output sensitivity with the same shape as $O$. Then

Apply the row formula to $D_P$ to obtain $D_S$. For the unmasked score product,

Masked score coordinates have zero sensitivity under a fixed mask. Check dimensions: $D_S$ is $n$-by-$m$, so $D_SK$ is $n$-by-$d_k$. A transposition error may survive square test cases; use unequal query/key counts and unequal $d_k,d_v$ in tests. These derivative formulas are taught here; the supplied lab intentionally implements only the forward pass. A backward implementation is an unassigned continuation, not an implied completed feature.

What a transformer adds, and what it does not establish#

An embedding table maps an integer token ID to a learned vector. The lookup is not an ordinal statement that token 200 is twice token 100. Positional information distinguishes sequence order; self-attention without position information has permutation symmetries that may be inappropriate for a sequence task. Learned projections construct queries, keys, and values from hidden states.

Multiple heads use separate projections, perform several attention operations, concatenate outputs, and project again. Residual paths add compatible-shaped representations; normalization and a position-wise feed-forward subnetwork provide further transformations. A decoder additionally enforces causal information flow. Exact normalization placement, positional schemes, and other details vary by architecture; the original paper is one specific design, not a specification of every current model.

Next-token training compares a model’s distribution at one position with a target token using a classification loss. Training on many such examples does not turn a high-probability continuation into established truth. Evaluation needs held-out data, contamination controls, clear task definitions, and failure analysis. Attention weights alone are not a causal explanation of the final answer: value vectors and later transformations can alter the consequence of a weight, and alternative parameterizations may change apparent importance.

Count computation and memory before claiming acceleration#

A straightforward self-attention score array for batch $B$, heads $H$, and sequence length $N$ has $BHN^2$ entries. At $B=2,H=8,N=1024$ with four-byte entries, that is $67{,}108{,}864$ bytes, or 64 MiB. Doubling $N$ makes this one array 256 MiB. This excludes queries, keys, values, gradients, activations, optimizer state, and temporary buffers. Efficient attention algorithms can avoid materializing some full intermediates; the quadratic storage statement describes the straightforward implementation, not an unavoidable lower bound for every exact implementation.

The dot products and weighted sums require work on the order of $N^2d_k+N^2d_v$ per head in the direct method. During autoregressive inference, a key/value cache reuses earlier projected states rather than recomputing them for every new token. It trades memory for work; it does not remove the need for the new query to interact with the allowed history under ordinary dense attention. This CPU lab is for correctness and dimensions, not GPU throughput or an inference benchmark.

Independent practice#

F foundation C standard course work T transfer S optional stretch L runnable lab or project. Attempt the complete set before opening its solutions.

Smallest useful hints#

For .01, write the inner dimensions of every matrix product. For .03, the first causal row has one allowed element, so its allowed weight must be one. For .05, MiB uses $2^{20}$ bytes, not $10^6$.

Complete solutions: G14#

Compare the first step where your work differs. A corrected answer is not yet an independent success.

Open every solution

"""G14: single-head scaled dot-product attention, not a trained transformer."""
import math


def softmax(scores, allowed=None):
    s = list(scores)
    mask = [True]*len(s) if allowed is None else list(allowed)
    if not s or len(mask) != len(s) or any(type(x) is not bool for x in mask) or not any(mask):
        raise ValueError("nonempty row with at least one allowed position required")
    if any(not math.isfinite(x) for x in s):
        raise ValueError("finite pre-mask scores required")
    top = max(v for v,ok in zip(s,mask) if ok)
    e = [math.exp(v-top) if ok else 0. for v,ok in zip(s,mask)]
    total = math.fsum(e)
    return [v/total for v in e]


def attention(Q,K,V,causal=False):
    n,m = len(Q),len(K)
    d = len(Q[0]) if n else 0
    dv = len(V[0]) if V else 0
    if (type(causal) is not bool or not 1 <= n <= 256 or not 1 <= m <= 256
            or not 1 <= d <= 128 or not 1 <= dv <= 128 or len(V) != m):
        raise ValueError("bounded nonempty matching matrices required")
    if causal and n != m:
        raise ValueError("this causal lab requires equal query/key lengths")
    for matrix,width in ((Q,d),(K,d),(V,dv)):
        if any(len(row) != width or any(not math.isfinite(x) for x in row) for row in matrix):
            raise ValueError("ragged or nonfinite matrix")
    out, weights = [],[]
    for i,q in enumerate(Q):
        scores = [math.fsum(a*b for a,b in zip(q,k))/math.sqrt(d) for k in K]
        p = softmax(scores,[j <= i for j in range(m)] if causal else None)
        weights.append(p)
        out.append([math.fsum(p[j]*V[j][v] for j in range(m)) for v in range(dv)])
    return out,weights


if __name__ == "__main__":
    print(attention([[0.],[0.]],[[1.],[2.]],[[2.],[8.]],causal=True))

Readiness and continuation#

Explain $Q$, $K$, $V$, scores, weights, and outputs as six distinct objects. Derive a backward shape on unequal lengths and demonstrate a masking counterexample. A next project can add the taught backward formulas and finite-difference tests, then a complete tiny decoder with leakage-free next-token evaluation. That additional project is not counted as completed work here.