Skip to content

mamba

PyTorch Code Documentation - Mamba

Overview

The Mamba model is designed for performing joint image and text processing. This documentation explains the purpose, functionality, usage, and core features of the Mamba class.

Purpose and Functionality

The Mamba model is designed to handle sequential processing tasks by combining information from text and images. The model employs a series of Mamba blocks to process the input data. The core functionality involves a forward propagation that processes the input and returns logits for text prediction. Key features of the Mamba model include the use of attention, layer normalization, and linear projection operations.

Class Definition

The Mamba class is defined with the following class signature and arguments: 

| Argument    | Type                      | Definition                                     | Default |
|-------------|---------------------------|------------------------------------------------|---------|
| vocab_size  | int                       | Size of the vocabulary                         | None    |
| dim         | int                       | Input dimension (for embedding)               | None    |
| depth       | int                       | Depth of the Mamba block                       | 5       |
| d_state     | int                       | State dimension                                 | 16      |
| expand      | int                       | Expansion factor                                | 2       |
| dt_rank     | Union[int, str]           | Rank of the temporal difference tensor         | "auto"  |
| d_conv      | int                       | Dimension of the convex kernel                 | 4       |

Functionality and Usage

The core functionality of the Mamba class is the forward pass, which processes the input and produces logits. The forward pass includes processing the input text and images, applying the Mamba blocks, and a final linear projection. The model is flexible to handle both image and text inputs. The Mamba model can be initialized with default parameters or with custom values during instantiation.

Examples

Example 1:

import torch

from zeta.nn import Mamba

x = torch.randint(0, 16, (1, 64))
model = Mamba(16, 64, 5, 16)
output = model(x)
print(output)

Example 2:

import torch

from zeta.nn import Mamba

x = torch.randint(0, 16, (1, 32))
img_features = torch.rand(1, 64)
model = Mamba(16, 32, 3, 16)
output = model(x, img_features)
print(output)

Example 3:

import torch

from zeta.nn import Mamba

x = torch.randint(0, 32, (1, 32))
model = Mamba(32, 32, 3, 16, 3, d_conv=8)
output = model(x)
print(output)

Additional Information

The Mamba model implementation adopts a mixed-type learning approach. It can handle both text and image inputs for generating context-aware predictions. Developers and data scientists may benefit from exploring the official GitHub repository for extended understanding and usage of this model.

References and Resources

This documentation provides detailed insights into the purpose, functionality, and usage of the Mamba class in PyTorch. By understanding core features, class definition, and usage scenarios, developers can effectively utilize the Mamba model for their specific applications.