Mamba Models

Configuration and a language-model constructor for selective Mamba-1 sequence models.

Each recurrent layer uses causal depthwise convolution, input-dependent time steps and B/C vectors, learned negative state rates, and a gated readout. It is built from generic differentiable operations shared by CPU and CUDA execution.

structure TorchLean.nn.models.Mamba.Config : Type

Configuration for the trainable one-hot-token Mamba language model.

• vocabularySize : ℕ

  Number of token categories accepted and predicted at each sequence position.

• modelWidths : List ℕ

  Feature widths of the Mamba layers. An empty list builds only the vocabulary projection.

• expansion : ℕ

  Expanded channels per model feature in the convolution and recurrent path.

• stateWidth : ℕ

  Diagonal recurrent states per expanded channel.

• kernelWidth : ℕ

  Newest-first taps in the causal depthwise convolution.

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

@[instance_reducible]

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

def TorchLean.nn.models.Mamba.Config.options (config : Config) : Runtime.Autograd.Model.Mamba.Options

Internal Mamba dimensions passed unchanged to the trainable layer.

def TorchLean.nn.models.Mamba.Config.validate (config : Config) (sequenceLength : ℕ) : Except String Unit

Validate model dimensions before allocating recurrent or projection parameters.

@[reducible, inline]

abbrev TorchLean.nn.models.Mamba.Config.inputShape (config : Config) (sequenceLength : ℕ) (batchShape : Shape := []) : Shape

One-hot input shape batchShape × sequenceLength × vocabularySize.

@[reducible, inline]

abbrev TorchLean.nn.models.Mamba.Config.outputShape (config : Config) (sequenceLength : ℕ) (batchShape : Shape := []) : Shape

Logit output shape batchShape × sequenceLength × vocabularySize.

def TorchLean.nn.models.Mamba.languageModel (config : Config) (sequenceLength : ℕ) (batchShape : Shape := []) : Builder (Sequential (config.inputShape sequenceLength batchShape) (config.outputShape sequenceLength batchShape))

Trainable selective Mamba-1 language model over one-hot token inputs.

The layers map each vocabularySize-wide token through modelWidths in order, and a final affine map produces vocabulary logits at every position. Each layer's internal width is expansion * modelWidth. Every sequence starts with zero hidden state and empty convolution history; each batch element has its own recurrence while sharing the eleven tensors per Mamba layer and the vocabulary projection. With no Mamba layers, the projection accepts empty sequences and ignores the unused core options.

The time-step projection is a dense matrix, matching Models.SelectiveMambaBlockSpec. The usual low-rank Mamba checkpoint stores two factors instead; their product matches a forward map here, but training a dense matrix gives a different parameterization. Model.Mamba.run exposes explicit state and convolution history for streaming computations.

def TorchLean.nn.models.Mamba.languageModel.buildLayers (config : Config) (sequenceLength : ℕ) (batchShape : Shape) (inputWidth : ℕ) (modelWidths : List ℕ) : Builder (Sequential (batchShape.concat [sequenceLength, inputWidth]) (config.outputShape sequenceLength batchShape))