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). Blocked entries receive zero numerator before row normalization. Over real scalars with a nonzero denominator, their attention weight is zero. This is the finite-scalar encoding of the PyTorch pattern scores.masked_fill(~mask, -torch.inf).

Rows with no allowed entries return zero weights directly, avoiding division by zero. This guarantees zero weights even for exceptional scores; the final multiplication by V still obeys the scalar backend's arithmetic. For IEEE scalars, 0 * ∞ can produce NaN. On a partially allowed row, exceptional allowed scores can also make the denominator and blocked weights NaN.

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 α] (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. The scalar context constructs the root from the count: binary backends can recover a finite root even when casting the dimension itself would overflow.

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 models the zero numerator contributed by exp(-∞). The zero-mass conclusion used in causal proofs requires the scalar division laws and a nonzero denominator. A row with no allowed positions returns zeros directly.

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 the LibTorch CUDA attention bridge (NN/Backend/LibTorch.lean).

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 α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {s : Shape} (weights dWeights : TorchLean.Tensor α s) (center : Bool := false) : 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 over real scalars with weights = 0 at blocked entries, the same formula gives zero gradient through blocked logits. Exceptional floating-point operands need not satisfy that zero-product conclusion:

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

Attention uses center := true: subtract the upstream coordinate at a maximal weight before taking the weighted sum. For normalized real weights this leaves the derivative unchanged; for an all-zero row both formulas return zero. In floating point it retains a small dominant-logit gradient after its probability rounds to one. The default retains the uncentered algebraic formula for arbitrary supplied weights, whose sum need not be one.

@[noinline]

def Spec.Internal.productCoefficients? {α : Type} [TorchLean.Storage α] [Context α] {m n p groups : ℕ} (left : Fin groups → TorchLean.Tensor α [m, n]) (right : Fin groups → TorchLean.Tensor α [n, p]) (scale : α) (row : Fin m) (column : Fin p) : Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))

Reconstruct a complete finite scaled contraction without rounding its coefficients.

@[noinline]

def Spec.Internal.replayProducts? {α : Type} [TorchLean.Storage α] [Context α] {m n p groups : ℕ} (left : Fin groups → TorchLean.Tensor α [m, n]) (right : Fin groups → TorchLean.Tensor α [n, p]) (scale ordinary : α) (row : Fin m) (column : Fin p) : Option α

Round a finite contraction only after reconstructing every original coefficient.

def Spec.Internal.scaledProducts {α : Type} [TorchLean.Storage α] [Context α] {m n p groups : ℕ} (left : Fin groups → TorchLean.Tensor α [m, n]) (right : Fin groups → TorchLean.Tensor α [n, p]) (scale : α) (scaleFirst : Bool := false) : TorchLean.Tensor α [m, p]

Scaled sums of matrix products used by attention's forward and differential rules.

One product serves the forward or VJP; two products retain the JVP's joint contraction. Within and between products, checked arithmetic follows the original left-fold schedule. scaleFirst keeps the VJP's input-scaling convention. Range loss or an absorbed nonzero term replays all original finite coefficient trees, including the scale, before rounding. Missing finite input encodings retain the ordinary result. Unsupported carriers use the original expression directly.

This is a reference scalar repair, not a numerical contract for an expanded IR or foreign kernel.

def Spec.Internal.scaledBatchProducts {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {m n p groups : ℕ} (left : Fin groups → TorchLean.Tensor α (batch.concat [m, n])) (right : Fin groups → TorchLean.Tensor α (batch.concat [n, p])) (scale : α) : TorchLean.Tensor α (batch.concat [m, p])

Scaled contractions over a shared batch prefix, without materializing an unscaled product.

def Spec.Internal.scaledMatmul {α : Type} [TorchLean.Storage α] [Context α] {batchA batchB batch : Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (left : TorchLean.Tensor α (batchA.concat [m, n])) (right : TorchLean.Tensor α (batchB.concat [n, p])) (scale : α) : TorchLean.Tensor α (batch.concat [m, p])

A scaled matrix product with broadcast batch prefixes and one retained contraction boundary.

def Spec.Internal.scaledMatmulJvp {α : Type} [TorchLean.Storage α] [Context α] {batchA batchB batch : Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (left deltaLeft : TorchLean.Tensor α (batchA.concat [m, n])) (right deltaRight : TorchLean.Tensor α (batchB.concat [n, p])) (scale : α) : TorchLean.Tensor α (batch.concat [m, p])

The two product-rule terms share one replay boundary, so large opposite tangents can cancel.

def Spec.Internal.scaledMatmulBackward {α : Type} [TorchLean.Storage α] [Context α] {batchA batchB batch : Shape} {m n p : ℕ} (broadcastA : batchA.CanBroadcastTo batch) (broadcastB : batchB.CanBroadcastTo batch) (left : TorchLean.Tensor α (batchA.concat [m, n])) (right : TorchLean.Tensor α (batchB.concat [n, p])) (gradient : TorchLean.Tensor α (batch.concat [m, p])) (scale : α) : TorchLean.Tensor α (batchA.concat [m, n]) × TorchLean.Tensor α (batchB.concat [n, p])

Cotangents of a scaled, broadcast matrix product, retaining its fixed scale in each product.

def Spec.Internal.replayPullbackCoefficientsRow? {α : Type} [TorchLean.Storage α] [Context α] {p : ℕ} (weights ordinary : TorchLean.Tensor α [p]) (coefficients : Fin p → Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))) : Option (TorchLean.Tensor α [p])

Replay one softmax row from retained finite cotangent coefficients.

def Spec.Internal.replayPullbackProductsRow? {α : Type} [TorchLean.Storage α] [Context α] {m n p groups : ℕ} (weights ordinary : TorchLean.Tensor α [p]) (left : Fin groups → TorchLean.Tensor α [m, n]) (right : Fin groups → TorchLean.Tensor α [n, p]) (scale : α) (row : Fin m) : Option (TorchLean.Tensor α [p])

Replay one softmax row from its original gradient contractions and saved weights.

def Spec.Internal.weightedProductCoefficients? {α : Type} [TorchLean.Storage α] [Context α] {m n p : ℕ} (left : TorchLean.Tensor α [m, n]) (right : TorchLean.Tensor α [n, p]) (factor : TorchLean.Tensor α [m, p]) (row : Fin m) (column : Fin p) : Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))

Multiply a retained contraction by one coordinate's factor before encoding its result.

def Spec.Internal.weightedValueCoefficients? {α : Type} [TorchLean.Storage α] [Context α] {rows keys width : ℕ} (weights multiplier : TorchLean.Tensor α [rows, keys]) (gradient : TorchLean.Tensor α [rows, width]) (key : Fin keys) (column : Fin width) : Option (TorchLean.Numeric.QuotientCoefficients (TorchLean.Numeric.QuotientArithmetic.depth α))

Retain both pointwise factors through the transposed value-gradient contraction.

def Spec.Internal.productsLoseRange {α : Type} [TorchLean.Storage α] [Context α] {m n p : ℕ} (left : TorchLean.Tensor α [m, n]) (right : TorchLean.Tensor α [n, p]) : Bool

Conservative range guard for an ordinary product whose coefficients must survive another map. Only flags escape this check; the caller keeps the original value-producing schedule.

def Spec.Internal.pullbackProducts {α : Type} [TorchLean.Storage α] [Context α] {m n p groups : ℕ} (weights gradient : TorchLean.Tensor α [m, p]) (left : Fin groups → TorchLean.Tensor α [m, n]) (right : Fin groups → TorchLean.Tensor α [n, p]) (scale : α) : TorchLean.Tensor α [m, p]

Retain the original gradient contractions through the centered softmax pullback.

The ordinary tensor calculation is unchanged unless an output coefficient loses range. Replaying the contraction and pullback together can recover a finite result even when the intervening gradient tensor cannot represent it. This uses the same coefficient replay as ordinary softmax.

def Spec.softmaxMatmulWeights {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {rows keys : ℕ} (scores : TorchLean.Tensor α (batch.concat [rows, keys])) (mask : Option (TorchLean.Tensor Bool [rows, keys]) := none) : TorchLean.Tensor α (batch.concat [rows, keys])

Row normalization followed by a differentiable multiplier and a value contraction. The forward schedule is deliberately the ordinary softmax, pointwise product, then matrix product. The compound boundary matters for derivatives, where the unweighted cotangent can overflow.

def Spec.Internal.softmaxMatmulFromWeights {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {rows keys width : ℕ} (weights : TorchLean.Tensor α (batch.concat [rows, keys])) (value : TorchLean.Tensor α (batch.concat [keys, width])) (multiplier : TorchLean.Tensor α (batch.concat [rows, keys])) : TorchLean.Tensor α (batch.concat [rows, width])

Mix value rows from saved weights, keeping pointwise multiplication before the contraction.

def Spec.softmaxMatmul {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {rows keys width : ℕ} (scores : TorchLean.Tensor α (batch.concat [rows, keys])) (value : TorchLean.Tensor α (batch.concat [keys, width])) (multiplier : TorchLean.Tensor α (batch.concat [rows, keys])) (mask : Option (TorchLean.Tensor Bool [rows, keys]) := none) : TorchLean.Tensor α (batch.concat [rows, width])

Forward value of the normalized matrix product, without changing its association.

def Spec.softmaxMatmulBackward {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {rows keys width : ℕ} (weights : TorchLean.Tensor α (batch.concat [rows, keys])) (value : TorchLean.Tensor α (batch.concat [keys, width])) (multiplier : TorchLean.Tensor α (batch.concat [rows, keys])) (gradient : TorchLean.Tensor α (batch.concat [rows, width])) : TorchLean.Tensor α (batch.concat [rows, keys]) × TorchLean.Tensor α (batch.concat [keys, width]) × TorchLean.Tensor α (batch.concat [rows, keys])

Cotangents of normalized scores, values and their differentiable multiplier. Finite coefficient replay keeps the original output cotangent and values through both the multiplier and centered softmax. Unsupported scalar carriers use the ordinary local formula.

def Spec.softmaxMatmulJvp {α : Type} [TorchLean.Storage α] [Context α] {batch : Shape} {rows keys width : ℕ} (weights : TorchLean.Tensor α (batch.concat [rows, keys])) (value : TorchLean.Tensor α (batch.concat [keys, width])) (multiplier deltaScores : TorchLean.Tensor α (batch.concat [rows, keys])) (deltaValue : TorchLean.Tensor α (batch.concat [keys, width])) (deltaMultiplier : TorchLean.Tensor α (batch.concat [rows, keys])) : TorchLean.Tensor α (batch.concat [rows, width])

Product-rule tangent of the normalized matrix product, using its saved row 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".

Over real scalars, each unmasked attention row sums to 1; a masked row with at least one allowed key has the same normalization. These equalities are not promised for arbitrary numeric contexts. A fully blocked row returns all-zero weights directly, avoiding a 0/0 result.

The ordinary matrix product and scaling order is retained. Finite-coefficient backends replay the complete scaled contraction after range loss or an absorbed nonzero term. This prevents an unscaled score overflow from poisoning otherwise representable scaled scores, without pre-scaling tiny queries to zero. Scalar conversion and native execution remain separate trust obligations.

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.

Over real scalars, masked calls use the VJP for hard masking: blocked logits have zero forward weight and therefore zero gradient. Finite-coefficient replay retains dOut and V through the centered pullback: their matrix product can overflow even when the score cotangent is finite. Nonfinite original operands retain the ordinary backend arithmetic.

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). Over real scalars, hard-masked blocked logits have zero forward weight, so their tangent contribution is zero in softmaxBackwardFromWeightsSpec. Upstream centering preserves that zero-row convention. The row-wise softmax Jacobian is symmetric, so the same centered formula serves as both VJP and JVP once the forward weights are known.

When a tangent contraction overflows before normalization, finite-coefficient replay carries both original product terms through the pullback. This also preserves zero tangents for a singleton softmax without rounding an unrepresentable intermediate tangent to infinity first.

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 α] {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 α] {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 beside the backward pass that produces it. 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).