ScaledDotProduct

End-to-end fderiv/backprop correctness for a scaled dot-product attention graph, built out of the proven tape nodes (matmul, matrixTranspose, scale, softmaxLast).

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

• This file matches the usual mathematical definition of scaled dot-product attention: softmax(c * Q Kᵀ) V. In PyTorch this corresponds to building the same computation with torch.matmul + torch.softmax, or using the dedicated helper: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

@[reducible, inline]

abbrev Proofs.Autograd.Attention.QKVShape (m d : ℕ) : Spec.Shape

Matrix shape for m×d Q/K/V inputs.

@[reducible, inline]

abbrev Proofs.Autograd.Attention.ΓQKV (m d : ℕ) : List Spec.Shape

Input context shapes: [Q, K, V], each m×d.

def Proofs.Autograd.Attention.idxQ {m d : ℕ} {ss : List Spec.Shape} : Idx (ΓQKV m d ++ ss) (QKVShape m d)

Context index of Q in ΓQKV.

def Proofs.Autograd.Attention.idxK {m d : ℕ} {ss : List Spec.Shape} : Idx (ΓQKV m d ++ ss) (QKVShape m d)

Context index of K in ΓQKV.

def Proofs.Autograd.Attention.idxV {m d : ℕ} {ss : List Spec.Shape} : Idx (ΓQKV m d ++ ss) (QKVShape m d)

Context index of V in ΓQKV.

@[reducible, inline]

abbrev Proofs.Autograd.Attention.ssScaledDotProduct (m d : ℕ) : List Spec.Shape

Saved tensors of the attention graph: Kᵀ, logits, scaled logits, probabilities, output.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.nodeKt (m d : ℕ) : Node (ΓQKV m d) (Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar))

Node 1: Kᵀ.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.nodeLogits (m d : ℕ) : Node (ΓQKV m d ++ [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar)]) (Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar))

Node 2: logits Q Kᵀ.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.nodeScaled (m d : ℕ) (c : ℝ) : Node (ΓQKV m d ++ [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]) (Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar))

Node 3: scaled logits c * (Q Kᵀ).

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.nodeProbs (m d : ℕ) : Node (ΓQKV m d ++ [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]) (Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar))

Node 4: row-wise softmax probabilities.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.nodeOut (m d : ℕ) : Node (ΓQKV m d ++ [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]) (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar))

Node 5: output probs * V.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.graphKt (m d : ℕ) : Graph (ΓQKV m d) [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar)]

Attention graph prefix through Kᵀ.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.graphLogits (m d : ℕ) : Graph (ΓQKV m d) [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]

Attention graph prefix through the logits.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.graphScaled (m d : ℕ) (c : ℝ) : Graph (ΓQKV m d) [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]

Attention graph prefix through the scaled logits.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.graphProbs (m d : ℕ) (c : ℝ) : Graph (ΓQKV m d) [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar)]

Attention graph prefix through the softmax probabilities.

noncomputable def Proofs.Autograd.Attention.scaledDotProductGraph {m d : ℕ} (c : ℝ) : Graph (ΓQKV m d) (ssScaledDotProduct m d)

Scaled dot-product attention as an explicit tape graph.

Computes Q K V ↦ softmax(c * (Q * Kᵀ)) * V and records intermediate values needed by backprop. The graph is a plain Graph.snoc chain so that Graph.evalVec can be computed node by node when relating the tape to the specification forward pass.

noncomputable def Proofs.Autograd.Attention.scaledDotProductDGraph {m d : ℕ} (c : ℝ) : DGraph (ΓQKV m d) (ssScaledDotProduct m d)

Scaled dot-product attention as a proved-correct DGraph.

Every node of scaledDotProductGraph carries its NodeFDerivCorrect certificate.

theorem Proofs.Autograd.Attention.backprop_eq_adjoint_fderiv_scaledDotProduct {m d : ℕ} (c : ℝ) (xV : CtxVec (ΓQKV m d)) (seedV : CtxVec (ΓQKV m d ++ [Spec.Shape.dim d (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim m Spec.Shape.scalar), Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar)])) : (scaledDotProductDGraph c).g.backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (scaledDotProductDGraph c).g.evalVec xV)) seedV

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

This is the formal statement that the tape reverse pass computes the VJP for the full attention computation.