Elamur Research⌕
☰
home / papers / 2024 / Zamba: A Compact 7B SSM Hybrid Model

Zamba: A Compact 7B SSM Hybrid Model

Paper · 2024 · Paolo Glorioso, Quentin Anthony, Yury Tokpanov et al.

Summary

Pure transformers pay a quadratic price for attention and need a KV cache that grows with context. State-space models such as Mamba run in linear time with fixed memory, but they still trail transformers, notably at in-context learning. Zamba, a 7B model from Zyphra, tries to combine the strengths of both. Its backbone is a stack of Mamba blocks. After every six of them, the network calls one attention-plus-MLP block whose weights are reused at each call. A separate, non-shared linear projection returns that block's output to the residual stream, and the original input embeddings are fed in alongside the current activations. Training used 1T tokens of openly available data in two stages: ordinary web pretraining, then an annealing stage on instruct and synthetic data with a rewarmed, exponentially decaying learning rate. The model nearly matches leading ~7B transformers trained on far more data. It lags slightly on reasoning and in-context learning but generates faster with less memory. All weights and intermediate checkpoints are public, which supports research into how non-transformer models learn.

Key points

  • Architecture: a Mamba backbone with one global shared self-attention (GSA) block, consisting of attention and MLP, inserted every 6 Mamba blocks. Sharing the parameters keeps both the parameter memory and the KV cache small. The design was inspired by how the cortex interacts with the hippocampus.
  • Training curriculum: about 1T tokens from open datasets. Phase 1 uses web data. The annealing phase uses high-quality instruct and synthetic data with a rewarmed learning rate and rapid exponential decay rather than linear decay. The annealing helped some downstream evaluations and had little effect on others.
  • Results: the authors report Zamba as the strongest SSM and dense SSM at the 7B scale. Among models with open checkpoints, it beats OLMo 1.7 and Pythia. It has lower forward-pass latency at 8k sequence length and lower memory use for generation than comparable transformers.
  • Released resources: the phase-1 model (Zyphra/Zamba-7B-v1-phase1), the annealed model (Zyphra/Zamba-7B-v1), and all training checkpoints from both phases, all on Hugging Face.
  • Configuration: 80 layers, hidden dimension 3712, state dimension 16, 16 attention heads, context length 4096, batch size 512, peak learning rate 1.5e-4, BF16 precision.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

Also on the byline: Zyphra.

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 19,503 worksArtificial Intelligence · 18,805 worksArts and Humanities · 17,959 worksLanguage and Linguistics · 17,789 worksComputer Vision and Pattern Recognition · 3,621 works

Cite

Cite as Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim and Beren Millidge. 2024. Zamba: A Compact 7B SSM Hybrid Model. arXiv:2405.16712.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2405.16712 · arXiv 2405.16712

@misc{glorioso2024zamba,
  title = {{Zamba: A Compact 7B SSM Hybrid Model}},
  author = {Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim and Beren Millidge},
  year = {2024},
  doi = {10.48550/arxiv.2405.16712},
  eprint = {2405.16712},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2405.16712}
}

References

42 references, 21 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Abdin, M., Jacobs, S. A., Awan, A. A., Aneja, J., Awadallah, A., Awadalla, H., Bach, N., Bahree, A., Bakhtiari, A., Behl, H., et al. (2024). Phi-3 technical report: A highly capable language model locally on your phone. arXiv preprint arXiv:2404.14219. AI2 (2024). OLMo 1.7–7B: A 24 point improvement…
  2. Anthony, Q., Tokpanov, Y., Glorioso, P., and Millidge, B. (2024). Blackmamba: Mixture of experts for state-space models. arXiv preprint arXiv:2402.01771. — [arXiv:2402.01771](https://arxiv.org/abs/2402.01771) find
  3. Bachmann, G., Anagnostidis, S., and Hofmann, T. (2023). Scaling mlps: A tale of inductive bias. Advances in Neural Information Processing Systems, 36. find
  4. Bengio, Y., Louradour, J., Collobert, R., and Weston, J. (2009). Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pages 41– 48.
  5. Biderman, S., Schoelkopf, H., Anthony, Q. G., Bradley, H., O’Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S., Raff, E., et al. (2023). Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pages 2397–2430… find
  6. Gunasekar, S., Zhang, Y., Aneja, J., Mendes, C. C. T., Del Giorno, A., Gopi, S., Javaheripi, M., Kauffmann, P., de Rosa, G., Saarikivi, O., et al. (2023). Textbooks are all you need. arXiv preprint arXiv:2306.11644. Gupta, K., Thérien, B., Ibrahim, A., Richter, M. L., Anthony, Q., Belilovsky, E., R… find
  7. Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, J. (2021). Measuring massive multitask language understanding.
  8. Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E., Rae, J. W., Vinyals, O., and Sifre, L.… — Training Compute-Optimal Large Language Models
  9. Hu, S., Tu, Y., Han, X., He, C., Cui, G., Long, X., Zheng, Z., Fang, Y., Huang, Y., Zhao, W., et al. (2024). Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.06395. Ibrahim, A., Thérien, B., Gupta, K., Richter, M. L., Anthony, Q.…
  10. Jacobs, R. A., Jordan, M. I., Nowlan, S. J., and Hinton, G. E. (1991). Adaptive mixtures of local experts. Neural computation, 3(1):79–87. find
  11. Jelassi, S., Brandfonbrener, D., Kakade, S. M., and Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying.
  12. Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. (2020). Scaling laws for neural language models. arXiv preprint arXiv:2001.08361.
  13. Krishna, K., Garg, S., Bigham, J. P., and Lipton, Z. (2023). Downstream datasets make surprisingly good pretraining corpora. In The 61st Annual Meeting Of The Association For Computational Linguistics. find
  14. Kwon, W., Li, Z., Zhuang, S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J. E., Zhang, H., and Stoica, I. (2023). Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles.
  15. Li, Y., Bubeck, S., Eldan, R., Del Giorno, A., Gunasekar, S., and Lee, Y. T. (2023). Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint arXiv:2309.05463.
  16. Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J., Dalmedigos, I., Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., Shwartz, E., Zusman, M., and Shoham, Y. (2024). Jamba: A hybrid t…
  17. Liu, C., He, S., Liu, K., and Zhao, J. (2018). Curriculum learning for natural answer generation. In Proceedings of the 27th International Joint Conference on Artificial Intelligence, pages 4223–4229. find
  18. Liu, Z., Qiao, A., Neiswanger, W., Wang, H., Tan, B., Tao, T., Li, J., Wang, Y., Sun, S., Pangarkar, O., et al. (2023). Llm360: Towards fully transparent open-source llms. arXiv preprint arXiv:2312.06550. find
  19. Maini, P., Seto, S., Bai, H., Grangier, D., Zhang, Y., and Jaitly, N. (2024). Rephrasing the web: A recipe for compute and data-efficient language modeling. arXiv preprint arXiv:2401.16380.
  20. NVIDIA (2023). Transformer Engine: A library for accelerating Transformer models on NVIDIA GPUs. https://github.com/NVIDIA/TransformerEngine. Accessed: May 28, 2024. ↗
  21. Park, J., Park, J., Xiong, Z., Lee, N., Cho, J., Oymak, S., Lee, K., and Papailiopoulos, D. (2024). Can mamba learn how to learn? a comparative study on in-context learning tasks.
  22. Parmar, J., Prabhumoye, S., Jennings, J., Patwary, M., Subramanian, S., Su, D., Zhu, C., Narayanan, D., Jhunjhunwala, A., Dattagupta, A., et al. (2024a). Nemotron-4 15b technical report. arXiv preprint arXiv:2402.16819.
  23. Parmar, J., Prabhumoye, S., Jennings, J., Patwary, M., Subramanian, S., Su, D., Zhu, C., Narayanan, D., Jhunjhunwala, A., Dattagupta, A., Jawa, V., Liu, J., Mahabaleshwarkar, A., Nitski, O., Brundyn, A., Maki, J., Martinez, M., You, J., Kamalu, J., LeGresley, P., Fridman, D., Casper, J., Aithal, A.,… — Nemotron-4 15B Technical Report
  24. Penedo, G., Malartic, Q., Hesslow, D., Cojocaru, R., Cappelli, A., Alobeidli, H., Pannier, B., Almazrouei, E., and Launay, J. (2023). The refinedweb dataset for falcon llm: outperforming curated corpora with web data, and web data only. arXiv preprint arXiv:2306.01116. find
  25. Peng, B., Alcaide, E., Anthony, Q., Albalak, A., Arcadinho, S., Biderman, S., Cao, H., Cheng, X., Chung, M., Grella, M., GV, K. K., He, X., Hou, H., Lin, J., Kazienko, P., Kocon, J., Kong, J., Koptyra, B., Lau, H., Mantri, K. S. I., Mom, F., Saito, A., Song, G., Tang, X., Wang, B., Wind, J. S., Wozn… find
  26. Peng, B., Goldstein, D., Anthony, Q., Albalak, A., Alcaide, E., Biderman, S., Cheah, E., Du, X., Ferdinan, T., Hou, H., Kazienko, P., GV, K. K., Kocoń, J., Koptyra, B., Krishna, S., au2, R. M. J., Muennighoff, N., Obeid, F., Saito, A., Song, G., Tu, H., Woźniak, S., Zhang, R., Zhao, B., Zhao, Q.,… find
  27. Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. (2020). Exploring the limits of transfer learning with a unified text-to-text transformer. Journal of machine learning research, 21(140):1–67.
  28. Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. (2022). Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale. In International conference on machine learning, pages 18332–18346.
  29. PMLR. Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., and Dean, J. (2016). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations. find
  30. Shen, Y., Guo, Z., Cai, T., and Qin, Z. (2024). Jetmoe: Reaching llama2 performance with 0.1 m dollars. arXiv preprint arXiv:2404.07413.
  31. Soldaini, L. and Lo, K. (2023). peS2o (Pretraining Efficiently on S2ORC) Dataset. Technical report, Allen Institute for AI. ODC-By, https://github.com/allenai/pes2o. ↗
  32. Sun, Y., Dong, L., Huang, S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, F. (2023). Retentive network: A successor to transformer for large language models. Team, G., Anil, R., Borgeaud, S., Wu, Y., Alayrac, J.-B., Yu, J., Soricut, R., Schalkwyk, J., Dai, A. M., Hauth, A., et al.
  33. (2023). Gemini: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805.
  34. Tirumala, K., Simig, D., Aghajanyan, A., and Morcos, A. (2024). D4: Improving llm pretraining via document deduplication and diversification. Advances in Neural Information Processing Systems, 36. find
  35. Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. (2023). Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288.
  36. Tscshantz, A., Millidge, B., Seth, A. K., and Buckley, C. L. (2023). Hybrid predictive coding: Inferring, fast and slow. PLoS Computational Biology, 19(8):e1011280. van Bergen, R. S. and Kriegeskorte, N. (2020). Going in circles is the way forward: the role of recurrence in visual inference. Current… find
  37. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017). Attention is all you need. Advances in neural information processing systems, 30.
  38. Whittington, J. C. R., Muller, T. H., Mark, S., Barry, C., Burgess, N., Behrens, T. E. E., Chen, G., Barry, C., Burgess, N., and Behrens, T. E. E. (2020). The Tolman-Eichenbaum Machine: Unifying Space and Relational Memory through Generalization in the Hippocampal Formation. Cell, 183(5):1249–1263.e… find
  39. Publisher: Elsevier Inc. Whittington, J. C. R., Warren, J., and Behrens, T. E. J. (2021). Relating transformers to models and neural representations of the hippocampal formation. International Conference on Learning Representations. find
  40. Xie, S. M., Pham, H., Dong, X., Du, N., Liu, H., Lu, Y., Liang, P. S., Le, Q. V., Ma, T., and Yu, A. W. (2024). Doremi: Optimizing data mixtures speeds up language model pretraining. Advances in Neural Information Processing Systems, 36. find
  41. Yang, Z., Zeng, X., Zhao, Y., and Chen, R. (2023). Alphafold2 and its applications in the fields of biology and medicine. Signal Transduction and Targeted Therapy, 8(1):115. find
  42. Zyphra (2024). Reproduction of Mamba-370M by Zyphra. https://huggingface.co/Zyphra/Mamba-370M. Accessed: May 28, 2024. ↗

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/glorioso-2024-zamba-compact-7b-ssm-hybrid/