Audit a Single-Head Self-Attention Layer

Audit self-attention in dependency order instead of trusting a plausible final tensor. An executable trace checks projections, scores, masks, softmax rows, forbidden weights, weighted values, output width, finite values, and common axis failures.

An implementation is easiest to trust when it reports the same intermediate objects as the equations. A final output with the expected shape is not enough: the wrong softmax axis or a late mask can still produce a finite tensor.

The audit order is

X,W_Q,W_K,W_V \rightarrow Q,K,V \rightarrow S \rightarrow S+M \rightarrow A \rightarrow H \rightarrow Y.

At each arrow, check both shape and meaning.

State the Contract Before the Code

This implementation uses one sequence rather than a batch so every matrix can be printed. It accepts

X:T\times d_{model},\quad W_Q,W_K:d_{model}\times d_k,\quad W_V:d_{model}\times d_v,\quad W_O:d_v\times d_{model}.

It must return $Y:T\times d_{model}$ . Its causal weight matrix must be finite, have row sums near 1, and contain exact zeroes above the diagonal.

Audit one causal self-attention layer

Run the verified trace, then change one parameter. The checks expose incompatible shapes, non-finite values, invalid row sums, and forbidden future weights.

The expected final rows are approximately

h_4=[0.592,-1.279],\qquad y_4=[1.871,-2.262].

Small rounding differences are normal. A different shape, a nonzero future weight, or a large numerical difference is not.

Five Failures and Their First Checks

Symptom	Likely error	First local check
score shape is $(d_k,d_k)$	position and feature axes were swapped	inspect Q and the transpose of K
columns sum to 1 instead of rows	softmax used the query axis	print row sums and column sums
future entries remain positive	mask was omitted or applied late	inspect masked scores before softmax
output width is $d_v$	output projection is missing or reversed	check $W_O:(d_v,d_{model})$
NaN or infinity appears	unsafe softmax, invalid inputs, or empty allowed row	find the first non-finite intermediate

The first non-finite or semantically wrong intermediate is usually more useful than the final exception. Debugging backward from $Y$ can hide the earliest failure behind several valid matrix operations.

Repair a Wrong Softmax Axis

Suppose a program has a score matrix $S:T\times T$ but normalizes each column. Every key column then distributes weight across different queries. That is not the intended operation: queries do not compete with one another for a key. Each query needs its own distribution over readable key-value positions, so the final axis must be normalized.

Q1. Diagnose a misleading success

An implementation returns a finite tensor with shape $(B,T,d_{model})$ , but each column of its attention matrix sums to 1 and most rows do not. What is the most direct repair?

Choose one

Normalize across the key-position axis within each query row

Transpose the output projection W_O

Remove the 1/sqrt(d_k) scale

Select one choice, then check.

Hint

Each query row must define one distribution over key positions.

Solution

Apply softmax across the final key-position axis. Then each query row, after masking, sums to 1 over the values it may read.

Preserve an Evidence Record

A useful attention audit records shapes, masks, row sums, forbidden weights, and selected numerical rows. It also separates observations from interpretations. “Weight 0.54 was assigned to position 4” is an observation; “position 4 caused the prediction” requires an intervention and evidence from the rest of the model.