Mamba架構(gòu)第一次做大！混合Transformer，打敗Transformer
取兩架構(gòu)之長，混合模型Jamba誕生
豐色 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
精彩精彩，第一個把爆火Mamba架構(gòu)真正擴(kuò)展到足夠大的工作來了。
520億參數(shù)，還是Mamba+Transformer混合架構(gòu)。
它的名字叫Jamba。
取兩種架構(gòu)之長，模型質(zhì)量和效率兼得，要吞吐量有吞吐量，要低內(nèi)存有低內(nèi)存。
初步跑分顯示：
- Jamba性能總體接近Mixtral 8x-7B，處理128k長上下文時吞吐量卻是其3倍。
- 一共支持256k上下文，而單張A100 GPU即可處理140k，直接拿下同等規(guī)模模型之最，高效又經(jīng)濟(jì)。
這項(xiàng)成果，來自以色列AI公司AI21labs。
Mamba原作者看了之后都激動轉(zhuǎn)發(fā)：
絕對的“大新聞”。
Mamba、Transformer，合體
由CMU和普林斯頓大學(xué)提出的Mamba，解決了Transformer的局限性（隨著推理上下文越長，模型內(nèi)存占用量越大，同時推理速度變慢，由此導(dǎo)致算力消耗巨大）。
但它也有自己的缺點(diǎn)——
在不關(guān)注整個上下文的情況下，Mamba的輸出質(zhì)量很差，尤其是在召回相關(guān)的任務(wù)上。
本著“既要也要”的原則，Jamba站出來提供兩全其美之作。
Jamba由Transformer、Mamba和MoE層組成，可同時優(yōu)化內(nèi)存、吞吐量和性能。
如下圖所示，為了集成兩種架構(gòu)，Jamba采用塊層（blocks-and-layers）組合的創(chuàng)新方法。
簡單來說，就是每個Jamba塊包含一個注意力層或一個Mamba層，再跟一個多層感知器MLP，總體比例保證為每八層一個Transformer層。
其次，Jamba利用MoE來增加模型參數(shù)的總量，同時簡化推理中使用的活動參數(shù)量。
最終模型容量高了，計算需求也沒有相應(yīng)的增加。
而為了在單張GPU（80GB）上最大限度地提高模型吞吐量，Jamba還優(yōu)化了所用MoE層和專家數(shù)量，最終為日常推理工作負(fù)載留出足夠內(nèi)存。
值得一提的是，在推理時，Jamba的MoE層僅需520億可用參數(shù)中的120億，就能同時保證比同等大小的僅Transformer模型更高效。
要知道，此前有人光是嘗試過擴(kuò)展Mamba，就沒能做到30億參數(shù)之上。
因此，除了成功合體Mamba和Transformer，Jamba也達(dá)成了第二大成就：
同類中第一個達(dá)到生產(chǎn)級規(guī)模和質(zhì)量的混合架構(gòu)（SSM混Transformer）（ps. Mamba就是一種狀態(tài)空間模型SSM）。
吞吐量和效率up
初步評估顯示，Jamba在吞吐量和效率等關(guān)鍵指標(biāo)上表現(xiàn)出色。
首先，Jamba可以在長上下文中提供3倍吞吐量，比Mixtral 8x7B等大小相當(dāng)?shù)腡ransformer模型都要高效。
如下圖所示，當(dāng)上下文窗口達(dá)到128k時，Jamba的每秒token數(shù)近乎1500，而此時表現(xiàn)最好的Mixtral 8x7B應(yīng)該才在500往上的樣子。
其次，在單張GPU上，Jamba最多可以容納140k上下文，經(jīng)濟(jì)又高效。
相比之下，Mixtral 8x7B為64k，Llama2 70B則僅為16k。
第三，Jamba的輸出質(zhì)量也得到了保證。
在如下一系列推理基準(zhǔn)上，4項(xiàng)中有3項(xiàng)它都拿下了SOTA。同時，在GSM8K等基準(zhǔn)上，Jamba即使沒有奪魁，也和SOTA模型打了個不相上下。
總體來說，Jamba的性能接近Mixtral 8x7B。
最后，作者提示，別忘了，這些都還只是初步改造后的結(jié)果，后續(xù)還有很多優(yōu)化空間（比如MoE并行、更快的Mamba實(shí)現(xiàn)）。所以到時性能會更強(qiáng)。
好消息：Jamba現(xiàn)在已經(jīng)上線Hugging Face，并且劃重點(diǎn)：采用apache-2.0許可。
（Jamba的指令版本則將很快通過AI21labs平臺上線。）
網(wǎng)友看完都感動哭了。
傳送門：
https://huggingface.co/ai21labs/Jamba-v0.1
參考鏈接：
[1]https://www.ai21.com/blog/announcing-jamba
[2]https://www.ai21.com/jamba
[3]https://twitter.com/AI21Labs/status/1773350888427438424?s=20
[4]https://twitter.com/tri_dao/status/1773418926518734957?s=20
相關(guān)閱讀
華為北大等聯(lián)手打造的Transformer竟在CV領(lǐng)域超過了CNN：多項(xiàng)底層視覺任務(wù)達(dá)到SOTA
多種任務(wù)共享一個Transformer 模塊。
