Why is Mamba creating waves? Is it a replacement for transformers?

Linear-Time Sequence Modeling with Selective State Spaces

ransformers do not scale very well to long sequence lengths largely because of quadratic self-attention complexity. In very simple terms, Quadratic self-attention complexity in transformers is how the computational requirements increase rapidly (in a quadratic manner) as the size of the input data grows. Specifically, in vanilla transformer models used in language processing (attention is all you need), the attention mechanism compares each element in the input sequence with every other element. This means if the sequence has 'n' elements, the attention step needs to perform calculations for 'n x n' pairs, which becomes very resource-intensive for long sequences.

This weekend, I read this very interesting paper in ICLR 2024 and explored its Github repository. It proposes a state space model architecture with an intriguing name “Mamba” to a model which stands out for those of us in tech who grapple with massive datasets. Traditionally, the vanilla transformer models have outperformed most models in performance.

While it can be seen in the comparison below that models like Linformer, Performer, and Synthesizer have pushed the envelope with improved processing speeds, the vanilla Transformer has remained competitive in performance since 2017 (which is quite a feat given the speed of change in the AI world) . It's not just about raw speed; the vanilla Transformer excels due to its deep and robust architecture, which allows it to understand complex patterns effectively. This is vividly captured in the attached image, where we see models like Reformer and Sinkhorn making strides in speed, but when balancing performance and speed, the Transformer consistently ranks competitively.

Performance [Long-Range Arena (LRA)] Vs Computational Speed

Mamba Architecture for Dummies

The paper's proposed models are like smart assistants for busy teachers grading papers. Imagine a stack of essays: the teacher needs to read through every page, but not every sentence is equally important. The new model, like our smart assistant, can quickly pinpoint which parts of the essay have the key points and which can be skimmed over. This way, the teacher focuses on the most important content, grading efficiently without getting bogged down by less relevant details. Just as this saves the teacher time while ensuring a fair grade, the new model processes data more efficiently by focusing only on crucial information.

To put it simply, the Mamba model proposed in the paper addresses the challenge of sifting through and prioritizing vast amounts of information:

• Selective Attention: Mamba can focus on the most important data, like a detective picking out vital clues from a pile of evidence.
• Smart Computation: It uses a clever method to process data that avoids overloading the computer's memory, much like a librarian organizes books to save space but keep them quickly accessible.
• Unified Design: The model combines the best parts of older models into one efficient design. Think of it as a Swiss Army knife for data analysis—compact yet powerful.
• Versatile and Fast: It's designed to be fast and accurate for different types of data, whether it's words in a book or notes in a song, up to very long sequences.
• Proven Effectiveness: Tests in the paper show that Mamba is good at tasks like copying text or understanding language, and it's even better than previous models at tasks involving audio and DNA sequences.Mamba zeroes in on essential data, achieving both speed and accuracy. This is crucial for analyzing complex data, whether predicting stock movements or decoding genetic information. I am yet to conduct some experiments on the real world datasets but I tried to explain everything that the paper offers in the most simplistic way for beginners and others who do not have too much time to read through the entire paper.

Installation and Usage

To set up Mamba for processing data, you need to install a couple of key components.

First, add an efficient causal Conv1d layer that's a part of the Mamba block using the pip command.

pip install causal-conv1d        

Then, install the main Mamba package with :

pip install mamba-ssm        

If you're developing directly from the source, use pip install . within the repository directory. Should there be any version conflicts with PyTorch, try using the --no-build-isolation option with pip.

Ensure you're running this setup PyTorch version 1.12 or newer, and CUDA 11.6 or later to ensure compatibility and optimal performance.

The central component of this repository is the Mamba block, which encompasses the selective state space model (SSM) as its core element.

import torch
from mamba_ssm import Mamba

batch, length, dim = 2, 64, 16
x = torch.randn(batch, length, dim).to("cuda")
model = Mamba(  # This module uses roughly 3 * expand * d_model^2 parameters
    d_model=dim, # Model dimension d_model
    d_state=16,  # SSM state expansion factor
    d_conv=4,    # Local convolution width
    expand=2,    # Block expansion factor
).to("cuda")
y = model(x)
assert y.shape == x.shape        

Mamba Language Model

Source: https://github.com/state-spaces/mamba/blob/main/mamba_ssm/models/mixer_seq_simple.py

Where can Mamba be used ?

There are some obvious use cases in daily life which come on top of my mind where we can use Mamba in place of vanilla transformers. In future (when I get time over the weekends), I would share some open-source example and compare one of the below use cases with both architectures:

1. Automated Customer Support: Utilizing the model to interpret customer queries and provide instant, accurate automated responses, reducing wait times for support.
2. Feedback Aggregation: Analyzing large sets of customer reviews and feedback to identify common issues or suggestions for product improvements.
3. Document Sorting and Organization: Classifying business documents by content, which streamlines finding and retrieving specific information when needed.
4. Email Filtering and Response: Managing high volumes of incoming emails by categorizing them and drafting preliminary responses, saving time for human reviewers.
5. Resume Screening: In recruitment, parsing through numerous job applications to quickly identify the most qualified candidates based on their submitted information.
6. Compliance Monitoring: Scanning communications and documents to ensure adherence to legal and regulatory standards, flagging potential non-compliant language or content.

Hope you found this article useful! We explored Mamba’s efficient selection mechanism, which smartly highlights vital data while discarding the irrelevant, and its unified design that merges the best of previous models. Practical applications, from automated customer support to compliance monitoring, were discussed, highlighting Mamba's versatility in various sectors. This architecture promises not only to revolutionize data analysis but also to provide tangible, real-world benefits across numerous industries. Stay tuned for more in-depth explorations and examples of Mamba in action.

The AI race is on. Do you have the right solution for your business? Multicloud4u Technologies stands at the forefront, offering cutting-edge Generative AI and Traditional ML solutions tailored to your business needs. Our expertise extends to Data Engineering and Cloud, ensuring scalable and efficient implementation. Don't miss the chance to transform your business with our innovative solutions. Book a quick appointment with us to discuss your business challenges and discover custom solutions: Schedule a Meeting or mail me at [email protected]

About the Author:

Bhaskar Tripathi is the Head of Data Science & Research Practices at Multicloud4U? Technologies and is a Ph.D. in Computational & Financial Mathematics. He is a leading open source contributor and creator of several popular open-source libraries on GitHub such as pdfGPT, text2diagram, sanitized gray wolf algorithm, tripathi-sharma low discrepancy sequence, TypeTruth AI Text Detector, HypothesisHub, Improved-CEEMDAN among many others.

1. Github: https://github.com/bhaskatripathi
2. Personal Website: https://www.bhaskartripathi.com
3. Google Scholar: Click Here