Residual direct-reshape attention

This file appends a residual addition to the graph in DirectReshapeAttention:

x ↦ x + MHA(x).

It can then be composed with post-norm LayerNorm:

LayerNorm(x + MultiHeadSelfAttention(x)).

The resulting reverse pass is the adjoint of this graph's Fréchet derivative. It retains the direct-reshape head split, which differs from Spec.splitHeadsSpec and the executable attention layout. Adding a residual does not establish correspondence with the public MHA or runtime Transformer; that requires a separate forward-layout bridge.

References:

• Vaswani et al., "Attention Is All You Need", NeurIPS 2017.
• PyTorch torch.nn.MultiheadAttention: https://pytorch.org/docs/stable/generated/torch.nn.MultiheadAttention.html
• PyTorch torch.nn.TransformerEncoderLayer: https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html

@[reducible, inline]

abbrev Proofs.Autograd.Transformer.ssMHAResidual (n dModel numHeads headDim : ℕ) : List Spec.Shape

Intermediate list for MHA followed by one residual-add output.

def Proofs.Autograd.Transformer.residualIdxX {n dModel numHeads headDim : ℕ} : Idx (DirectReshapeAttention.ΓMHA n dModel numHeads headDim ++ DirectReshapeAttention.ssMHA n dModel numHeads headDim) (DirectReshapeAttention.XShape n dModel)

Original sequence input x, weakened into the context after the MHA intermediates.

def Proofs.Autograd.Transformer.residualIdxAttnOut {n dModel numHeads headDim : ℕ} : Idx (DirectReshapeAttention.ΓMHA n dModel numHeads headDim ++ DirectReshapeAttention.ssMHA n dModel numHeads headDim) (DirectReshapeAttention.XShape n dModel)

The final output of mhaDGraph, i.e. the projected attention result.

The literal index is intentional: DirectReshapeAttention.ssMHA has 14 saved values, and the final entry is the attention output with the same shape as the input sequence.

noncomputable def Proofs.Autograd.Transformer.mhaResidualDGraph {n dModel numHeads headDim : ℕ} (c : ℝ) : DGraph (DirectReshapeAttention.ΓMHA n dModel numHeads headDim) (ssMHAResidual n dModel numHeads headDim)

Proof-carrying graph for x + MHA(x), using the direct-reshape attention variant.

Context layout is inherited from MHA: [x, Wq, Wk, Wv, Wo].

theorem Proofs.Autograd.Transformer.mhaResidual_backpropVec_eq_adjoint_fderiv {n dModel numHeads headDim : ℕ} (c : ℝ) (xV : CtxVec (DirectReshapeAttention.ΓMHA n dModel numHeads headDim)) (seedV : CtxVec (DirectReshapeAttention.ΓMHA n dModel numHeads headDim ++ ssMHAResidual n dModel numHeads headDim)) : (mhaResidualDGraph c).g.backpropVec xV seedV = (ContinuousLinearMap.adjoint (fderiv ℝ (mhaResidualDGraph c).g.evalVec xV)) seedV

End-to-end VJP theorem for the residual-attention sublayer x + MHA(x).

This graph is composed with LayerNorm in the post-norm proof modules. LayerNorm itself has its own current-spec VJP theorem in NN.Proofs.Autograd.Tape.Ops.Norm.LayerNorm; the composed post-norm attention sublayer and the two-sublayer post-norm bridge are packaged in NN.Proofs.Autograd.Tape.Ops.Transformer.PostNorm.