Selective Mamba

The Mamba-1 recurrence, expressed through the same differentiable operations used by other TorchLean layers. A token determines its convolution feature, time step, and input/output state vectors. The previous hidden state enters only the affine state update.

Parameter matrices use the orientation of Models.SelectiveMambaBlockSpec: a row vector multiplies a matrix on the right. The convolution kernel is newest-first. We store logA and compute A = exp(logA) before the recurrence; over the reals this makes every continuous-time rate -A negative throughout training. Floating-point overflow and underflow still follow the chosen backend's arithmetic.

step and runArray carry both the hidden state and projected-token history. Keeping references in that cache lets a continued computation retain gradients through earlier chunks. Callers that want truncated backpropagation can detach those references explicitly.

This implementation follows the Spec's dense time-step projection. It does not use a fused variable-coefficient scan or the low-rank time-step parameterization of the authors' default block.

structure Runtime.Autograd.Model.Mamba.Options : Type

Dimensions inside a selective Mamba layer, independently of its input and output widths.

• expansion : ℕ

  Expanded channels per output feature: innerWidth = expansion * outputWidth.

• stateWidth : ℕ

  Number of diagonal recurrent states carried by each expanded channel.

• kernelWidth : ℕ

  Number of newest-first taps in the causal depthwise convolution.

def Runtime.Autograd.Model.Mamba.instReprOptions.repr : Options → ℕ → Std.Format

@[instance_reducible]

instance Runtime.Autograd.Model.Mamba.instReprOptions : Repr Options

@[instance_reducible]

instance Runtime.Autograd.Model.Mamba.instDecidableEqOptions : DecidableEq Options

def Runtime.Autograd.Model.Mamba.instDecidableEqOptions.decEq (x✝ x✝¹ : Options) : Decidable (x✝ = x✝¹)

def Runtime.Autograd.Model.Mamba.Options.validate (options : Options) : Except String Unit

Reject empty internal axes before a layer allocates parameters or records operations.

structure Runtime.Autograd.Model.Mamba.Parameters (T : Spec.Shape → Type) (inputWidth innerWidth stateWidth outputWidth kernelWidth : ℕ) : Type

The eleven trainable tensors, in Spec matrix orientation.

The content, gate, B, C, and output projections have no bias. The convolution and time-step projection each have a bias. logA parameterizes positive rate magnitudes, and dSkip multiplies the activated convolution feature directly.

• xProj : T [inputWidth, innerWidth]

  Project raw tokens into the content path carried by the convolution cache.

• zProj : T [inputWidth, innerWidth]

  Project raw tokens into the SiLU gate, independently of the convolution path.

• convKernel : T [kernelWidth, innerWidth]

  Depthwise taps: row zero multiplies the current projected token.

• convBias : T [innerWidth]

  Additive channel bias before the content SiLU.

• dtProj : T [innerWidth, innerWidth]

  Dense map from activated content to per-channel time steps before softplus.

• dtBias : T [innerWidth]

  Time-step bias before softplus, stored separately from the dense projection.

• logA : T [innerWidth, stateWidth]

  Logarithms of rate magnitudes; transitions use exp(-delta * exp(logA)).

• bProj : T [innerWidth, stateWidth]

  Produce a token's B vector, shared by the expanded channels.

• cProj : T [innerWidth, stateWidth]

  Produce a token's C vector for reading out the updated state.

• dSkip : T [innerWidth]

  Per-channel coefficient of the direct D * u readout path.

• outProj : T [innerWidth, outputWidth]

  Project gated expanded channels back to the requested output width.

structure Runtime.Autograd.Model.Mamba.State (T : Spec.Shape → Type) (innerWidth stateWidth : ℕ) : Type

State at a chunk boundary.

history[0] is the most recent projected content token, before convolution and SiLU. A step prepends the new projection and retains at most kernelWidth entries. Missing entries mean zero padding; the hidden state has one stateWidth vector for each expanded channel.

• hidden : T [innerWidth, stateWidth]

  Diagonal state after the most recently processed token.

• history : Array (T [innerWidth])

  Projected content tokens in newest-first order.

def Runtime.Autograd.Model.Mamba.Internal.project {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {input output : ℕ} (x : Ref [input]) (weight : Ref [input, output]) : m (Ref [output])

Right-multiply a vector by a Spec-oriented projection without introducing a bias parameter.

def Runtime.Autograd.Model.Mamba.Internal.stepWithRates {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {inputWidth innerWidth stateWidth outputWidth kernelWidth : ℕ} (parameters : Parameters Ref inputWidth innerWidth stateWidth outputWidth kernelWidth) (rates : Ref [innerWidth, stateWidth]) (state : State Ref innerWidth stateWidth) (x : Ref [inputWidth]) : m (State Ref innerWidth stateWidth × Ref [outputWidth])

One step with rate magnitudes already computed.

Sharing exp(logA) across a sequence avoids recording the same exponential for every token. Its reference still participates in each transition, so reverse accumulation sums all rate gradients.

def Runtime.Autograd.Model.Mamba.zeroState {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {innerWidth stateWidth : ℕ} : m (State Ref innerWidth stateWidth)

Start a fresh sequence with zero recurrent state and zero-padded convolution history.

def Runtime.Autograd.Model.Mamba.step {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {inputWidth innerWidth stateWidth outputWidth kernelWidth : ℕ} (parameters : Parameters Ref inputWidth innerWidth stateWidth outputWidth kernelWidth) (state : State Ref innerWidth stateWidth) (x : Ref [inputWidth]) : m (State Ref innerWidth stateWidth × Ref [outputWidth])

Advance one token using the selective Mamba-1 equations.

With u = SiLU(causalConv(x @ xProj)), the update is h'[d,n] = exp(-softplus(u @ dtProj + dtBias)[d] * exp(logA[d,n])) * h[d,n] + (delta[d] * (u @ bProj)[n]) * u[d]. The readout contracts h' with u @ cProj, adds dSkip * u, gates by SiLU(x @ zProj), and applies outProj.

def Runtime.Autograd.Model.Mamba.runArray {α : Type} [TorchLean.Storage α] [Context α] {m : Type → Type} [Monad m] [Ops m α] {inputWidth innerWidth stateWidth outputWidth kernelWidth : ℕ} (parameters : Parameters Ref inputWidth innerWidth stateWidth outputWidth kernelWidth) (state : State Ref innerWidth stateWidth) (xs : Array (Ref [inputWidth])) : m (State Ref innerWidth stateWidth × Array (Ref [outputWidth]))

Run a chunk and return the complete cache needed by the next chunk, together with its outputs.

An empty chunk preserves both pieces of state. Chunk boundaries introduce no detach operation: using the returned references in a later chunk keeps gradients through the earlier computation.