TorchLean NN: Attention

def Runtime.Autograd.Model.Layers.Attention.probabilities {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch heads n : ℕ} (scores : Ref [batch, heads, n, n]) (mask : Option (TorchLean.Tensor Bool [n, n])) : m (Ref [batch, heads, n, n])

Softmax over the visible keys of each attention row.

The mask uses true for a visible key. Gathering those keys before softmax keeps masked scores out of both the row maximum and its denominator. An entirely masked row stays zero. In particular, we do not approximate exclusion with a large negative constant, whose behavior depends on the score scale and cannot represent an empty row.

def Runtime.Autograd.Model.Layers.Attention.forward {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {batch n modelWidth heads headWidth : ℕ} (queryWeight keyWeight valueWeight : Ref [modelWidth, heads * headWidth]) (outputWeight : Ref [heads * headWidth, modelWidth]) (queryBias keyBias valueBias : Option (Ref [heads * headWidth])) (outputBias : Option (Ref [modelWidth])) (probability : Option (Ref Spec.Shape.scalar)) (input : Ref [batch, n, modelWidth]) (mask : Option (TorchLean.Tensor Bool [n, n])) (seed : ℕ) (training : Bool) : m (Ref [batch, n, modelWidth])

Attention with optional affine projection biases and dropout on softmax probabilities.

The operation order is projection, head splitting, scaled dot products, masked softmax, probability dropout, value aggregation, and output projection. Dropout therefore removes individual query/key contributions before values are mixed. The existing fused attention path remains available to the bias-free constructors without probability dropout.

def Runtime.Autograd.Model.Layers.Attention.optionalShapes (enabled : Bool) (shape : Spec.Shape) : List Spec.Shape

Include one state slot only when the corresponding option is enabled.

def Runtime.Autograd.Model.Layers.Attention.optionalState (enabled : Bool) {shape : Spec.Shape} (value : TorchLean.Tensor Float shape) : TorchLean.TensorPack Float (optionalShapes enabled shape)

Initial value for an optional state slot, with no placeholder when it is disabled.

def Runtime.Autograd.Model.Layers.Attention.optionalInit (enabled : Bool) (shape : Spec.Shape) (init : Module.RuntimeInit.FloatInit) : Module.RuntimeInit.Plan (optionalShapes enabled shape)

Native initialization follows the same optional shape list as the ordinary state pack.

def Runtime.Autograd.Model.Layers.Attention.takeOptional {Ref : Spec.Shape → Type} (enabled : Bool) (shape : Spec.Shape) {rest : List Spec.Shape} : RefList Ref (optionalShapes enabled shape ++ rest) → Option (Ref shape) × RefList Ref rest

Read one optional state slot and return the remaining, still shape-indexed references.

def Runtime.Autograd.Model.Layers.multiHeadAttention (batchSize sequenceLength modelWidth headCount headWidth : ℕ) {sequenceLengthNonzero : sequenceLength ≠ 0} (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed : ℕ := 0) (weightInitialization? outputWeightInitialization? : Option Torch.Init.Scheme := none) (mask : Option (TorchLean.Tensor Bool [sequenceLength, sequenceLength]) := none) : Layer [batchSize, sequenceLength, modelWidth] [batchSize, sequenceLength, modelWidth]

Multi-head self-attention layer for a sequence (sequenceLength × modelWidth) → (sequenceLength × modelWidth).

This layer packs the four projection matrices (Wq, Wk, Wv, Wo) and calls the TorchLean attention primitive. An optional boolean mask of shape (sequenceLength × sequenceLength) can be provided, for example for causal masking.

PyTorch analogy: torch.nn.MultiheadAttention(embed_dim=modelWidth, num_heads=headCount) in self-attention mode.

def Runtime.Autograd.Model.Layers.multiHeadAttentionOutputBias (batchSize sequenceLength modelWidth headCount headWidth : ℕ) {sequenceLengthNonzero : sequenceLength ≠ 0} (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed : ℕ := 0) (weightInitialization? outputWeightInitialization? : Option Torch.Init.Scheme := none) (mask : Option (TorchLean.Tensor Bool [sequenceLength, sequenceLength]) := none) : Layer [batchSize, sequenceLength, modelWidth] [batchSize, sequenceLength, modelWidth]

Multi-head self-attention with a trainable bias on the final output projection.

The Q/K/V projections remain bias-free. This is the parameterization used in Karpathy's educational GPT implementation: the three per-head projections are linear maps without bias, while the projection applied after concatenating the heads is affine.

def Runtime.Autograd.Model.Layers.multiHeadAttentionConfigured (batchSize sequenceLength modelWidth headCount headWidth : ℕ) {sequenceLengthNonzero : sequenceLength ≠ 0} (queryWeightSeed keyWeightSeed valueWeightSeed outputWeightSeed : ℕ := 0) (weightInitialization? outputWeightInitialization? : Option Torch.Init.Scheme := none) (mask : Option (TorchLean.Tensor Bool [sequenceLength, sequenceLength]) := none) (inputBias outputBias : Bool := false) (dropout? : Option Float := none) (dropoutSeed : ℕ := 0) : Layer [batchSize, sequenceLength, modelWidth] [batchSize, sequenceLength, modelWidth]

Configurable affine self-attention with optional probability dropout.

The first four state slots remain queryWeight, keyWeight, valueWeight, outputWeight. Enabling input bias appends queryBias, keyBias, valueBias; enabling output bias appends outputBias; enabling dropout appends its non-trainable scalar probability. Disabled options allocate no state slots. All biases start at zero and consume no initialization seeds.