MultiHeadSelfAttention

End-to-end fderiv/backprop correctness for a Multi-Head Self-Attention graph, decomposed into proven tape nodes:

• linear projections via matmul,
• head split/merge via reshape and coordinate reindexing,
• attention core via batched matmul, coordinate reindexing, scaling, and row-wise softmax.

This is spec-level over ℝ. It is a corollary of the general graph theorem once each node used by the graph has a NodeFDerivCorrect instance.

PyTorch correspondence / citations

• The construction matches the usual “project → split heads → scaled dot-product attention → concat heads → output projection” pipeline used by torch.nn.MultiheadAttention. https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html
• The core attention step corresponds to torch.nn.functional.scaled_dot_product_attention. https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.XShape (n dModel : ℕ) : Spec.Shape

Sequence input shape n×dModel.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.BigShape (n numHeads headDim : ℕ) : Spec.Shape

Concatenated-head representation n×(numHeads*headDim).

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.HeadsShape (n numHeads headDim : ℕ) : Spec.Shape

Split-head representation (numHeads)×n×headDim.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.KtShape (n numHeads headDim : ℕ) : Spec.Shape

Key-transposed shape (numHeads)×headDim×n used for Q Kᵀ.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ScoresShape (n numHeads : ℕ) : Spec.Shape

Attention scores shape (numHeads)×n×n.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.SwappedShape (n numHeads headDim : ℕ) : Spec.Shape

Intermediate shape after swapping axes for concatenation n×numHeads×headDim.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ssMHA (n dModel numHeads headDim : ℕ) : List Spec.Shape

Intermediate node output shapes (tape “saved tensors”) for the MHA graph.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ssMHAProjections (n numHeads headDim : ℕ) : List Spec.Shape

Saved tensors produced while projecting the input into query, key, and value heads.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ssMHAScores (n numHeads headDim : ℕ) : List Spec.Shape

Saved tensors through scaled query-key scores, before softmax is applied.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ssMHAAttention (n numHeads headDim : ℕ) : List Spec.Shape

Saved tensors through the head-wise attention output, before heads are merged.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.WqShape (dModel numHeads headDim : ℕ) : Spec.Shape

Projection weight shape dModel×(numHeads*headDim) (used for Q/K/V).

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.WoShape (dModel numHeads headDim : ℕ) : Spec.Shape

Output projection weight shape (numHeads*headDim)×dModel.

@[reducible, inline]

abbrev Proofs.Autograd.MultiHeadAttention.ΓMHA (n dModel numHeads headDim : ℕ) : List Spec.Shape

Input context shapes: [x, Wq, Wk, Wv, Wo].

def Proofs.Autograd.MultiHeadAttention.idxX {n dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMHA n dModel numHeads headDim ++ ss) (XShape n dModel)

Context index of the sequence input x in ΓMHA.

def Proofs.Autograd.MultiHeadAttention.idxWq {n dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMHA n dModel numHeads headDim ++ ss) (WqShape dModel numHeads headDim)

Context index of Wq in ΓMHA.

def Proofs.Autograd.MultiHeadAttention.idxWk {n dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMHA n dModel numHeads headDim ++ ss) (WqShape dModel numHeads headDim)

Context index of Wk in ΓMHA.

def Proofs.Autograd.MultiHeadAttention.idxWv {n dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMHA n dModel numHeads headDim ++ ss) (WqShape dModel numHeads headDim)

Context index of Wv in ΓMHA.

def Proofs.Autograd.MultiHeadAttention.idxWo {n dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMHA n dModel numHeads headDim ++ ss) (WoShape dModel numHeads headDim)

Context index of Wo in ΓMHA.

theorem Proofs.Autograd.MultiHeadAttention.size_big_to_heads (n numHeads headDim : ℕ) : (BigShape n numHeads headDim).size = (HeadsShape n numHeads headDim).size

Reshaping between the flat [n, numHeads * headDim] view and the split [n, numHeads, headDim] view preserves the element count.

This is what makes the head split a pure reinterpretation: no data moves, only the shape changes, so the reshape node is the identity on the underlying vector.

theorem Proofs.Autograd.MultiHeadAttention.size_swap_to_concat (n numHeads headDim : ℕ) : (Spec.Shape.dim n (Spec.Shape.dim numHeads (Spec.Shape.dim headDim Spec.Shape.scalar))).size = (BigShape n numHeads headDim).size

The same count equality in the direction used when concatenating heads back together.

noncomputable def Proofs.Autograd.MultiHeadAttention.appendScaleLast {Γ ss : List Spec.Shape} {s : Spec.Shape} (dg : DGraph Γ (ss ++ [s])) (c : ℝ) : DGraph Γ (ss ++ [s] ++ [s])

Append a constant scaling node for the most recently saved tensor.

noncomputable def Proofs.Autograd.MultiHeadAttention.appendBatchedSoftmaxLast {Γ ss : List Spec.Shape} {h m n : ℕ} (dg : DGraph Γ (ss ++ [Spec.Shape.dim h (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))])) : DGraph Γ (ss ++ [Spec.Shape.dim h (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))] ++ [Spec.Shape.dim h (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar))])

Append a row-wise softmax node for the most recently saved batched matrix.

noncomputable def Proofs.Autograd.MultiHeadAttention.mhaProjectionDGraph {n dModel numHeads headDim : ℕ} : DGraph (ΓMHA n dModel numHeads headDim) (ssMHAProjections n numHeads headDim)

Project the sequence input into query, transposed-key, and value head tensors.

noncomputable def Proofs.Autograd.MultiHeadAttention.mhaScoresDGraph {n dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (ΓMHA n dModel numHeads headDim) (ssMHAScores n numHeads headDim)

Form and scale the query-key score matrices for every head.

noncomputable def Proofs.Autograd.MultiHeadAttention.mhaAttentionDGraph {n dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (ΓMHA n dModel numHeads headDim) (ssMHAAttention n numHeads headDim)

Normalize the score rows and combine the resulting probabilities with the value heads.

noncomputable def Proofs.Autograd.MultiHeadAttention.mhaDGraph {n dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (ΓMHA n dModel numHeads headDim) (ssMHA n dModel numHeads headDim)

Multi-head self-attention as a proof-carrying graph.

This implements x Wq Wk Wv Wo ↦ Wo (concat_heads (softmax(c * (Q Kᵀ)) V)), with Q, K, and V projected from x. Reshaping and axis swaps model the usual runtime head split and merge.

theorem Proofs.Autograd.MultiHeadAttention.mha_backpropVec_eq_adjoint_fderiv {n dModel numHeads headDim : ℕ} (c : ℝ) (xV : CtxVec (ΓMHA n dModel numHeads headDim)) (seedV : CtxVec (ΓMHA n dModel numHeads headDim ++ ssMHA n dModel numHeads headDim)) : (mhaDGraph c).g.backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (mhaDGraph c).g.evalVec xV)) seedV

Corollary of the general DAG theorem: backprop equals (fderiv eval)† for the MHA graph.

This is the formal “VJP correctness” statement for the full MHA computation (as laid out by mhaDGraph).