Jamba: Hybrid Transformer-Mamba Language Models
language models;state-space models;mamba;hybrid architecture;foundation models
Yes
No ethics review needed.
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.
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.
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.
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.
Yes
No ethics review needed.
- 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)?
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.
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.
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.
Yes
No ethics review needed.
Add more comparsion against the other hybrid attention-SSM models under a fair setting.
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.
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.
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."
Yes
No ethics review needed.
- 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?
- 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
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.
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
Hybrid Transformer-Mamba models give better throughput and long-context capabilities at large scale
@inproceedings{
anonymous2024jamba,
title={Jamba: Hybrid Transformer-Mamba Language Models},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=JFPaD7lpBD},
note={under review}
}
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.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
language models state-space models mamba hybrid architecture foundation models
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/0baf09cce88ac1d047c42083ef6546ff4bcd7dcf.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Jamba: Hybrid Transformer-Mamba Language Models
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission