TorchLean API

NN.Runtime.Autograd.Model.Layers.Mamba

Mamba Layer #

The trainable layer around the selective recurrence in Model.Mamba. Every call starts with zero hidden state and empty convolution history. Streaming callers can use Mamba.runArray with an explicit cache instead; the ordinary layer has no mutable sequence state.

def Runtime.Autograd.Model.Layers.mamba (sequenceLength inputWidth hiddenWidth : ) (inputWeightSeed stateWeightSeed gateWeightSeed : := 0) (options : Mamba.Options := { }) :
Layer [sequenceLength, inputWidth] [sequenceLength, hiddenWidth]

Selective Mamba-1 layer on a time-major sequence.

hiddenWidth is the output width. The convolution and recurrent path use innerWidth = options.expansion * hiddenWidth channels, each with options.stateWidth diagonal states. The kernel contains options.kernelWidth newest-first taps. Outputs depend only on the current token and its prefix.

State order is xProj, zProj, convKernel, convBias, dtProj, dtBias, logA, bProj, cProj, dSkip, outProj. Matrix orientations match Models.SelectiveMambaBlockSpec; the positive Spec rate tensor is exp(logA). This is the dense time-step parameterization: importing a factorized reference checkpoint requires multiplying its two time-step projection matrices.

Projection and kernel tensors use Xavier initialization, convolution bias starts at zero, time-step bias gives an initial step of 0.01 when the projected feature is zero, the diagonal rate magnitudes start at 1, ..., stateWidth, and the skip coefficients start at one. The three seed arguments control content/convolution, time-step/B/C, and gate/output initialization.

The recurrence is composed from generic differentiable operations, including the convolution. Eager and typed graph execution therefore differentiate the same program. The CUDA variable-coefficient scan binding is not used here.

Instances For