Authors:
(1) Opher Lieber, with Equal contribution; (2) Barak Lenz, with Equal contribution; (3) Hofit Bata; (4) Gal Cohen; (5) Jhonathan Osin; (6) Itay Dalmedigos; (7) Erez Safahi; (8) Shaked Meirom; (9) Yonatan Belinkov; (10) Shai Shalev-Shwartz; (11) Omri Abend; (12) Raz Alon; (13) Tomer Asida; (14) Amir Bergman; (15) Roman Glozman; (16) Michael Gokhman; (17) Avashalom Manevich; (18) Nir Ratner; (19) Noam Rozen; (20) Erez Shwartz; (21) Mor Zusman; (22) Yoav Shoham. Table of Links Part 1 Part 2 Part 3 Part 4 Part 5 Part 6 6. Ablations and Insights This section discusses ablation experiments we ran for different design choices in our implementation of the Jamba architecture. First we show the benefit of combining attention and Mamba layers, at which ratio they should be combined, and how to interleave them. We investigate cases where pure Mamba fails, suggesting that it struggles to develop in-context learning capabilities, while the Attention–Mamba hybrid exhibits in-context learning similar to vanilla Transformers. Then we show the benefit of adding MoE on top of a hybrid Attention–Mamba model. Finally, we share two additional learnings that we found useful: explicit positional information is not needed in Jamba, and Mamba layers necessitate special normalization to stabilize training at large scale.[4] For these ablations, we report the following measures, which exhibit informative performance even at small data or model scale. • Academic benchmarks: HellaSwag (10-shot) [47], WinoGrande (5-shot) [37], Natural Questions closed-book (NQ; 5-shot) [26]. • HuggingFace OpenLLM leaderboard (OLLM) [11]: a summary statistic of several datasets. We report results with our reproduction. • Perplexity evaluations: we report log-prob (per byte) on texts from three domains: C4, Books, and code. 6.1 Benefits of combining Attention and Mamba We first investigate the ratio of Attention to Mamba layers (a : m), with 1.3B parameters models trained for 250B tokens. As Table 4 shows, the hybrid Jamba model outperforms the pure attention or Mamba models. The ratio of attention-to-Mamba layers may be 1:3 or 1:7 with virtually no performance difference. Figure 5 shows the training loss of these models, where Jamba exhibits improved loss during training. Given that a 1:7 ratio is more compute-efficient and shows similar performance, we opt for it in our larger-scale experiments. Next, we compare performance of vanilla Transformer, vanilla Mamba, and Attention-Mamba hybrid models, at 7B model size, after training on 50B tokens. As Table 5 shows, the pure Mamba layer is quite competitive, but lags slightly behind pure Attention. The hybrid Attention-Mamba (without MoE) outperforms the pure models while obtaining better throughput than the vanilla Transformer (Section 3.2). Figure 6 shows the training loss of the three architectures. While the pure Transformer and Mamba models have a similar convergence, the hybrid Jamba (no MoE) has a lower loss throughout this run. This paper is available on arxiv under CC BY-SA 4.0 DEED license. [4] In all the ablation experiments, “pure Mamba” refers to models with Mamba layers interleaved with MLP layers. Authors: (1) Opher Lieber, with Equal contribution; (2) Barak Lenz, with Equal contribution; (3) Hofit Bata; (4) Gal Cohen; (5) Jhonathan Osin; (6) Itay Dalmedigos; (7) Erez Safahi; (8) Shaked Meirom; (9) Yonatan Belinkov; (10) Shai Shalev-Shwartz; (11) Omri Abend; (12) Raz Alon; (13) Tomer Asida; (14) Amir Bergman; (15) Roman Glozman; (16) Michael Gokhman; (17) Avashalom Manevich; (18) Nir Ratner; (19) Noam Rozen; (20) Erez Shwartz; (21) Mor Zusman; (22) Yoav Shoham. Authors: Authors: (1) Opher Lieber, with Equal contribution; (2) Barak Lenz, with Equal contribution; (3) Hofit Bata; (4) Gal Cohen; (5) Jhonathan Osin; (6) Itay Dalmedigos; (7) Erez Safahi; (8) Shaked Meirom; (9) Yonatan Belinkov; (10) Shai Shalev-Shwartz; (11) Omri Abend; (12) Raz Alon; (13) Tomer Asida; (14) Amir Bergman; (15) Roman Glozman; (16) Michael Gokhman; (17) Avashalom Manevich; (18) Nir Ratner; (19) Noam Rozen; (20) Erez Shwartz; (21) Mor Zusman; (22) Yoav Shoham. Table of Links Part 1 Part 1 Part 2 Part 2 Part 3 Part 3 Part 4 Part 4 Part 5 Part 5 Part 6 Part 6 6. Ablations and Insights This section discusses ablation experiments we ran for different design choices in our implementation of the Jamba architecture. First we show the benefit of combining attention and Mamba layers, at which ratio they should be combined, and how to interleave them. We investigate cases where pure Mamba fails, suggesting that it struggles to develop in-context learning capabilities, while the Attention–Mamba hybrid exhibits in-context learning similar to vanilla Transformers. Then we show the benefit of adding MoE on top of a hybrid Attention–Mamba model. Finally, we share two additional learnings that we found useful: explicit positional information is not needed in Jamba, and Mamba layers necessitate special normalization to stabilize training at large scale.[4] For these ablations, we report the following measures, which exhibit informative performance even at small data or model scale. • Academic benchmarks: HellaSwag (10-shot) [47], WinoGrande (5-shot) [37], Natural Questions closed-book (NQ; 5-shot) [26]. • HuggingFace OpenLLM leaderboard (OLLM) [11]: a summary statistic of several datasets. We report results with our reproduction. • Perplexity evaluations: we report log-prob (per byte) on texts from three domains: C4, Books, and code. 6.1 Benefits of combining Attention and Mamba We first investigate the ratio of Attention to Mamba layers (a : m), with 1.3B parameters models trained for 250B tokens. As Table 4 shows, the hybrid Jamba model outperforms the pure attention or Mamba models. The ratio of attention-to-Mamba layers may be 1:3 or 1:7 with virtually no performance difference. Figure 5 shows the training loss of these models, where Jamba exhibits improved loss during training. Given that a 1:7 ratio is more compute-efficient and shows similar performance, we opt for it in our larger-scale experiments. Next, we compare performance of vanilla Transformer, vanilla Mamba, and Attention-Mamba hybrid models, at 7B model size, after training on 50B tokens. As Table 5 shows, the pure Mamba layer is quite competitive, but lags slightly behind pure Attention. The hybrid Attention-Mamba (without MoE) outperforms the pure models while obtaining better throughput than the vanilla Transformer (Section 3.2). Figure 6 shows the training loss of the three architectures. While the pure Transformer and Mamba models have a similar convergence, the hybrid Jamba (no MoE) has a lower loss throughout this run. This paper is available on arxiv under CC BY-SA 4.0 DEED license. This paper is available on arxiv under CC BY-SA 4.0 DEED license. available on arxiv [4] In all the ablation experiments, “pure Mamba” refers to models with Mamba layers interleaved with MLP layers.

Part of HackerNoon's growing list of open-source research papers, promoting free access to academic material.

HuggingFace

Analyzing the Impact of Model Scaling on Long-Form Factuality

256K Tokens on One GPU? Jamba’s Engineering Magic Explained

LanguageModels.TECH

Read My Stories

Too Long; Didn't Read

Make resilience your competitive advantage

Inside Jamba’s Architecture: Mamba Layers, MoE, and the Future of AI Models

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

Untitled Story

TnT-LLM Implementation Details: Pipeline Design, Robustness, and Efficiency

Why Jamba Is the First Truly Scalable Hybrid LLM for Long Contexts

What Jamba’s Benchmark Wins Tell Us About the Power of Hybrid LLMs

How Jamba Combines Transformers and Mamba to Build Smarter Language Models

256K Tokens on One GPU? Jamba’s Engineering Magic Explained

TnT-LLM Implementation Details: Pipeline Design, Robustness, and Efficiency

Why Jamba Is the First Truly Scalable Hybrid LLM for Long Contexts

What Jamba’s Benchmark Wins Tell Us About the Power of Hybrid LLMs

How Jamba Combines Transformers and Mamba to Build Smarter Language Models

256K Tokens on One GPU? Jamba’s Engineering Magic Explained

Light-Mode

Classic

Newspaper

Dark-Mode

Neon Noir

Minty

HN StartUps