Post-Norm Transformer Encoder Block

This module builds one concrete SSA graph for a post-norm encoder block and proves its end-to-end VJP theorem:

x ↦ LN₂(FFN(LN₁(MHA(x) + x)) + LN₁(MHA(x) + x)).

The two sublayers already have graph-level VJP theorems in PostNorm.lean. encoderBlockGraph assembles them into a single graph over the context [x, Wq, Wk, Wv, Wo, gamma₁, beta₁, gamma₂, beta₂] with VJP theorem encoderBlock_backpropVec_eq_adjoint_fderiv_at, and twoSublayerPostNormBlock_hasFDerivAt gives map-level composition for a model proof that needs the whole encoder block as one differentiable map.

Attention uses the direct-reshape variant from DirectReshapeAttention, whose head split differs from the executable attention layout. The FFN maps and biases are fixed inputs to this construction, not trainable entries in its context. The theorem proves this exact-real graph's derivative and reverse pass, not correspondence with the public runtime encoder.

Concrete SSA encoder block

The definitions below assemble one mathematical proof graph for a post-norm encoder block. The context is:

[x, Wq, Wk, Wv, Wo, gamma₁, beta₁, gamma₂, beta₂].

The FFN affine maps are supplied as fixed sequence-level linear maps, matching the interface used by seqFfnResidualDGraph.

@[reducible, inline]

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

Full encoder-block context: MHA parameters, first LayerNorm affine parameters, second LayerNorm affine parameters.

@[reducible, inline]

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

Saved tensors for one post-norm encoder block.

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

First LayerNorm scale parameter in the full encoder-block context.

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

First LayerNorm shift parameter in the full encoder-block context.

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

Second LayerNorm scale parameter in the full encoder-block context.

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

Second LayerNorm shift parameter in the full encoder-block context.

noncomputable def Proofs.Autograd.Transformer.encoderMhaResidualDGraph {seqLen dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (ΓEncoderBlock seqLen dModel numHeads headDim) (ssMHAResidual seqLen dModel numHeads headDim)

Residual-attention prefix while carrying both LayerNorm parameter pairs.

def Proofs.Autograd.Transformer.idxEncoderMhaResidual {seqLen dModel numHeads headDim : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim) (LayerNorm.MatShape seqLen dModel)

Residual stream x + MHA(x) after the attention prefix.

def Proofs.Autograd.Transformer.encoderNorm1Inputs {seqLen dModel numHeads headDim : ℕ} : LayerNorm.Inputs (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim) seqLen dModel

First LayerNorm input triple in the concrete encoder block.

noncomputable def Proofs.Autograd.Transformer.encoderAfterNorm1Graph {seqLen dModel numHeads headDim : ℕ} (c ε₁ : ℝ) : Graph (ΓEncoderBlock seqLen dModel numHeads headDim) (ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel])

Encoder graph through the first post-norm attention sublayer.

def Proofs.Autograd.Transformer.idxEncoderNorm1Out {seqLen dModel numHeads headDim : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

First post-norm sublayer output.

def Proofs.Autograd.Transformer.idxEncoderNorm1OutAfterFfn {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF, SeqFFNModelShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

First post-norm output weakened through the FFN intermediates.

def Proofs.Autograd.Transformer.idxEncoderFfnHiddenPre {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF]) (SeqFFNHiddenShape seqLen dFF)

First FFN affine output in the full encoder graph.

def Proofs.Autograd.Transformer.idxEncoderFfnHiddenAct {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF]) (SeqFFNHiddenShape seqLen dFF)

FFN activation output in the full encoder graph.

def Proofs.Autograd.Transformer.idxEncoderFfnProjected {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ [SeqFFNHiddenShape seqLen dFF, SeqFFNHiddenShape seqLen dFF, SeqFFNModelShape seqLen dModel]) (SeqFFNModelShape seqLen dModel)

FFN projection output in the full encoder graph.

def Proofs.Autograd.Transformer.idxEncoderFfnResidual {seqLen dModel numHeads headDim dFF : ℕ} : Idx (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF) (SeqFFNModelShape seqLen dModel)

FFN residual output after the second sublayer residual add.

noncomputable def Proofs.Autograd.Transformer.encoderFfnResidualGraph {seqLen dModel numHeads headDim dFF : ℕ} (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 ε₁ : ℝ) : Graph (ΓEncoderBlock seqLen dModel numHeads headDim) (ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF)

Graph through the FFN residual part after the first LayerNorm.

def Proofs.Autograd.Transformer.encoderNorm2Inputs {seqLen dModel numHeads headDim dFF : ℕ} : LayerNorm.Inputs (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssMHAResidual seqLen dModel numHeads headDim ++ [LayerNorm.MatShape seqLen dModel] ++ ssSeqFFNResidual seqLen dModel dFF) seqLen dModel

Second LayerNorm input triple in the concrete encoder block.

noncomputable def Proofs.Autograd.Transformer.encoderBlockGraph {seqLen dModel numHeads headDim dFF : ℕ} (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 ε₁ ε₂ : ℝ) : Graph (ΓEncoderBlock seqLen dModel numHeads headDim) (ssEncoderBlock seqLen dModel numHeads headDim dFF)

Concrete SSA graph for one full post-norm Transformer encoder block.

noncomputable def Proofs.Autograd.Transformer.encoderFfnResidualGraphFDerivCorrectAt {seqLen dModel numHeads headDim dFF : ℕ} (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 ε₁ : ℝ) (xV : CtxVec (ΓEncoderBlock seqLen dModel numHeads headDim)) (hε₁ : 0 < ε₁) : GraphFDerivCorrectAt (encoderFfnResidualGraph fc1 b1 fc2 b2 c ε₁) xV

Pointwise analytic correctness for the graph through the FFN residual.

noncomputable def Proofs.Autograd.Transformer.encoderBlockGraphFDerivCorrectAt {seqLen dModel numHeads headDim dFF : ℕ} (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 ε₁ ε₂ : ℝ) (xV : CtxVec (ΓEncoderBlock seqLen dModel numHeads headDim)) (hε₁ : 0 < ε₁) (hε₂ : 0 < ε₂) : GraphFDerivCorrectAt (encoderBlockGraph fc1 b1 fc2 b2 c ε₁ ε₂) xV

Pointwise analytic correctness for the complete concrete encoder-block graph.

theorem Proofs.Autograd.Transformer.encoderBlock_backpropVec_eq_adjoint_fderiv_at {seqLen dModel numHeads headDim dFF : ℕ} (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 ε₁ ε₂ : ℝ) (xV : CtxVec (ΓEncoderBlock seqLen dModel numHeads headDim)) (seedV : CtxVec (ΓEncoderBlock seqLen dModel numHeads headDim ++ ssEncoderBlock seqLen dModel numHeads headDim dFF)) (hε₁ : 0 < ε₁) (hε₂ : 0 < ε₂) : (encoderBlockGraph fc1 b1 fc2 b2 c ε₁ ε₂).backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (encoderBlockGraph fc1 b1 fc2 b2 c ε₁ ε₂).evalVec xV)) seedV

End-to-end VJP theorem for the concrete post-norm Transformer encoder-block graph.