Scaled dot-product attention: tape graph versus specification

NN.Proofs.Autograd.Tape.Ops.Attention.ScaledDotProduct proves backprop = (fderiv eval)† for a tape graph built from proven nodes. NN.Spec.Layers.Attention defines Spec.scaledDotProductAttention directly on tensors. This file connects the two:

• get_sdpaOutIdx_evalVec: the output block of the graph evaluation is exactly the vectorized specification forward pass softmax(Q Kᵀ / √d) V;
• hasFDerivAt_sdpaSpecForwardVec: the specification forward pass is Fréchet-differentiable, with derivative read off the graph;
• backpropVec_eq_adjoint_fderiv_scaledDotProductAttention: the tape reverse pass seeded on the output block is the vector-Jacobian product of Spec.scaledDotProductAttention;
• backpropVec_eq_adjoint_fderiv_scaledDotProductAttention_allTrueMask: the same statement for the masked code path with the all-true Boolean mask.

Blocks of the evaluated graph are read with the Graph.evalVec lemmas of NN.Proofs.Autograd.Tape.Ops.Norm.CtxVecEval, and matrix coordinates are read through TapeNodes.Matmul.tensorToVec_idxMN: attention is proved in the flat Vec world where Mathlib's calculus lives, and stated about Tensor ℝ [m, n]. The first half relates the flattened node functions (matmulVec, transposeVec, forwardMN) to Spec.matMulSpec, swapAdjacentAxes, scaleSpec, and Activation.softmaxSpec 1; the second half walks the attention graph node by node.

Flattened matrices

theorem Proofs.Autograd.tensorToVec_scaleSpec {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) (c : ℝ) : tensorToVec (A.scaleSpec c) = c • tensorToVec A

Vectorization is homogeneous: scaling a matrix scales its vector. This is what lets the 1/√d factor in scaled dot-product attention be handled as ordinary scalar multiplication.

theorem Proofs.Autograd.transposeVec_idxMN {m n : ℕ} (a : Vec (TapeNodes.Matmul.matSize m n)) (i : Fin n) (j : Fin m) : (TapeNodes.MatTranspose.transposeVec a).ofLp (TapeNodes.Matmul.idxMN i j) = a.ofLp (TapeNodes.Matmul.idxMN j i)

Transposition on flat vectors swaps the row and column index, as expected.

theorem Proofs.Autograd.tensorToVec_swapAdjacentAxes {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) : tensorToVec (A.swapAdjacentAxes 0) = TapeNodes.MatTranspose.transposeVec (tensorToVec A)

Swapping the two axes of a matrix tensor corresponds to the flat transpose.

Attention transposes the key matrix, so without this the Kᵀ in Q Kᵀ would have to be reasoned about at the tensor level and at the vector level separately.

theorem Proofs.Autograd.castVec_tensorToVec_finProdFinEquiv {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) (hsz : (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)).size = m * n) (i : Fin m) (j : Fin n) : (castVec hsz (tensorToVec A)).ofLp (finProdFinEquiv (i, j)) = Spec.get2 A i j

Reading a vectorized matrix through the Fin m × Fin n product equivalence, after the size cast, gives the matrix entry. This is TapeNodes.Matmul.tensorToVec_idxMN in the indexing the softmax development uses.

theorem Proofs.Autograd.rows_castVec_tensorToVec {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) (hsz : (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)).size = m * n) (i : Fin m) : TapeNodes.SoftmaxLastAxis.rows (castVec hsz (tensorToVec A)) i = getScalarE (Spec.get A i)

Row i of a vectorized matrix is the vectorization of row i of the tensor.

theorem Proofs.Autograd.castVec_tensorToVec_softmaxSpec_one {m n : ℕ} (S : TorchLean.Tensor ℝ [m, n]) (hsz : (Spec.Shape.dim m (Spec.Shape.dim n Spec.Shape.scalar)).size = m * n) : castVec hsz (tensorToVec (Activation.softmaxSpec 1 S)) = TapeNodes.SoftmaxLastAxis.forwardMN (castVec hsz (tensorToVec S))

Row-wise softmax on tensors agrees with the flat forwardMN softmax on vectors.

The temperature is fixed to 1 because that is the only case attention needs; the scaling is already folded into the scores by tensorToVec_scaleSpec.

Attention graph evaluation

def Proofs.Autograd.Attention.sdpaOutIdx (m d : ℕ) : Idx (ΓQKV m d ++ ssScaledDotProduct m d) (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar))

Index of the attention output block in the saved context.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.ctxQ {m d : ℕ} (xV : CtxVec (ΓQKV m d)) : TorchLean.Tensor ℝ [m, d]

The Q tensor stored in a vectorized attention context.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.ctxK {m d : ℕ} (xV : CtxVec (ΓQKV m d)) : TorchLean.Tensor ℝ [m, d]

The K tensor stored in a vectorized attention context.

@[reducible, inline]

noncomputable abbrev Proofs.Autograd.Attention.ctxV {m d : ℕ} (xV : CtxVec (ΓQKV m d)) : TorchLean.Tensor ℝ [m, d]

The V tensor stored in a vectorized attention context.

noncomputable def Proofs.Autograd.Attention.attentionForwardSpec {m d : ℕ} (c : ℝ) (Q K V : TorchLean.Tensor ℝ [m, d]) : TorchLean.Tensor ℝ [m, d]

Tensor-level attention forward pass with an explicit scale factor c.

theorem Proofs.Autograd.Attention.scaledDotProductAttention_eq_attentionForwardSpec {m d : ℕ} (hm : m ≠ 0) (Q K V : TorchLean.Tensor ℝ [m, d]) : Spec.scaledDotProductAttention { Q := Q, K := K, V := V, mask := none } = attentionForwardSpec (1 / Spec.attentionScaleDenom d) Q K V

Unmasked Spec.scaledDotProductAttention is attentionForwardSpec at the scale 1 / √d.

theorem Proofs.Autograd.Attention.get_sdpaOutIdx_evalVec {m d : ℕ} (c : ℝ) (xV : CtxVec (ΓQKV m d)) : CtxVec.get (sdpaOutIdx m d) ((scaledDotProductGraph c).evalVec xV) = tensorToVec (attentionForwardSpec c (ctxQ xV) (ctxK xV) (ctxV xV))

The output block of the attention graph is the vectorized tensor forward pass.

noncomputable def Proofs.Autograd.Attention.sdpaSpecForwardVec {m d : ℕ} (c : ℝ) (xV : CtxVec (ΓQKV m d)) : Vec (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar)).size

The specification forward pass as a function on vectorized contexts.

theorem Proofs.Autograd.Attention.sdpaSpecForwardVec_eq_getCLM_evalVec {m d : ℕ} (c : ℝ) : sdpaSpecForwardVec c = fun (xV : CtxVec (ΓQKV m d)) => (CtxVec.getCLM (sdpaOutIdx m d)) ((scaledDotProductGraph c).evalVec xV)

The specification forward pass is the output projection of the graph evaluation.

theorem Proofs.Autograd.Attention.hasFDerivAt_sdpaSpecForwardVec {m d : ℕ} (c : ℝ) (xV : CtxVec (ΓQKV m d)) : HasFDerivAt (sdpaSpecForwardVec c) (CtxVec.getCLM (sdpaOutIdx m d) ∘SL fderiv ℝ (scaledDotProductGraph c).evalVec xV) xV

The specification attention forward pass is Fréchet-differentiable; its derivative is the output projection of the graph derivative.

theorem Proofs.Autograd.Attention.backpropVec_eq_adjoint_fderiv_sdpaSpec {m d : ℕ} (c : ℝ) (xV : CtxVec (ΓQKV m d)) (δ : Vec (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar)).size) : (scaledDotProductGraph c).backpropVec xV (CtxVec.single (sdpaOutIdx m d) δ) = (ContinuousLinearMap.adjoint (fderiv ℝ (sdpaSpecForwardVec c) xV)) δ

Tape backprop seeded on the output block is the VJP of the specification forward pass.

theorem Proofs.Autograd.Attention.get_sdpaOutIdx_evalVec_scaledDotProductAttention {m d : ℕ} (hm : m ≠ 0) (xV : CtxVec (ΓQKV m d)) : CtxVec.get (sdpaOutIdx m d) ((scaledDotProductGraph (1 / Spec.attentionScaleDenom d)).evalVec xV) = tensorToVec (Spec.scaledDotProductAttention { Q := ctxQ xV, K := ctxK xV, V := ctxV xV, mask := none })

The output block of the attention graph at scale 1 / √d is the vectorized Spec.scaledDotProductAttention forward pass.

theorem Proofs.Autograd.Attention.backpropVec_eq_adjoint_fderiv_scaledDotProductAttention {m d : ℕ} (hm : m ≠ 0) (xV : CtxVec (ΓQKV m d)) (δ : Vec (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar)).size) : (scaledDotProductGraph (1 / Spec.attentionScaleDenom d)).backpropVec xV (CtxVec.single (sdpaOutIdx m d) δ) = (ContinuousLinearMap.adjoint (fderiv ℝ (fun (xV : CtxVec (ΓQKV m d)) => tensorToVec (Spec.scaledDotProductAttention { Q := ctxQ xV, K := ctxK xV, V := ctxV xV, mask := none })) xV)) δ

Tape backprop is the vector-Jacobian product of unmasked Spec.scaledDotProductAttention.

theorem Proofs.Autograd.Attention.backpropVec_eq_adjoint_fderiv_scaledDotProductAttention_allTrueMask {m d : ℕ} (hm : m ≠ 0) (xV : CtxVec (ΓQKV m d)) (δ : Vec (Spec.Shape.dim m (Spec.Shape.dim d Spec.Shape.scalar)).size) : (scaledDotProductGraph (1 / Spec.attentionScaleDenom d)).backpropVec xV (CtxVec.single (sdpaOutIdx m d) δ) = (ContinuousLinearMap.adjoint (fderiv ℝ (fun (xV : CtxVec (ΓQKV m d)) => tensorToVec (Spec.scaledDotProductAttention { Q := ctxQ xV, K := ctxK xV, V := ctxV xV, mask := some (Spec.allTrueMask m m) })) xV)) δ

With the all-true Boolean mask, the masked code path of Spec.scaledDotProductAttention has the same vector-Jacobian product, computed by the same tape backprop.