Jamba: Hybrid Transformer-Mamba Language Models

Authors: Barak Lenz, Opher Lieber, Alan Arazi, Amir Bergman, Avshalom Manevich, Barak Peleg, Ben Aviram, Chen Almagor, Clara Fridman, Dan Padnos, Daniel Gissin, Daniel Jannai, Dor Muhlgay, Dor Zimberg, Edden M. Gerber, Elad Dolev, Eran Krakovsky, Erez Safahi, Erez Schwartz, Gal Cohen, Gal Shachaf, Haim Rozenblum, Hofit Bata, Ido Blass, Inbal Magar, Itay Dalmedigos, Jhonathan Osin, Julie Fadlon, Maria Rozman, Matan Danos, Michael Gokhman, Mor Zusman, Naama Gidron, Nir Ratner, Noam Gat, Noam Rozen, Oded Fried, Ohad Leshno, Omer Antverg, Omri Abend, Or Dagan, Orit Cohavi, Raz Alon, Ro'i Belson, Roi Cohen, Rom Gilad, Roman Glozman, Shahar Lev, Shai Shalev-Shwartz, Shaked Haim Meirom, Tal Delbari, Tal Ness, Tomer Asida, Tom Ben Gal, Tom Braude, Uriya Pumerantz, Josh Cohen, Yonatan Belinkov, Yuval Globerson, Yuval Peleg Levy, Yoav Shoham
Primary area: foundation or frontier models, including LLMs
Keywords: language models, state-space models, mamba, hybrid architecture, foundation models

Abstract:
We present Jamba, a novel hybrid Transformer-Mamba mixture-of-experts (MoE) architecture. Jamba interleaves blocks of Transformer and Mamba layers, enjoying the benefits of both model families. MoE is added in some of these layers to increase model capacity while keeping active parameter usage manageable. This flexible architecture allows resource- and objective-specific configurations. We implement two configurations: Jamba-1.5-Large, with 94B active parameters, and Jamba-1.5-mini, with 12B active parameters. Built at large scale, Jamba models provide high throughput and small memory footprint compared to vanilla Transformers, especially at long-context tasks, with an effective context length of 256K tokens, the largest amongst open-weight models. At the same time, they are also competitive on standard language modeling and chatbot benchmarks. We study various architectural decisions, such as how to combine Transformer and Mamba layers, and how to mix experts, and show that some of them are crucial in large scale modeling. To support cost-effective inference, we introduce ExpertsInt8, a novel quantization technique that allows fitting Jamba-1.5-Large on a machine with 8 80GB GPUs when processing 256K-token contexts without loss of quality. We also describe several interesting properties of this architecture that the training and evaluation of Jamba have revealed. The model weights are publicly available.

Decision: accept (poster)
Reviewer scores: 8, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730160266500

thread_type:
official_review

<thread_object>
creator: reviewer_pw77

<content>
summary:
The paper introduces a model family of Transformer-Mamba models, augmented with Mixture of Expert (MoE) layers. Due to the Mamba and MoE components, the model has throughput and memory benefits on long context tasks. The authors describe a quantization method for the MoE layers which further enhances the inference performance. Finally, the authors describe some evaluations for long context, chat, and natural language benchmarks.

strengths:
- The paper proposes a large-scale Hybrid of Mamba, Transformers, and MoE layers. The large-scale aspect of it with respect to model size is certainly novel.
- Although quantizing the Expert layers on MoE models has been explored before [1] [2] the quantization on Int8 is novel and interesting.

[1] https://neurips2023-enlsp.github.io/papers/paper_81.pdf 

[2] https://arxiv.org/abs/2406.08155

weaknesses:
Unfortunately there are many key aspects missing from all the sections of the paper which makes it hard to quantify if the proposed approach is better than the state of the art in an apples to apples setup, or for future reproductions and/or comparisons. The most important weakness is the lack of disclosing the pre-training size (or FLOPs budget in general), which is known to be indicative of scaling trends and allows for comparing models in an standardised way [5]. While I believe the proposed protocol is novel and interesting enough, I think the following weaknesses should be addressed so that this paper is useful for the ICLR community.

**Efficient Inference details**

- Section 4.1 and 5 focus on the latency aspects of the proposed quantization method, but there is no empirical evidence to quantify the loss in model performance (if any).
- Figure 2 from section 4.1 seems focused on latency and missing throughput, which is an important consideration to compare to the other approaches.
- Relevant very similar work in MoE quantization on the experts [2] is missing from section 4.1, there should be at least a discussion, at best a latency comparison.
- Section 5 lacks important details of what software stack was used to run the latency and throughput comparisons. It is not possible to assess to what degree the comparison is fair, or if it paints a good or bad picture for the proposed method. Concretely: what quantization is used in the baseline models if at all? What’s the software stack used for running these comparisons?

**Pre-training details**

- There are no meaningful details about the composition of the pre-training dataset or the ablations conducted, which makes it hard for future work to compare to this model (e.g. was Jamba trained strongly on X or Y language, or domain).
- There is no indication on section 6 or the rest of the paper about either the pre-training dataset size, or the amount of tokens the model saw during pre-training. Without this information, it is simply not possible to understand if the proposed model outperforms in a FLOPs apples-to-apples comparison with other models. This makes it hard for future work to make a call of whether or not it is worth it to implement the Jamba architecture with a given FLOPs budget — which is a tradeoff often times researchers in Academia and Industry must think about.
- MoE models are known to be finicky to pre-train [3], there’s no details about the MoE pre-training protocol. Such as for example, routing loss function and/or hyperparameters. Broadly speaking, there’s no discussion about pre-training hyperparameters at all, which makes it hard to reproduce Jamba in the future, or to run sensible comparable ablations.
- There’s no information about the tokenizer training either: the mixture of corpus size used, or software stack which is known to matter in practice [4].
- There is no mention at all about the pre-training effective throughput and FLOPs per step. This makes it hard for future work and practitioners to decide whether the proposed approach scales once the model performance is viewed as a function of FLOPs and wall-clock time.
- [1] is a very close recent work but is not discussed or ablated throughout the paper.

**Post-training details**

- There are no details at all about what algorithm was used during Post-training and what are their hyper-parameters and compute budget, which makes it hard to reproduce and to compare to other models.
- There’s a small mention of synthetic data, but there are no details about what model was used to bootstrap the synthetic data, the size, and under what protocol. This makes it hard to assess to what degree this model can be a distillation of other models.

**Evaluation**

- Section 7 has this statement “We mainly compare with recent open-weight models of the same size range” . This is problematic because the delta in performance is dictated by both model size **and** data budget [5]. Since the full compute budget is not disclosed, it’s hard to draw conclusions that can inform future research ideas and questions (Eg is the Jamba architecture better than transformers for a given compute budget).
- There are no details about the inference used to compute the self-reported numbers in table 4. This is important since some models/tasks are sensible to the use of different floating point precisions and quantization schemes. This is also important to guarantee a fair comparison.
- A seemingly arbitrary set of (model, tasks) are self-reported, whereas others are drawn from previous papers (without citation) or from the Huggingface OpenLLM Leaderboard. It is not clear why this is the case, which makes it harder for future work to reference this table or to compare to it.
- Since some numbers are self-reported and for the rest there’s no citation of where they come from (except for the leaderboard) it is not clear at all if this is a fair comparison under the same setup: prompts, split sets, same shot number. While this is a very challenging area of the LLM literature, there should be at least an explanation of the setup used to evaluate Jamba so that future work is somewhat comparable.
- The claim on Figure 5 seems overreaching, since Jamba 1.5-large is both larger in total and active parameters than llama 3.1 70b, and is well above taking into account error bars, so is not either of a similar size (the other models are smaller) or having competitive performance.

**Typos**

- quantization technique that **allows**
- Section 7.1.1: The results for key **proprietary**

[1] https://arxiv.org/abs/2402.01771

[2] https://neurips2023-enlsp.github.io/papers/paper_81.pdf

[3] https://arxiv.org/abs/2202.08906

[4] https://aclanthology.org/2024.findings-naacl.247/

[5] https://proceedings.neurips.cc/paper_files/paper/2022/file/c1e2faff6f588870935f114ebe04a3e5-Paper-Conference.pdf

[6] https://arxiv.org/abs/2404.06654

[7] https://arxiv.org/abs/2403.05530

questions:
- Is there a regression downstream performance when using ExpertsInt8 Quantization?
- Why adding Gemma on Table 3 if it does not support the context window for this task?
- Section 7.2 begins with “While not our main focus” — what’s the main focus? This seems important to contextualize to what degree this model is applicable for future work.
- What is the rationale to self-report some metrics and models in table 4?
- What does strict/flexible mean in table 4?
- Why are jamba models having higher error bars than the other models in figure 5?
- There's a couple of ablations in the appendix comparing Jamba to vanilla transformers (with and without MoE) but it lacks the details to assess if it was a fair comparison. For example, a transformer and mamba models are notably different architectures, are the optimization parameters tuned for each model in figure 6 and 7? (Learning rate, scheduler, beta values if Adam, etc). Also, are the two models parameter count equivalent? If so, how is this compensated?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730281728499

thread_type:
official_review

<thread_object>
creator: reviewer_7okn

<content>
summary:
This paper presents Jamba, a hybrid MoE model with interleaved Attention and Mamba layers. They open-source Jamba-1.5-Large (94B active parameters) and Jamba-1.5-Mini (12B active parameters). The authors conduct extensive experiments on various design choices, including different ways to interleave attention and Mamba layers and mix experts. Jamba demonstrates strong results and inference efficiency on long-context tasks, and is competitive on standard language tasks (e.g., MMLU). They introduce ExpertsInt8 to save time and memory of model loading for MoE models.

strengths:
1)	The proposed model, Jamba, is open sourced for research use.
2)	Jamba achieves strong results on long-context tasks, while offering better inference efficiency on throughput and memory footprint.
3)	The authors provide detailed ablations and insights on different ways of combining attention and Mamba layers, including the ratio of attention and Mamba layers, using Mamba-1 or Mamba-2 layers.

weaknesses:
Lack of a fair comparison on the effectiveness of Jamba against the other hybrid attention-SSM architectures. Although Jamba achieved promising results, the paper lacks a comparison with other hybrid attention-SSM models, such as YOCO [1] and Samba [2], using the same training data and parameter scale, particularly on the long-context and standard language tasks (e.g., MMLU).


[1] Sun, Yutao, et al. "You only cache once: Decoder-decoder architectures for language models."
[2] Ren, Liliang, et al. "Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling."

questions:
Add more comparsion against the other hybrid attention-SSM models under a fair setting.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730561323991

thread_type:
official_review

<thread_object>
creator: reviewer_x4g4

<content>
summary:
The paper introduces Jamba-1.5-{Large, Mini} models based on a novel hybrid of Transformer, Mamba, and MoE architectures. Combining the architectural design and the novel ExpertsInt8 quantization technique, the models are efficient in serving in terms of latency, throughput, and memory footprint. Evaluations on long-context benchmarks demonstrate that the model has strong performance in its effective 256K context length. Academic and chat benchmark evaluations indicate that Jamba-1.5 models perform similarly to SOTA public models.

strengths:
1. Serving the proposed Jamba architecture is efficient in terms of low latency, high throughput, and reduced memory footprint, especially in longer context length. 
2. The paper shares some insights found in the pre-training process, which can be beneficial to the community.
3. The performance of Jamba-1.5 models is close to other public models of similar activated parameters.

weaknesses:
1. The details of model training are not revealed, including data source, data mixture ratio, number of tokens trained, context length, and post-training techniques. This makes the training process unclear and not reproducible.
2.  A direct comparison between Jamba, Llama, and Mistral is insufficient to demonstrate the effectiveness of the proposed Jamba architecture. 
	- The training of Jamba, Llama, and Mistral models are different in data, model size, and computation, making it difficult to attribute the evaluation results. A comparison between architectures (like Jamba, Mamba, and Attention; with and without MoE) under the same pre/mid/post-training condition (like with the same data and same computation or number of tokens) can better reveal the effectiveness. 
	- Although Appendix B provides some ablation studies in architecture, the training computation, model size, and reported evaluation datasets are rather limited, and evaluation after mid/post-training is not included. This is insufficient to demonstrate the architectural effectiveness at scale and in comprehensive downstream tasks.
3. It would be better to demonstrate the scaling ability of the Jamba architecture. For instance, the relationship of LM loss/downstream task performance wrt. training computation.

questions:
- Is it possible to provide details on training procedures? (e.g., pre-training: data source, number of tokens, training throughput, training time, training context length, training batch size; post-training: strategies used and tokens trained)
- Is it possible to provide experimental results of comparison concerning the Jamba/Mamba/Attention/MoE architectures under the same training setting at scale (like more tokens and larger model size) and for more comprehensive tasks (like all tasks in Section 7)?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695621442

thread_type:
official_review

<thread_object>
creator: reviewer_bbms

<content>
summary:
The paper introduces Jamba, a hybrid language model combining Transformer, Mamba, and Mixture-of-Experts (MoE) layers to enhance memory efficiency, throughput, and long-context handling, supporting up to 256K tokens. With ExpertsInt8 quantization, Jamba achieves scalable, cost-effective deployment on 8-GPU setups. Experimental results demonstrate competitive performance across benchmarks in language modeling, chatbot, and multilingual tasks, highlighting its adaptability to various hardware and resource constraints.

strengths:
1.Hybrid Architecture with Enhanced Flexibility : Jamba’s combination of Transformer layers, Mamba layers, and Mixture-of-Experts (MoE) modules offers a unique balance of efficiency and performance, addressing the limitations of each component individually. This hybrid design allows for configurable trade-offs between memory usage, throughput, and model capacity, making Jamba adaptable to diverse hardware configurations.

2. Long-Context Capabilities: Jamba supports an effective context length of 256K tokens, one of the longest among open-weight models. This is particularly advantageous for long-context tasks, as demonstrated in benchmarks like RULER and ∞BENCH, where Jamba-1.5 models perform competitively, underscoring their suitability for tasks requiring extensive memory.

3. Cost-Effective Inference via ExpertsInt8 Quantization
The introduction of ExpertsInt8 quantization allows Jamba-1.5-Large to fit on hardware with 8 80GB GPUs without compromising quality, even for long-context processing. This quantization technique shows latency advantages on A100 GPUs, where FP8 isn’t available, thereby providing a cost-efficient alternative to traditional approaches.

4.Competitive Performance Across Benchmarks
Jamba models perform comparably to state-of-the-art models on standard language modeling, chatbot evaluations, and multilingual benchmarks. This includes maintaining high throughput and latency efficiency, particularly at large context lengths, further demonstrating the model’s real-world applicability.

weaknesses:
1. Minimal Discussion of MoE and Mamba Layer Interaction : While the hybrid design leverages both MoE and Mamba layers, the paper provides limited analysis of how these components interact to optimize performance. An in-depth exploration of how each component contributes to throughput, especially in long-context scenarios, would clarify the architectural benefits and limitations.

2. Sparse Performance Analysis on Edge and CPU: Although Jamba is designed to balance memory and compute efficiency, its performance on edge devices and CPUs is not evaluated. Given the trend toward deploying models in low-resource environments, this data would offer practical insights into the model’s viability beyond high-end GPU setups.

questions:
In Section 3.2, the authors mention that “Jamba leverages a balanced mix of Transformer, Mamba, and Mixture-of-Experts (MoE) layers to optimize both throughput and model capacity, particularly in long-context scenarios.” It would be helpful if the authors could provide further insights into how the specific ratio of these layers was determined. An ablation study examining the impact of varying the number or sequence of each layer type (Transformer, Mamba, MoE) on key performance metrics such as latency, memory usage, and accuracy would lend greater clarity to the efficacy of this particular configuration and substantiate its contribution to Jamba’s overall efficiency.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732629029393

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response; clarifying the contributions

comment:
We thank the reviewers for their reviews and for engaging in discussion. The main weakness, pointed out by two reviewers, is the missing details on pre/post-training, such as data mix and training budget. We appreciate this perspective and have clarified as many details as possible during the rebuttal. We’re pleased we were able to mitigate some of the concern as reflected in the discussion and revised evaluation by one of the reviewers. However, we believe that these details are mostly **orthogonal to the main contributions of the paper**. We wish to emphasize the key contributions of this work:

1. We introduce a **and novel hybrid Transformer-Mamba-MoE architecture and demonstrate via careful and extensive ablation experiments its benefits over vanilla Transformer or Mamba architectures**. These ablations are done at large data and model scale, which required substantial resources and incurred high costs. Importantly, **the ablations are done in a fair comparison** using exactly the same training data and the same computation resources. This is a feature of our work, as many papers that introduce novel architectures do not report extensive ablations as we do (some don’t report any ablations). Moreover, we introduce a novel quantization technique to support efficient serving of large Jamba models and demonstrate its superiority over other quantizations with no quality degradation. In addition, we describe a novel loss we introduce to support efficient inference of Jamba models at large scale. Finally, we report throughput and latency measurements, which show the benefit of the Jamba architecture with long contexts. All these investigations and innovations are a beneficial contribution to the academic community. 
2. We provide **implementations of Jamba as open-weight models of two large sizes**, along with an extensive evaluation on common benchmarks. This evaluation demonstrates the excellent performance of the released models in long context benchmarks, with high-quality on mainstream short-context benchmarks. Thus the released models are a useful artifact for the academic community to use and study. 

We find the **requirement to disclose details such as the training data mix overly restrictive**. Open-weight models (Mistral, LLaMA) are the de facto standard for open models in the academic literature. However, papers describing such leading models do not disclose such details, not to mention closed API models (GPT, Claude). Indeed, there are numerous ICLR papers using LLaMA, Mistral, and similar models, although none of them were subject to the same requirements that our work is facing. Even the GPT3 paper, which won best paper at NeurIPS, did not disclose such information. At present there are only very few truly open-source models that disclose such information (e.g., OLMo) and they lag behind Jamba and other open-weight models, and are much less frequently used for academic work than, e.g., LLaMa or Mistral. 

Moreover, most companies releasing frontier open-weight models **do not even bother to engage with the academic community and submit to peer review**, and instead only post technical reports on their websites. In contrast, we believe it’s important for papers about novel architectures and frontier open models to go through the peer-review process and engage with the academic discourse. 

Penalizing this paper on the grounds of not disclosing the training information is not only orthogonal to the paper’s main contributions, but in fact means that **ICLR does not accept work describing novel architectures implemented as frontier models at large scale**. Subsequently, ICLR may be restricted to papers using existing released models, rather than introducing new ones. 

In summary, we therefore contend that the position reflected in two of the reviews greatly misses the role of the ICLR community as we understand it.
</content>
</thread_object>

</thread>


Metareview:
metareview:
> The paper introduces Jamba, a hybrid language model combining Transformer, Mamba, and Mixture-of-Experts (MoE) layers to enhance memory efficiency, throughput, and long-context handling, supporting up to 256K tokens. With ExpertsInt8 quantization, Jamba achieves scalable, cost-effective deployment on 8-GPU setups. Experimental results demonstrate competitive performance across benchmarks in language modeling, chatbot, and multilingual tasks, highlighting its adaptability to various hardware and resource constraints.

Overall the results are good, but the experimental details are scarce and the technical decisions are not all backed by data. *I lean accept due to the performance but I am not sure it matches the scholarship required for ICLR*. Sadly this is a problem with the domain more than this single paper.

Reviewer x4G4 notes:
> The training of Jamba, Llama, and Mistral models are different in data, model size, and computation, making it difficult to attribute the evaluation results. A comparison between architectures (like Jamba, Mamba, and Attention; with and without MoE) under the same pre/mid/post-training condition (like with the same data and same computation or number of tokens) can better reveal the effectiveness.

And while this is true, a strict comparison is not really feasible in practice.

Reviewers x4G4 also discusses (in the rebuttal period):

> However, for a paper introducing a new model architecture, the lack of scaling experiments and the inability to disclose pre-training and post-training details for the two main models are significant drawbacks.

And the lack of scaling experiments is potentially the main drawback to propose a new architecture and convince the community.