Submitted on 1st December, 2023 on arXiv, the paper titled “Mamba: Linear-Time Sequence Modeling with Selective State Spaces” proposed an interesting approach to sequence modeling. The authors — Albert Gu, Tri Dao — introduced, ‘Mamba’ that utilized ‘selective’ state space models (SSM) to attain results that compete with the performance of the, now ubiquitous, Transformer model.

Transformers have seen recent popularity with the rise of Large Language Models (LLMs) like LLaMa-2, GPT-4, Claude, Gemini, etc., but it surely suffers from the issue of context window. The problem with transformers lies in it’s core, the multi head-attention mechanism.

The predominant issue with multi-head attention sprouts from the indisputable fact that for input sequence length n, the time complexity and space complexity scales by O(n²). This limits the length of the context window of an LLM. Because, to extend it by 10x, we want to scale the hardware requirement (most notably GPU VRAM) by 100x.

Mamba, alternatively, scales by O(n)!, i.e., Linearly.

This linear scaling is what has taken wind for researchers to take a position that Mamba is perhaps the longer term of sequence modeling.

The core of the Mamba model comes from the concept of State Space Models. State Space Models, like Transformers and RNN, process sequences of data, like text, audio signals, video frames, DNA sequences, etc.

State Space Models come from an idea of describing a physical system as a set of input, outputs, and variables. These variables are: A, B, C, D. The technique of SSM involves calculation of an internal state vector h(t), given an input x(t). Then, we do a weighted sum of h(t) and x(t) where the weights are A, B, C, & D. In the best form (continuous time-invariant), the method formulation looks like:

h(t) is usually called the ‘hidden’ or the ‘latent’ state, I can be sticking to calling it the ‘hidden’ state for higher clarity. It will be important to notice that A, B, C, and D are learnt parameters in SSM.

What are the variables?

The variables, A, B, C & D, are learnt parameters, they usually could be described as:

• A: How much should the previous hidden state (h) be considered to calculate the brand new hidden state
• B: How much should the input (x) be consider to calculate the brand new hidden state.
• C: How much should the brand new hidden state be considered in calculating the output (y).
• D: How much should the input (x) be consider in calculating the output (y).

D comes in the long run of the computations and doesn’t affect how the hidden state is calculated. Hence, it will likely be considered outside of ssm, and could be considered a skip connection.

Going from continuous spaces to discrete spaces

The above formulation applies to a system where the input and output belong to a continuous space. But in cases, like language modeling, where the input and output belong to discrete spaces (token values in a vocabulary). Also, finding h(t) is analytically difficult. This could be achieved by performing a Zero-order hold.

In a zero-order hold, each time an input is received, the model holds its value till the following input is received. This results in a continuous input space.

This length of ‘hold’ is set by a brand new parameter called, step size ∆. It might be considered the resolution of the input. Ideally, ∆ must be infinitesimal.

Mathematically, Zero-order hold could be described as:

Finally, we are able to create a discrete SSM, as:

Since, D is used with a skip connection outside of SSM, the output could be reduced to:

In SSMs, the hidden state is carried over to when the following input is received. This is comparable to how Recurrent Neural Networks function.

This recurrent format of SSM could be unwrapped, identical to RNNs. But unlike RNNs, that are iterative and slow, SSM can process the input sequence in parallel (identical to transformers) and this makes the training processes faster.

Note that ‘D’ is utilized in a skip connection, which is outside of SSM.

The important thing insight in how SSM make training fast is to make use of the variables A, B, C in a pre-computed convolutional kernel. Maarten Grootendorst wrote a extremely good explanation on how this canonical ‘convolutional’ kernel is constructed. But here’s an easy mathematical explanation.

Consider the output y. For a sequence length of k, the output for y(k) can be represented (assuming h0 = zero):

Similarly, y3 could be represented as:

Extrapolating the pattern, yk could be represented as:

This formulation could be further reduced to:

The funny looking multiplication symbol represents a convolution operation, where the convolution kernel is K. Notice that K shouldn’t be depending on x, hence K could be pre-computed right into a convolutional kernel, which makes the method faster.

Pretty much as good because the computational capability of SSM sounds, it seems to be pretty meh in metrics like accuracy in comparison with Transformers.

The core issue lies with the variables, ∆, A, B, & C. Seems that since we apply the identical matrices to each input, they can’t really process the context of the sequence.

SSMs are inflexible in the way in which they process data[4]

So what’s so special about Mamba? In mamba, we use a process called ‘selective’ SSM, where the variables, ∆, B, & C, are computed based on the input. 🤔. We do that by passing the present input through Linear layers, and take the output to be the ∆, B, & C.

But then this makes ∆, B, & C input dependent, hence meaning that they can’t be pre-computed 😢, fast convolution isn’t going to work here. But, the authors discuss a way, which relies on parallel associative scan.

Parallel Associative Scan

Parallel associative scan is a strong technique utilized in parallel computing to perform a prefix sum operation, which is a cumulative operation on a sequence of numbers. This operation is “associative”, meaning the way in which numbers are grouped within the operation doesn’t change the result.

Within the context of the Mamba model, by defining an associative operator, elements and associative operators for a parallel associative scan operation are obtained. This enables for solving problems on the entire time interval in parallel, leading to logarithmic time complexity within the variety of sub-intervals.

Hardware aware algorithm

Together with associative scan, the authors also propose a hardware aware algorithm, where they use the quirks inside Nvidia GPUs related to the speed of HBM and SRAM. They argue that the computation of SSM states could be sped up by:

• keeping the hidden state and A within the faster but less capability SRAM,
• while computing ∆, B, & C, within the slower but larger capability HBM.
• They then transfer ∆, B, & C to the SRAM, compute the brand new hidden state inside SRAM.
• After which write ∆, B & C back to HBM.

Within the implementation section, I won’t be discussing on the right way to work with the hardware aware algorithm, somewhat I can be only using parallel associative scan.

With all of this in mind, let’s explore and implement the Mamba architecture using Keras and TensorFlow.

The Mamba architecture, after reading the paper and evaluation of the code, could be broken into a number of key components that are connected as:

The Mamba architecture consists of multiple stacked layers of ‘Mamba blocks’. Which, judging from the above illustration, consists of quite a number of components. One other essential thing to notice is that the authors add the output from Selective SSM to the unique input after which apply a normalization layer to it. This normalization could be either a Layer normalization or an RMS normalization.

Lets start with coding a part of Mamba. We are going to using the next dependencies:

tensorflow[and-cuda]==2.15.0.post1 # if you ought to use GPU or
tensorflow==2.15.0.post1 # if you ought to only use CPU
transformers==4.36.2 # for using the bert tokenizer
einops==0.7.0 # useful to make matrix manipulation faster
datasets==2.16.1 # to load datasets
# all other modules (like numpy) can be auto installed

Imports:

import tensorflow_datasets as tfds
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, Model
from dataclasses import dataclass
from einops import rearrange, repeat
from typing import Union
from transformers import AutoTokenizer
import datasets
import math
import numpy as np

To make the modeling argument processing easier, let’s create an easy ModelArgs dataclass as a config class. This enables us to only pass the dataclass variable within the arguments after we are initializing the model.

@dataclass
class ModelArgs:
model_input_dims: int = 64
model_states: int = 64
projection_expand_factor: int = 2
conv_kernel_size: int = 4
delta_t_min: float = 0.001
delta_t_max: float = 0.1
delta_t_scale: float = 0.1
delta_t_init_floor: float = 1e-4
conv_use_bias: bool = True
dense_use_bias: bool = False
layer_id: int = -1
seq_length: int = 128
num_layers: int = 5
dropout_rate: float = 0.2
use_lm_head: float = False
num_classes: int = None
vocab_size: int = None
final_activation = None
loss:Union[str, keras.losses.Loss] = None
optimizer: Union[str, keras.optimizers.Optimizer] = keras.optimizers.AdamW()
metrics = ['accuracy']
def __post_init__(self):
self.model_internal_dim: int = int(self.projection_expand_factor * self.model_input_dims)
self.delta_t_rank = math.ceil(self.model_input_dims/16)
if self.layer_id == -1:
self.layer_id = np.round(np.random.randint(0, 1000), 4)
if self.vocab_size == None:
raise ValueError("vocab size can't be none")
if self.use_lm_head:
self.num_classes=self.vocab_size
else:
if self.num_classes == None:
raise ValueError(f'num classes can't be {self.num_classes}')
if self.num_classes == 1:
self.final_activation = 'sigmoid'
else:
self.final_activation = 'softmax'
if self.loss == None:
raise ValueError(f"loss can't be {self.loss}")

Load the bert-base-uncased tokenizer:

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
vocab_size = tokenizer.vocab_size

Before we implement our Mamba and SSM classes, we want to implement the parallel associative scan, the code looks like this:

def selective_scan(u, delta, A, B, C, D):
# first step of A_bar = exp(ΔA), i.e., ΔA
dA = tf.einsum('bld,dn->bldn', delta, A) 
dB_u = tf.einsum('bld,bld,bln->bldn', delta, u, B)
dA_cumsum = tf.pad(
dA[:, 1:], [[0, 0], [1, 1], [0, 0], [0, 0]])[:, 1:, :, :]
dA_cumsum = tf.reverse(dA_cumsum, axis=[1])  # Flip along axis 1
# Cumulative sum along all of the input tokens, parallel prefix sum, 
# calculates dA for all of the input tokens parallely
dA_cumsum = tf.math.cumsum(dA_cumsum, axis=1)  
# second step of A_bar = exp(ΔA), i.e., exp(ΔA)
dA_cumsum = tf.exp(dA_cumsum)  
dA_cumsum = tf.reverse(dA_cumsum, axis=[1])  # Flip back along axis 1
x = dB_u * dA_cumsum
# 1e-12 to avoid division by 0
x = tf.math.cumsum(x, axis=1)/(dA_cumsum + 1e-12) 
y = tf.einsum('bldn,bln->bld', x, C)
return y + u * D 

With this, we are able to implement the MambaBlock:

class MambaBlock(layers.Layer):
def __init__(self, modelargs: ModelArgs, *args, **kwargs):
super().__init__(*args, **kwargs)
self.args = modelargs
args = modelargs
self.layer_id = modelargs.layer_id
self.in_projection = layers.Dense(
args.model_internal_dim * 2, 
input_shape=(args.model_input_dims,), use_bias=False)
self.conv1d = layers.Conv1D(
filters=args.model_internal_dim,
use_bias=args.conv_use_bias,
kernel_size=args.conv_kernel_size,
groups=args.model_internal_dim,
data_format='channels_first',
padding='causal'
)
# this layer takes in current token 'x' 
# and outputs the input-specific Δ, B, C (in accordance with S6)
self.x_projection = layers.Dense(args.delta_t_rank + args.model_states * 2, use_bias=False)
# this layer projects Δ from delta_t_rank to the mamba internal 
# dimension
self.delta_t_projection = layers.Dense(args.model_internal_dim, 
input_shape=(args.delta_t_rank,), use_bias=True)
self.A = repeat(
tf.range(1, args.model_states+1, dtype=tf.float32), 
'n -> d n', d=args.model_internal_dim)
self.A_log = tf.Variable(
tf.math.log(self.A), 
trainable=True, dtype=tf.float32, 
name=f"SSM_A_log_{args.layer_id}")
self.D = tf.Variable(
np.ones(args.model_internal_dim), 
trainable=True, dtype=tf.float32, 
name=f"SSM_D_{args.layer_id}")
self.out_projection = layers.Dense(
args.model_input_dims, 
input_shape=(args.model_internal_dim,), 
use_bias=args.dense_use_bias)
def call(self, x):
"""Mamba block forward. This looks the identical as Figure 3 in Section 3.4 within the Mamba pape.
Official Implementation:
class Mamba, https://github.com/state-spaces/mamba/blob/predominant/mamba_ssm/modules/mamba_simple.py#L119
mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/predominant/mamba_ssm/ops/selective_scan_interface.py#L311
"""
(batch_size, seq_len, dimension) = x.shape
x_and_res = self.in_projection(x) # shape = (batch, seq_len, 2 * model_internal_dimension)
(x, res) = tf.split(x_and_res, 
[self.args.model_internal_dim, 
self.args.model_internal_dim], axis=-1)
x = rearrange(x, 'b l d_in -> b d_in l')
x = self.conv1d(x)[:, :, :seq_len]
x = rearrange(x, 'b d_in l -> b l d_in')
x = tf.nn.swish(x)
y = self.ssm(x)
y = y * tf.nn.swish(res)
return self.out_projection(y)
def ssm(self, x):
"""Runs the SSM. See:
- Algorithm 2 in Section 3.2 within the Mamba paper
- run_SSM(A, B, C, u) in The Annotated S4
Official Implementation:
mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/predominant/mamba_ssm/ops/selective_scan_interface.py#L311
"""
(d_in, n) = self.A_log.shape
# Compute ∆ A B C D, the state space parameters.
#     A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A is not selective)
#     ∆, B, C are input-dependent (this can be a key difference between Mamba and the linear time invariant S4,
#                                  and is why Mamba known as **selective** state spaces)
A = -tf.exp(tf.solid(self.A_log, tf.float32)) # shape -> (d_in, n)
D = tf.solid(self.D, tf.float32)
x_dbl = self.x_projection(x) # shape -> (batch, seq_len, delta_t_rank + 2*n)
(delta, B, C) = tf.split(
x_dbl, 
num_or_size_splits=[self.args.delta_t_rank, n, n], 
axis=-1) # delta.shape -> (batch, seq_len) & B, C shape -> (batch, seq_len, n)
delta = tf.nn.softplus(self.delta_t_projection(delta)) # shape -> (batch, seq_len, model_input_dim)
return selective_scan(x, delta, A, B, C, D)

Finally, a residual block to implement the external skip connection.

class ResidualBlock(layers.Layer):
def __init__(self, modelargs: ModelArgs, *args, **kwargs):
super().__init__(*args, **kwargs)
self.args = modelargs
self.mixer = MambaBlock(modelargs)
self.norm = layers.LayerNormalization(epsilon=1e-5)
def call(self, x):
"""
Official Implementation:
Block.forward(), https://github.com/state-spaces/mamba/blob/predominant/mamba_ssm/modules/mamba_simple.py#L297
Note: the official repo chains residual blocks that appear to be
[Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> ...
where the primary Add is a no-op. This is solely for performance reasons as this
allows them to fuse the Add->Norm.
We as an alternative implement our blocks because the more familiar, simpler, and numerically equivalent
[Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> ....
"""
return self.mixer(self.norm(x)) + x

With this, we are able to initialize our model. In this instance, I can be demonstrating the right way to use the Mamba block to create an easy classification model, but it may well be easily modified to turn into a language model. Let’s load the IMDB reviews dataset for an easy sentiment classifier.

from datasets import load_dataset
from tqdm import tqdm
dataset = load_dataset("ajaykarthick/imdb-movie-reviews")

First we create a function that may take the model args and return a model.

def init_model(args: ModelArgs):
input_layer = layers.Input(shape=(args.seq_length,), name='input_ids')
x = layers.Embedding(
args.vocab_size, 
args.model_input_dims, 
input_length=args.seq_length)(input_layer)
for i in range(args.num_layers):
x = ResidualBlock(args, name=f"Residual_{i}")(x)
x = layers.Dropout(args.dropout_rate)(x) # for regularization
x = layers.LayerNormalization(epsilon=1e-5)(x) # normalization layer
# use flatten provided that we are usually not using the model as an LM
if not args.use_lm_head: 
x = layers.Flatten()(x)
x = layers.Dense(1024, activation=tf.nn.gelu)(x)
output_layer = layers.Dense(
args.num_classes, 
activation=args.final_activation)(x)
model = Model(
inputs=input_layer, 
outputs=output_layer, name='Mamba_ka_Mamba')
model.compile(
loss=args.loss,
optimizer=args.optimizer,
metrics=args.metrics
)
return model

Now we are able to initialize our model, and summarize it:

args = ModelArgs(
model_input_dims=128,
model_states=32,
num_layers=12,
dropout_rate=0.2,
vocab_size=vocab_size,
num_classes=1,
loss='binary_crossentropy',
)
model = init_model(args)
model.summary()

Model: "Mamba_ka_Mamba"
_________________________________________________________________
Layer (type)                Output Shape              Param #   
=================================================================
input_ids (InputLayer)      [(None, 128)]             0         
embedding_2 (Embedding)     (None, 128, 128)          3906816   
Residual_0 (ResidualBlock)  (None, 128, 128)          129024    
dropout_24 (Dropout)        (None, 128, 128)          0         
Residual_1 (ResidualBlock)  (None, 128, 128)          129024    
dropout_25 (Dropout)        (None, 128, 128)          0
... (I even have shrinked this to make it more readable)
dropout_35 (Dropout)        (None, 128, 128)          0         
layer_normalization_38 (La  (None, 128, 128)          256       
yerNormalization)                                               
flatten_2 (Flatten)         (None, 16384)             0         
dense_148 (Dense)           (None, 1024)              16778240  
dense_149 (Dense)           (None, 1)                 1025      
=================================================================
Total params: 22234625 (84.82 MB)
Trainable params: 22234625 (84.82 MB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________

For easier processing, lets pre-tokenize our data right into a numpy arrays, then convert them into tf.data.Dataset objects:

train_labels, test_labels = [], []
train_ids = np.zeros((len(dataset['train']), args.seq_length))
test_ids = np.zeros((len(dataset['test']), args.seq_length))
for i, item in enumerate(tqdm(dataset['train'])):
text = item['review']
train_ids[i, :] = tokenizer.encode_plus(
text, 
max_length=args.seq_length, 
padding='max_length', 
return_tensors='np')['input_ids'][0][:args.seq_length]
train_labels.append(item['label'])
for i, item in enumerate(tqdm(dataset['test'])):
text = item['review']
test_ids[i, :] = tokenizer.encode_plus(
text, 
max_length=args.seq_length, 
padding='max_length', 
return_tensors='np')['input_ids'][0][:args.seq_length]
test_labels.append(item['label'])
del dataset # delete the unique dataset to save lots of some memory
BATCH_SIZE = 32
train_dataset = tf.data.Dataset.from_tensor_slices((train_ids, train_labels)).batch(BATCH_SIZE).shuffle(1000)
test_dataset = tf.data.Dataset.from_tensor_slices((test_ids, test_labels)).batch(BATCH_SIZE).shuffle(1000)

Now the model could be trained:

history = model.fit(train_dataset, validation_data=test_dataset, epochs=10)

You possibly can mess around with the inference algorithm:

def infer(text: str, model: Model, tokenizer):
tokens = tokenizer.encode(
"Hello what's up", 
max_length=args.seq_length, 
padding='max_length', return_tensors='np')
output = model(tokens)[0, 0]
return output

This model could be converted right into a language model and algorithms like beam search, top-k sampling, greedy sampling, etc. could be used to generate language.

This code could be found on my Github.

Quite a lot of the code is inspired from the mamba’s official implementation[2] and one other pytorch implementation called ‘mamba-tiny’[3]

Thanks for reading.