7B新王登基！Zamba 2完勝同級模型，推理效率比Llama 3提升20%，內存用量更少

作者：新智元 2024-10-25 10:57:11

Zamba2-7B是一款小型語言模型，在保持輸出質量的同時，通過創新架構實現了比同類模型更快的推理速度和更低的內存占用，在圖像描述等任務上表現出色，能在各種邊緣設備和消費級GPU上高效運行。

除了不斷增加語言模型的尺寸來提升性能外，小語言模型（SLM）賽道也是越來越卷，研究人員在保證輸出質量盡量不變的情況下，不斷降低模型尺寸，減少內存占用量，提升推理效率，從而能夠在各種邊緣計算設備和消費級GPU上部署使用。

最近，Zyphra發布Zamba2-7B模型，在質量和性能上都優于Mistral、Google的Gemma和Meta的Llama3系列同尺寸小語言模型；在推理效率上，與 Llama3-8B 等模型相比，第一個token的時間縮短了 25%，每秒token數量提高了 20%，并且內存使用量顯著減少。

Instruct下載鏈接：https://huggingface.co/Zyphra/Zamba2-7B-Instruct

base下載鏈接：https://huggingface.co/Zyphra/Zamba2-7B

相對于上一代Zamba1-7B，新模型在架構上的改進包括：

1. 把Mamba1塊已替換為Mamba2塊；

2. 把單個共享注意力塊增加為兩個共享注意力塊，在整個網絡中以 ABAB 模式交錯排列，增強了網絡對信息的處理能力；

3. 為每個共享的多層感知機（MLP）模塊應用了一個LoRA投影器，可以讓網絡在每次調用共享層時，針對不同層次專門化（specialize）MLP模塊以適應不同層次的數據處理，從而提高網絡的性能和效率。

模型質量

Zamba2在標準語言建模評估集上表現非常出色，考慮到延遲和生成速度，其在小語言模型（≤8B）中，在質量和性能上都處于領先地位，主要原因如下：

1. 新型共享注意力架構可以讓更多的參數分配到Mamba2智能體的骨干網絡中，從而讓共享的Transformer模塊就能夠保留注意力計算中豐富的跨序列依賴性；新架構通過優化參數分配，使得智能體在處理圖像描述等任務時，能夠更好地理解和利用數據中的復雜關系。

2. 預訓練數據集達到了3萬億個token的規模，混合了Zyda數據和公開可用的數據，經過了嚴格的過濾和去重處理，確保了數據的高質量，在與現有的頂級開源預訓練數據集的比較中，數據處理也達到了最先進的水平。

3. 模型中還有一個特別的「退火」（annealing）預訓練階段，在處理100B個高質量token的過程中快速降低學習率，其中退火數據集是精心策劃和從多個高質量來源整理而來的，以確保質量。智能體在這個階段通過處理大量高質量的數據，能夠更快地學習和適應，從而提高其在圖像描述等任務上的表現。

由于我們的預訓練和退火數據集的卓越質量，Zamba2-7B智能體在每個訓練token上的表現非常出色，輕松超越了競爭對手模型的性能曲線。

Zamba2-7B智能體利用并擴展了初代的Zamba混合SSM-注意力架構，核心的Zamba架構由Mamba層構成的骨干網絡與一個或多個共享注意力層交錯組成（Zamba1有一個共享注意力層，Zamba2有兩個），注意力機制的權重共享，以最小化模型的參數成本。

研究人員發現，將輸入的原始模型嵌入與這個注意力塊進行連接可以提高性能，很可能是因為可以更好地保持了信息在網絡深度上的傳遞。

Zamba2架構還對共享的MLP應用了LoRA投影矩陣，以在每個模塊中獲得一些額外的表達能力，并允許每個共享模塊稍微專門化，以適應其獨特的位置，同時保持額外的參數開銷很小。

類似于在智能體的「大腦」中添加了一種特殊的「眼鏡」，使其能夠更清晰地看到每個數據點的獨特之處，同時保持整體的簡潔和高效。

通過這種方式，Zamba2-7B智能體在處理圖像描述等任務時，能夠更加精準地理解和生成內容。

模型實現了最先進的推理效率，包括延遲、吞吐量和內存使用，主要原因如下：

1. Mamba2模塊的效率極高，其吞吐量大約是同等參數Transformer模塊的4倍，也就意味著Mamba2模塊在處理數據時更快，能夠更迅速地完成圖像描述等智能體任務。

2. Mamba模塊只需要存儲較小的隱藏狀態，并且不需要KV緩存，所以只需要為共享注意力模塊的調用存儲KV狀態，就好像智能體在記憶信息時，不需要記住每一個細節，而是只記住最關鍵的部分，既節省了空間，也提高了效率。

3. 選擇的模型尺寸非常適合在現代硬件上進行并行處理（例如，GPU上的多個流式多處理器，CPU上的多個核心），像是在工廠里使用多條生產線同時工作，可以大大提高生產速度和效率。

這些設計使得該智能體在處理圖像描述等任務時，不僅速度快，而且資源消耗少，為用戶提供了高效且流暢的體驗。

使用基于Megatron-LM開發的內部訓練框架，在128個H100 GPU上進行了訓練了大約50天，表明即使在70億參數的規模上，前沿技術仍然是可及且可以超越的，即使是小團隊和適度預算也能實現。

Zamba2-7B智能體的開源許可證允許研究人員、開發者和公司使用。

今年5月，Zamba發布，開創性地結合了Mamba骨干網絡和單一共享注意力模塊的獨特架構，以最小的參數成本，保持了注意力機制的優勢，實現了比同類的Transformer模型更高的推理效率，并且在生成長序列時所需的內存量也大大減少。

論文鏈接：https://arxiv.org/pdf/2405.16712

Zamba的預訓練分為兩個階段：

1. 基于現有的網絡數據集預訓練；

2. 退火階段包括在高質量的指導性和合成數據集上對模型進行退火處理，其特點是學習率快速衰減。

Zamba智能體的架構設計上，由一系列標準的Mamba模塊構成骨干網絡，并與一個共享的注意力和多層感知機（MLP）模塊相連，其中共享模塊每6個Mamba模塊重復一次，但參數是共享的，使得Mamba能夠在相同的內存成本下利用更多的浮點運算（FLOPs）來提升性能。

輸入的嵌入始終與殘差流一起連接到共享注意力模塊，為模型提供了一個額外的路徑來記住輸入信息；在模塊處理完畢后，用一個可學習的線性投影將輸出映射回殘差流。

在推理和生成效率方面，Zamba智能體表現出色，雖然參數共享機制導致每個參數使用的FLOPs更多，但Zamba智能體的前向傳遞速度明顯快于7B規模的競品模型，隨著序列長度的增加，優勢更加明顯。

由于Zamba智能體的SSM骨干網絡，Mamba所需的KV緩存內存比其他類似規模的模型減少了很多，從而使Zamba智能體能夠更有效地生成內容，并在單個設備上實現更長的上下文。

責任編輯：張燕妮來源：新智元