Additive-Bias Multi-Head Attention Core

This module proves the head-wise fixed-score-bias attention core. The surrounding projection/split/merge graph lives in MultiHeadSelfAttention.lean; the theorem here is the reusable additive-bias replacement for the score/probability/value part of that graph:

softmax(c · QKᵀ + bias) V.

The bias has shape (heads × seq × seq). This is not the boolean causal-mask semantics: boolean attention masks in the spec/runtime path use hard masking, where blocked entries contribute zero softmax numerator. This file is for intentional finite additive score biases.

@[reducible, inline]

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

Context for a fixed-bias attention core: [Q_heads, Kᵀ_heads, V_heads].

@[reducible, inline]

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

Saved tensors for the fixed-bias multi-head attention core.

def Proofs.Autograd.DirectReshapeAttention.idxMaskedCoreQ {n numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMaskedCore n numHeads headDim ++ ss) (HeadsShape n numHeads headDim)

Query-head index in the fixed-bias attention context.

def Proofs.Autograd.DirectReshapeAttention.idxMaskedCoreKt {n numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMaskedCore n numHeads headDim ++ ss) (KtShape n numHeads headDim)

Transposed-key index in the fixed-bias attention context.

def Proofs.Autograd.DirectReshapeAttention.idxMaskedCoreV {n numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓMaskedCore n numHeads headDim ++ ss) (HeadsShape n numHeads headDim)

Value-head index in the fixed-bias attention context.

noncomputable def Proofs.Autograd.DirectReshapeAttention.maskedCoreDGraph {n numHeads headDim : ℕ} (c : ℝ) (bias : Vec (ScoresShape n numHeads).size := 0) : DGraph (ΓMaskedCore n numHeads headDim) (ssMaskedCore n numHeads headDim)

Proof-carrying fixed-bias attention core.

The fixed bias is added after scaling the score tensor and before the row-wise softmax.

theorem Proofs.Autograd.DirectReshapeAttention.maskedCore_backpropVec_eq_adjoint_fderiv {n numHeads headDim : ℕ} (c : ℝ) (bias : Vec (ScoresShape n numHeads).size := 0) (xV : CtxVec (ΓMaskedCore n numHeads headDim)) (seedV : CtxVec (ΓMaskedCore n numHeads headDim ++ ssMaskedCore n numHeads headDim)) : (maskedCoreDGraph c bias).g.backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (maskedCoreDGraph c bias).g.evalVec xV)) seedV

Reverse-mode theorem for the fixed-bias multi-head attention core.

theorem Proofs.Autograd.DirectReshapeAttention.maskedCoreAfterProjection_hasFDerivAt {E : Type u} [NormedAddCommGroup E] [NormedSpace ℝ E] {n numHeads headDim : ℕ} (c : ℝ) (bias : Vec (ScoresShape n numHeads).size := 0) (projectPack : E → CtxVec (ΓMaskedCore n numHeads headDim)) (DprojectPack : E →L[ℝ] CtxVec (ΓMaskedCore n numHeads headDim)) (x : E) (hProject : HasFDerivAt projectPack DprojectPack x) : HasFDerivAt (fun (z : E) => (maskedCoreDGraph c bias).g.evalVec (projectPack z)) (fderiv ℝ (maskedCoreDGraph c bias).g.evalVec (projectPack x) ∘SL DprojectPack) x

Composition theorem for the additive-bias core after a projection/split stage.

projectPack is the mathematical interface exposed by a full MHA front half: it builds [Q_heads, Kᵀ_heads, V_heads] from an outer context. The theorem composes that front half with the proved additive-bias attention core.

theorem Proofs.Autograd.DirectReshapeAttention.projectedMaskedAttention_hasFDerivAt {E : Type u} [NormedAddCommGroup E] [NormedSpace ℝ E] {F : Type v} [NormedAddCommGroup F] [NormedSpace ℝ F] {n numHeads headDim : ℕ} (c : ℝ) (bias : Vec (ScoresShape n numHeads).size := 0) (projectPack : E → CtxVec (ΓMaskedCore n numHeads headDim)) (DprojectPack : E →L[ℝ] CtxVec (ΓMaskedCore n numHeads headDim)) (mergePack : CtxVec (ΓMaskedCore n numHeads headDim ++ ssMaskedCore n numHeads headDim) → F) (DmergePack : CtxVec (ΓMaskedCore n numHeads headDim ++ ssMaskedCore n numHeads headDim) →L[ℝ] F) (x : E) (hProject : HasFDerivAt projectPack DprojectPack x) (hMerge : HasFDerivAt mergePack DmergePack ((maskedCoreDGraph c bias).g.evalVec (projectPack x))) : HasFDerivAt (fun (z : E) => mergePack ((maskedCoreDGraph c bias).g.evalVec (projectPack z))) (DmergePack ∘SL fderiv ℝ (maskedCoreDGraph c bias).g.evalVec (projectPack x) ∘SL DprojectPack) x

Fixed-bias attention composition contract.

The theorem composes:

• a differentiable front half that projects/splits inputs into Q, Kᵀ, and V;
• the proved fixed-bias split-head attention core;
• a differentiable back half that merges/projects outputs or packages residual data for the caller.

Instantiating projectPack and mergePack with the concrete projection/split/merge graphs gives the corresponding fixed-bias block's differentiability statement. The theorem assumes the two outer derivative certificates; it does not construct them or prove Boolean causal masking.