27 · Content revision 2

Compute one unmasked self-attention head

Learn how one unmasked Transformer self-attention head scores queries against keys, normalizes each row, and mixes values with inspectable Rust evidence.

Predict which value each query will retrieve

Continue directly from Chapter 26. Its three projections produced

Q=\begin{bmatrix}0&3\\2&-1\end{bmatrix},\quad K=\begin{bmatrix}3&0\\-1&2\end{bmatrix},\quad V=\begin{bmatrix}3&-3\\1&3\end{bmatrix}.

There is one batch with two token positions. The query/key feature width and value width are both two, so $Q,K,V\in\mathbb{R}^{1\times2\times2}$ .

Before calculating a softmax, predict which value each query will retrieve. The first query produces the two dot products

q_0K^\top=[0,6],

so it should favor $v_1=[1,3]$ . The second produces

q_1K^\top=[6,-4],

so it should strongly favor $v_0=[3,-3]$ . The signs and relative gaps matter; normalization will turn each score row into retrieval weights.

Score, scale, normalize, and mix

One unmasked scaled dot-product attention head follows

A=\operatorname{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right),\quad O=AV

For this worked example, the score matrices before and after scaling are

QK^\top=\begin{bmatrix}0&6\\6&-4\end{bmatrix},\qquad \frac{QK^\top}{\sqrt{2}}\approx \begin{bmatrix}0&4.242641\\4.242641&-2.828427\end{bmatrix}.

Softmax runs across key positions independently for each query:

A\approx \begin{bmatrix} 0.014166&0.985834\\ 0.999151&0.000849 \end{bmatrix},\qquad \sum_j A_{ij}=1.

The first weighted value mixture is

o_0 =0.014166[3,-3]+0.985834[1,3] \approx[1.028332,2.915004].

The second is

o_1 =0.999151[3,-3]+0.000849[1,3] \approx[2.998303,-2.994908].

Thus

O\approx \begin{bmatrix} 1.028332&2.915004\\ 2.998303&-2.994908 \end{bmatrix}.

Keep query rows and key columns separate

$Q$ contains one query row per token position: what should this position retrieve?
$K$ contains the candidate key rows: how can each position be matched?
$d_k$ is the shared query/key width.
$A$ contains one normalized retrieval row for every query.
$V$ contains the value rows whose content is mixed.
$O$ contains one resulting mixture per query position.
$B$ is the batch size, $T$ the token count, and $d_v$ the value width.

The complete shape rule is

Q,K\in\mathbb{R}^{B\times T\times d_k},\quad V\in\mathbb{R}^{B\times T\times d_v},\quad A\in\mathbb{R}^{B\times T\times T},\quad O\in\mathbb{R}^{B\times T\times d_v}.

For fixed batch $b$ and query position $i$ , the key position $j$ varies inside one probability row:

S_{bij}=\frac{q_{bi}\cdot k_{bj}}{\sqrt{d_k}},\qquad A_{bij}=\frac{\exp(S_{bij})}{\sum_{r=0}^{T-1}\exp(S_{bir})}.

That row sum makes $A_{bij}$ a retrieval weight. It is not a calibrated probability that a token or claim is correct.

Vaswani et al. motivate the denominator under a specific assumption: if query and key components are independent with mean zero and variance one, their dot product has variance $d_k$ . Under those assumptions, dividing by $\sqrt{d_k}$ keeps the variance of the softmax input from growing with $d_k$ . This is not a universal theorem, an overflow guarantee, or proof that the scale is an optimal temperature.

A two-coordinate probe makes the effect concrete. For the scores $[1,0]$ , the favored weight without scaling is

p_{\mathrm{unscaled}}=\frac{e^1}{e^1+e^0}\approx0.731059.

With $d_k=2$ , the same weight after square-root scaling is

p_{\mathrm{scaled}} =\frac{e^{1/\sqrt{2}}}{e^{1/\sqrt{2}}+e^0} \approx0.669762.

Scaling therefore softens this particular distribution; the assumptions above explain why the denominator is useful, not how sharp every attention row must be.

From recurrent context to all-position retrieval

This is neural-attention history on the road to modern LLMs, not programming- language history.

A basic recurrent encoder-decoder can force an entire source sentence through one fixed-size vector, while additive attention still advances recurrently and computes a new source alignment at each output step.

Bahdanau, Cho, and Bengio, Neural Machine Translation by Jointly Learning to Align and Translate introduced the relevant bridge. Bahdanau, Cho, and Bengio describe the possible fixed-length-vector bottleneck of a basic encoder-decoder and compute each new context as a softmax-weighted sum of encoder annotations scored against the previous decoder state.

Calling that learned scoring function additive attention follows the later retrospective classification used by Vaswani and colleagues. The earlier paper does not use this chapter’s query, key, and value notation, scaled dot products, or same-sequence score grid.

The Transformer combines many queries into matrices and uses scaled dot-product self-attention to compare positions in the same available sequence, allowing one layer to form its full score grid with batched matrix operations.

Vaswani et al., Attention Is All You Need provide that later step. Vaswani et al. define scaled dot-product attention as a softmax-normalized matrix of scaled query-key dot products applied to values, combine simultaneous queries in a matrix, and define self-attention as relating positions within one sequence.

Each decoder self-attention head turns learned queries and keys into row-normalized attention weights, then mixes learned values; a causal decoder additionally masks future key positions before normalization.

Forming that grid does not make total work constant, and it does not make autoregressive token generation parallel. The important historical change is structural: recurrent alignment computes a new context as the decoder advances, while self-attention relates all available positions within one layer.

The executable contrast below isolates the change in which positions supply the query side and the key/value side:

Contrast encoder-decoder alignment sources with same-sequence query-key pairs rust/demos/ch27-self-attention/src/lib.rs#historical-attention-contrast

fn attention_pairs<'a>(queries: &'a [&'a str], keys: &'a [&'a str]) -> Vec<(&'a str, &'a str)> {
    queries
        .iter()
        .flat_map(|query| keys.iter().map(move |key| (*query, *key)))
        .collect()
}

fn historical_attention_contrast() -> HistoryEvidence {
    let decoder_states = ["decoder-state-0", "decoder-state-1"];
    let encoder_annotations = ["encoder-annotation-0", "encoder-annotation-1"];
    let hidden_sequence = ["hidden-position-0", "hidden-position-1"];

    let encoder_decoder_alignment = attention_pairs(&decoder_states, &encoder_annotations);
    let self_attention = attention_pairs(&hidden_sequence, &hidden_sequence);
    assert!(
        encoder_decoder_alignment
            .iter()
            .all(|(query, key)| query.starts_with("decoder") && key.starts_with("encoder"))
    );
    assert!(
        self_attention
            .iter()
            .all(|(query, key)| query.starts_with("hidden") && key.starts_with("hidden"))
    );

    HistoryEvidence {
        earlier: "recurrent-fixed-context",
        bridge: "additive-encoder-decoder-alignment",
        transformer: "scaled-dot-product-self-attention",
        comparison: "all-sequence-positions",
    }
}

Compose the head from differentiable tensor operations

scaled_dot_product_self_attention accepts three rank-three TensorValue objects. It composes the existing transpose, batched matmul, scalar multiply, stable log_softmax, exp, and second matmul operations. It exposes raw scores, scaled scores, probabilities, output, scale, key width, and value width:

Compose one inspectable unmasked scaled dot-product attention head rust/crates/llm-from-scratch/src/attention/self_attention.rs#self-attention-forward

/// Shared, validated score preparation for unmasked and causally masked heads.
#[derive(Clone, Debug)]
pub(crate) struct ScaledSelfAttentionScores {
    pub(crate) raw_scores: TensorValue,
    pub(crate) scaled_scores: TensorValue,
    pub(crate) scale: f64,
    pub(crate) key_width: usize,
    pub(crate) value_width: usize,
}

/// Inspectable evidence from one unmasked attention head.
#[derive(Clone, Debug)]
pub struct SelfAttentionForward {
    raw_scores: TensorValue,
    scaled_scores: TensorValue,
    weights: TensorValue,
    output: TensorValue,
    scale: f64,
    key_width: usize,
    value_width: usize,
}

impl SelfAttentionForward {
    /// The unnormalized matrix `Q K^T` with shape `[batch, tokens, tokens]`.
    pub fn raw_scores(&self) -> &TensorValue {
        &self.raw_scores
    }

    /// Alias that emphasizes that each raw cell is one query-key dot product.
    pub fn dot_products(&self) -> &TensorValue {
        &self.raw_scores
    }

    /// The raw scores divided by the square root of the query/key width.
    pub fn scaled_scores(&self) -> &TensorValue {
        &self.scaled_scores
    }

    /// Row-normalized probabilities over key positions.
    pub fn weights(&self) -> &TensorValue {
        &self.weights
    }

    /// Alias for the row-normalized attention weights.
    pub fn probabilities(&self) -> &TensorValue {
        &self.weights
    }

    /// The weighted value rows with shape `[batch, tokens, value_width]`.
    pub fn output(&self) -> &TensorValue {
        &self.output
    }

    /// The fixed score multiplier `1 / sqrt(key_width)` used by this pass.
    pub const fn scale(&self) -> f64 {
        self.scale
    }

    pub const fn key_width(&self) -> usize {
        self.key_width
    }

    pub const fn value_width(&self) -> usize {
        self.value_width
    }

    pub fn into_output(self) -> TensorValue {
        self.output
    }

    pub fn into_parts(self) -> (TensorValue, TensorValue, TensorValue, TensorValue) {
        (
            self.raw_scores,
            self.scaled_scores,
            self.weights,
            self.output,
        )
    }
}

/// Computes one unmasked scaled dot-product self-attention head.
///
/// Q, K, and V must describe the same batch and token positions. Q and K share
/// one nonzero feature width; V may use a different nonzero output width.
pub fn scaled_dot_product_self_attention(
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
) -> Result<SelfAttentionForward, SelfAttentionError> {
    let prepared = scaled_self_attention_scores(query, key, value)?;
    let log_weights = prepared
        .scaled_scores
        .log_softmax(2)
        .map_err(autodiff_error(SelfAttentionStage::LogSoftmax))?;
    let weights = log_weights
        .exp()
        .map_err(autodiff_error(SelfAttentionStage::Probabilities))?;
    let output = weights
        .matmul(value)
        .map_err(autodiff_error(SelfAttentionStage::ValueMixture))?;

    Ok(SelfAttentionForward {
        raw_scores: prepared.raw_scores,
        scaled_scores: prepared.scaled_scores,
        weights,
        output,
        scale: prepared.scale,
        key_width: prepared.key_width,
        value_width: prepared.value_width,
    })
}

pub(crate) fn scaled_self_attention_scores(
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
) -> Result<ScaledSelfAttentionScores, SelfAttentionError> {
    let query_shape = query.shape();
    let key_shape = key.shape();
    let value_shape = value.shape();

    for (input, shape) in [
        (SelfAttentionInput::Query, query_shape.as_slice()),
        (SelfAttentionInput::Key, key_shape.as_slice()),
        (SelfAttentionInput::Value, value_shape.as_slice()),
    ] {
        if shape.len() != 3 {
            return Err(SelfAttentionError::InputRank {
                input,
                rank: shape.len(),
            });
        }
    }

    if query_shape[0] != key_shape[0] || query_shape[0] != value_shape[0] {
        return Err(SelfAttentionError::BatchMismatch {
            query: query_shape[0],
            key: key_shape[0],
            value: value_shape[0],
        });
    }
    if query_shape[1] != key_shape[1] || query_shape[1] != value_shape[1] {
        return Err(SelfAttentionError::TokenMismatch {
            query: query_shape[1],
            key: key_shape[1],
            value: value_shape[1],
        });
    }
    if query_shape[1] == 0 {
        return Err(SelfAttentionError::EmptyTokens);
    }
    if query_shape[2] == 0 {
        return Err(SelfAttentionError::EmptyFeatureWidth {
            input: SelfAttentionInput::Query,
        });
    }
    if key_shape[2] == 0 {
        return Err(SelfAttentionError::EmptyFeatureWidth {
            input: SelfAttentionInput::Key,
        });
    }
    if query_shape[2] != key_shape[2] {
        return Err(SelfAttentionError::QueryKeyWidthMismatch {
            query: query_shape[2],
            key: key_shape[2],
        });
    }
    if value_shape[2] == 0 {
        return Err(SelfAttentionError::EmptyFeatureWidth {
            input: SelfAttentionInput::Value,
        });
    }

    let key_transposed = key
        .transpose(1, 2)
        .map_err(autodiff_error(SelfAttentionStage::KeyTranspose))?;
    let raw_scores = query
        .matmul(&key_transposed)
        .map_err(autodiff_error(SelfAttentionStage::RawScores))?;

    let scale = 1.0 / (query_shape[2] as f64).sqrt();
    let scale_tensor =
        Tensor::from_vec(Vec::new(), vec![scale]).map_err(|source| SelfAttentionError::Tensor {
            stage: SelfAttentionStage::ScaleTensor,
            source,
        })?;
    let scale_value = TensorValue::constant(scale_tensor)
        .map_err(autodiff_error(SelfAttentionStage::ScaleTensor))?;
    let scaled_scores = raw_scores
        .mul(&scale_value)
        .map_err(autodiff_error(SelfAttentionStage::ScaledScores))?;
    Ok(ScaledSelfAttentionScores {
        raw_scores,
        scaled_scores,
        scale,
        key_width: query_shape[2],
        value_width: value_shape[2],
    })
}

Query, key, and value inputs must have rank three and matching batch and token axes. The token axis and all feature axes must be nonempty. Query and key widths must agree, but $d_v$ may differ from $d_k$ . An empty batch is valid. Typed errors preserve input-rank, batch, token, empty-token, feature-width, and forward-stage context:

Reject invalid self-attention inputs before exposing a partial result rust/crates/llm-from-scratch/src/attention/self_attention.rs#self-attention-errors

/// One of the three inputs to a self-attention head.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SelfAttentionInput {
    Query,
    Key,
    Value,
}

impl fmt::Display for SelfAttentionInput {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::Query => "query",
            Self::Key => "key",
            Self::Value => "value",
        })
    }
}

/// The forward stage at which a cumulative tensor operation failed.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SelfAttentionStage {
    KeyTranspose,
    RawScores,
    ScaleTensor,
    ScaledScores,
    LogSoftmax,
    Probabilities,
    ValueMixture,
}

impl fmt::Display for SelfAttentionStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::KeyTranspose => "key transpose",
            Self::RawScores => "raw query-key scores",
            Self::ScaleTensor => "score scale",
            Self::ScaledScores => "scaled query-key scores",
            Self::LogSoftmax => "row log-softmax",
            Self::Probabilities => "attention probabilities",
            Self::ValueMixture => "weighted value mixture",
        })
    }
}

/// A rejected Q/K/V shape or cumulative tensor operation.
#[derive(Clone, Debug, PartialEq)]
pub enum SelfAttentionError {
    InputRank {
        input: SelfAttentionInput,
        rank: usize,
    },
    BatchMismatch {
        query: usize,
        key: usize,
        value: usize,
    },
    TokenMismatch {
        query: usize,
        key: usize,
        value: usize,
    },
    QueryKeyWidthMismatch {
        query: usize,
        key: usize,
    },
    EmptyTokens,
    EmptyFeatureWidth {
        input: SelfAttentionInput,
    },
    Tensor {
        stage: SelfAttentionStage,
        source: TensorError,
    },
    Autodiff {
        stage: SelfAttentionStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for SelfAttentionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InputRank { input, rank } => write!(
                formatter,
                "self-attention {input} must have rank three [batch, tokens, features], got rank {rank}"
            ),
            Self::BatchMismatch { query, key, value } => write!(
                formatter,
                "self-attention batch sizes must match, got query {query}, key {key}, value {value}"
            ),
            Self::TokenMismatch { query, key, value } => write!(
                formatter,
                "unmasked self-attention token counts must match, got query {query}, key {key}, value {value}"
            ),
            Self::QueryKeyWidthMismatch { query, key } => write!(
                formatter,
                "self-attention query and key widths must match, got query {query}, key {key}"
            ),
            Self::EmptyTokens => formatter.write_str(
                "unmasked self-attention needs at least one token so every probability row has a key",
            ),
            Self::EmptyFeatureWidth { input } => write!(
                formatter,
                "self-attention {input} needs a nonzero feature width"
            ),
            Self::Tensor { stage, source } => {
                write!(formatter, "self-attention {stage}: {source}")
            }
            Self::Autodiff { stage, source } => {
                write!(formatter, "self-attention {stage}: {source}")
            }
        }
    }
}

impl Error for SelfAttentionError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Tensor { source, .. } => Some(source),
            Self::Autodiff { source, .. } => Some(source),
            _ => None,
        }
    }
}

fn autodiff_error(
    stage: SelfAttentionStage,
) -> impl FnOnce(TensorAutodiffError) -> SelfAttentionError {
    move |source| SelfAttentionError::Autodiff { stage, source }
}

For the reverse example, bars denote reverse-mode gradients:

\bar X=\frac{\partial L}{\partial X}.

Choose the upstream seed

\bar O=\begin{bmatrix}1&0\\0&1\end{bmatrix},

which defines the scalar objective

L=\langle O,\bar O\rangle=O_{00}+O_{11}\approx-1.966576.

First the output matrix multiplication gives

\bar V_b=A_b^{\mathsf T}\bar O_b,\qquad \bar A_b=\bar O_bV_b^{\mathsf T}.

For each query row, the softmax backward rule is

\bar S_{bij} =A_{bij}\left(\bar A_{bij}-\sum_r A_{bir}\bar A_{bir}\right).

The scaled dot products then send gradients to queries and keys:

\bar Q_b=\frac{\bar S_bK_b}{\sqrt{d_k}},\qquad \bar K_b=\frac{\bar S_b^{\mathsf T}Q_b}{\sqrt{d_k}}.

Here $\mathsf T$ transposes the final two axes separately inside each batch.

For the worked inputs, these equations give

\bar Q\approx \begin{bmatrix} 0.079000&-0.039500\\ -0.014389&0.007195 \end{bmatrix},

\bar K\approx \begin{bmatrix} -0.007195&0.062847\\ 0.007195&-0.062847 \end{bmatrix},

\bar V\approx \begin{bmatrix} 0.014166&0.999151\\ 0.985834&0.000849 \end{bmatrix}.

All four coordinates of each input agree with central differences using step $10^{-6}$ and tolerance $2\times10^{-6}$ . Equal keys produce equal weights, a single token gives its only value weight one, matching permutations of $Q$ , $K$ , and $V$ produce the same output permutation, and separate batch elements remain independent. An empty batch with $Q,K$ shaped $[0,2,2]$ and $V$ shaped $[0,2,3]$ returns $A$ shaped $[0,2,2]$ and $O$ shaped $[0,2,3]$ . The value width may differ from the query/key width: values shaped $[1,2,1]$ produce an output shaped $[1,2,1]$ .

The complete evidence builder keeps these forward, reverse, scale, shape, boundary, and replay checks together:

Build the worked self-attention evidence from the cumulative tensor operations rust/demos/ch27-self-attention/src/lib.rs#self-attention-fixture

pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let primary = primary_once()?;
    let replay = primary_once()?;
    let (query_checks, key_checks, value_checks, gradcheck_passed) = gradient_evidence(&primary)?;
    Ok(LearnerEvidence {
        replay_bitwise: primary == replay,
        scale: scale_evidence()?,
        single_token: single_token_evidence()?,
        shapes: shape_evidence()?,
        errors: error_evidence()?,
        history: historical_attention_contrast(),
        primary,
        query_checks,
        key_checks,
        value_checks,
        gradcheck_passed,
    })
}

Run the complete unmasked self-attention example rust/demos/ch27-self-attention/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let evidence = ch27_self_attention::learner_evidence()?;
    print!("{}", ch27_self_attention::render_report(&evidence));
    Ok(())
}

Run cargo run --quiet --locked -p ch27-self-attention to inspect the same scores, probabilities, mixtures, gradients, shapes, and boundary results.

Trace every attention row from scores to output

The diagram follows the exact inputs through raw and scaled scores, probability rows, weighted value terms, outputs, gradients, shapes, rejected inputs, and the historical transition to self-attention:

Follow every score into a weighted value mixture

Follow exact query, key, value, score, probability, mixture, gradient, shape, history, and rejected-boundary evidence through one unmasked attention head.

Solid query border
Dashed key border
Double value border
Dotted score cue
Double-line probability cue

Trace one complete unmasked attention calculation

Rows and columns preserve the two-token calculation exactly; borders and text carry every distinction without color.

Start with query, key, and value rows

Query rows: what should each position retrieve?

$Q$

$0$ $1$

$q_0$ $0.000000$ $3.000000$
$q_1$ $2.000000$ $-1.000000$

Shape: $[1,2,2]$
Key rows: how can each position be matched?

$K$

$0$ $1$

$k_0$ $3.000000$ $0.000000$
$k_1$ $-1.000000$ $2.000000$

Shape: $[1,2,2]$
Value rows: what content can each position contribute?

$V$

$0$ $1$

$v_0$ $3.000000$ $-3.000000$
$v_1$ $1.000000$ $3.000000$

Shape: $[1,2,2]$

	$0$	$1$
$q_0$	$0.000000$	$3.000000$
$q_1$	$2.000000$	$-1.000000$

	$0$	$1$
$k_0$	$3.000000$	$0.000000$
$k_1$	$-1.000000$	$2.000000$

	$0$	$1$
$v_0$	$3.000000$	$-3.000000$
$v_1$	$1.000000$	$3.000000$

Compare every query with every key

$QK^\top$

	$k_0$	$k_1$
$q_0$	$0.000000$	$6.000000$
$q_1$	$6.000000$	$-4.000000$

Scale the score grid

$S=QK^\top/\sqrt{d_k}$

	$k_0$	$k_1$
$q_0$	$0.000000$	$4.242641$
$q_1$	$4.242641$	$-2.828427$

Attention scale: $1/\sqrt{d_k}=0.707107$
Shape: $[1,2,2]$

Normalize each query row over keys

$q_0$

$A_{0,:}=[0.014166,0.985834]$

Probability-row check

$\sum_j A_{0j}=1.000000$

Normalization axis

key
$q_1$

$A_{1,:}=[0.999151,0.000849]$

Probability-row check

$\sum_j A_{1j}=1.000000$

Normalization axis

key

Mix value rows into outputs

$o_0$

Normalize each query row over keys

$[0.014166,0.985834]$

Already-weighted value terms

$[0.042498,-0.042498]$ $+[0.985834,2.957502]$

Output row

$o_0=[1.028332,2.915004]$
$o_1$

Normalize each query row over keys

$[0.999151,0.000849]$

Already-weighted value terms

$[2.997454,-2.997454]$ $+[0.000849,0.002546]$

Output row

$o_1=[2.998303,-2.994908]$

Check gradients, shapes, and rejected boundaries

The gradients, shapes, and boundary cases show what the same attention operation preserves and rejects.

Reverse evidence

$\bar O$	$[1.000000,0.000000,0.000000,1.000000]$
$\bar Q$	$[0.079000,-0.039500,-0.014389,0.007195]$
$\bar K$	$[-0.007195,0.062847,0.007195,-0.062847]$
$\bar V$	$[0.014166,0.999151,0.985834,0.000849]$

Batch isolation and shapes

$Q,K,V$: $[2,2,2]$
$A$: $[2,2,2]$
$O$: $[2,2,2]$
Batch independence: Verified

One token has no competing key

$A=[1.000000]$

$O=[5.000000,-2.000000]$

Query gradient: $\bar Q=0$

Key gradient: $\bar K=0$

Visibility boundary

Every key position is visible

$d_k=2,\quad d_v=2$

Rejected input boundaries

Rejected Error kind: input-rank The query input must expose batch, token, and feature axes. Rejected shape evidence: operand=query|rank=2
Rejected Error kind: batch-mismatch Query, key, and value tensors must have the same batch size. Rejected shape evidence: query=1|key=2|value=1
Rejected Error kind: token-mismatch Query, key, and value tensors must have the same token count. Rejected shape evidence: query=2|key=3|value=2
Rejected Error kind: empty-token-axis Attention needs at least one key position to normalize each row. Rejected shape evidence: tokens=0
Rejected Error kind: query-key-width-mismatch Query and key feature widths must match for their dot products. Rejected shape evidence: query=2|key=3

Numerical checks

Probability-row check: $0.000000000001$
Gradient check: gradcheck=true
Gradient coordinates checked: $4+4+4$
Gradient-check tolerance: $0.000002$
Deterministic replay: replay=bitwise

Follow the neural-attention path toward modern LLMs

The comparison describes model structure, not programming-language history or a hardware benchmark.

Fixed recurrent context

One fixed-size source vector is reused across the recurrent decoder steps.
Additive encoder-decoder alignment

Each decoder step retrieves a new weighted context from encoder annotations.
Scaled dot-product self-attention

One layer forms a score for every query-key position pair in the available sequence.

Every available key position is permitted for every query in this unmasked head.

Read each probability row across key positions before following it into the matching weighted value terms. Solid, dashed, double, and dotted borders keep query, key, value, score, and probability roles distinct without relying on color. Every position is still visible; Chapter 28 adds the missing causal boundary.

Predict before reading the evidence

Compute all four entries of $QK^\top$ for the worked example.
Predict which key each query favors before applying softmax.
Explain why each row of $A$ sums to one instead of each column.
Predict the output shape for $Q,K\in\mathbb{R}^{4\times7\times3}$ and $V\in\mathbb{R}^{4\times7\times5}$ .
Predict the attention probability and output when there is one token with value $[5,-2]$ .
Decide what happens if the same two token positions are swapped in $Q$ , $K$ , and $V$ .
Predict the probabilities when the two key rows are equal.
Identify which future access would leak target information during causal decoder training.
Contrast recurrent alignment with the score-grid structure of self-attention.

Check the predictions

The score rows are $[0,6]$ and $[6,-4]$ .
Query zero favors key one; query one favors key zero.
Each query chooses among key positions, so softmax normalizes its key axis.
The output shape is $[4,7,5]$ .
The only probability is $[1]$ , and the output is $[5,-2]$ .
Jointly permuting the rows of $Q$ , $K$ , and $V$ applies the same permutation to the output rows. Position-dependent inputs or a mask must be analyzed separately.
Equal scores produce the uniform row $[0.5,0.5]$ .
A query can use later target positions because this head is unmasked; that future access would leak target information during autoregressive training.
Recurrent alignment computes a new context as the decoder advances; self-attention forms relationships among all available positions within one layer. This does not make total work constant or autoregressive generation parallel.

Mask future keys next

The cumulative decoder can now turn one projected query/key/value triplet into the output of an unmasked attention head. Chapter 28 will exclude future key positions before each score row is normalized.

The head already preserves batch and query-position axes and returns both the inspectable probability matrix and mixed values. It is not yet safe for autoregressive decoding: every query can read every key. Causal masking is the next boundary; positional information, multiple heads, output projection, residual wrapping, and cached decoding remain later chapters.