Post-Norm Decoder Core

This module builds a decoder-core SSA graph with a fixed additive score bias and proves its VJP theorem (decoderCore_backpropVec_eq_adjoint_fderiv_at). Map-level composition uses twoSublayerPostNormBlock_hasFDerivAt from PostNorm.lean.

The concrete graph starts with already projected head tensors. Its merge and FFN affine maps are fixed arguments, not trainable entries in the differentiated context. The score bias is finite real data; it does not implement Boolean hard or causal masking. These are exact-real derivative theorems, not a correspondence theorem for the executable decoder or LibTorch.

The map-level composition theorem accepts differentiable residual-pack maps. The attention composition theorem in NN.Proofs.Autograd.Tape.Ops.Attention.MaskedMultiHeadSelfAttention can supply such a map when the caller provides the projection and merge derivative certificates. It does not construct a full decoder from a token/parameter context automatically.

Concrete fixed-bias decoder-core SSA graph

The concrete graph below starts after the Q/K/V projection split:

[Q_heads, Kᵀ_heads, V_heads, residual_stream, gamma₁, beta₁, gamma₂, beta₂].

It then adds the fixed bias to the scaled attention scores, merges the head output through a supplied affine map, adds the residual stream, and applies the two post-norm sublayers.

@[reducible, inline]

abbrev Proofs.Autograd.Transformer.ΓDecoderCore (seqLen dModel numHeads headDim : ℕ) : List Spec.Shape

Concrete decoder-core context: head tensors, residual stream, and two LayerNorm parameter pairs.

@[reducible, inline]

abbrev Proofs.Autograd.Transformer.ssDecoderCore (seqLen dModel numHeads headDim dFF : ℕ) : List Spec.Shape

Saved tensors for the concrete fixed-bias decoder-core block.

def Proofs.Autograd.Transformer.idxDecoderResidualInput {seqLen dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ ss) (LayerNorm.MatShape seqLen dModel)

Residual-stream input in the decoder-core context.

def Proofs.Autograd.Transformer.idxDecoderNorm1Gamma {seqLen dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ ss) (LayerNorm.VecShape dModel)

First LayerNorm scale parameter in the decoder-core context.

def Proofs.Autograd.Transformer.idxDecoderNorm1Beta {seqLen dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ ss) (LayerNorm.VecShape dModel)

First LayerNorm shift parameter in the decoder-core context.

def Proofs.Autograd.Transformer.idxDecoderNorm2Gamma {seqLen dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ ss) (LayerNorm.VecShape dModel)

Second LayerNorm scale parameter in the decoder-core context.

def Proofs.Autograd.Transformer.idxDecoderNorm2Beta {seqLen dModel numHeads headDim : ℕ} {ss : List Spec.Shape} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ ss) (LayerNorm.VecShape dModel)

Second LayerNorm shift parameter in the decoder-core context.

noncomputable def Proofs.Autograd.Transformer.decoderMaskedCoreDGraph {seqLen dModel numHeads headDim : ℕ} (c : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) : DGraph (ΓDecoderCore seqLen dModel numHeads headDim) (DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim)

Fixed-bias attention core while carrying residual and LayerNorm parameters.

def Proofs.Autograd.Transformer.idxDecoderHeadOut {seqLen dModel numHeads headDim : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim) (DirectReshapeAttention.HeadsShape seqLen numHeads headDim)

Head output after the fixed-bias attention core.

def Proofs.Autograd.Transformer.idxDecoderMergedAttention {seqLen dModel numHeads headDim : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel]) (LayerNorm.MatShape seqLen dModel)

Merged attention output after the supplied output projection.

def Proofs.Autograd.Transformer.idxDecoderResidualInputAfterMerge {seqLen dModel numHeads headDim : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel]) (LayerNorm.MatShape seqLen dModel)

Residual input weakened past the merged attention output.

def Proofs.Autograd.Transformer.idxDecoderAttentionResidual {seqLen dModel numHeads headDim : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel]) (LayerNorm.MatShape seqLen dModel)

Sum of the residual stream and projected attention output before the first LayerNorm.

def Proofs.Autograd.Transformer.decoderNorm1Inputs {seqLen dModel numHeads headDim : ℕ} : LayerNorm.Inputs (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel]) seqLen dModel

First LayerNorm input triple for the concrete decoder-core graph.

noncomputable def Proofs.Autograd.Transformer.decoderAfterNorm1Graph {seqLen dModel numHeads headDim : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (c ε₁ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) : Graph (ΓDecoderCore seqLen dModel numHeads headDim) (DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel])

Decoder graph through the first post-norm attention sublayer.

def Proofs.Autograd.Transformer.idxDecoderNorm1Out {seqLen dModel numHeads headDim : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

First decoder post-norm output.

def Proofs.Autograd.Transformer.idxDecoderNorm1OutAfterFfn {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF, SeqFFNModelShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

First decoder post-norm output weakened through FFN intermediates.

def Proofs.Autograd.Transformer.idxDecoderFfnHiddenPre {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF]) (SeqFFNHiddenShape seqLen dFF)

First FFN affine output in the concrete decoder graph.

def Proofs.Autograd.Transformer.idxDecoderFfnHiddenAct {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF]) (SeqFFNHiddenShape seqLen dFF)

FFN activation output in the concrete decoder graph.

def Proofs.Autograd.Transformer.idxDecoderFfnProjected {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF, SeqFFNModelShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

FFN projection output in the concrete decoder graph.

def Proofs.Autograd.Transformer.idxDecoderFfnResidual {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF) (SeqFFNModelShape seqLen dModel)

FFN residual output before the second decoder LayerNorm.

noncomputable def Proofs.Autograd.Transformer.decoderFfnResidualGraph {seqLen dModel numHeads headDim dFF : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (fc1 : Vec (SeqFFNModelShape seqLen dModel).size →L[ℝ] Vec (SeqFFNHiddenShape seqLen dFF).size) (b1 : Vec (SeqFFNHiddenShape seqLen dFF).size) (fc2 : Vec (SeqFFNHiddenShape seqLen dFF).size →L[ℝ] Vec (SeqFFNModelShape seqLen dModel).size) (b2 : Vec (SeqFFNModelShape seqLen dModel).size) (c ε₁ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) : Graph (ΓDecoderCore seqLen dModel numHeads headDim) (DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF)

Decoder graph through the FFN residual.

def Proofs.Autograd.Transformer.decoderNorm2Inputs {seqLen dModel numHeads headDim dFF : ℕ} : LayerNorm.Inputs (ΓDecoderCore seqLen dModel numHeads headDim ++ DirectReshapeAttention.ssMaskedCore seqLen numHeads headDim ++ [LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel, LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF) seqLen dModel

Second LayerNorm input triple for the concrete decoder-core graph.

noncomputable def Proofs.Autograd.Transformer.decoderCoreGraph {seqLen dModel numHeads headDim dFF : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (fc1 : Vec (SeqFFNModelShape seqLen dModel).size →L[ℝ] Vec (SeqFFNHiddenShape seqLen dFF).size) (b1 : Vec (SeqFFNHiddenShape seqLen dFF).size) (fc2 : Vec (SeqFFNHiddenShape seqLen dFF).size →L[ℝ] Vec (SeqFFNModelShape seqLen dModel).size) (b2 : Vec (SeqFFNModelShape seqLen dModel).size) (c ε₁ ε₂ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) : Graph (ΓDecoderCore seqLen dModel numHeads headDim) (ssDecoderCore seqLen dModel numHeads headDim dFF)

Concrete SSA graph for one additive-bias decoder-core block.

noncomputable def Proofs.Autograd.Transformer.decoderFfnResidualGraphFDerivCorrectAt {seqLen dModel numHeads headDim dFF : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (fc1 : Vec (SeqFFNModelShape seqLen dModel).size →L[ℝ] Vec (SeqFFNHiddenShape seqLen dFF).size) (b1 : Vec (SeqFFNHiddenShape seqLen dFF).size) (fc2 : Vec (SeqFFNHiddenShape seqLen dFF).size →L[ℝ] Vec (SeqFFNModelShape seqLen dModel).size) (b2 : Vec (SeqFFNModelShape seqLen dModel).size) (c ε₁ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) (xV : CtxVec (ΓDecoderCore seqLen dModel numHeads headDim)) (hε₁ : 0 < ε₁) : GraphFDerivCorrectAt (decoderFfnResidualGraph merge mergeBias fc1 b1 fc2 b2 c ε₁ bias) xV

Pointwise analytic correctness for the decoder graph through the FFN residual.

noncomputable def Proofs.Autograd.Transformer.decoderCoreGraphFDerivCorrectAt {seqLen dModel numHeads headDim dFF : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (fc1 : Vec (SeqFFNModelShape seqLen dModel).size →L[ℝ] Vec (SeqFFNHiddenShape seqLen dFF).size) (b1 : Vec (SeqFFNHiddenShape seqLen dFF).size) (fc2 : Vec (SeqFFNHiddenShape seqLen dFF).size →L[ℝ] Vec (SeqFFNModelShape seqLen dModel).size) (b2 : Vec (SeqFFNModelShape seqLen dModel).size) (c ε₁ ε₂ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) (xV : CtxVec (ΓDecoderCore seqLen dModel numHeads headDim)) (hε₁ : 0 < ε₁) (hε₂ : 0 < ε₂) : GraphFDerivCorrectAt (decoderCoreGraph merge mergeBias fc1 b1 fc2 b2 c ε₁ ε₂ bias) xV

Pointwise analytic correctness for the complete concrete decoder-core graph.

theorem Proofs.Autograd.Transformer.decoderCore_backpropVec_eq_adjoint_fderiv_at {seqLen dModel numHeads headDim dFF : ℕ} (merge : Vec (DirectReshapeAttention.HeadsShape seqLen numHeads headDim).size →L[ℝ] Vec (LayerNorm.MatShape seqLen dModel).size) (mergeBias : Vec (LayerNorm.MatShape seqLen dModel).size) (fc1 : Vec (SeqFFNModelShape seqLen dModel).size →L[ℝ] Vec (SeqFFNHiddenShape seqLen dFF).size) (b1 : Vec (SeqFFNHiddenShape seqLen dFF).size) (fc2 : Vec (SeqFFNHiddenShape seqLen dFF).size →L[ℝ] Vec (SeqFFNModelShape seqLen dModel).size) (b2 : Vec (SeqFFNModelShape seqLen dModel).size) (c ε₁ ε₂ : ℝ) (bias : Vec (DirectReshapeAttention.ScoresShape seqLen numHeads).size := 0) (xV : CtxVec (ΓDecoderCore seqLen dModel numHeads headDim)) (seedV : CtxVec (ΓDecoderCore seqLen dModel numHeads headDim ++ ssDecoderCore seqLen dModel numHeads headDim dFF)) (hε₁ : 0 < ε₁) (hε₂ : 0 < ε₂) : (decoderCoreGraph merge mergeBias fc1 b1 fc2 b2 c ε₁ ε₂ bias).backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (decoderCoreGraph merge mergeBias fc1 b1 fc2 b2 c ε₁ ε₂ bias).evalVec xV)) seedV

End-to-end VJP theorem for the concrete additive-bias decoder-core graph.