MatMamba: A Matryoshka State Space Model

We thank the reviewer for their response. We appreciate that they recognize the close-to-linear downscaling capability of MatMamba when compared to other methods. We address the reviewer’s questions and concerns below:

It is a combination of existing techniques.

We would like to highlight the rapidly growing adoption of the Mamba (and other linear attention / efficient recurrent style models) in both research and production environments. We believe that combining the adaptivity of Matryoshka-style learning and the efficiency of state space models (SSMs) like Mamba is a novel contribution.

Can you give a citation / relevant prior work for the claim in L.456f ?

The scaling laws series of works (e.g. Kaplan et al. 2020 and related follow ups) offer evidence to support the claim. However, we agree that presenting downstream eval results gives a clearer picture of the model performance. We have amended line 456 to no longer contain the original claim, and we also present detailed results for 13 downstream tasks (LAMBADA, Hellaswag, PIQA, ARC etc) using the standard lm-evaluation-harness settings. The results show a similar trend to the validation loss in which performance improves with model size for certain tasks (e.g. PIQA, LAMBADA etc.).

There is a typo in L.522.

Fixed

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?

We think that this is unclear and is up for further exploration. The language models we trained are larger than the image models (and the ImageNet dataset is also balanced/easier to learn similar representations for 1000 classes, whereas the LM vocab size is approx. 50k and needs a larger number of training steps). We observed that for the earlier stages of training, the Mix’n’Match LM curves also lie 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). This needs more rigorous understanding.

As to whether this behavior is specific to MatMamba, we point the reader to figure 2 (a) of the MatFormer paper (https://arxiv.org/abs/2310.07707) which also has a similar trend (albeit slightly less pronounced). This suggests that the behavior is likely not unique to MatMamba.

Why is g set to 4? Have you explored other values? How does the method scale when g is lower/higher?

g=4 is an empirical choice (can be adjusted as needed), although choosing D, D/2, D/4, and D/8 as the four granularities was intentional in order to keep the training cost theoretically within the ceiling of 2X FLOPs compared to just training a single model (given that N + N/2 + N/4 + … = 2N). This also made it easy to compare our results with MatFormer. In some smaller scale experiments, we did observe that choosing extremely low values like D/32 can be counterproductive because it increases the overhead while also being potentially noisy/unstable for representation learning in the most important first few dimensions.

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.

Another great question! While we do not explicitly do super long-context evals in this work (beyond processing images at higher resolutions), there is strong evidence that Mamba2 works very well on long context tasks (when compared to the original Mamba or prior SSM’s that struggled with these). Please see the MQAR results in Figure 8 of the Mamba2 paper (https://arxiv.org/abs/2405.21060), on which the original Mamba (and even Attention) has much worse performance. Moreover, (Mat)Mamba models are significantly faster at longer context lengths for both training and inference. In future work, we plan to apply MatMamba to long-form video.

Once again, thank you for the detailed review with several great observations which we think will help improve the work!

We thank the reviewer for their response. We are glad that they enjoyed reading the paper! We address the reviewer’s questions and concerns below:

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?

This is a fantastic question. We appreciate the reviewer’s attention to detail :)

For a given parameter budget, one could do a number of things including reducing all dimensions proportionally, reducing some dimensions a lot but keeping others the same, or various combinations of these that fit the particular deployment hardware or other constraints.

In the work, we empirically found that for image models, a pyramid/inverted pyramid type of dimension selection (e.g. 512-D in first 12 layers, then 1024-D in next 12 layers; or the other way round) worked better for ImageNet models, whereas proportionally reducing all dimensions worked better for language models (e.g. for 50% of parameter count, we choose all dimensions to be 1024 in a base 2048-D model).

The convenient thing is that we can choose all of these dynamically on-the-fly, so just a few quick validation set passes with a number of different combinations can help tune this for the task at hand. We are studying more systematic (instead of just empirical) methods for this process in future work.

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.

We believe that combining the adaptivity of Matryoshka-style learning and the efficiency of state space models (SSMs) like Mamba is a novel contribution (2 very different types of algorithmic efficiency being demonstrated to work well together). Also, MatMamba has several advantages over MatFormer (e.g. >=95% of parameter count being in the projection layers, which makes tensor parallelism efficient and easier without needing FlashAttention etc, or that we do the Matryoshka process over the whole network instead of just the MLP layers which leads to a larger parameter reduction).

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.

We expect the scaling trends to continue. The strongest intuition we have here is: given that there exists a duality between Mamba2 and Transformers (see detailed discussion in the Mamba2 paper https://arxiv.org/pdf/2405.21060), and that scaling trends continue on Transformers, we would expect them to also continue on Mamba2. We will add more detailed discussion on (expected) scaling trends.

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).

Thank you for the valuable suggestion (and for the specific pointer to lm-evaluation-harness). We have now reported results for 13 downstream tasks (LAMBADA, Hellaswag, PIQA, ARC etc) using the standard lm-evaluation-harness settings. We have also amended line 456 to no longer contain the original claim, which could be confusing to readers.

We thank the reviewer for their response. We are glad that they appreciated the adaptive inference, scalability, and performance of MatMamba, and that they noted the positive effect that the consistent metric space has on tasks like retrieval. We address the reviewer’s concerns and questions below:

(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.

While we acknowledge the slight degradation of validation loss for the untrained granularities in MatMamba-LM, we would like to highlight that in the MatMamba-Vision models, the untrained granularities are not only at par, but often exceed the expected performance on the performance-compute curve. The retrieval results also clearly indicate that the models preserve the metric space, which is an extremely useful property (even in cases where the validation loss slightly suffers in the untrained granularities).

Additionally, Mix’n’Match is very useful in helping pick near pareto-optimal models for free. Often, this is not feasible without additional training. If someone does use a Mix’n’Match model from an untrained granularity that is slightly lesser performant, even that can offer an extremely strong starting point for finetuning with a small number of steps.

(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.

Self-distillation techniques (e.g. DistillBERT) are complementary to Matryoshka-style training (e.g. one could use the output of the largest model to train the smaller ones). However, these necessitate additional training runs and need more compute. We will explore these ideas in-depth in future work.

(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.

Thank you for the valuable suggestion. We have now added results for downstream LM evals in Tables 5,6,7, and 8. The results show a similar trend to the validation loss in which performance improves with model size for certain tasks (e.g. PIQA, LAMBADA etc.). Note however, that downstream results at this scale are fairly noisy, whereas validations loss/scaling laws are more predictive of model performance.

We thank the reviewer for their response. We are glad that the reviewer agrees that MatMamba validates the scalability/effectiveness of Matryoshka learning when applied to SSMs, and that the problem and our methodology are well motivated. We address the reviewers concerns and questions below:

For the MatMamba-LM model family, only evaluation loss is reported. The submission would be strengthened by including evaluations on downstream tasks as well.

We appreciate the reviewer’s suggestion. We have now included results for MatMamba-LM on a number of downstream tasks (in Tables 5,6,7,8 in the Appendix). The results show a similar trend to the validation loss in which performance improves with model size for certain tasks (e.g. PIQA, LAMBADA etc.).

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.

A crucial difference between MatFormer and MatMamba is that MatFormer only does nesting on the MLP portion of the Transformer block, whereas MatMamba does the nesting on all applicable parts of the model. Thus, the parameter count (and theoretical inference speed assuming optimal hardware utilization) reduces linearly in MatMamba with nesting, whereas in MatFormer the attention parts stay constant. Follow-up works like Flextron have extended the idea to attention layers as well, but contain additional design elements (e.g. a surrogate model and a latency-based loss) which we did not study in the context of Mamba-like models in this work.

Another interesting property of the Mamba2 model is that >=95% of the typical parameter count lies in the input and output projection layers alone (whereas in a transformer it is split 40-60 between attention and MLP portions). This helps us use tensor parallelism easily and effectively.

While we did not encounter any other major challenges, we are eager to build upon this work and further optimize the nesting mechanism, which is an open problem.