Multi-head attention with a direct reshape split

End-to-end fderiv/backprop correctness for a multi-head attention variant with a direct reshape split, 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 a graph-level result over ℝ. It is a corollary of the general graph theorem once each node used by the graph has a NodeFDerivCorrect instance.

The Q/K/V split here directly reshapes [n, numHeads * headDim] to [numHeads, n, headDim], preserving flat coordinate order. In contrast, Spec.splitHeadsSpec first reshapes to [n, numHeads, headDim] and swaps the token/head axes; the executable attention paths use that permutation. For n = numHeads = headDim = 2, flat projection entries 0,1,2,3,4,5,6,7 stay in that order in this graph, while the spec split produces 0,1,4,5,2,3,6,7.

The theorem below proves the derivative of this declared graph. It does not establish equality with Spec.MultiHeadAttention.forward or the executable attention implementation.

PyTorch correspondence / citations

• The construction has the “project → split heads → scaled dot-product attention → concat heads → output projection” stages of torch.nn.MultiheadAttention, with the different split layout described above. 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.DirectReshapeAttention.XShape (n dModel : ℕ) : Spec.Shape

Sequence input shape n×dModel.

@[reducible, inline]

abbrev Proofs.Autograd.DirectReshapeAttention.FlatHeadsShape (n numHeads headDim : ℕ) : Spec.Shape

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

@[reducible, inline]

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

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

@[reducible, inline]

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

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

@[reducible, inline]

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

Attention scores shape (numHeads)×n×n.

@[reducible, inline]

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

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

@[reducible, inline]

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

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

@[reducible, inline]

abbrev Proofs.Autograd.DirectReshapeAttention.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.DirectReshapeAttention.ssMHAScores (n numHeads headDim : ℕ) : List Spec.Shape

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

@[reducible, inline]

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

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

@[reducible, inline]

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

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

@[reducible, inline]

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

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

@[reducible, inline]

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

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

def Proofs.Autograd.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.size_flatHeads_to_heads (n numHeads headDim : ℕ) : (FlatHeadsShape n numHeads headDim).size = (HeadsShape n numHeads headDim).size

Reshaping [n, numHeads * headDim] directly to [numHeads, n, headDim] preserves the element count.

The graph's reshape node preserves flat coordinate order. This count equality does not provide the token/head axis permutation used by Spec.splitHeadsSpec.

theorem Proofs.Autograd.DirectReshapeAttention.size_swap_to_concat (n numHeads headDim : ℕ) : (SwappedShape n numHeads headDim).size = (FlatHeadsShape n numHeads headDim).size

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

noncomputable def Proofs.Autograd.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.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.DirectReshapeAttention.mhaDGraph {n dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (ΓMHA n dModel numHeads headDim) (ssMHA n dModel numHeads headDim)

Multi-head attention with a direct reshape split as a proof-carrying graph.

This implements x Wq Wk Wv Wo ↦ concat_heads (softmax(c * (Q Kᵀ)) V) * Wo, with Q, K, and V projected from x. Each projection is reshaped directly to [numHeads, n, headDim] without permuting its flat entries. The output merge swaps the head/token axes before flattening. These are the graph's declared operations; its input split differs from Spec.splitHeadsSpec.

theorem Proofs.Autograd.DirectReshapeAttention.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).