Zamba: コンパクトな7B SSMハイブリッドモデル

Zamba: A Compact 7B SSM Hybrid Model

May 26, 2024

著者: Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, Beren Millidge

cs.AI

要旨

本技術レポートでは、Zambaという新たな7B SSM-トランスフォーマーハイブリッドモデルを紹介します。Zambaは、同等規模の主要なオープンウェイトモデルと競合する性能を達成しています。Zambaは、公開されているデータセットから1兆トークンを学習し、この規模において最良の非トランスフォーマーモデルです。Zambaは、Mambaバックボーンと単一の共有アテンションモジュールを組み合わせた独自のアーキテクチャを採用し、最小限のパラメータコストでアテンションの利点を得ています。このアーキテクチャにより、Zambaは同等のトランスフォーマーモデルよりも推論が大幅に高速で、長いシーケンスの生成に必要なメモリも大幅に削減されます。Zambaは2段階で事前学習されています：第1段階は既存のウェブデータセットに基づいており、第2段階は高品質な指示データセットと合成データセットを用いたアニーリングで、学習率の急速な減衰が特徴です。Zambaの重みとすべてのチェックポイントを、第1段階とアニーリング段階の両方を通じてオープンソースとして公開します。

English

In this technical report, we present Zamba, a novel 7B SSM-transformer hybrid model which achieves competitive performance against leading open-weight models at a comparable scale. Zamba is trained on 1T tokens from openly available datasets and is the best non-transformer model at this scale. Zamba pioneers a unique architecture combining a Mamba backbone with a single shared attention module, thus obtaining the benefits of attention at minimal parameter cost. Due to its architecture, Zamba is significantly faster at inference than comparable transformer models and requires substantially less memory for generation of long sequences. Zamba is pretrained in two phases: the first phase is based on existing web datasets, while the second one consists of annealing the model over high-quality instruct and synthetic datasets, and is characterized by a rapid learning rate decay. We open-source the weights and all checkpoints for Zamba, through both phase 1 and annealing phases.