Long Short-Term Memory Models

Higher‑level LSTM architectures built from module specs (Spec.Module.Chain), including:

• sequence‑to‑sequence outputs,
• classifier heads (many‑to‑one),
• multi‑layer compositions.

Cell equations are in NN/Spec/Layers/Lstm.lean; this file focuses on composing modules.

References (math + PyTorch behavior):

• Hochreiter and Schmidhuber (1997), "Long Short-Term Memory" (original LSTM): https://www.bioinf.jku.at/publications/older/2604.pdf
• PyTorch nn.LSTM docs: https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html
• PyTorch nn.LSTMCell docs: https://pytorch.org/docs/stable/generated/torch.nn.LSTMCell.html

Model and gradient types

The LSTM model layer exposes first-class model objects with:

• a forward pass,
• a standard training objective, and
• an explicit reverse-mode / BPTT backward pass producing parameter gradients.

The backward functions reuse the gate-aware implementation in NN.Spec.Layers.Lstm.

Gradient records

structure Spec.Lstm.Grads (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Parameter gradients for Lstm.Model.

This bundles the LSTM cell gradients and the time-distributed linear head gradients.

• cell : LSTMGateGradients α inputSize hiddenSize

  Gradients of the recurrent cell.

• output : LinearParameterGradients α hiddenSize outputSize

  Gradients of the output projection.

structure Spec.Lstm.ClassifierGrads (α : Type) [TorchLean.Storage α] (inputSize hiddenSize numClasses : ℕ) : Type

Parameter gradients for an LSTM classifier.

• cell : LSTMGateGradients α inputSize hiddenSize

  Gradients of the recurrent cell.

• classifier : LinearParameterGradients α hiddenSize numClasses

  Gradients of the classifier head.

def Spec.Lstm.sequence {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize outputSize : ℕ} (lstmSpec : LSTMSpec α inputSize hiddenSize) (linearSpec : LinearSpec α hiddenSize outputSize) : Module.Chain α [seqLen, inputSize] [seqLen, outputSize]

Sequence-to-sequence LSTM model as a Spec.Module.Chain: LSTM over time, then a per-timestep linear head.

PyTorch analogue: nn.LSTM producing an output sequence, followed by nn.linear applied at each time step.

def Spec.Lstm.classifier {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize numClasses : ℕ} (lstmSpec : LSTMSpec α inputSize hiddenSize) (classifierHead : LinearSpec α hiddenSize numClasses) (h : seqLen ≠ 0) : Module.Chain α [seqLen, inputSize] [numClasses]

Many-to-one LSTM classifier as a Spec.Module.Chain.

This runs an LSTM over the sequence and applies a linear classifier head to the final hidden state. PyTorch analogue: nn.LSTM + nn.linear, taking the last output/hidden.

def Spec.Lstm.stacked {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize outputSize : ℕ} (firstSpec : LSTMSpec α inputSize hiddenSize) (secondSpec : LSTMSpec α hiddenSize hiddenSize) (linearSpec : LinearSpec α hiddenSize outputSize) : Module.Chain α [seqLen, inputSize] [seqLen, outputSize]

Two-layer LSTM stack (sequence-to-sequence), followed by a per-timestep linear head.

The second LSTM consumes the hidden stream produced by the first.

def Spec.Lstm.languageModel {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen vocabularySize hiddenSize : ℕ} (embeddingSpec : LinearSpec α vocabularySize hiddenSize) (lstmSpec : LSTMSpec α hiddenSize hiddenSize) (outputSpec : LinearSpec α hiddenSize vocabularySize) : Module.Chain α [seqLen, vocabularySize] [seqLen, vocabularySize]

Simple LSTM language-model pipeline as a Spec.Module.Chain: embedding, LSTM core, and output projection.

In this spec layer we represent the embedding/projection as LinearSpecs (often used with one-hot token vectors). PyTorch analogue: nn.Embedding (conceptually) + nn.LSTM + nn.linear.

def Spec.Lstm.bidirectionalClassifier {α : Type} [TorchLean.Storage α] [Context α] [DecidableRel fun (x1 x2 : α) => x1 > x2] {seqLen inputSize hiddenSize numClasses : ℕ} (forwardSpec backwardSpec : LSTMSpec α inputSize hiddenSize) (classifierHead : LinearSpec α (hiddenSize + hiddenSize) numClasses) (h : seqLen ≠ 0) : Module.Chain α [seqLen, inputSize] [numClasses]

Bidirectional LSTM followed by a classifier on the final concatenated state.

structure Spec.Lstm.Model (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Bundle of parameters for a single-layer LSTM model with a linear output head.

This is a direct record representation (as opposed to the Spec.Module.Chain representation above).

• lstm : LSTMSpec α inputSize hiddenSize

  Recurrent cell parameters.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

structure Spec.Lstm.StackedModel (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize numLayers : ℕ) : Type

Bundle of parameters for a multi-layer LSTM model with a linear output head.

The first layer consumes inputSize, and all subsequent layers consume hiddenSize.

• firstLayer : LSTMSpec α inputSize hiddenSize

  First recurrent layer, whose input may differ from the hidden width.

• hiddenLayers : Fin (numLayers - 1) → LSTMSpec α hiddenSize hiddenSize

  Remaining recurrent layers.

• outputLayer : LinearSpec α hiddenSize outputSize

  Linear output projection.

structure Spec.Lstm.Classifier (α : Type) [TorchLean.Storage α] (inputSize hiddenSize numClasses : ℕ) : Type

Bundle of parameters for a many-to-one LSTM classifier.

The classifier head is applied to the final hidden state.

• lstm : LSTMSpec α inputSize hiddenSize

  Recurrent cell parameters.

• classifier : LinearSpec α hiddenSize numClasses

  Linear classifier head.

structure Spec.Lstm.Generator (α : Type) [TorchLean.Storage α] (vocabularySize hiddenSize : ℕ) : Type

Bundle of parameters for a many-to-many LSTM generator (language-model style).

This includes an (embedding) linear map, recurrent core, and output projection back to vocabulary.

• embedding : LinearSpec α vocabularySize hiddenSize

  Token projection used by this one-hot specification.

• lstm : LSTMSpec α hiddenSize hiddenSize

  Recurrent cell parameters.

• outputProjection : LinearSpec α hiddenSize vocabularySize

  Projection from hidden states to vocabulary logits.

structure Spec.Lstm.BidirectionalModel (α : Type) [TorchLean.Storage α] (inputSize hiddenSize outputSize : ℕ) : Type

Bundle of parameters for a bidirectional LSTM model with an output head.

The head consumes the concatenation of forward and backward hidden states. PyTorch analogue: nn.LSTM(..., bidirectional=true) plus a linear projection.

• forwardLstm : LSTMSpec α inputSize hiddenSize

  Recurrent cell for the original sequence order.

• backwardLstm : LSTMSpec α inputSize hiddenSize

  Recurrent cell for the reversed sequence order.

• outputLayer : LinearSpec α (hiddenSize + hiddenSize) outputSize

  Projection from concatenated forward and backward states.

structure Spec.Lstm.LanguageModel (α : Type) [TorchLean.Storage α] (vocabularySize hiddenSize : ℕ) : Type

Bundle of parameters for a stacked LSTM language model with deterministic dropout.

This model uses an array of LSTM layers (all with hiddenSize input/output) and applies an evaluation-mode dropout step between the recurrent stack and the output projection.

• embedding : LinearSpec α vocabularySize hiddenSize

  Token projection used by this one-hot specification.

• layers : Array (LSTMSpec α hiddenSize hiddenSize)

  Recurrent layers, ordered from input to output.

• outputProjection : LinearSpec α hiddenSize vocabularySize

  Projection from hidden states to vocabulary logits.

• dropoutRate : α

  Dropout probability used between the recurrent stack and output projection.

def Spec.Lstm.Model.forward {α : Type} [TorchLean.Storage α] [Context α] {inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (input : TorchLean.Tensor α [inputSize]) (state : LSTMState α hiddenSize) : TorchLean.Tensor α [outputSize] × LSTMState α hiddenSize

One-step forward pass for Lstm.Model.

Given an input vector and the previous LSTM state (hidden, cell), compute (output, new_state). PyTorch analogue: nn.LSTMCell step followed by a nn.linear head.

def Spec.Lstm.Model.forwardSequence {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) : TorchLean.Tensor α [seqLen, outputSize] × LSTMState α hiddenSize

Sequence forward pass for Lstm.Model.

Runs the LSTM over all timesteps (time-major), applies the output head to each hidden state, and returns (outputs, final_state).

Backward pass (BPTT) for the simple LSTM sequence model

This is the model-level analogue of Spec.lstmSequenceBackwardSpec. The only extra work we do here is to backprop through the per-timestep output projection and feed its gradient into the LSTM sequence backward pass.

def Spec.Lstm.Model.backward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) (outputGrad : TorchLean.Tensor α [seqLen, outputSize]) : Grads α inputSize hiddenSize outputSize × TorchLean.Tensor α [seqLen, inputSize] × LSTMState α hiddenSize

Backward pass for Lstm.Model.forwardSequence.

Returns:

• parameter gradients (Lstm.Grads)
• gradient w.r.t. input sequence (dInputs)
• gradient w.r.t. initial recurrent state (dInitialState)

def Spec.Lstm.Model.mseLoss {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (targets : TorchLean.Tensor α [seqLen, outputSize]) (initialState : LSTMState α hiddenSize) : α

MSE loss for the simple LSTM sequence model.

This runs Lstm.Model.forwardSequence and compares the predicted output sequence against targets using mseSpec.

def Spec.Lstm.Model.mseGrad {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (targets : TorchLean.Tensor α [seqLen, outputSize]) (initialState : LSTMState α hiddenSize) : α × Grads α inputSize hiddenSize outputSize

Compute (loss, grads) for the simple LSTM sequence model under MSE.

This is the “full training API” building block: an optimizer (SGD/Adam) can consume these grads.

def Spec.Lstm.Classifier.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (model : Classifier α inputSize hiddenSize numClasses) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) : TorchLean.Tensor α [numClasses]

Forward pass for an Lstm.Classifier (many-to-one).

This uses the final hidden state of the LSTM sequence as the classifier input.

Backward for the classifier head (many-to-one)

The classifier only consumes the final hidden state. We express that by feeding a gradient sequence that is zero everywhere except the last timestep.

def Spec.Lstm.Classifier.backward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (model : Classifier α inputSize hiddenSize numClasses) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialState : LSTMState α hiddenSize) (logitGrad : TorchLean.Tensor α [numClasses]) : ClassifierGrads α inputSize hiddenSize numClasses × TorchLean.Tensor α [seqLen, inputSize] × LSTMState α hiddenSize

Backward pass for an Lstm.Classifier (many-to-one).

This backprops through the classifier head, then runs an LSTM sequence backward pass where the hidden-state gradient is zero at all timesteps except the last.

def Spec.Lstm.Generator.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (model : Generator α vocabularySize hiddenSize) (inputTokens : TorchLean.Tensor α [seqLen, vocabularySize]) (initialState : LSTMState α hiddenSize) : TorchLean.Tensor α [seqLen, vocabularySize] × LSTMState α hiddenSize

Forward pass for an Lstm.Generator (many-to-many).

This applies an embedding linear map to each token vector, runs the LSTM, and projects each hidden state back into vocabulary space.

def Spec.Lstm.BidirectionalModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : BidirectionalModel α inputSize hiddenSize outputSize) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (forwardState backwardState : LSTMState α hiddenSize) : TorchLean.Tensor α [seqLen, outputSize]

Forward pass for a bidirectional LSTM model (time-major).

This runs a forward LSTM on the sequence, a backward LSTM on the reversed sequence, concatenates the two hidden streams per timestep, and applies an output head.

def Spec.Lstm.StackedModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize numLayers : ℕ} (model : StackedModel α inputSize hiddenSize outputSize numLayers) (inputs : TorchLean.Tensor α [seqLen, inputSize]) (initialStates : Fin numLayers → LSTMState α hiddenSize) (hLayers : 0 < numLayers) : TorchLean.Tensor α [seqLen, outputSize] × (Fin numLayers → LSTMState α hiddenSize)

Forward pass for a Lstm.StackedModel.

This runs the first layer on the input sequence, then threads the resulting hidden stream through each additional hidden layer, and finally applies the output head per timestep.

@[irreducible]

def Spec.Lstm.StackedModel.forward.processHiddenLayers {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize numLayers : ℕ} (model : StackedModel α inputSize hiddenSize outputSize numLayers) (hLayers : 0 < numLayers) (layer : ℕ) (layerInput : TorchLean.Tensor α [seqLen, hiddenSize]) (states : Fin numLayers → LSTMState α hiddenSize) : TorchLean.Tensor α [seqLen, hiddenSize] × (Fin numLayers → LSTMState α hiddenSize)

def Spec.Lstm.LanguageModel.forward {α : Type} [TorchLean.Storage α] [Context α] {seqLen vocabularySize hiddenSize : ℕ} (model : LanguageModel α vocabularySize hiddenSize) (inputTokens : TorchLean.Tensor α [seqLen, vocabularySize]) (initialStates : Array (LSTMState α hiddenSize)) : Option (TorchLean.Tensor α [seqLen, vocabularySize] × Array (LSTMState α hiddenSize))

Forward pass for Lstm.LanguageModel (teacher forcing, time-major).

This runs the embedding, then a stack of LSTM layers with provided initial states, applies evaluation-mode dropout (dropoutInferenceSpec), and projects to vocabulary logits.

def Spec.Lstm.LanguageModel.forward.processLayers {α : Type} [TorchLean.Storage α] [Context α] {seqLen hiddenSize : ℕ} (layers : List (LSTMSpec α hiddenSize hiddenSize)) (states : List (LSTMState α hiddenSize)) (layerInput : TorchLean.Tensor α [seqLen, hiddenSize]) : Option (TorchLean.Tensor α [seqLen, hiddenSize] × List (LSTMState α hiddenSize))

def Spec.Lstm.Model.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : Model α inputSize hiddenSize outputSize) : Module α [seqLen, inputSize] [seqLen, outputSize]

Package Lstm.Model as a shape-indexed module.

The Python expression records the intended runtime analogue; forward remains the mathematical meaning of the module.

def Spec.Lstm.Classifier.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize numClasses : ℕ} (model : Classifier α inputSize hiddenSize numClasses) : Module α [seqLen, inputSize] [numClasses]

Package Lstm.Classifier as an Spec.Module.

PyTorch analogue: nn.LSTM feeding a nn.linear classifier head.

def Spec.Lstm.BidirectionalModel.toModule {α : Type} [TorchLean.Storage α] [Context α] {seqLen inputSize hiddenSize outputSize : ℕ} (model : BidirectionalModel α inputSize hiddenSize outputSize) : Module α [seqLen, inputSize] [seqLen, outputSize]

Package Lstm.BidirectionalModel as an Spec.Module.

PyTorch analogue: nn.LSTM(..., bidirectional=true) feeding a per-timestep linear head.