Causal Transformer Models

This module defines the model structure shared by TorchLean's causal language-model examples. It supports one-hot inputs and integer token tensors, with either an independent vocabulary head or a projection tied to the token-embedding table. Tokenization and checkpoint formats live in their respective API modules.

structure TorchLean.nn.models.CausalTransformerConfig : Type

Configuration shared by TorchLean's GPT-style causal language models.

The model has the common GPT-2 “shape”:

embedding → learned positional embedding → (masked self-attention + FFN)×layers → LayerNorm → linear

The configuration is independent of how token ids enter the model. One-hot, integer-token, and pre-embedded constructors reuse the same Transformer width, depth, and output vocabulary.

• batch : ℕ
• seqLen : ℕ
• vocab : ℕ
• numHeads : ℕ
• headDim : ℕ
• ffnHidden : ℕ
• layers : ℕ
• activation : blocks.Activation

  Feed-forward activation used in every Transformer block.

• dropout? : Option Float

  Dropout probability for attention and feed-forward outputs.

• normFirst : Bool

  Use pre-normalized Transformer blocks.

• attentionOutputBias : Bool

  Add a trainable bias after each attention output projection.

• parameterInit? : Option Runtime.Autograd.Torch.Init.Scheme

  Shared initialization for embedding and projection weights. none keeps layer defaults.

• residualProjectionInit? : Option Runtime.Autograd.Torch.Init.Scheme

  Initialization for projections whose outputs are added to residual streams.

  GPT-2 scales these weights by network depth. Keeping the setting explicit lets other causal Transformers use their own residual initialization without changing the block implementation.

• seedStride : ℕ

  Seed stride used when initializing repeated blocks.

def TorchLean.nn.models.instReprCausalTransformerConfig.repr : CausalTransformerConfig → ℕ → Std.Format

@[implicit_reducible]

instance TorchLean.nn.models.instReprCausalTransformerConfig : Repr CausalTransformerConfig

def TorchLean.nn.models.CausalTransformerConfig.dModel (cfg : CausalTransformerConfig) : ℕ

Transformer width implied by numHeads * headDim.

@[reducible, inline]

abbrev TorchLean.nn.models.causalVocabularyShape (cfg : CausalTransformerConfig) : Shape

Vocabulary-grid shape (batch × seqLen × vocab) used by one-hot inputs and output logits.

@[reducible, inline]

abbrev TorchLean.nn.models.causalEmbeddingShape (cfg : CausalTransformerConfig) : Shape

Embedded-token tensor shape (batch × seqLen × dModel).

def TorchLean.nn.models.causalTransformerHiddenFromEmbeddings (cfg : CausalTransformerConfig) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg))

Causal Transformer hidden-state stack after token embeddings have been computed.

The stack adds learned positions, applies causally masked Transformer blocks, and finishes with LayerNorm. It deliberately has no vocabulary projection. Language models can therefore choose an independent output head or reuse their token-embedding matrix.

def TorchLean.nn.models.causalTransformerFromEmbeddings (cfg : CausalTransformerConfig) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg))

GPT-style causal Transformer body with an independent affine vocabulary head.

Use causalTransformerHiddenFromEmbeddings when the caller needs hidden states or a tied token-embedding/output matrix.

def TorchLean.nn.models.causalTransformerOneHot (cfg : CausalTransformerConfig) (h_seqLen : cfg.seqLen ≠ 0 := by decide) (h_dModel : cfg.dModel ≠ 0 := by decide) : M (Sequential (causalVocabularyShape cfg) (causalVocabularyShape cfg))

Build a GPT-2-style causal language model over one-hot tokens.

This is the shared constructor used by the runnable GPT-2 examples. It stays in nn.M so it composes with the rest of the API-layer model-building interface.

@[reducible, inline]

abbrev TorchLean.nn.models.causalTokenShape (cfg : CausalTransformerConfig) : Shape

Flattened shape of one batch of token ids or row weights.

@[reducible, inline]

abbrev TorchLean.nn.models.causalTransformerTokenParamShapes (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) : List Shape

Parameter shapes for an indexed-token model with an independent vocabulary head.

def TorchLean.nn.models.causalTransformerTokenForward {α : Type} [Context α] [DecidableEq Shape] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) (params : Runtime.Autograd.Torch.RefList (Runtime.Autograd.TorchLean.NN.Seq.RefT m α) (causalTransformerTokenParamShapes cfg body)) (tokens : Runtime.Autograd.Torch.NatTensorRef (causalTokenShape cfg)) : m (Runtime.Autograd.TorchLean.NN.Seq.RefT m α (causalVocabularyShape cfg))

Run indexed-token embedding and a causal Transformer body in any TorchLean backend.

Tokens remain Nat tensors throughout the call. Only the embedding table and Transformer parameters use the model scalar type, so the program cannot differentiate with respect to token ids or reinterpret them as floating-point data.

def TorchLean.nn.models.causalTransformerTokenProgramWithMode (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.ProgramWithNatInputs α (causalTransformerTokenParamShapes cfg body) [causalTokenShape cfg] (causalVocabularyShape cfg)

Backend-polymorphic indexed-token forward program.

def TorchLean.nn.models.causalTransformerTokenProgram (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.ProgramWithNatInputs α (causalTransformerTokenParamShapes cfg body) [causalTokenShape cfg] (causalVocabularyShape cfg)

Evaluation-mode indexed-token forward program.

@[reducible, inline]

abbrev TorchLean.nn.models.causalTransformerTiedTokenParamShapes (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) : List Shape

Parameter shapes for an indexed-token model whose embedding and output projection are tied.

def TorchLean.nn.models.causalTransformerTiedTokenForward {α : Type} [Context α] [DecidableEq Shape] {m : Type → Type} [Monad m] [Runtime.Autograd.Torch.Ops m α] (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) (params : Runtime.Autograd.Torch.RefList (Runtime.Autograd.TorchLean.NN.Seq.RefT m α) (causalTransformerTiedTokenParamShapes cfg body)) (tokens : Runtime.Autograd.Torch.NatTensorRef (causalTokenShape cfg)) : m (Runtime.Autograd.TorchLean.NN.Seq.RefT m α (causalVocabularyShape cfg))

Run a causal Transformer whose token lookup and vocabulary projection share one matrix.

The matrix has shape (vocab, dModel). The forward pass gathers its rows to construct the input embeddings, runs the hidden-state Transformer, and multiplies the final hidden states by its transpose. Consequently, gradients from both lookup and prediction accumulate into the same parameter.

def TorchLean.nn.models.causalTransformerTiedTokenProgramWithMode (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.ProgramWithNatInputs α (causalTransformerTiedTokenParamShapes cfg body) [causalTokenShape cfg] (causalVocabularyShape cfg)

Backend-polymorphic tied-token forward program.

def TorchLean.nn.models.causalTransformerTiedTokenProgram (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) {α : Type} [Context α] [DecidableEq Shape] : Runtime.Autograd.TorchLean.ProgramWithNatInputs α (causalTransformerTiedTokenParamShapes cfg body) [causalTokenShape cfg] (causalVocabularyShape cfg)

Evaluation-mode tied-token forward program.

def TorchLean.nn.models.causalTransformerTokenScalarModuleDefWithMode (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) (reduction : Loss.Reduction := Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (causalTransformerTokenParamShapes cfg body) [] [causalTokenShape cfg, causalTokenShape cfg]

Scalar loss for causal language modeling with integer token ids.

The public one-hot constructor above is useful for small teaching examples because the input is an ordinary Float tensor. File-backed tokenized datasets use the representation found in language-model training systems: token ids are Nats, the embedding table is a trainable Float parameter, and the loss gathers the target classes directly instead of building one-hot targets.

tokens and targets are flattened (batch * seqLen) vectors. This matches the backend gather ops and keeps dataset storage simple; the embedding helper reshapes gathered rows back to (batch, seqLen, dModel) before running the Transformer body.

def TorchLean.nn.models.causalTransformerTokenScalarModuleDef (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalVocabularyShape cfg)) (reduction : Loss.Reduction := Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (causalTransformerTokenParamShapes cfg body) [] [causalTokenShape cfg, causalTokenShape cfg]

Training-mode wrapper for integer-token causal language modeling.

def TorchLean.nn.models.causalTransformerTiedTokenScalarModuleDefWithMode (mode : Runtime.Autograd.TorchLean.NN.Mode) (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) (reduction : Loss.Reduction := Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (causalTransformerTiedTokenParamShapes cfg body) [] [causalTokenShape cfg, causalTokenShape cfg]

Scalar causal-language-model loss for a tied token embedding and output projection.

def TorchLean.nn.models.causalTransformerTiedTokenScalarModuleDef (cfg : CausalTransformerConfig) (body : Sequential (causalEmbeddingShape cfg) (causalEmbeddingShape cfg)) (reduction : Loss.Reduction := Loss.Reduction.mean) : Runtime.Autograd.TorchLean.ScalarModuleDef (causalTransformerTiedTokenParamShapes cfg body) [] [causalTokenShape cfg, causalTokenShape cfg]

Training-mode wrapper for tied-token causal language modeling.