Causal Transformer Architecture

Configuration, shape families, and graph constructors for GPT-style causal Transformers.

structure TorchLean.nn.models.CausalTransformer.Config : 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)×layerCount → LayerNorm → linear

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

• sequenceLength : ℕ

  Number of token positions processed in one model invocation. Must be positive.

• vocabularySize : ℕ

  Number of token categories accepted by the input and vocabulary head. Must be positive.

• headCount : ℕ

  Number of parallel self-attention heads in each Transformer block. Must be positive.

• headWidth : ℕ

  Width of each attention head. Must be positive.

• feedForwardWidth : ℕ

  Hidden width of each block's position-wise feed-forward network.

• layerCount : ℕ

  Number of stacked causal Transformer blocks.

• activation : Activation.Kind

  Feed-forward activation used in every Transformer block.

• dropout? : Option Float

  Dropout probability for attention and feed-forward outputs.

• attentionDropout? : Option Float

  Drop attention probabilities before value aggregation, independently of residual dropout.

• feedForwardDropout? : Option Float

  Drop activated FFN hidden units before their output projection.

• attentionInputBias : Bool

  Enable separate query, key, and value biases; false keeps the existing parameter layout.

• normalizeFirst : Bool

  Use pre-normalized Transformer blocks, as in GPT-2.

• attentionOutputBias : Bool

  Add a trainable bias after each attention output projection, as in GPT-2.

• parameterInitialization? : Option Init.Scheme

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

• residualProjectionInitialization? : Option 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.

@[instance_reducible]

instance TorchLean.nn.models.CausalTransformer.instReprConfig : Repr Config

def TorchLean.nn.models.CausalTransformer.instReprConfig.repr : Config → ℕ → Std.Format

def TorchLean.nn.models.CausalTransformer.Config.modelWidth (config : Config) : ℕ

Transformer width implied by headCount * headWidth.

def TorchLean.nn.models.CausalTransformer.Config.validateBody (config : Config) : Except String Unit

Validate the hidden Transformer independently of its token-vocabulary boundary.

def TorchLean.nn.models.CausalTransformer.Config.validate (config : Config) : Except String Unit

Validate the complete language-model configuration before allocating any parameters.

@[reducible, inline]

abbrev TorchLean.nn.models.CausalTransformer.Config.tokenShape (config : Config) (batchShape : Shape := []) : Shape

Bounded token-id tensor shape.

@[reducible, inline]

abbrev TorchLean.nn.models.CausalTransformer.Config.vocabularyShape (config : Config) (batchShape : Shape := []) : Shape

Per-token vocabulary tensor shape used by one-hot inputs and output logits.

@[reducible, inline]

abbrev TorchLean.nn.models.CausalTransformer.Config.embeddingShape (config : Config) (batchShape : Shape := []) : Shape

Embedded-token tensor shape.

def TorchLean.nn.models.CausalTransformer.hidden (config : Config) (batchShape : Shape := []) : Builder (Sequential (config.embeddingShape batchShape) (config.embeddingShape batchShape))

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.CausalTransformer.fromEmbeddings (config : Config) (batchShape : Shape := []) : Builder (Sequential (config.embeddingShape batchShape) (config.vocabularyShape batchShape))

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

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

def TorchLean.nn.models.CausalTransformer.oneHot (config : Config) (batchShape : Shape := []) : Builder (Sequential (config.vocabularyShape batchShape) (config.vocabularyShape batchShape))

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.Builder so it composes with the rest of the API-layer model-building interface.