Evaluation

1. “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 two kinds of evaluations and they serve different purposes. The evaluations reported in Section 7 are of a fully trained model. Here, we compare with other open-weight models, which are also similarly released as is. The purpose of these evaluations are to provide users and researchers with information when they choose which model to use as is or as a starting point.

A separate kind of evaluation is to answer the question of which architecture to use. Here, indeed, it is important to compare under a given compute budget, which is precisely what the ablations in Appendix C (previously Appendix B) do. As explained above, these ablations support the benefits of Jamba over alternative configurations under a given compute budget.

2+3+4. “ no details about the inference used to compute the self-reported numbers in table 4.” and “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” and “not clear why this is the case” and “not clear at all if this is a fair comparison under the same setup”

Table 4 contains symbols indicating which evaluations were run by us and which were taken from the OpenLLM leaderboard. Results with competitors that are not marked are taken from the papers of these models. Whenever possible, we used previously reported numbers. When those were not available, we calculated them ourselves.

We have made every effort to ensure the comparison is fair by using the same setup: prompts, official splits, and number of shots (which is indicated in the table). When evaluating Jamba and self-reporting results of other models, we always used the official repositories. We mainly used the lm-evaluation-harness (https://github.com/EleutherAI/lm-evaluation-harness) whenever possible. Regarding evaluations not in lm-evaluation-harness: For HumanEval, BFCL, and IFEval, we used the official dataset and metrics, as mentioned in their official repos. For RealToxicity, we used the official dataset and used their API for judgment. As you rightfully note, this is a very challenging area, but our comparison is as fair as possible.

In terms of the inference for evaluation, we have used vLLM for all self-reported scores, no quantization for all models except for ExpertsInt8 quantization for Jamba-Large. If you have specific questions here, we would be happy to provide more information.

5. “The claim on Figure 5 seems overreaching”

We agree that LLaMA-3.1-70B is better than Jamba-Large, but note that Jamba-Mini is better than LLaMA-3.1-8B. We have revised the statement to make it more fair.

Typos: fixed, thank you.

Thank you for your detailed review. We’re glad you appreciated the “memory and throughput benefits on long context tasks”, the large-scale aspect of our work, and that you found the quantization “novel and interesting”.

Concerning the main weakness: “many key aspects missing from all the sections of the paper” and “most important weakness is the lack of disclosing the pre-training size (or FLOPs budget in general)”

We agree that our inability to disclose full details about model training is a limitation of this work. We have added a Limitations section at the top of the appendix that acknowledges this. Nevertheless, we emphasize that the novel aspects of this work, namely the new Jamba architecture and its efficient performance at long-contexts, are independent of these questions. We show evidence for this by providing ablations (Appendix B) of different model configurations when pre-training on exactly the same data with exactly the same compute and evaluation budget. That said, following your question, we have added additional details (Section 6.2 and 6.3) about the pre-training and post-training procedures, including the major pre-training data types and the pre-training context length. We have also added details about the post-training process.

We believe our results and analyses provide valuable insights for the ICLR community, which is why we have decided to go through the peer-review process, in contrast to other work at this scale, which does not bother to engage with the ICLR community.

Below we address the points you raised one by one.

Efficient inference details

1. “no empirical evidence to quantify the loss in model performance” of the proposed quantization method

As mentioned in Section 4.1, our quantization technique incurs no loss in model quality. Following your question, we have added in Appendix C.7 empirical evidence supporting this statement.

2. “Figure 2 from section 4.1 seems focused on latency and missing throughput”

Throughput and latency are naturally inversely related. Following your question, we have run additional throughput measurements. As Appendix C.7 now reports, we observe similar trends in throughput, supporting the benefits of ExpertsInt8.

3. “Relevant very similar work in MoE quantization on the experts [2] is missing”

Thank you for this reference. MoQE is focused on very low-bit quantization, which incurs a loss in quality (they report BLEU on machine translation). The other work you cited (Li et al.), which is based on GPTQ, is also focused on very low-bit quantization and doesn’t compare with a non-quantized model. It’s also from June 2024, which was too recent for us to directly compare with. Importantly, we provide a public implementation of our quantization directly in vLLMs’ fused_moe kernel. We have added mention of these two papers in footnote 5.

4. “Section 5 lacks important details of what software stack was used to run the latency and throughput comparisons” and “what quantization was used in the baseline models if at all”

All cases were tested using the same stack of software (vLLM image == 0.5.5) , same hardware (H100/A100, depending on the chart), and same code for testing (vLLM latency test). We did not quantize any of the models, except for Jamba-Large (with ExpertsInt8) and LLaMA-3.1-405B (with FP8), since both do not fit on a single 8xA100 node without quantization. We managed running LLaMA-3.1-405B on A100 with FP8, but reached only a sequence length of ~100K tokens, hitting out-of-memory after that. In contrast, Jamba-Large can reach 256K tokens.

Pre-training details

1+2. “no meaningful details about the composition of the pre-training dataset” and “pre-training dataset size, or the amount of tokens the model saw during pre-training”

As mentioned above, we acknowledge this limitation of our work, and have added some details following your question.

However, we disagree with your statement that “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”. As clearly shown in the ablations in Appendix C (previously Appendix B), training the Jamba architecture with a given FLOPs budget leads to strictly better results than various alternative configurations (vanilla Mamba or Transformer, no-MoE, other Attention-to-Mamba ratios). Thus, our ablations provide valuable evidence for the ICLR community in designing and implementing future hybrid models.

3. “no details about the MoE pre-training protocol.”

We follow standard protocols for MoE training, as mentioned in Appendix B (previously Appendix A). In particular, we used standard router load-balancing loss to balance the routing decisions and a router zloss to avoid large router logits. We’ve added this information to Appendix B. The number of total and active experts is specified in Section 3.

4. “No information about the tokenizer training”

We used a similar mixture corpus for training the tokenizer as in our pre-training mixture. We trained with the same software and very similar settings as the LLaMA/Mistral tokenizers, namely, SentencePiece BPE, which was shown [4] to be better than HuggingFace BPE. We’ve added this information to Appendix B (previously Appendix A). As noted in the Appendix, we set numbers to single digits as is common in recent models.

5. “pre-training effective throughput and FLOPs per step”

The important question to answer is how efficient the Mamba and Attention layers are during pre-training. We evaluated this with the models FLOPs utilization metric (MFU), proposed in the PaLM paper. In our measurements, the MFU with vanilla attention drops from 60% to 20% at 32K token length, while the Mamba MFU remains close to 60%. In longer contexts, the difference would be even bigger.

6. “[1] is a very close recent work”

Thank you for this reference to BlackMamba. We were not aware of this work and have added mention of this to Section 1. One small difference is that BlackMamba applied MoE to MLP, while we apply it only at every other layer. And, of course, BlackMamba is not a hybrid Attention-Mamba model.

Post-training details

1. “no details at all about what algorithm was used during Post-training and what are their hyper-parameters and compute budget”

As Section 6.3 and Appendix C.2 (previously Appendix B.2) indicate, our training algorithm was supervised fine-tuning. While we are not able to provide details on hyper-parameters and compute budget, we share interesting observations in Appendix C.2, which we hope will be useful for the community.

2. “no details about what model was used to bootstrap the synthetic data, the size, and under what protocol”

Appendix D (previously Appendix C) provides some information about the synthetic data generation protocol. We have now added additional details on our general strategy of sampling and filtering.

Questions

1. “Is there a regression downstream performance when using ExpertsInt8 Quantization?”

No. See new results in Appendix C.7.

2. “Why adding Gemma on Table 3 if it does not support the context window for this task?”

This was done just as a matter of style to conform with the other tables.

3. “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.”

The main focus is to provide long context capabilities with improved throughput, which is why Section 7.1 appears first. We provide standard academic benchmarks in Section 7.2 as a useful reference to compare with other models.

4. “What is the rationale to self-report some metrics and models in table 4?”

As explained above, we always used previously reported scores whenever they were available. We only self-report results that are missing from prior studies.

4. “What does strict/flexible mean in table 4?”

Strict and flexible are two common ways to evaluate on GSM8K. Since LLaMA-3.1 models perform poorly with the standard strict evaluation, we also reported for them the flexible metric, which allows for higher results. Strict means that the answer must follow “####” while flexible searches for the numbers in the completion. See the LM Evaluation Harness for the regexes: https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/gsm8k/gsm8k.yaml.

5. “Why are jamba models having higher error bars than the other models in figure 5?”

Jamba models are more recent than other models, so they have fewer votes on the ChatbotArena and thus higher error bars.

6. “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.”

We have tuned each hyper-parameter separately for each architecture in an effort to get the best possible model. The models are also using a very similar parameter count. In general it’s not possible to use exactly the same parameter counts due to difference in architectures, but our models are very similar. Importantly, the boost in performance from pure Mamba to the hybrid Attention-Mamba comes with the hybrid having fewer parameters.

We hope we were able to answer as many of your comments as possible, and that you now agree that the contributions of this work are valuable to the ICLR community.

While I think it's unsatisfactory that the authors are not able to disclose what open models were used in Appendix D, general pre-training hyperparameters, and I share the concerns of reviews x4G4 (hope the AC weighs in), I think the authors have made a notable effort to improve the manuscript and are genuinely engaged on sharing as much as they can -- an increasingly challenging point in industry research -- . Hence, I'm increasing my score. Thanks to the authors for engaging.

Thank you for your helpful review. We’re glad you appreciated our “extensive experiments” with “strong results on long-context tasks” and “better inference efficiency”. We’re also happy you appreciate that Jamba is publicly available, and our “detailed ablations and insights”.

We would like to address the weakness and question you mentioned.

“Lack of a fair comparison on the effectiveness of Jamba against the other hybrid attention-SSM architectures” and “comparsion against the other hybrid attention-SSM models under a fair setting”:

As discussed in the Introduction, all prior attempts with hybrid models are either of smaller scale or perform worse than vanilla Mamba or vanilla Transformers. Concerning YOCO and Samba, they were posted on arxiv in May and June 2024, which was too recent for us to compare with under the same setting. We have now added mention of these recent attempts to footnote 2 in the Introduction. We did try the main aspect that Samba notes, namely the sliding-window attention. In our preliminary experiments with a 1.3B-parameter model, combining Mamba with sliding window attention performed worse than Mamba with attention without sliding window. (Note that the YOCO and Samba models themselves appear weaker than Jamba (e.g., on MMLU), but that is not a fair comparison, so we are not stressing this point.)

We hope we have addressed your concern. Given the positive note of your review, we would appreciate it if you can consider increasing your score. Please let us know if you have any other questions.

Dear reviewer, we hope our response and new revision has addressed the points you have raised. Could you let us know if you have any other comments or questions? Assuming we have answered you concerns, we would appreciate it if you can consider revising your evaluation.

Thank you for your helpful review. We’re glad you noted the work’s strengths: the model’s “strong performance in its effective 256K context length”, “serving of the proposed Jamba architecture [...] efficient”, and the “insights shared during the pre-training process”.

We address the weaknesses and questions below.

1. “The details of model training are not revealed” and the question on “details on training procedures”

We agree that our inability to disclose full details about model training is a limitation of this work. We have added a Limitations section at the top of the appendix that acknowledges this. Nevertheless, we emphasize that the novel aspects of this work, namely the new Jamba architecture and its efficient performance at long-contexts, are independent of these questions. We show evidence for this by providing ablations (Appendix C, previously Appendix B) of different model configurations when pre-training on exactly the same data with exactly the same compute and evaluation budget.

That said, following your question, we have added additional details (in Section 6.2 and 6.3) about the pre-training and post-training procedures, including the major data types in pre-training and the pre-training context length. We have also added details about the types of data used in post-training. Refer also to Appendix D (previously Appendix C), where we have provided information and insights about post-training strategies.

2. “A direct comparison between Jamba, Llama, and Mistral is insufficient to demonstrate the effectiveness of the proposed Jamba architecture.” and the question about “experimental results of comparison…”

We agree that comparing different released models is insufficient to make claims about a proposed architecture. For this reason, we provide the ablation studies in Appendix C (previously Appendix B), as you acknowledge. These studies clearly justify our design choices in terms of the advantage of hybrid over vanilla models, which attention-to-Mamba ratio to use, and the benefit of MoE. The takeaways are summarized at the top of Section 3.

We emphasize that these experiments were conducted under exactly the same data and same computation (number of tokens) for each compared variant. Thus, the comparison is fair. Moreover, the ablations are conducted at a fairly large scale: 1.3B/7B parameter models trained for 250B/50B tokens. This is a large scale for ablation studies that already required substantial compute and expenses.

Finally, it is impossible to perform architectural ablations at mid/post-training, since this would require full pre-pretraining of multiple configurations, which is generally unrealistic.

Moreover, the ablations in Appendix C (previously Appendix B) already include results from 3 academic benchmarks, the OpenLLM leaderboard summary (which is an aggregate metric over 7 datasets), and log-probs with 3 data sources of different domains. These were chosen because we found them informative for ablations at a smaller scale than full pre-training, since not all benchmarks provide informative results at this scale. Nevertheless, given your question, we have added additional ablation results to Appendix C on more tasks. Specifically, we added results on BoolQ, NarrativeQA, ARC-Challenge, and IMDB. While some of them are unstable (which we expanded upon in Appendix C.3), the results are in line with our prior observations.

3. “demonstrate the scaling ability of the Jamba architecture”

Thank you for this suggestion. Conducting a full scaling study is a complicated procedure taking many resources, which is beyond our scope. However, we are looking into providing measurements at multiple scales and will get back to you ASAP.

We hope our response has helped answer your concerns. If you have any additional questions, please let us know. We would appreciate it if you can revise your evaluation accordingly.

Dear reviewer, we hope our response and the new paper revision have addressed your concerns. Please let us know if you have any remaining questions. We would appreciate it if you can go over our response and revision, and consider revising your evaluation assuming we have answered your concerns.

Thank you for your helpful review. We’re glad you appreciated the benefits provided by our “Hybrid Architecture with Enhanced Flexibility”, which “offers a unique balance of efficiency and performance”. We’re happy you noted the strengths of this architecture in “Long-Context Capabilities” and the “Cost-Effective Inference via ExpertsInt8 Quantization”, as well as the “Competitive Performance Across Benchmarks”.

Below we address the weaknesses noted.

1. “Minimal Discussion of MoE and Mamba Layer Interaction”

Table 10 in Appendix B.4 provides a comparison of MoE and no-MoE hybrid models in terms of quality. The table shows improvements with MoE across the board.

Considering the effect of throughput, we have conducted measurements and have added a discussion in the new version of the paper (Appendix C.4). The conclusion is that the effect of MoE on throughput/latency is linear in sequence length, as expected. Interestingly, compared to regular MLP layers, MoE has only about a 1.3/1.6x effect on throughput/latency in Jamba-Mini/Large, even though 2 experts are used (meaning 2x active parameters). This indicates that MLP/MoE layers are I/O bounded and since 2 experts can be computed independently, batching and fusing save time. In general increasing MLP compute via MoE helps reduce the relative effect of Attention layers, which are quadratic in sequence length. In Jamba-Mini, Attention layers become dominant around context length of 540K tokens, until then MOE is dominant. Attention starts to have a greater effect on throughput/latency compared to Mamba at ~300K tokens. In Jamba-Large, when the matrices are larger, I/O is less of a problem, and MoE is the dominant factor.

2. “Sparse Performance Analysis on Edge and CPU”

Thank you for raising this point. We find the prospects of the Jamba architecture for edge and CPU devices exciting. While implementing this would face some challenges (especially with custom Mamba kernels), we think this is a great avenue for future work.

Regarding your question about “how the specific ratio of these layers was determined” and "An ablation study […] on key performance metrics such as latency, memory usage, and accuracy”:

• Appendix C (previously Appendix B) provides a detailed ablation study in terms of quality. Appendix C.1 compares hybrid Attention-Mamba models to vanilla Attention or Mamba, including different ratios of Attention-to-Mamba. Appendix C.4 compares hybrid models with and without MoE.
• In terms of latency and memory, we conducted detailed measurements and have added a new Appendix C.8 that reports our findings. Thank you for this suggestion.

Thank you once again for your review and please let us know if you have any other questions or comments.