MatMamba: A Matryoshka State Space Model

• URL: http://arxiv.org/abs/2410.06718v1
• Date: Wed, 9 Oct 2024 09:41:34 GMT
• Title: MatMamba: A Matryoshka State Space Model
• Authors: Abhinav Shukla, Sai Vemprala, Aditya Kusupati, Ashish Kapoor,
• Abstract summary: MatMamba is a state space model which combines Matryoshka-style learning with Mamba2. MatMamba allows for efficient and adaptive deployment across various model sizes. We train language and image models at a variety of parameter sizes from 35M to 1.4B.
• Score: 24.85566171753877
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 \url{https://github.com/ScaledFoundations/MatMamba}

Related papers

• Starbucks: Improved Training for 2D Matryoshka Embeddings [32.44832240958393]
  We propose Starbucks, a new training strategy for Matryoshka-like embedding models. For the fine-tuning phase, we provide a fixed list of layer-dimension pairs, from small size to large sizes. We also introduce a new pre-training strategy, which applies masked autoencoder language modelling to sub-layers and sub-dimensions.
  arXiv  Detail & Related papers  (2024-10-17T05:33:50Z)
• Scalable Autoregressive Image Generation with Mamba [23.027439743155192]
  We introduce AiM, an autoregressive (AR) image generative model based on Mamba architecture. Mamba is a novel state-space model characterized by its exceptional performance for long-sequence modeling with linear time. We provide AiM models in various scales, with parameter counts ranging from 148M to 1.3B.
  arXiv  Detail & Related papers  (2024-08-22T09:27:49Z)
• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models [92.36510016591782]
  We present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs) Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture using only 3B tokens and a hybrid version (Hybrid Phi-Mamba) using 5B tokens. Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models.
  arXiv  Detail & Related papers  (2024-08-19T17:48:11Z)
• LaMamba-Diff: Linear-Time High-Fidelity Diffusion Models Based on Local Attention and Mamba [54.85262314960038]
  Local Attentional Mamba blocks capture both global contexts and local details with linear complexity. Our model exhibits exceptional scalability and surpasses the performance of DiT across various model scales on ImageNet at 256x256 resolution. Compared to state-of-the-art diffusion models on ImageNet 256x256 and 512x512, our largest model presents notable advantages, such as a reduction of up to 62% GFLOPs.
  arXiv  Detail & Related papers  (2024-08-05T16:39:39Z)
• CLIP-Mamba: CLIP Pretrained Mamba Models with OOD and Hessian Evaluation [18.383760896304604]
  This report introduces the first attempt to train a Mamba model utilizing contrastive technical-image pretraining (CLIP) A Mamba model 67 million parameters is on par with a 307 million- parameters Vision Transformer (ViT) model in zero-shot classification tasks.
  arXiv  Detail & Related papers  (2024-04-30T09:40:07Z)
• The Hidden Attention of Mamba Models [54.50526986788175]
  The Mamba layer offers an efficient selective state space model (SSM) that is highly effective in modeling multiple domains. We show that such models can be viewed as attention-driven models. This new perspective enables us to empirically and theoretically compare the underlying mechanisms to that of the self-attention layers in transformers.
  arXiv  Detail & Related papers  (2024-03-03T18:58:21Z)
• BlackMamba: Mixture of Experts for State-Space Models [10.209192169793772]
  State-space models (SSMs) have recently demonstrated competitive performance to transformers at large-scale language modeling benchmarks. MoE models have shown remarkable performance while significantly reducing the compute and latency costs of inference. We present BlackMamba, a novel architecture that combines the Mamba SSM with MoE to obtain the benefits of both.
  arXiv  Detail & Related papers  (2024-02-01T07:15:58Z)
• MatFormer: Nested Transformer for Elastic Inference [94.1789252941718]
  MatFormer is a nested Transformer architecture designed to offer elasticity in a variety of deployment constraints. We show that a 2.6B decoder-only MatFormer language model (MatLM) allows us to extract smaller models spanning from 1.5B to 2.6B. We also observe that smaller encoders extracted from a universal MatFormer-based ViT (MatViT) encoder preserve the metric-space structure for adaptive large-scale retrieval.
  arXiv  Detail & Related papers  (2023-10-11T17:57:14Z)
• UnIVAL: Unified Model for Image, Video, Audio and Language Tasks [105.77733287326308]
  UnIVAL model goes beyond two modalities and unifies text, images, video, and audio into a single model. Our model is efficiently pretrained on many tasks, based on task balancing and multimodal curriculum learning. Thanks to the unified model, we propose a novel study on multimodal model merging via weight generalization.
  arXiv  Detail & Related papers  (2023-07-30T09:48:36Z)
• eP-ALM: Efficient Perceptual Augmentation of Language Models [70.47962271121389]
  We propose to direct effort to efficient adaptations of existing models, and propose to augment Language Models with perception. Existing approaches for adapting pretrained models for vision-language tasks still rely on several key components that hinder their efficiency. We show that by freezing more than 99% of total parameters, training only one linear projection layer, and prepending only one trainable token, our approach (dubbed eP-ALM) significantly outperforms other baselines on VQA and Captioning.
  arXiv  Detail & Related papers  (2023-03-20T19:20:34Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.