Mamba Models

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

The recurrent core 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.

• modelWidth : ℕ

  Output feature width of the Mamba block, before 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.

@[instance_reducible]

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

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

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 block maps each vocabularySize-wide token to modelWidth features, and a final affine map produces vocabulary logits at every position. Its 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 Mamba tensors and vocabulary projection.

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.runArray exposes explicit state and convolution history for streaming computations.