Transformer (spec model)

This file defines Transformer-style spec components in a way that matches the usual PyTorch mental model:

• encoder layers (self-attention + FFN, each wrapped in residual + LayerNorm),
• decoder layers (masked self-attention, cross-attention, FFN, each wrapped in residual + LayerNorm),
• an encoder-decoder wrapper (Transformer),
• spec-level backward passes for the encoder stack.

Sinusoidal positional encodings live in NN.Spec.Layers.PositionalEncoding as Spec.sinusoidalPositionalEncodingSpec; causal masks are Spec.causalMask in NN.Spec.Layers.Attention.

Shapes follow the common convention:

• sequence tensors are (seqLen × embedDim),
• attention is "last-axis softmax" over the key dimension.

PyTorch analogy:

• TransformerEncoderLayer.forward corresponds to the core of torch.nn.TransformerEncoderLayer (ignoring dropout and some configuration knobs),
• TransformerEncoder.forward corresponds to torch.nn.TransformerEncoder.
• TransformerDecoderLayer.forward corresponds to the core of torch.nn.TransformerDecoderLayer,
• TransformerDecoder.forward corresponds to torch.nn.TransformerDecoder,
• Transformer.forward is similar in spirit to torch.nn.Transformer (but simplified).

References:

• Vaswani et al., "Attention Is All You Need" (2017).
• Ba et al., "Layer Normalization" (2016).
• He et al., "Deep Residual Learning for Image Recognition" (2015) for the residual/skip-connection pattern.

PyTorch docs (for API shape intuition, not semantics):

• torch.nn.TransformerEncoderLayer: https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoderLayer.html
• torch.nn.TransformerEncoder: https://pytorch.org/docs/stable/generated/torch.nn.TransformerEncoder.html
• torch.nn.TransformerDecoderLayer: https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoderLayer.html
• torch.nn.TransformerDecoder: https://pytorch.org/docs/stable/generated/torch.nn.TransformerDecoder.html
• torch.nn.Transformer: https://pytorch.org/docs/stable/generated/torch.nn.Transformer.html

Gradient containers

To keep the backward pass readable (and easy to reuse from downstream models like ViT/Seq2Seq), we bundle parameter gradients into records that mirror the parameter records.

structure Spec.FeedForwardGrads (embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Gradients for a FeedForward block (field-for-field).

This container is used by downstream models that want a readable backward pass.

• inputWeight : TorchLean.Tensor α [embedDim, hiddenDim]

  Gradient of the input projection weight.

• outputWeight : TorchLean.Tensor α [hiddenDim, embedDim]

  Gradient of the output projection weight.

• inputBias : TorchLean.Tensor α [hiddenDim]

  Gradient of the input projection bias.

• outputBias : TorchLean.Tensor α [embedDim]

  Gradient of the output projection bias.

structure Spec.TransformerEncoderLayerGrads (headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Gradients for a TransformerEncoderLayer (field-for-field).

This container is intended to keep the backward pass readable by mirroring the parameter layout.

• mha : MultiHeadAttentionParameterGradients headCount embedDim (embedDim / headCount) α

  Gradients for the self-attention block.

• ffn : FeedForwardGrads embedDim hiddenDim α

  Gradients for the feedforward block.

• norm1Scale : TorchLean.Tensor α [embedDim]

  Gradient of LayerNorm 1 gamma (attention "Add & Norm").

• norm1Bias : TorchLean.Tensor α [embedDim]

  Gradient of LayerNorm 1 beta (attention "Add & Norm").

• norm2Scale : TorchLean.Tensor α [embedDim]

  Gradient of LayerNorm 2 gamma (FFN "Add & Norm").

• norm2Bias : TorchLean.Tensor α [embedDim]

  Gradient of LayerNorm 2 beta (FFN "Add & Norm").

structure Spec.FeedForward (embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

2-layer position-wise feedforward network used inside Transformer layers.

Semantics (per token): ffn(x) = (relu(x * W1 + b1) * W2) + b2.

PyTorch analogue: the linear1 / linear2 submodule in torch.nn.TransformerEncoderLayer.

• inputWeight : TorchLean.Tensor α [embedDim, hiddenDim]

  First linear layer weights (embedDim -> hiddenDim).

• outputWeight : TorchLean.Tensor α [hiddenDim, embedDim]

  Second linear layer weights (hiddenDim -> embedDim).

• inputBias : TorchLean.Tensor α [hiddenDim]

  First layer bias (length hiddenDim).

• outputBias : TorchLean.Tensor α [embedDim]

  Second layer bias (length embedDim).

def Spec.FeedForward.forward {α : Type} [TorchLean.Storage α] [Context α] {embedDim hiddenDim seqLen : ℕ} (ffn : FeedForward embedDim hiddenDim α) (x : TorchLean.Tensor α [seqLen, embedDim]) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for FeedForward.

Shape convention: inputs and outputs are (seqLen × embedDim); the feedforward operates independently on each sequence position.

structure Spec.TransformerEncoderLayer (headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Transformer encoder layer (post-norm).

This follows the common "Add & Norm" structure:

1. Self-attention, residual add, LayerNorm
2. Feedforward, residual add, LayerNorm

PyTorch analogue: torch.nn.TransformerEncoderLayer with norm_first=False (post-norm), ignoring dropout and other configuration knobs.

• mha : MultiHeadAttention α headCount embedDim (embedDim / headCount)

  Multi-head self-attention block.

• ffn : FeedForward embedDim hiddenDim α

  Position-wise feedforward block.

• norm1Scale : TorchLean.Tensor α [embedDim]

  LayerNorm 1 gamma (attention "Add & Norm").

• norm1Bias : TorchLean.Tensor α [embedDim]

  LayerNorm 1 beta (attention "Add & Norm").

• norm2Scale : TorchLean.Tensor α [embedDim]

  LayerNorm 2 gamma (FFN "Add & Norm").

• norm2Bias : TorchLean.Tensor α [embedDim]

  LayerNorm 2 beta (FFN "Add & Norm").

def Spec.TransformerEncoderLayer.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {headCount embedDim hiddenDim seqLen : ℕ} (layer : TransformerEncoderLayer headCount embedDim hiddenDim α) (x : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for a post-norm TransformerEncoderLayer.

Input/output shape: (seqLen × embedDim). The proofs h1/h2 are used by layerNorm to justify nondegenerate normalization.

structure Spec.TransformerEncoder (numLayers headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Transformer encoder: a stack of TransformerEncoderLayers.

PyTorch analogue: torch.nn.TransformerEncoder (a list of layers composed sequentially).

• layers : TorchLean.Tensor (TransformerEncoderLayer headCount embedDim hiddenDim α) [numLayers]

  Encoder layers, with the stack depth recorded in the type.

def Spec.TransformerEncoder.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {numLayers headCount embedDim hiddenDim seqLen : ℕ} (encoder : TransformerEncoder numLayers headCount embedDim hiddenDim α) (x : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for TransformerEncoder (left-fold over layers).

Input/output shape: (seqLen × embedDim).

Decoder notes

We include a small Transformer-style decoder layer for completeness:

• self-attention over the decoder sequence,
• cross-attention where queries come from the decoder and keys/values come from the encoder,
• then the same feedforward block.

PyTorch analogy: this corresponds to the core of torch.nn.TransformerDecoderLayer (ignoring dropout and a few configuration knobs).

Cross-attention helper

The attention layer provides MultiHeadAttention.forward for the common self-attention case (Q=K=V=x). A decoder block also needs cross-attention, where Q comes from the decoder stream and K,V come from the encoder stream.

We keep the helper here small and explicit by following the same structure as the self-attention definition: project, split into heads, run scaled dot-product attention per head, combine heads, then project with Wo.

def Spec.multiHeadCrossAttention {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {headCount embedDim nQ nK : ℕ} (hQ : nQ ≠ 0) (hK : nK ≠ 0) (mha : MultiHeadAttention α headCount embedDim (embedDim / headCount)) (qInput : TorchLean.Tensor α [nQ, embedDim]) (kvInput : TorchLean.Tensor α [nK, embedDim]) (mask : Option (TorchLean.Tensor Bool [nQ, nK])) : TorchLean.Tensor α [nQ, embedDim]

Cross-attention forward pass using a MultiHeadAttention parameter record.

This is the decoder-specific variant of MultiHeadAttention.forward:

• queries come from qInput (decoder stream),
• keys/values come from kvInput (encoder stream),
• an optional boolean mask of shape (nQ × nK) can be applied.

Shape conventions:

• qInput : (nQ × embedDim),
• kvInput : (nK × embedDim),
• output : (nQ × embedDim).

PyTorch analogue: the cross-attention inside torch.nn.TransformerDecoderLayer, typically implemented via torch.nn.MultiheadAttention with separate query and key/value inputs.

structure Spec.TransformerDecoderLayer (headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Transformer decoder layer (post-norm).

This mirrors the standard structure:

1. Self-attention (decoder stream), residual add, LayerNorm
2. Cross-attention (queries from decoder, keys/values from encoder), residual add, LayerNorm
3. Feedforward, residual add, LayerNorm

PyTorch analogue: torch.nn.TransformerDecoderLayer with norm_first=False (post-norm), ignoring dropout and a few configuration knobs.

• selfAttn : MultiHeadAttention α headCount embedDim (embedDim / headCount)

  Self-attention block over the decoder sequence.

• crossAttn : MultiHeadAttention α headCount embedDim (embedDim / headCount)

  Cross-attention block (decoder queries, encoder keys/values).

• ffn : FeedForward embedDim hiddenDim α

  Position-wise feedforward block.

• norm1Scale : TorchLean.Tensor α [embedDim]

  Scale of the layer normalization after self-attention.

• norm1Bias : TorchLean.Tensor α [embedDim]

  Bias of the layer normalization after self-attention.

• norm2Scale : TorchLean.Tensor α [embedDim]

  Scale of the layer normalization after cross-attention.

• norm2Bias : TorchLean.Tensor α [embedDim]

  Bias of the layer normalization after cross-attention.

• norm3Scale : TorchLean.Tensor α [embedDim]

  Scale of the layer normalization after the feed-forward block.

• norm3Bias : TorchLean.Tensor α [embedDim]

  Bias of the layer normalization after the feed-forward block.

def Spec.TransformerDecoderLayer.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {headCount embedDim hiddenDim seqLen : ℕ} (layer : TransformerDecoderLayer headCount embedDim hiddenDim α) (x encoderOutput : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) (targetMask memoryMask : Option (TorchLean.Tensor Bool [seqLen, seqLen]) := none) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for a post-norm TransformerDecoderLayer.

Input/output shape: (seqLen × embedDim). This spec uses the same seqLen for encoder and decoder streams for simplicity (cross-attention uses nQ = nK = seqLen).

targetMask controls decoder self-attention, and memoryMask controls attention to the encoder output. In both masks, true allows an attention edge and false blocks it. Pass some (causalMask seqLen) as targetMask for autoregressive decoding: position i then uses only target positions 0, ..., i. The source sequence remains available through cross-attention unless memoryMask restricts it separately. Both arguments default to none for unmasked attention.

structure Spec.TransformerDecoder (numLayers headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

Transformer decoder: a stack of TransformerDecoderLayers.

PyTorch analogue: torch.nn.TransformerDecoder (a list of decoder layers composed sequentially).

• layers : TorchLean.Tensor (TransformerDecoderLayer headCount embedDim hiddenDim α) [numLayers]

  Decoder layers, with the stack depth recorded in the type.

def Spec.TransformerDecoder.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {numLayers headCount embedDim hiddenDim seqLen : ℕ} (decoder : TransformerDecoder numLayers headCount embedDim hiddenDim α) (x encoderOutput : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) (targetMask memoryMask : Option (TorchLean.Tensor Bool [seqLen, seqLen]) := none) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for TransformerDecoder (left-fold over layers).

Input/output shape: (seqLen × embedDim). Every layer receives the same target and memory masks. A causal target mask must reach the whole stack: unmasked attention in any layer would let future target positions enter the earlier residual streams.

structure Spec.Transformer (numLayers headCount embedDim hiddenDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

End-to-end encoder-decoder Transformer (spec model).

This is a seq2seq Transformer wrapper built out of the encoder and decoder stacks above. It models the core tensor algebra of torch.nn.Transformer while making the proof-relevant choices explicit:

• embeddings are modeled as explicit linear projections,
• sequence length is shared between source and target streams,
• we omit dropout, caching, and most configuration knobs.

Shape convention: all activations in this file use (seqLen × embedDim). In a full implementation, outputProjection would usually map to a vocabulary size; here it is kept as an embedDim -> embedDim projection to stay in the "core tensor algebra" setting.

• encoder : TransformerEncoder numLayers headCount embedDim hiddenDim α

  Encoder stack.

• decoder : TransformerDecoder numLayers headCount embedDim hiddenDim α

  Decoder stack.

• inputEmbedding : TorchLean.Tensor α [embedDim, embedDim]

  Source/input embedding projection matrix.

• outputEmbedding : TorchLean.Tensor α [embedDim, embedDim]

  Target embedding projection matrix.

• outputProjection : TorchLean.Tensor α [embedDim, embedDim]

  Final output projection matrix (here embedDim -> embedDim).

def Spec.Transformer.forward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {numLayers headCount embedDim hiddenDim seqLen : ℕ} (transformer : Transformer numLayers headCount embedDim hiddenDim α) (input target : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) (targetMask memoryMask : Option (TorchLean.Tensor Bool [seqLen, seqLen]) := none) : TorchLean.Tensor α [seqLen, embedDim]

Forward pass for Transformer.

Runs:

1. source embedding projection,
2. encoder stack,
3. target embedding projection,
4. decoder stack (with cross-attention to the encoder output),
5. output projection.

All tensors in this simplified spec have shape (seqLen × embedDim). targetMask and memoryMask are passed to every decoder layer. For next-token prediction, supply shifted target inputs and targetMask := some (causalMask seqLen); the mask permits the current input token while blocking later inputs. Source encoding stays bidirectional.

def Spec.FeedForward.backward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {embedDim hiddenDim seqLen : ℕ} (ffn : FeedForward embedDim hiddenDim α) (x outputGrad : TorchLean.Tensor α [seqLen, embedDim]) (h_seq : seqLen > 0) (_h_embed : embedDim > 0) : FeedForwardGrads embedDim hiddenDim α × TorchLean.Tensor α [seqLen, embedDim]

Backward pass for FeedForward.forward.

Given the input x and an upstream gradient outputGrad = dL/dy (w.r.t. the FFN output), returns:

• parameter gradients (as FeedForwardGrads),
• the gradient w.r.t. the input x.

This is a spec-level backward that reconstructs the forward intermediates (pre-activations and ReLU mask) instead of relying on a mutable tape, similar to the math underlying PyTorch autograd.

def Spec.TransformerEncoderLayer.backward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {headCount embedDim hiddenDim seqLen : ℕ} (layer : TransformerEncoderLayer headCount embedDim hiddenDim α) (x outputGrad : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) : TransformerEncoderLayerGrads headCount embedDim hiddenDim α × TorchLean.Tensor α [seqLen, embedDim]

Backward pass for TransformerEncoderLayer.forward.

Inputs:

• x: the layer input (seqLen × embedDim),
• outputGrad: upstream gradient w.r.t. the layer output.

Outputs:

• parameter gradients (TransformerEncoderLayerGrads),
• gradient w.r.t. x.

The implementation mirrors the forward pass structure (residuals + LayerNorm) and uses layerNormBackward and multiHeadAttentionBackward as its core primitives.

Backward pass for an encoder stack

The encoder is a fixed-length vector of layers applied sequentially. To compute gradients we:

1. re-run the forward pass to collect each layer's input (a small "cache"),
2. traverse layers in reverse, applying TransformerEncoderLayer.backward,
3. return per-layer parameter gradients plus the gradient w.r.t. the encoder input.

This is purely a spec (no mutation, no state), so we do the simplest thing: recompute.

def Spec.TransformerEncoder.backward {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {numLayers headCount embedDim hiddenDim seqLen : ℕ} (encoder : TransformerEncoder numLayers headCount embedDim hiddenDim α) (x outputGrad : TorchLean.Tensor α [seqLen, embedDim]) (h1 : seqLen > 0) (h2 : embedDim > 0) : TorchLean.Tensor (TransformerEncoderLayerGrads headCount embedDim hiddenDim α) [numLayers] × TorchLean.Tensor α [seqLen, embedDim]

Backward pass for TransformerEncoder.forward (a sequential stack of layers).

Returns:

• one parameter-gradient record per layer, in the same order as encoder.layers,
• the gradient w.r.t. the encoder input x.

Because this is a pure spec, we recompute forward intermediates (each layer input) instead of storing a mutable cache.