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. The first half of this file 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.val_finProdFinEquiv {m k : ℕ} (i : Fin m) (j : Fin k) : ↑(finProdFinEquiv (i, j)) = ↑j + k * ↑i

theorem Proofs.Autograd.val_idxMN {m n : ℕ} (i : Fin m) (j : Fin n) : ↑(TapeNodes.Matmul.idxMN i j) = ↑j + n * ↑i

The flattened index of entry (i, j) is j + n * i: row-major layout, as in PyTorch.

theorem Proofs.Autograd.exists_idxMN {m n : ℕ} (ip : Fin (TapeNodes.Matmul.matSize m n)) : ∃ (i : Fin m) (j : Fin n), TapeNodes.Matmul.idxMN i j = ip

Every flattened matrix index is the image of a row/column pair.

theorem Proofs.Autograd.tensorToVec_idxMN {m n : ℕ} (A : TorchLean.Tensor ℝ [m, n]) (i : Fin m) (j : Fin n) : (tensorToVec A).ofLp (TapeNodes.Matmul.idxMN i j) = Spec.get2 A i j

Coordinate idxMN i j of a vectorized matrix is the matrix entry A i j.

This is the bridge the whole file is named for. Attention is proved in the flat Vec world, where Mathlib's calculus lives, and stated about Tensor ℝ [m, n]; this lemma is what connects the two without ever unfolding the flattening.

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 Qᵀ 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 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.adjoint_getCLM {Γ : List Spec.Shape} {s : Spec.Shape} (idx : Idx Γ s) (v : Vec s.size) : (ContinuousLinearMap.adjoint (CtxVec.getCLM idx)) v = CtxVec.single idx v

The adjoint of block projection is block injection.

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.