Vision Transformer

Patch embedding is an arbitrary-dimensional convolution. The spatial output is flattened into a token axis before the Transformer block, so the construction applies equally to one-dimensional signals, images, volumes, and higher-dimensional grids.

inductive TorchLean.nn.models.ViT.Pooling : Type

How a vision Transformer turns encoded tokens into one vector per sample.

• mean : Pooling

  Average every encoded patch token.

• cls : Pooling

  Prepend a learned class slot and select it after the encoder.

@[instance_reducible]

instance TorchLean.nn.models.ViT.instDecidableEqPooling : DecidableEq Pooling

def TorchLean.nn.models.ViT.instReprPooling.repr : Pooling → ℕ → Std.Format

@[instance_reducible]

instance TorchLean.nn.models.ViT.instReprPooling : Repr Pooling

structure TorchLean.nn.models.ViT.EncoderConfig (d : ℕ) : Type

Configuration for a Transformer encoder over patches from a d-dimensional spatial domain.

• inputChannels : ℕ

  Number of channels in each input sample.

• spatial : Tensor ℕ [d]

  Extent of each spatial axis.

• patchEmbedding : Convolution.Config d

  Convolution that extracts and embeds patches.

• headCount : ℕ

  Number of attention heads. Must be positive.

• headWidth : ℕ

  Width of each attention head. Must be positive.

• feedForwardWidth : ℕ

  Width of the feed-forward sublayer.

• layerCount : ℕ

  Number of Transformer encoder blocks.

• dropout? : Option Float

  Dropout probability for attention and feed-forward outputs.

• attentionDropout? : Option Float

  Drop attention probabilities before value aggregation; independent of residual dropout.

• feedForwardDropout? : Option Float

  Drop activated FFN hidden units before their output projection.

• attentionInputBias : Bool

  Enable trainable query, key, and value biases while preserving the legacy default layout.

• pooling : Pooling

  Classifier readout; .cls also prepends a learned class token.

structure TorchLean.nn.models.ViT.Config (d : ℕ) extends TorchLean.nn.models.ViT.EncoderConfig d : Type

Configuration for a vision Transformer classifier.

• inputChannels : ℕ
• spatial : Tensor ℕ [d]
• patchEmbedding : Convolution.Config d
• headCount : ℕ
• headWidth : ℕ
• feedForwardWidth : ℕ
• layerCount : ℕ
• dropout? : Option Float
• attentionDropout? : Option Float
• feedForwardDropout? : Option Float
• attentionInputBias : Bool
• pooling : Pooling
• classCount : ℕ

  Number of classifier outputs per sample.

def TorchLean.nn.models.ViT.EncoderConfig.classifier {d : ℕ} (config : EncoderConfig d) (classCount : ℕ) : Config d

Attach a classifier output width to reusable encoder settings.

def TorchLean.nn.models.ViT.Config.encoder {d : ℕ} (config : Config d) : EncoderConfig d

Encoder settings embedded in a classifier configuration.

def TorchLean.nn.models.ViT.EncoderConfig.validate {d : ℕ} (config : EncoderConfig d) (kind : String := "ViT") : Except String Unit

Validate patch extraction and the complete Transformer template before allocating parameters.

def TorchLean.nn.models.ViT.Config.validate {d : ℕ} (config : Config d) : Except String Unit

Validate both the reusable encoder and classifier head before construction.

def TorchLean.nn.models.ViT.EncoderConfig.patchGrid {d : ℕ} (config : EncoderConfig d) : Tensor ℕ [d]

Grid produced by patch embedding.

def TorchLean.nn.models.ViT.EncoderConfig.patchCount {d : ℕ} (config : EncoderConfig d) : ℕ

Number of patch tokens.

def TorchLean.nn.models.ViT.EncoderConfig.prefixLength {d : ℕ} (config : EncoderConfig d) : ℕ

Number of slots inserted before the patch sequence.

def TorchLean.nn.models.ViT.EncoderConfig.sequenceLength {d : ℕ} (config : EncoderConfig d) : ℕ

Number of tokens passed through the Transformer encoder.

def TorchLean.nn.models.ViT.EncoderConfig.flattenedWidth {d : ℕ} (config : EncoderConfig d) : ℕ

Number of scalar features in the complete encoded token sequence.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.EncoderConfig.inputShape {d : ℕ} (config : EncoderConfig d) (batchShape : Shape := []) : Shape

Input shape for any caller-supplied batch shape.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.EncoderConfig.patchShape {d : ℕ} (config : EncoderConfig d) (batchShape : Shape := []) : Shape

Patch tensor produced by the embedding convolution.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.EncoderConfig.patchTokenShape {d : ℕ} (config : EncoderConfig d) (batchShape : Shape := []) : Shape

Patch tokens after moving the embedding width to the final axis.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.EncoderConfig.outputShape {d : ℕ} (config : EncoderConfig d) (batchShape : Shape := []) : Shape

Tokens consumed and produced by the Transformer stack.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.Config.inputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Classifier input shape for any caller-supplied batch shape.

@[reducible, inline]

abbrev TorchLean.nn.models.ViT.Config.outputShape {d : ℕ} (config : Config d) (batchShape : Shape := []) : Shape

Classifier output shape for the same batch shape as the input.

def TorchLean.nn.models.Internal.patchesToTokens {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Layer (config.patchShape batchShape) (config.patchTokenShape batchShape)

Implementation layer that turns a patch grid into a token sequence.

def TorchLean.nn.models.Internal.prependClassToken {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Sequential (config.patchTokenShape batchShape) (batchShape.concat [1 + config.patchCount, config.patchEmbedding.outChannels])

Prepend one shared, trainable class token to every independently mapped sequence.

def TorchLean.nn.models.Internal.tokensToChannels {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Sequential (config.outputShape batchShape) (batchShape.concat [config.patchEmbedding.outChannels, config.sequenceLength])

Implementation layer that moves the embedding width before the token axis.

def TorchLean.nn.models.Internal.meanVitTokens {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Builder (Sequential (config.outputShape batchShape) (batchShape.appendDim config.patchEmbedding.outChannels))

Implementation layer for mean-pool classification.

def TorchLean.nn.models.Internal.firstVitToken {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Sequential (config.outputShape batchShape) (batchShape.appendDim config.patchEmbedding.outChannels)

Implementation layer for class-token classification.

def TorchLean.nn.models.vitEncoder {d : ℕ} (config : ViT.EncoderConfig d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build the patch and Transformer portion of a vision transformer.

Mean readout leaves the patch sequence unchanged. Class-token readout prepends one learned token before the Transformer. Classification, reconstruction, and other tasks can attach their own heads without rebuilding the patch pipeline.

def TorchLean.nn.models.vit {d : ℕ} (config : ViT.Config d) (batchShape : Shape := []) : Builder (Sequential (config.inputShape batchShape) (config.outputShape batchShape))

Build a vision Transformer encoder followed by a linear classifier.