Block future keys with a causal mask

Predict the visible triangle

Extend Chapter 27 with a third position:

There is one batch, three token positions, a query/key width of two, and a value width of two. Thus

The unmasked raw score rows are $[0,6,3]$, $[6,-4,3]$, and $[3,1,3]$. Before running the fixture, mark the cells each query may use. Query $0$ keeps key $0$; query $1$ keeps keys $0$ and $1$; query $2$ keeps all three keys. That is six allowed cells and three blocked cells, including all three diagonal cells among the allowed six.

Mask before softmax

For query row $i$ and key column $j$, use the additive mask

The diagonal is deliberately allowed. During decoder training, target inputs are shifted by one position. The input representation on the diagonal therefore contains an earlier known token; the prediction does not read its own target. The shift and mask work together to preserve autoregressive conditioning.

Let the scaled score tensor be

Add the visibility rule before normalization, then mix values:

For this example,

Every blocked probability is exactly zero. Each allowed prefix keeps the full unit mass:

The resulting rows are

Zeroing future probabilities after an ordinary all-key softmax is not equivalent: the remaining values would retain a denominator that included the blocked keys. A large finite negative sentinel can approximate the ideal mask, but it is not mathematically identical to $-\infty$.

Keep visibility separate from position

• $S$ is the scaled query-key score tensor before masking.
• $M$ is the additive visibility mask.
• $i$ indexes a query position and $j$ indexes a key position.
• $A$ contains attention probabilities after the future keys are excluded.
• $Q$, $K$, and $V$ retain their Chapter 27 query, key, and value roles.
• $O$ contains one visible-prefix mixture per query position.
• $B$ is the batch size, $T$ the token count, $d_k$ the query/key width, and $d_v$ the value width.

The complete shape rule remains

For an allowed cell, the row normalization is

and for a blocked cell,

This rule supplies visibility, not position. Rows still need a separate absolute or relative position signal to distinguish order. Padding masks, variable lengths, multiple heads, output projection, and key/value caching are also separate concerns.

From recurrent prefix state to an explicit decoder mask

Graves, Generating Sequences With Recurrent Neural Networks trains next-element prediction one sequence element at a time. During generation, each sampled element becomes the next recurrent input, so future elements do not yet exist. This prefix boundary comes from sequential recurrence, whose state must advance one step at a time. The claim concerns the paper’s text-generation path, not its separately conditioned handwriting-synthesis model.

Vaswani et al., Attention Is All You Need pack queries into matrices and mask decoder self-attention so a row cannot attend to subsequent positions. Illegal pre-softmax scores are set to $-\infty$; together with output embeddings shifted by one position, row $i$ depends only on known outputs before $i$.

This lets the masked attention rows for known target positions be evaluated together during training, while ordinary autoregressive decoding still appends one token at a time. A decoder-only Transformer applies this causal boundary in each self-attention layer. The mask controls visibility; it does not encode position.

Keep the mask inspectable and recorded values finite

causal_additive_mask constructs a plain $[T,T]$ tensor. Allowed cells contain $0$ and future cells contain $-\infty$:

/// Builds an additive square mask with zero for `key <= query` and negative
/// infinity for future keys.
pub fn causal_additive_mask(tokens: usize) -> Result<Tensor, CausalMaskingError> {
    let elements = tokens
        .checked_mul(tokens)
        .ok_or(CausalMaskingError::MaskTensor(TensorError::ShapeOverflow))?;
    let mut values = Vec::new();
    values
        .try_reserve_exact(elements)
        .map_err(|_| CausalMaskingError::MaskAllocationFailed { elements })?;
    for query in 0..tokens {
        for key in 0..tokens {
            values.push(if key <= query { 0.0 } else { f64::NEG_INFINITY });
        }
    }
    Tensor::from_vec(vec![tokens, tokens], values).map_err(CausalMaskingError::MaskTensor)
}

A differentiable TensorValue rejects nonfinite leaf data. Its causal_softmax operation implements the same mathematical boundary without storing $-\infty$: it reads only $j\le i$, subtracts the maximum of that allowed prefix, normalizes the allowed cells, and writes the exact floating-point value $+0.0$ to every blocked cell.

fn causal_softmax_forward(input: &Tensor) -> Result<Tensor, TensorAutodiffError> {
    if input.rank() < 2 {
        return Err(ModelOpError::CausalSoftmaxRank { rank: input.rank() }.into());
    }
    let queries = input.shape()[input.rank() - 2];
    let keys = input.shape()[input.rank() - 1];
    if queries != keys {
        return Err(ModelOpError::CausalSoftmaxNonSquare { queries, keys }.into());
    }
    if queries == 0 {
        return Err(ModelOpError::CausalSoftmaxEmptyTokens.into());
    }

    let mut probabilities = zeros(input.shape())?;
    let grids = input.len() / (queries * keys);
    for grid in 0..grids {
        for query in 0..queries {
            let row_start = (grid * queries + query) * keys;
            let allowed = &input.as_slice()[row_start..=row_start + query];
            let maximum = allowed.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            let mut exponential_tail = 0.0;
            let mut skipped_one_maximum = false;
            for &score in allowed {
                let shifted = score - maximum;
                if shifted == 0.0 && !skipped_one_maximum {
                    skipped_one_maximum = true;
                } else {
                    exponential_tail += shifted.exp();
                }
            }
            debug_assert!(skipped_one_maximum);
            let denominator = 1.0 + exponential_tail;
            for (key, &score) in allowed.iter().enumerate() {
                let probability = (score - maximum).exp() / denominator;
                probabilities.as_mut_slice()[row_start + key] =
                    if probability == 0.0 { 0.0 } else { probability };
            }
        }
    }
    Ok(probabilities)
}

The operation accepts square score tensors of rank two or greater. This lets a future score shape such as $[B,H,T,T]$ preserve its independent leading axes. It rejects a rank below two, unequal final axes, an empty token axis, or a released operand.

causal_scaled_dot_product_self_attention reuses the Chapter 27 score construction, applies causal_softmax, and multiplies by $V$:

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

impl CausalSelfAttentionForward {
    pub fn raw_scores(&self) -> &TensorValue {
        &self.raw_scores
    }

    pub fn dot_products(&self) -> &TensorValue {
        &self.raw_scores
    }

    pub fn scaled_scores(&self) -> &TensorValue {
        &self.scaled_scores
    }

    /// The plain additive mask. It is intentionally not a tape value because
    /// its blocked cells contain negative infinity.
    pub const fn additive_mask(&self) -> &Tensor {
        &self.additive_mask
    }

    pub fn weights(&self) -> &TensorValue {
        &self.weights
    }

    pub fn probabilities(&self) -> &TensorValue {
        &self.weights
    }

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

    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
    }
}

/// Computes one scaled self-attention head with an inclusive-prefix mask.
pub fn causal_scaled_dot_product_self_attention(
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
) -> Result<CausalSelfAttentionForward, CausalMaskingError> {
    let prepared = scaled_self_attention_scores(query, key, value)?;
    let tokens = query.shape()[1];
    let additive_mask = causal_additive_mask(tokens)?;
    let weights = prepared
        .scaled_scores
        .causal_softmax()
        .map_err(autodiff_error(CausalMaskingStage::MaskedSoftmax))?;
    let output = weights
        .matmul(value)
        .map_err(autodiff_error(CausalMaskingStage::ValueMixture))?;

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

For an allowed score, reverse mode follows

while a blocked score receives

All six coordinates of each of $Q$, $K$, and $V$ agree with central differences using step $10^{-6}$ and tolerance $4\times10^{-6}$. The checked boundary also includes a single token, empty batches, rank-two and rank-four score grids, independent leading axes, extreme blocked scores, typed failures, released tape values, and bitwise replay.

The executable prints the checked mask, probabilities, outputs, prefix-invariance result, gradients, and boundary cases:

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

Run cargo run --quiet --locked -p ch28-causal-masking to inspect the complete worked result.

Follow the lower triangle through attention

The triangular tables show which scores survive before normalization. Solid, dashed, and double borders distinguish allowed, blocked, and diagonal cells without relying on color.

Changing only

leaves $o_0$ and $o_1$ bitwise unchanged, while

Predict before opening the answers

1. Write the allowed key-index set for query rows $0$, $1$, and $2$.
2. Predict whether changing only $k_2$ and $v_2$ can change $o_0$ or $o_1$.
3. Explain why all three diagonal cells are allowed without letting a prediction read its own target.
4. Predict the query and key gradients for a one-token causal self-attention head.
5. Explain why setting future probabilities to zero after an ordinary softmax does not preserve a unit row sum.
6. Decide whether the mask makes known-target training, autoregressive generation, both, or neither parallel across token positions.
7. Identify which separate mechanism Chapter 29 must add without changing the lower-triangular boundary.

Preserve the prefix boundary as the decoder grows

The cumulative decoder now produces an output at position $i$ using only keys and values through position $i$. If a loss uses only the first two positions, the future suffix receives exact zero gradient:

That is the information boundary required by an autoregressive decoder. Chapter 29 adds relative position information without widening it.