MatMamba: A Matryoshka State Space Model

## Abstract

Abstract State Space Models (SSMs) like Mamba2 are a promising alternative to Transformers, with faster theoretical training and inference times – especially for long context lengths. Recent work on Matryoshka Representation Learning – and its application to Transformer backbones in works like MatFormer – showed how to introduce nested granularities of smaller submodels in one universal elastic model. In this work, we present MatMamba: a state space model which combines Matryoshka-style learning with Mamba2, by modifying the block to contain nested dimensions to enable joint training and adaptive inference. MatMamba allows for efficient and adaptive deployment across various model sizes. We train a single large MatMamba model and are able to get a number of smaller nested models for free – while maintaining or improving upon the performance of a baseline smaller model trained from scratch. We train language and image models at a variety of parameter sizes from 35M to 1.4B. Our results on ImageNet and FineWeb show that MatMamba models scale comparably to Transformers, while having more efficient inference characteristics. This makes MatMamba a practically viable option for deploying large-scale models in an elastic way based on the available inference compute. Code and models are open sourced at https://github.com/ScaledFoundations/MatMamba

## 1 Introduction

Deep learning practitioners often train different sizes of the same kind of model to facilitate deployment in a variety of ranges of available inference compute. For example, the Llama 3.2 (Dubey et al., ) series has 1B, 3B, 11B, and 90B variations. These models are extremely powerful individually – but due to independent training do not necessarily share the same metric space – a property which can be extremely useful for inference applications like speculative decoding (Leviathan et al., ), hybrid cloud-edge inference, or just general input or compute adaptive processing.
Moreover, because training these models is expensive, we typically see only a few chosen sizes trained. This is not desirable in situations where the deployment setup can optimally support an intermediate model (e.g. a 2B model), but has to settle for the less accurate 1B model instead.

Techniques like model compression and distillation aim to address these issues, but require additional training (for which data may not be available), and can sometimes drop accuracy (Jaiswal et al., ). Thus, methods that offer adaptive inference out of the box at intermediate granularities are extremely useful. This has been explored for Transformers (Devvrit et al., ; Cai et al., ) and ConvNets (Yu & Huang, ; Cai et al., ). The core focus of this work is to try to enable out of the box adaptive inference in a newer architecture: Mamba2 (Dao & Gu, ).

State Space Models like Mamba2 (Dao & Gu, ) and a number of other related newer architectures (see Section [2](#S2)) have shown tremendous potential as they try to improve on the efficiency of Transformers, while maintaining their potency as accurate and general sequence processing architectures. Mamba2 has comparable scaling properties to Transformers, while being significantly faster at longer context lengths.

In this work, we introduce MatMamba, a nested Matryoshka structure (Kusupati et al., ) within a Mamba2 block (Dao & Gu, ). MatMamba enables the extraction of hundreds of nested submodels from the same set of weights, without requiring any additional training during deployment. MatMamba is a general-purpose sequence processing architecture that can be applied to any type of model (encoder/decoder), modality (language/vision/sound/actions), loss function, or learning algorithm compatible with a Transformer or Mamba2 layer.

The philosophically closest work to MatMamba is MatFormer (Devvrit et al., ) – which imposes a nested structure on the FFN block in a Transformer layer. We use the same concept to impose a nested structure on any learnable parameter in a Mamba2 block that depends upon the hidden dimensionality of the block. Formally, a MatMamba block consists of a nested combination of $g$ Mamba2 blocks $M_{i}$, such that $M_{1}\subset M_{2}\subset...\subset M_{g}$, where $M_{i}\subset M_{j}$ means that all the parameters of a sub-block $M_{i}$ are present in $M_{j}$. We train the model using $g$ forward passes with gradient accumulation followed by a single backward pass for parameter updates (see Figure [1](#S1.F1)).

By jointly training all $g$ granularities, the smallest sub-blocks are incentivized to represent the most important information, like in Matryoshka Representation Learning (Kusupati et al., ). We can now use any of the $g$ nested sub-blocks $M_{i}$ flexibly. Additionally, we can flexibly slice the block along any dimensionality (even beyond the $g$ explicitly optimized granularities). Using Mix’n’Match (Section [3.4](#S3.SS4)), we can perform this operation over multiple layers at varying granularities to flexibly extract a combinatorially large number of models from the single larger model. We observe that these extracted models preserve the metric space of the larger model, and are accurate across a variety of tested tasks – effectively allowing us to choose a tradeoff between model performance and compute.

We train MatMamba-based vision models (MatMamba-Vision), and find that: (a) MatMamba-Vision models scale as well as baseline Mamba2 based models at all $g=4$ granularities; (b) Using Mix’n’Match, we can flexibly extract submodels between the explicitly optimized granularities. The submodels span (and sometimes exceed) the pareto optimal accuracy-vs-compute curve; (c) MatMamba-Vision models are significantly faster at higher resolutions than ViTs, making them promising candidates for long-form and high resolution visual tasks, while enabling adaptive visual processing with the nested submodels (see Section [4.1.1](#S4.SS1.SSS1)).

Furthermore, MatMamba-Vision models can act as elastic image encoders for adaptive image retrieval. We can encode visual datasets with the largest model, and because the smaller submodels share its metric space, we can use them as query encoders, needing drastically lower compute with minimal loss in accuracy (see Section [4.1.2](#S4.SS1.SSS2)).

We also train MatMamba-based decoder language models (MatMamba-LM) at various sizes from 130M-1.4B parameters, and at $g=4$ granularities. We make similar observations here too, that MatMamba-LM models scale as well as Mamba2 baselines with the same architecture for all nested granularities. We also observe interesting homogenous scaling behaviour between the nested granularities for different models (see Section [4.2](#S4.SS2)).

Through MatMamba, for the first time, we bring together the adaptivity of Matryoshka-style learning and the efficiency of state space models (SSMs) like Mamba2 (Dao & Gu, ).

We make the following research contributions:

- 1.
We introduce MatMamba, which imposes a nested Matryoshka structure on a Mamba2 state space model. We jointly optimize all nested granularities to train a single elastic model.
- 2.
We show that MatMamba models scale as well as the baseline Mamba2 models for a variety of model sizes from 35M-1.4B parameters on language and vision tasks.
- 3.
Using Mix’n’Match with MatMamba allows the flexible extraction of hundreds of submodels to perform adaptive inference. These submodels preserve the metric space of the original model.
- 4.
MatMamba-Vision models are comparably accurate and significantly faster at higher resolutions than ViTs, making them well suited for long-form/high resolution and adaptive visual processing.

Figure: Figure 1: MatMamba introduces a nested Matryoshka (Kusupati et al., ) structure in a Mamba2 (Dao & Gu, ) block. We jointly train a few chosen granularities to get a single model from which we can flexibly extract a large number of nested submodels for adaptive inference based on the available deployment compute.
Refer to caption: /html/2410.06718/assets/x1.png

## 2 Related Work

The ever growing demand of AI models across various accuracy and resource constraints makes it infeasible to train a different model for each use case. Instead, these adaptive deployment needs are often solved through introducing elasticity in models (Kusupati, ). Work on slimmable networks (Yu et al., ; Yu & Huang, ) and once-for-all networks (Cai et al., ) brought the idea of training multiple submodels present within one universal model. Nested dropout (Rippel et al., ) generalizes this idea to learn ordered representations which further extended to enable elasticity at each dense vector embeddings through Matryoshka Representation Learning (MRL) (Kusupati et al., ). MRL simplifies the training process to induce elasticity with a small set of nested granularities (hence the name Matryoshka), exponentially separated in size, all optimized with the same target loss function as the full vector. MRL further smoothly interpolates to the granularities not seen during training, thus allowing for complete elasticity to extract sub-vectors based on the requirements.

Matryoshka information packing and learning has been widely adopted in bringing adaptivity not only in output space, but also in input (Beyer et al., ) and model weights (Devvrit et al., ; Cai et al., ; Valipour et al., ). MatFormer (Devvrit et al., ) is a direct translation of MRL to every hidden activation vector of a MLP sub-block within a Transformer layer (Vaswani et al., ). MatFormer showed scaling trends similar to Transformer, while also providing the capability to adaptively extract submodels that fall on the accuracy-vs-compute pareto curve. More recent works (Cai et al., ; Jain et al., ) developed dynamic routing on top of the conditional computation enabled by MatFormer to realize performance gains in deployment. Further, matryoshka packing was also used for flexible tokenization (Cai et al., ; Hu et al., ) as well as diffusion models (Gu et al., ).

Transformers (Vaswani et al., ) have been fundamental sequence processing blocks in neural networks for the past few years. There has been a recent wave of work on efficient sequence processing architectures that aim to be faster and equally performant alternatives to Transformers.
Mamba (Gu & Dao, ) and Mamba2 (Dao & Gu, ) are the most relevant to this work, with other very closely related works like Linear Attention (Katharopoulos et al., ), Test-time training (Sun et al., ), RWKV (Peng et al., ), Griffin (De et al., ), Jamba (Lieber et al., ), xLSTM (Beck et al., ), HGRN2 (Qin et al., ), RetNet (Sun et al., ), RecurrentGemma (Botev et al., ).
Waleffe et al. () present a detailed study of how to train large-scale Mamba-based language models. Works like MambaVision (Hatamizadeh & Kautz, ), MambaND (Li et al., ), Vision Mamba (Zhu et al., ),
VideoMamba (Li et al., ), and Sonic (CartesiaAI, ) have all shown how a Mamba layer can process visual data and other modalities. Liu et al. () present a detailed survey of Mamba-based vision models.

## 3 MatMamba

### 3.1 Mamba2 Preliminaries

MatMamba is based on Mamba2. We make simple modifications to the Mamba2 block to impose the Matryoshka structure. A detailed description of the internals of Mamba2 can be found in the original paper Dao & Gu (). However for the purposes of this work, we treat the Mamba2 block as a combination of an input linear projection ($W_{in}$, which can be broken down into $W_{z}$, $W_{x}$, $W_{B}$, $W_{C}$, $W_{dt}$), a causal 1D convolution layer with kernel size 4 (with weights that are a concatenation of $W_{conv_{x}}$, $W_{conv_{B}}$, and $W_{conv_{C}}$ applied in groups), a chunk + selective scan operation (SSM), and an output projection layer ($W_{out}$). Similar to a Transformer, this block takes in an $(b,l,d)$ shaped tensor – $b$ is batch size, $l$ is sequence length, and $d$ is the dimensionality – and produces a $(b,l,d)$ shaped output after a sequence transformation. For an input tensor $u$, the Mamba2 block $M(u)$ consists of the following steps:

$$ $\textit{XBC}(u)=\sigma(\textit{Conv}(W_{conv_{x}}{{}^{\frown}}W_{conv_{B}}{{}^{\frown}}W_{conv_{C}},W_{x}.u{{}^{\frown}}W_{B}.u{{}^{\frown}}W_{C}.u))$ (1) $$

$$ $Y(u)=\textit{SSM}(\textit{XBC}(u),W_{dt}.u,A,D)$ (2) $$

$$ $M(u)=\textit{Norm}(Y(u).{\sigma}(W_{z}.u)).W_{out}^{T}$ (3) $$

where ^⌢ is the concatenation operation, $\textit{Conv}(k,s)$ applies a 1-D causal convolution with weights $k$ (applied in $len(k)$ groups) on a sequence $s$, and $A$ and $D$ are learnable SSM parameters. $\sigma$ is a nonlinearity which we set to SiLU (Elfwing et al., ), and Norm is a layer norm function which we set to RMSNorm (Zhang & Sennrich, ).

### 3.2 MatMamba Block

A MatMamba block also has both input and output shapes as $(b,l,d)$. It is defined as a nested combination of $g$ Mamba2 blocks $M_{i}$, such that $M_{1}\subset M_{2}\subset...\subset M_{g}$, where $M_{i}\subset M_{j}$ means that all the parameters of a sub-block $M_{i}$ are present in $M_{j}$. Works like MatFormer (Devvrit et al., ), OFA (Cai et al., ), and Flextron (Cai et al., ) all share similar designs in which the largest model $M_{g}$ is the single universal base model from which numerous smaller submodels $M_{i}$ can be flexibly extracted. In MatMamba, we impose the nested structure along the dimensions of the model parameters. Specifically for a sub-block $M_{i}$ with expansion factor $e=\frac{d_{inner}}{d_{model}}$, we choose a Matryoshka dimension $m_{i}$, such that $0<m_{i}<d_{model}$, which results in an inner slice dimension $d_{i}=e\times m_{i}$ and number of heads $h_{i}=\frac{d_{i}}{d_{head}}$, subject to $d_{i}\mod d_{head}=0$. For example, parameters like $W_{x}$ have a shape of $(d_{inner},d_{model})$. For the $M_{i}$ sub-block, it will become $W_{x}[0:d_{i}]$ by slicing it along the $d_{inner}$ dimension. Similarly for parameters like A which have a shape of $(n_{heads})$, it will become $A[0:h_{i}]$. Concretely, the MatMamba block $M_{i}(u)$ when applied to an input tensor $u$ is these steps:

$$ $\displaystyle\textit{XBC}_{i}(u)=\sigma(\textit{Conv}(W_{conv_{x}}[0:d_{i}]{{}^{\frown}}W_{conv_{B}}{{}^{\frown}}W_{conv_{C}},W_{x}[0:d_{i}].u{{}^{\frown}}W_{B}.u{{}^{\frown}}W_{C}.u))$ (4) $$

$$ $Y_{i}(u)=\textit{SSM}(\textit{XBC}_{i}(u),W_{dt}[0:h_{i}].u,A[0:h_{i}],D[0:h_{i}])$ (5) $$

$$ $M_{i}(u)=\textit{Norm}(Y_{i}(u).{\sigma}(W_{z}[0:d_{i}].u)).W_{out}[0:d_{i}]^{T}$ (6) $$

In practice, $W_{z}$, $W_{x}$, $W_{B}$, $W_{C}$, and $W_{dt}$ are implemented as a single input projection layer with tensor parallelism, with appropriate rearranging of dimensions depending on $m_{i}$. Figure [1](#S1.F1) illustrates the MatMamba block. We also provide PyTorch-style pseudocode for the block in Appendix [A](#A1), to provide a clearer understanding of our implementation.

Compared to MatFormer (Devvrit et al., ), where the Matryoshka structure is only applied on the MLP subblock of the Transformer block, MatMamba applies nesting to the entire block wherever the inner dimension plays a role. This leads to a nearly linear reduction in total parameter count (and also a nearly linear reduction in flop count due to the nature of Mamba2). Also, typically $>95\%$ of the parameter count in a MatMamba block is in the input and output projections, which can be converted into nested layers while maintaining the well understood systems characteristics of projection layers. See Appendix [A](#A1) for a detailed example of parameter count reduction.

We can stack $L$ such MatMamba blocks to create a MatMamba model. For a given $m_{i}$ and nested blocks $M_{1}\subset M_{2}\subset...\subset M_{g}$, we can create a MatMamba model $f_{i}$ with $L$ layers, and $g$ nested models $f_{1}\subset f_{2}\subset...\subset f_{g}$. Each $f_{i}$ is formed by stacking $M_{i}$ $L$ times. Like Mamba2, the MatMamba backbone is a general purpose sequence processing architecture, which with an appropriate tokenizer and output head can process a variety of modalities.

### 3.3 Training

To train a model comprised of MatMamba blocks for $g$ chosen granularities, we perform $g$ forward passes to calculate a joint loss function. For an input $x$, model $f$, target $y$ and loss function $\mathcal{L}$:

$$ $\mathcal{L}_{joint}(x,y)=\sum_{i=1}^{g}\lambda_{i}.\mathcal{L}(f_{i}(x),y)$ (7) $$

where $\lambda_{i}$ is the weight of the $i$-th nested submodel’s loss. In this work, we train $g=4$ nested submodels with a uniform $\lambda_{i}=1/g=0.25$ for each submodel. As shown in Figure [1](#S1.F1), during each forward pass, we accumulate gradients. The parameter update is done with a single backward pass. During the whole process, the model and the weights are the same, thereby also making memory usage the same as a regular Mamba2 block. In this work, we train MatMamba models with $g=4$ nested granularities, with the corresponding list of $m_{i}$’s being $[d_{model},d_{model}/2,d_{model}/4,d_{model}/8]$, i.e. a halving of dimensionality for every sub-model.
Like MatFormer (Devvrit et al., ) and Flextron (Cai et al., ), we note that it is also possible to finetune an existing pretrained model to produce a nested structure. However, in this work, we focus on training from scratch to study the scaling characteristics of MatMamba models.

### 3.4 Mix’n’Match

We can apply the Mix’n’Match strategy from MatFormer (Devvrit et al., ) to flexibly extract any submodel from MatMamba for inference. Concretely, for a model $f$ with $L$ layers, we need to choose a dimensionality $m_{i}$ at each layer $i$. Note that $m_{i}$ can be either one of the explicitly optimized $g$ granularites (e.g. picking from one of [1024, 512, 256, 128] from a 135M-MatMamba-Vision model, see section [4.1](#S4.SS1)), or we can choose interpolated dimensionalities that were not explicitly optimized for (e.g. picking any random valid value like 768 or 384 that was not explicitly trained). For instance, we could choose $m_{1}=256$ (25% size) in layer 1, $m_{2}=1024$ (100% size) in layer 2, $m_{3}=768$ (75% size) in layer 3, and so on. The only constraint on $m_{i}$ in MatMamba is that it needs to lead to an integer number of heads, or that $(e\times{m_{i}})\mod d_{head}=0$, where $e=\frac{d_{inner}}{d_{model}}$. This leads to a combinatorially large number of possible submodels (beyond the $g$ explicitly optimized granularities) that can be flexibly extracted – all from the same set of base model weights – as shown in Figure [1](#S1.F1). Due to the Matryoshka structure, the first few dimensions (that are shared among all the nested submodels) are incentivized to learn the strongest representations.

### 3.5 Elastic Inference

When deploying a MatMamba model for inference, we typically need to store the single universal model $f_{g}$ in memory. If compute is not constrained (or if the inference workload is predictable), then we can use the full model to get the most accurate results. However, depending on dynamic constraints (e.g. available inference compute, energy usage, system load, desired accuracy etc.), we can perform a forward pass on a chosen slice of the network on the fly.

There are exciting possibilities like combining cloud and edge inference – we could store a smaller model $f_{i}$ on the edge device and when necessary, use the larger model $f_{j}$ on the cloud, or using a smaller model to act as a draft model for speculative decoding (Leviathan et al., ) with a larger verifier model. We could also potentially do input-adaptive sub-model selection (e.g. use a larger model for a more difficult input). All of these are possible only because MatMamba has a consistent and nested Matryoshka structure, in which all the sub-models share the same metric space.

## 4 Experiments

In this section, we demonstrate the effectiveness of MatMamba-based models across two modalities: vision (MatMamba-Vision) and language (MatMamba-LM). For vision, we show results for image classification (Section [4.1.1](#S4.SS1.SSS1)) and adaptive image retrieval (Section [4.1.2](#S4.SS1.SSS2)). For language, we train decoder language models (Section [4.2](#S4.SS2)). We train models at a variety of scales from 35M to 1.4B parameters.
For a fair comparison, we also independently train baseline Mamba2 models which have the same architecture as the submodels of each MatMamba granularity. Please note that we do not aim to achieve state-of-the-art results in this work on either language or vision for the chosen model sizes. We instead focus on properties like nested structure consistency, parameter reduction, inference speedups/memory usage for submodels, and scaling of simple networks built using the MatMamba block.

### 4.1 MatMamba-Vision

Figure: Figure 2: MatMamba layers for vision tasks. Similar to a ViT (Dosovitskiy, ), we convert an image into a tensor of embedded patches. Because of the causal nature of the Mamba2 block, we suffix the [CLS] token. We intentionally keep the design simple to better study the properties of the MatMamba block.
Refer to caption: /html/2410.06718/assets/x2.png

**Table 1: Base model architectures for MatMamba-Vision (35M and 135M) with the explicitly optimized submodels for $g=4$ nested granularities.**
|  |  |  |  | Parameters |  |
| --- | --- | --- | --- | --- | --- |
| Base Model | Layers | $m_{i}$ | $h_{i}$ | Patch embed | MatMamba Layers |
| 135M-1024D | 20 | 1024 | 32 | 787,456 | 132,739,840 |
| 20 | 512 | 16 | 787,456 | 69,004,160 |  |
| 20 | 256 | 8 | 787,456 | 37,136,320 |  |
| 20 | 128 | 4 | 787,456 | 21,202,400 |  |
| 35M-512D | 20 | 512 | 32 | 393,728 | 34,927,360 |
| 20 | 256 | 16 | 393,728 | 18,787,200 |  |
| 20 | 128 | 8 | 393,728 | 10,717,120 |  |
| 20 | 64 | 4 | 393,728 | 6,682,080 |  |

MatMamba-Vision (Figure [2](#S4.F2)) contains a patch embedding followed by $L$ MatMamba blocks with a unidirectional SSM scan. One crucial design choice we make is to use the [CLS] token as a suffix instead of the conventional prefix. This allows it to attend to information from the entire sequence. We find that this simple architecture works effectively on both image classification and adaptive retrieval. We train two model variations (35M with $d_{model}=512$ and 135M with $d_{model}=1024$, see Table [1](#S4.T1)) with patch size 16 and $L=20$ layers on ImageNet-1k Deng et al. () which has 1.28M training images and 50k validation images.
Compared to other recent work on SSM’s for vision tasks like MambaVision (Hatamizadeh & Kautz, ), MambaND (Li et al., ), and Vision Mamba (Zhu et al., ) – all of which have major design changes on top of Mamba layers like bidirectional scan with additional projections, varying order of scans, or combining SSM layers with attention and convolution layers – we keep our network architecture as simple as possible.
We use FFCV (Leclerc et al., ) dataloaders for efficient training. We apply augmentations like RandAug (Cubuk et al., ), Random Erasing (Zhong et al., ), Mixup (Zhang, ), Cutmix (Yun et al., ), and a number of other settings following DEiT-3 (Touvron et al., ), AugReg (Steiner et al., ), and Better ViT Baselines (Beyer et al., ).
The exact detailed experimental settings can be seen in Appendix [A](#A1).

Figure: (a) 35M model
Refer to caption: /html/2410.06718/assets/x3.png

#### 4.1.1 Image Classification

In Figure [3](#S4.F3), we see that for both the 35M and 135M MatMamba-Vision models, the explicitly optimized submodels closely match the 4 independently trained baseline models with the same architecture as the nested submodel. However, instead of needing four separate models, we can get all levels of performance/parameter counts flexibly in a single model.

Figure: (a) MatMamba-Vision FPS vs. Image Size
Refer to caption: /html/2410.06718/assets/x5.png

Adaptive Inference using Mix’n’Match: Additionally (Figure [3](#S4.F3)), using Mix’n’Match at a variety of combined granularities yields models that smoothly interpolate (and sometimes exceed) the accuracy on the line joining the explicitly optimized granularities. This points towards powerful adaptivity, because we can extract a combinatorially large number of submodels along the accuracy-compute curve. We can optimize submodel selection for deployment constraints flexibly, all while only using the weights of a single nested universal model.

Inference Speeds at Higher Resolutions: In Figure [4](#S4.F4), we also study the inference speed tradeoffs of nested granularities of MatMamba-Vision models when compared with each other and a ViT-B/16 model. We find that at or below 512px resolution, the sequence length is low enough for the ViT to be the fastest model (due to GPU parallelism and optimizations like FlashAttention). However, as we increase the resolution to 1024px and beyond, Mamba-style models start outperforming ViT in both througput and latency. We also study inference memory usage, and find that MatMamba-Vision scales slightly better than an optimized ViT-B/16 as the resolution increases. Both of these observations offer promising evidence that MatMamba-based models can be suitable for processing longer sequences of visual data at higher resolutions on a single accelerator (as opposed to scaling context length in Transformers using methods like RingAttention Liu et al. () which needs multiple interconnected accelerators for a single forward pass at long sequence lengths).

Figure: (a) 1-NN Retrieval for 35M model
Refer to caption: /html/2410.06718/assets/x7.png

#### 4.1.2 Adaptive Image Retrieval

Image retrieval aims to locate semantically similar images using representations generated by a pretrained encoder (Chen et al., ). The standard method involves encoding both database and query images with the same encoder and then performing nearest neighbor retrieval. While using a powerful encoder for database images is feasible, the query encoder must be efficient for real-time applications. Moreover, query encoding scenarios can vary, such as on-device versus cloud processing and varying query load and complexity. Existing solutions with fixed encoders often compromise accuracy or cost in different settings.

Due to its flexibility, MatMamba-Vision is a promising candidate for query encoding. However, retrieval also requires that submodels maintain distance relationships between fixed database (encoded with a larger encoder) and query embeddings across various granularities. Using smaller baseline Mamba2 models solely for query encoding can lead to significant distance preservation issues and poor retrieval accuracy (as illustrated in Figure [5](#S4.F5)).

We evaluated both the baseline and MatMamba-Vision encoders on ImageNet-1K for image retrieval at 35M and 135M parameter scales. Using the [CLS] token representation, we calculated 1-nearest neighbor (NN) accuracy. Figure [5](#S4.F5) demonstrates that submodels extracted from MatMamba can effectively preserve distances and offer greater flexibility. For example, MatMamba-Vision-135M can reduce compute cost by 55% with a minimal accuracy loss of less than 0.5%. While causal models with suffix [CLS] token might not be as accurate as bi-directional encoders for retrieval, this is a promising start towards better long-context encoders while enabling adaptive query processing.

### 4.2 MatMamba-LM

**Table 2: Base model architectures for MatMamba-LM**
| Base Model | Layers | $d_{model}$ | $d_{head}$ | Embed params | Non-embed params | Tokens |
| --- | --- | --- | --- | --- | --- | --- |
| 130M | 24 | 768 | 24 | 38,615,040 | 90,368,448 | 62.9B |
| 370M | 48 | 1024 | 32 | 51,486,720 | 316,851,712 | 125.8B |
| 790M | 48 | 1536 | 48 | 77,230,080 | 702,918,912 | 125.8B |
| 1.4B | 48 | 2048 | 64 | 102,973,440 | 1,240,767,488 | 251.6B |

We train decoder language models using the MatMamba block (MatMamba-LM). The models closely follow the training procedure and hyperparameters of llm.c (Karpathy, ). We use the GPT-2 (Radford et al., ) tokenizer with a padded vocabulary size of 50,280. We use the FineWeb (Penedo et al., ) dataset to train all models. We train 4 separate models (with base model parameter sizes 130M, 370M, 790M, and 1.4B). For each of these base models, we optimize $g=4$ nested granularities $[d_{model},d_{model}/2,d_{model}/4,d_{model}/8]$. For baselines, we train vanilla Mamba2 models with the same architecture as the nested submodels. Table [2](#S4.T2) shows the exact configurations for each model.

Figure: (a) 1.4B model
Refer to caption: /html/2410.06718/assets/x9.png

MatMamba-LM scales as well as Mamba2:
In Figure [6](#S4.F6), we see that MatMamba-LM models scale with training tokens as well as Mamba2 models for the largest granularity. In Figure [7](#S4.F7), we also see that for all granularities, the final trained models of every granularity scale as well as the baseline model trained with the same architecture. Furthermore, we observe that the validation loss in Figure [6](#S4.F6) for every nested granularity is at a similar distance (usually a delta of 0.4 in val loss) between the largest model ($m_{i}=d_{model}$) and the smallest model ($m_{i}=d_{model}/8$), with consistent gaps for the intermediate models. Given that validation loss on a large and diverse dataset is the strongest proxy for language model performance (as opposed to few-shot evals on noisy datasets), these scaling trends offer very promising evidence that a single nested MatMamba-LM model can be used in a variety of deployments instead of training 4 separate models independently.

In Figure [7](#S4.F7), we show results for adaptive inference using Mix’n’Match on all 4 MatMamba-LM variants. We see a smooth interpolation between the $d_{model}/2$ and $d_{model}$ granularities (e.g. between $d_{model}/8$ and $d_{model}/4$). However, for the lower granularities, even though the explicitly optimized granularities scale as well as expected, the Mix’n’Match models that have not been explicitly trained suffer a slight performance degradation. We observed that during earlier stages of training, the Mix’n’Match trends for all granularities were exactly on the performance-compute curve. However, towards the later stages, the explicitly optimized granularities improve faster than the Mix’n’Match granularities (almost like anchor points). There are mechanisms that can potentially fix this: like a self-distillation loss with the output of the largest submodel, training with more than $g=4$ granularities, or the surrogate model structure used in Flextron (Cai et al., ), that should make the Mix’n’Match trend smooth. However, this requires more rigorous understanding, and we leave deeper exploration to future work.

Figure: (a) 1.4B model

## 5 Conclusions

In this work, we presented MatMamba, which is a way to impose a nested Matryoshka strcuture on a Mamba2 state space model. It brings together the best of both Mamba-style models (faster inference times, especially for longer sequences) and Matryoshka-style learning. A single MatMamba model contains hundreds of nested and accurate submodels that can be flexibly extracted for inference. MatMamba-Vision and MatMamba-LM models match the performance and accuracy of the independently trained Mamba2 baselines. MatMamba models allow us to choose a desired performance-compute tradeoff, all while being a single Matryoshka-style model instead of multiple different models for specific scenarios. This enables interesting use cases like speculative decoding using a smaller draft model and a larger verifier model, input-adaptive submodel selection, and hybrid cloud-edge inference with the same model based on available compute.

## Appendix A Appendix

Figure: Figure 8: Pytorch-style pseudocode for a MatMamba block

**Table 3: Training Configuration for ImageNet runs**
|  | MatMamba-Vision |  |
| --- | --- | --- |
| Procedure | 135M | 35M |
| Model Dim. | 1024 | 512 |
| Layers | 20 | 20 |
| Batch Size | 4096 | 8192 |
| Training Steps | 249,600 | 124,800 |
| Optimizer | AdamW | AdamW |
| LR | 0.005 | 0.005 |
| LR Decay | Cosine | Cosine |
| Weight decay | 0.1 | 0.1 |
| Warmup steps | 10,000 | 10,000 |
| Label smoothing eps. | 0.1 | 0.1 |
| Dropout | 0.1 | 0.1 |
| Stochastic depth | 0.1 | 0.1 |
| Repeated Aug | Yes | Yes |
| Gradient clip | 1.0 | 1.0 |
| Horizontal flip | Yes | Yes |
| Random Resized Crop | Yes | Yes |
| RandAugment | (2,9) | (2,9) |
| MixUp Alpha | 0.8 | 0.8 |
| CutMix Alpha | 1.0 | 1.0 |
| RandomErase prob. | 0.3 | 0.3 |
| ColorJitter | 0.3 | 0.3 |
| Test crop ratio | 1.0 | 1.0 |

**Table 4: Learnable parameters (without biases) in a MatMamba layer, with example parameter reduction from a Mamba2 layer for $d_{model}=1024$, $d_{head}=32$, $d_{inner}=2\times 1024=2048$ (expand factor 2), $d_{state}=128$, $m_{i}=512$, and $h_{i}=16$ (half of the original dimensions and half of original heads being used inside the model).**
| Parameter | Mamba Shape | MatMamba Shape | Reduction Fraction |
| --- | --- | --- | --- |
| $W_{z}$ | $d_{inner}\times d_{model}$ | $(2\times m_{i})\times d_{model}$ | 0.5x |
| $2048\times 1024$ | $(2\times 512)\times 1024$ |  |  |
| 2,097,152 | 1,048,576 |  |  |
| $W_{x}$ | $d_{inner}\times d_{model}$ | $(2\times m_{i})\times d_{model}$ | 0.5x |
| $2048\times 1024$ | $(2\times 512)\times 1024$ |  |  |
| 2,097,152 | 1,048,576 |  |  |
| $W_{B}$ | $d_{state}\times d_{model}$ | $d_{state}\times d_{model}$ | 1x |
| $128\times 1024$ | $128\times 1024$ |  |  |
| 131,072 | 131,072 |  |  |
| $W_{C}$ | $d_{state}\times d_{model}$ | $d_{state}\times d_{model}$ | 1x |
| $128\times 1024$ | $128\times 1024$ |  |  |
| 131,072 | 131,072 |  |  |
| $W_{dt}$ | $n_{heads}\times d_{model}$ | $h_{i}\times d_{model}$ | 0.5x |
| $32\times 1024$ | $16\times 512$ |  |  |
| 32,768 | 16,384 |  |  |
| $D$ | $n_{heads}$ | $h_{i}$ | 0.5x |
| $32$ | $16$ |  |  |
| $A$ | $n_{heads}$ | $h_{i}$ | 0.5x |
| $32$ | $16$ |  |  |
| $W_{conv_{x}}$ | $d_{inner}\times 1\times 4$ | $2\times m_{i}\times 1\times 4$ | 0.5x |
| $2048\times 1\times 4$ | $(2\times 512)\times 1\times 4$ |  |  |
| 8,192 | 4,096 |  |  |
| $W_{conv_{BC}}$ | $(2\times d_{state})\times 1\times 4$ | $(2\times d_{state})\times 1\times 4$ | 1x |
| $(2\times 128)\times 1\times 4$ | $(2\times 128)\times 1\times 4$ |  |  |
| 256 | 256 |  |  |
| $W_{out}$ | $d_{model}\times d_{inner}$ | $d_{model}\times(2\times m_{i})$ | 0.5x |
| $1024\times 2048$ | $1024\times(2\times 512$) |  |  |
| 2,097,152 | 1,048,576 |  |  |
| Total | 6,594,880 | 3,428,640 | 0.519x |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: - What is the main difference you observed between elastic inference on Matformer and MatMamba during the experiment? Since MatMamba essentially follows the approach used in Matformer, I'm curious if any unique challenges were encountered specific to the Mamba architecture.
soundness: 4
strengths: - This work validates the scalability of Matryoshka Representation Learning by demonstrating its applicability to state-space models (SSMs). It considers both Mamba-vision and Mamba-language models, confirming effectiveness across both tasks. 
- The methodology and problem setting are well-established and well-motivated.
summary: This work extends Matryoshka Representation Learning to Mamba2, a representative architecture for state-space models, proposing MatMamba. MatMamba introduces a training pipeline that generates numerous sub-models, each offering different performance-efficiency trade-offs, within a single training run. The authors analyze the effectiveness and scalability of the proposed method.
weaknesses: - For the MatMamba-LM model family, only evaluation loss is reported. The submission would be strengthened by including evaluations on downstream tasks as well.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: See the weaknesses
soundness: 3
strengths: (1) Adaptive Inference Resource: MatMamba allows the extraction of multiple nested submodels from one single model, enabling versatile deployment options, from edge devices to cloud settings, without the need for retraining. By using the Mix’n’Match approach, MatMamba optimizes for various performance-compute trade-offs, allowing resource allocation adjustments based on current needs.
(2) Scalability and Performance: The model exhibits scalability on par with transformers and baseline Mamba2 models in both image and language tasks. The author provides multiple model sizes ranging from 30M parameters to 1.4B parameters, which demonstrates scalability.
(3) The nested structure maintains a consistent metric space, ensuring that different granularities yield compatible outputs, which is beneficial for downstream applications like retrieval.
summary: This paper introduces MatMamba, a novel approach to State Space Models (SSMs) that combines Mamba model with Matryoshka-style learning. The proposed model embeds nested sub-models within a single super-architecture, allowing for dynamic inference across different granularities without the need for separate training. This setup facilitates scalable and efficient deployment in both vision and language models by dynamically adjusting model sizes based on compute availability. Through extensive testing on tasks such as image classification and language modeling, MatMamba demonstrates comparable or superior performance to baseline Mamba2 models while maintaining flexibility and efficiency.
weaknesses: (1) Dependency on Explicitly Trained Granularities for Optimal Performance: While MatMamba’s Mix’n’Match approach offers flexibility, untrained granularities do not perform as effectively as explicitly optimized ones, showing degradation in performance and accuracy.
(2) Limited Exploration of Self-Distillation Techniques: The paper mentions that self-distillation or other techniques could further enhance interpolation accuracy in untrained granularities, suggesting room for improvement in the training setup.
(3) Limited Evaluation on Language Tasks: The author trained a large model up to 1.4B parameters, but did not conduct more experiments on text-generation benchmarks used to evaluate LLMs. I reckon that it would be better if the author could run the proposed MatMamba through datasets like MMLU, HellaSwag, GSM8K, ARC and etc.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: * Why is g set to 4? Have you explored other values? How does the method scale when g is lower/higher?
* How is MatMamba's long context performance? SSM-based architectures can perform poorly on long-context tasks, and I’d be curious to see if MatMamba maintains this trend, or perhaps improves it. MatMamba’s superior performance at higher resolutions in the vision domain provide somewhat of an indication here.
soundness: 2
strengths: * The paper is very well-written, with the main ideas and results presented clearly. It was a pleasure to read.
* Elastic networks have shown significant promise for Transformer networks, since they can be used to generate tailored networks for specific deployment scenarios at no additional training cost (once elastic training is complete). The paper makes a contribution in this relevant and promising direction.
* Evaluation and results across vision and language looks strong overall.
summary: This paper presents MatMamba, which introduces a nested Matryoshka structure on a Mamba2 state space model. The resulting elastic network architecture can be used to generate multiple smaller sub-networks at deployment time without any additional training, similar to prior work in this area such as Matformer and Flextron. The training process involves gradient accumulation on 4 subnetworks with a single backward pass for parameter updates. MatMamba is evaluated on vision and language tasks with elastic network sizes ranging from 35M to 1.4B.
weaknesses: * There doesn’t seem to be a systematic way of selecting smaller sub-networks given a deployment constraint. For a given parameter or latency budget, should all dimensions be proportionally reduced? How is this proportion decided?
* Novelty: this work appears to be a straightforward application of Matryoshka principles to the SSM domain; apart from the obvious changes needed for SSMs, the training algorithm remains nearly identical to the Matformer work.
* It’s not clear if the scaling trends continue beyond the ~1.5B parameter mark. Do the authors have any insights here? I understand that training larger networks can be quite resource-intensive, but at least a discussion on expected scaling trends would be useful.
* I don’t agree with the claim made on line 456. Downstream task performance must be reported for MatMamba-LM using a standard framework like LM-Eval-Harness. While validation loss is a good starting point, performance on a variety of downstream tasks can provide a more complete picture of LLM performance (which is why most work dealing with LLMs reports these numbers).

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Can you give a citation / relevant prior work for the claim in L.456f ?
There is a typo in L.522.
Do you think that the different Mix'n'Match behavior for language and image modeling is intrinsic to the modality or might be specific to MatMamba?
soundness: 3
strengths: The authors show that Matryoshka-style model tying/training can be applied to Linear Transformers/Matrix SSMs/Fast-Weight Programmers.
They show results on how the Mix'n'Match interpolation degrades performance across modalities.
summary: The paper presents MatMamba the application of Matryoshka techniques (Kusupati 2022) to models based on the Mamba2 architecture.
The authors apply the Matroyshka technique to most layers, leading to a close-to linear downscaling capability. 
They can show that jointly learning one model at four granularities matches the performance of models trained at the single scales.
For vision experiments, the Mix'n'Match technique using layer dimensionalities in between granularities matches/surpasses an interpolated 
performance, whereas in Language Modeling there is a degradation using this technique. The presented technique binds representations at 
different model scales, making potential for e.g. more efficient speculative decoding techniques.
weaknesses: It is a combination of existing techniques.

Meta-review: The submission presents MatMamba, a novel application of Matryoshka-style representation learning to state-space models (SSMs), enabling flexible nested submodels within a single architecture. While the paper provides a thorough exposition of the method and demonstrates its scalability and efficiency across language and vision tasks, it struggles to establish sufficient novelty beyond prior work, such as Matformer and Flextron. The approach appears to be a relatively straightforward extension of existing techniques, and the core contributions are incremental. The dependence on explicit training for specific granularities and limited exploration of broader applications or benchmarks further weakens its impact. Additionally, while some improvements over baseline models are observed, the paper fails to provide insights into challenges unique to MatMamba or its broader implications for SSM-based architectures.

Despite the strong presentation and interesting results, the lack of a systematic approach for submodel selection, as well as underexplored areas such as self-distillation and long-context performance, limits the paper’s contribution to the field. Furthermore, the paper would benefit from additional evaluations on more diverse and challenging language modeling benchmarks to substantiate its claims.

Decision: Reject