Attention (spec layer)

This file defines the standard scaled dot-product attention primitive and a simple multi-head wrapper.

Attention(Q,K,V) = softmax(Q Kᵀ / √d) V

TorchLean goal here is to mirror the math you see in deep learning libraries (especially PyTorch), but keep everything as pure functions on TorchLean.Tensor so the same definitions can be reused for:

• proofs (e.g. reasoning about shapes and gradients),
• reference implementations (runtime extraction),
• verification backends (e.g. interval semantics).

Shapes and conventions

We model the "single batch element" case. Batched attention is obtained by prepending [B] and mapping over it.

Core shapes:

• Q : (nQ × d) queries
• K : (nK × d) keys
• V : (nK × dV) values

In many transformer blocks dV = d, and this file uses that common choice for simplicity.

The optional Boolean mask has shape (nQ × nK). In the main spec, masks use the true -∞ semantics: blocked entries receive zero numerator before row normalization, so their attention weight is definitionally zero. This is the finite-scalar encoding of the PyTorch pattern scores.masked_fill(~mask, -torch.inf).

Rows with no allowed entries evaluate to the zero vector. This total convention agrees with the native TorchLean and SDPA paths and avoids the undefined 0 / 0 normalization of an empty row.

PyTorch analogy:

• scaledDotProductAttention corresponds to torch.nn.functional.scaled_dot_product_attention (no dropout), with Boolean masks interpreted as true -∞ masks.
• MultiHeadAttention.forward corresponds to the core computation inside nn.MultiheadAttention / transformer blocks, ignoring biases and dropout.

Scaled Dot-Product Attention

We separate out the single-head primitive (scaledDotProductAttention) because:

• it is the core mathematical object, reused in multi-head attention,
• it is a good target for proofs and for "spec vs runtime" comparisons.

Boolean masks

TorchLean uses the same boolean mask convention as PyTorch SDPA:

• true means a key/value position is allowed to be attended to,
• false means it is blocked (its softmax numerator is exactly zero).

If an entire row is false, every output weight in that row is zero.

PyTorch reference: torch.nn.functional.scaled_dot_product_attention uses the same convention for boolean attn_mask entries: True entries are included, and False entries are blocked.

def Spec.allTrueMask (nQ nK : ℕ) : TorchLean.Tensor Bool [nQ, nK]

A (nQ × nK) mask where every position is allowed (true).

def Spec.allFalseMask (nQ nK : ℕ) : TorchLean.Tensor Bool [nQ, nK]

A (nQ × nK) mask where every position is blocked (false).

def Spec.causalMask (n : ℕ) : TorchLean.Tensor Bool [n, n]

Causal (lower-triangular) self-attention mask of shape (n, n).

mask[i,j] = true iff j ≤ i, i.e. each query position can attend to itself and past positions.

def Spec.futureMask (n : ℕ) : TorchLean.Tensor Bool [n, n]

Future-only (upper-triangular) self-attention mask of shape (n, n).

This is the (strict) complement of causalMask: mask[i,j] = true iff i < j.

structure Spec.AttentionContext (α : Type) [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] (nQ nK dModel : ℕ) (h1 : nQ ≠ 0) (h2 : nK ≠ 0) : Type

Bundled inputs and mask needed for scaled dot-product attention.

• Q : TorchLean.Tensor α [nQ, dModel]
• K : TorchLean.Tensor α [nK, dModel]
• V : TorchLean.Tensor α [nK, dModel]
• mask : Option (TorchLean.Tensor Bool [nQ, nK])

def Spec.attentionScaleDenom {α : Type} [Context α] (dModel : ℕ) : α

Denominator used by scaled dot-product attention.

Standard attention requires a positive feature dimension and divides scores by sqrt(dModel). TorchLean's tensor shapes also admit zero dimensions. In that degenerate case the result has no feature coordinates, so choosing denominator 1 gives the unique empty-feature result without introducing a division by zero.

Exact hard masking

TorchLean encodes the usual "true -∞ before softmax" behavior without requiring the tensor scalar type itself to contain infinities. Instead of replacing blocked logits by a finite sentinel, we form stable softmax numerators directly. If rowMax is the greatest allowed score in the row, then

numerator_j = if mask_j then exp(score_j - rowMax) else 0.

This is exactly what exp(-∞)=0 contributes to softmax. Blocked positions therefore have exactly zero attention mass, which is the property causal proofs need. A row with no allowed positions is defined to contain only zeros.

def Spec.hardMaskedMax? {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n : ℕ} (scores : TorchLean.Tensor α [n]) (mask : TorchLean.Tensor Bool [n]) : Option α

Maximum allowed score in one hard-masked row, or none when every entry is blocked.

def Spec.hardMaskedSoftmaxVecSpec {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n : ℕ} (scores : TorchLean.Tensor α [n]) (mask : TorchLean.Tensor Bool [n]) : TorchLean.Tensor α [n]

Hard-masked softmax on one vector.

mask[j] = false makes the j-th numerator exactly zero before normalization. This is the ordinary finite-scalar encoding of softmax with true -∞ masked logits.

The maximum and denominator are computed only over allowed entries. Subtracting the allowed-row maximum gives the usual numerically stable softmax formula. If every mask entry is false, the result is the zero vector, matching PyTorch SDPA and TorchLean's native CUDA providers.

def Spec.hardMaskedSoftmaxSpec {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Shape} : TorchLean.Tensor α s → TorchLean.Tensor Bool s → TorchLean.Tensor α s

Hard-masked softmax along the innermost axis of an arbitrary tensor.

def Spec.softmaxBackwardFromWeightsSpec {α : Type} [TorchLean.Storage α] [Context α] {s : Shape} : TorchLean.Tensor α s → TorchLean.Tensor α s → TorchLean.Tensor α s

VJP/JVP helper for a softmax-like row-normalization when the forward weights are already known.

For ordinary softmax, weights = softmax(scores). For hard-masked softmax, blocked entries have weights = 0, and the same formula gives zero gradient through blocked logits:

dScores = weights ⊙ (dWeights - Σⱼ dWeightsⱼ * weightsⱼ).

def Spec.scaledDotProductAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) : TorchLean.Tensor α [nQ, dModel]

Scaled dot-product attention (forward).

Given:

• Q : (nQ × d), K : (nK × d), V : (nK × d),

we compute:

1. scores S = Q Kᵀ with shape (nQ × nK)
2. scaled scores S' = S / √d
3. (optional) mask: for each (i,j), if mask[i,j] = false, its softmax numerator is exactly zero (the finite-scalar encoding of true -∞ masking)
4. attention weights A by row normalization over the last axis
5. output Out = A V with shape (nQ × d)

Mask convention:

mask[i,j] = true means "this key position is allowed", and false means "mask it out".

For unmasked attention, each attention row sums to 1. A masked row with at least one allowed key has the same normalization. A fully blocked row is defined to have all-zero weights, matching PyTorch SDPA and avoiding a 0/0 result.

PyTorch analogy: torch.softmax(scores.masked_fill(~mask, -torch.inf), dim=-1) row-wise, then a final matrix multiply by V.

def Spec.scaledDotProductAttentionBackward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) (dOut : TorchLean.Tensor α [nQ, dModel]) : TorchLean.Tensor α [nQ, dModel] × TorchLean.Tensor α [nK, dModel] × TorchLean.Tensor α [nK, dModel]

Backward/VJP for scaled dot-product attention.

Returns (dQ, dK, dV) given an upstream gradient dOut.

We recompute the forward intermediates locally so this spec stays self-contained and does not rely on a global tape.

For masked calls, this is the VJP for true hard masking. Blocked logits have zero forward weight, and softmaxBackwardFromWeightsSpec therefore gives zero gradient through those blocked positions.

def Spec.scaledDotProductAttentionJvp {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {nQ nK dModel : ℕ} {h1 : nQ ≠ 0} {h2 : nK ≠ 0} (ctx : AttentionContext α nQ nK dModel h1 h2) (dQ : TorchLean.Tensor α [nQ, dModel]) (dK dV : TorchLean.Tensor α [nK, dModel]) : TorchLean.Tensor α [nQ, dModel]

Forward-mode JVP for scaled dot-product attention.

This differentiates the pure attention equation

Out = softmax(mask(Q Kᵀ / sqrt(d))) V

in the direction (dQ,dK,dV). For hard-masked calls, blocked logits have zero forward weight, so their tangent contribution is zero in softmaxBackwardFromWeightsSpec. The row-wise softmax Jacobian is symmetric, so the same formula serves as both VJP and JVP once the forward weights are known.

structure Spec.MultiHeadAttention (α : Type) [TorchLean.Storage α] (numHeads dModel headDim : ℕ) : Type

Multi-head attention parameters (projection matrices).

PyTorch analogy: this corresponds to the four linear maps used in attention blocks:

• Wq, Wk, Wv project dModel -> (numHeads * headDim)
• Wo projects (numHeads * headDim) -> dModel

This spec keeps them as explicit matrices, without bias terms, so the parameterization remains visible in statements about the forward and derivative maps.

• queryWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Query projection from dModel to all attention heads.

• keyWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Key projection from dModel to all attention heads.

• valueWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Value projection from dModel to all attention heads.

• outputWeight : TorchLean.Tensor α [numHeads * headDim, dModel]

  Projection from the concatenated heads back to dModel.

def Spec.splitHeadsSpec {α : Type} [TorchLean.Storage α] [Context α] {n dModel : ℕ} (x : TorchLean.Tensor α [n, dModel]) (numHeads headDim : ℕ) (h : dModel = numHeads * headDim) : TorchLean.Tensor α [numHeads, n, headDim]

Split (n, dModel) into (numHeads, n, headDim).

We store heads as the outermost axis so that "per-head computation" is just a Tensor.dim over Fin numHeads.

The feature coordinate is interpreted as (head, coordinate-within-head): first reshape to (n, numHeads, headDim), then swap the token and head axes. Reshaping directly to (numHeads, n, headDim) would preserve the wrong row-major coordinate order.

def Spec.combineHeadsSpec {α : Type} [TorchLean.Storage α] [Context α] {n numHeads headDim : ℕ} (heads : TorchLean.Tensor α [numHeads, n, headDim]) : TorchLean.Tensor α [n, numHeads * headDim]

Combine a tensor-of-heads back into a single (n, numHeads*headDim) tensor.

Implementation detail:

1. Tensor.swapAdjacentAxes exchanges the head and token axes.
2. reshapeSpec flattens the final two axes into (n, numHeads * headDim).

def Spec.MultiHeadAttention.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {numHeads dModel headDim : ℕ} (n : ℕ) (h1 : n ≠ 0) (mha : MultiHeadAttention α numHeads dModel headDim) (x : TorchLean.Tensor α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n])) : TorchLean.Tensor α [n, dModel]

Multi-head attention forward pass (self-attention when mask is square).

High-level structure:

1. project x into Q,K,V
2. split the projection dimension into heads
3. run scaled dot-product attention per head (sharing the same mask)
4. combine heads back and project with Wo

structure Spec.MultiHeadAttentionParameterGradients (numHeads dModel headDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Parameter gradients for multi-head attention: one per projection matrix.

The record lives here, beside the backward pass that produces it, rather than in the transformer model file where it was originally declared. All four fields are named after the corresponding field of MultiHeadAttention, so { queryWeight, keyWeight, valueWeight, outputWeight } works with the anonymous constructor at every call site.

PyTorch analogue: the .grad of nn.MultiheadAttention.in_proj_weight split into its three blocks, plus out_proj.weight.grad.

• queryWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Gradient of the query projection matrix.

• keyWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Gradient of the key projection matrix.

• valueWeight : TorchLean.Tensor α [dModel, numHeads * headDim]

  Gradient of the value projection matrix.

• outputWeight : TorchLean.Tensor α [numHeads * headDim, dModel]

  Gradient of the output projection matrix.

structure Spec.MultiHeadAttentionGradients (n numHeads dModel headDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Everything multiHeadAttentionBackward sends backwards: the four parameter gradients and the gradient with respect to the attended sequence.

• parameters : MultiHeadAttentionParameterGradients numHeads dModel headDim α

  Gradients for the four projection matrices.

• input : TorchLean.Tensor α [n, dModel]

  Gradient with respect to the input sequence x.

def Spec.multiHeadAttentionBackward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (mha : MultiHeadAttention α numHeads dModel headDim) (x : TorchLean.Tensor α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n])) (gradOutput : TorchLean.Tensor α [n, dModel]) : MultiHeadAttentionGradients n numHeads dModel headDim α

Multi-head attention backward pass.

Returns gradients for input x and all projection matrices (Wq,Wk,Wv,Wo). The forward intermediates are recomputed locally instead of relying on a global tape.

def Spec.multiHeadAttentionJvp {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n numHeads dModel headDim : ℕ} (h1 : n ≠ 0) (mha dmha : MultiHeadAttention α numHeads dModel headDim) (x dx : TorchLean.Tensor α [n, dModel]) (mask : Option (TorchLean.Tensor Bool [n, n])) : TorchLean.Tensor α [n, dModel]

Forward-mode JVP for multi-head attention.

The rule follows the same computational graph as MultiHeadAttention.forward:

1. project tangents through Q/K/V,
2. split primal and tangent projections into heads,
3. apply scaledDotProductAttentionJvp head-wise,
4. combine head tangents, then differentiate the final output projection.

Attention forward-mode AD is explicit at the specification layer rather than hidden behind a runtime-only implementation.

def Spec.selfAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n dModel projDim : ℕ} (x : TorchLean.Tensor α [n, dModel]) (Wq Wk Wv : TorchLean.Tensor α [dModel, projDim]) (Wo : TorchLean.Tensor α [projDim, dModel]) (h1 : n ≠ 0) : TorchLean.Tensor α [n, dModel]

Self-attention on a single sequence.

This uses the same input x for Q/K/V, runs scaled dot-product attention, then applies the output projection Wo.

PyTorch analogue: the core of nn.MultiheadAttention / TransformerEncoderLayer (ignoring the batch axis).

def Spec.crossAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n1 n2 dModel projDim : ℕ} (query : TorchLean.Tensor α [n1, dModel]) (key value : TorchLean.Tensor α [n2, dModel]) (Wq Wk Wv : TorchLean.Tensor α [dModel, projDim]) (Wo : TorchLean.Tensor α [projDim, dModel]) (h1 : n1 ≠ 0) (h2 : n2 ≠ 0) : TorchLean.Tensor α [n1, dModel]

Cross-attention between two sequences.

query is length n1 and attends to key/value of length n2.

PyTorch analogue: the attention block in a Transformer decoder layer (nn.MultiheadAttention with distinct query and key/value inputs).

def Spec.sparseAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {n dModel projDim : ℕ} (x : TorchLean.Tensor α [n, dModel]) (sparsityPattern : TorchLean.Tensor Bool [n, n]) (Wq Wk Wv : TorchLean.Tensor α [dModel, projDim]) (Wo : TorchLean.Tensor α [projDim, dModel]) (h1 : n ≠ 0) : TorchLean.Tensor α [n, dModel]

Sparse attention using a Boolean attention pattern.