之前引爆了AI圈的Mamba架構(gòu)，今天又推出了一版超強(qiáng)變體！

人工智能獨(dú)角獸AI21 Labs剛剛開(kāi)源了Jamba，世界上第一個(gè)生產(chǎn)級(jí)的Mamba大模型！

Jamba在多項(xiàng)基準(zhǔn)測(cè)試中表現(xiàn)亮眼，與目前最強(qiáng)的幾個(gè)開(kāi)源Transformer平起平坐。

特別是對(duì)比性能最好的、同為MoE架構(gòu)的Mixtral 8x7B，也互有勝負(fù)。

具體來(lái)說(shuō)它——

• 是基于全新SSM-Transformer混合架構(gòu)的首個(gè)生產(chǎn)級(jí)Mamba模型
• 與Mixtral 8x7B相比，長(zhǎng)文本處理吞吐量提高了3倍
• 實(shí)現(xiàn)了256K超長(zhǎng)上下文窗口
• 是同等規(guī)模中，唯一一個(gè)能在單張GPU上處理140K上下文的模型
• 以Apache 2.0開(kāi)源許可協(xié)議發(fā)布，開(kāi)放權(quán)重

之前的Mamba因?yàn)楦鞣N限制，只做到了3B，還被人質(zhì)疑能否接過(guò)Transformer的大旗，而同為線性RNN家族的RWKV、Griffin等也只擴(kuò)展到了14B。

——Jamba這次直接干到52B，讓Mamba架構(gòu)第一次能夠正面硬剛生產(chǎn)級(jí)別的Transformer。

Jamba在原始Mamba架構(gòu)的基礎(chǔ)上，融入了Transformer的優(yōu)勢(shì)來(lái)彌補(bǔ)狀態(tài)空間模型（SSM）的固有局限性。

可以認(rèn)為，這實(shí)際上是一種新的架構(gòu)——Transformer和Mamba的混合體，最重要的是，它可以在單張A100上運(yùn)行。

它提供了高達(dá)256K的超長(zhǎng)上下文窗口，單個(gè)GPU就可以跑140K上下文，而且吞吐量是Transformer的3倍！

與Transformer相比，看Jamba如何擴(kuò)展到巨大的上下文長(zhǎng)度，非常震撼

Jamba采用了MoE的方案，52B中有12B是活躍參數(shù)，目前模型在Apache 2.0下開(kāi)放權(quán)重，可以在huggingface上下載。

模型下載：https://huggingface.co/ai21labs/Jamba-v0.1

LLM新里程碑

Jamba的發(fā)布標(biāo)志著LLM的兩個(gè)重要里程碑：

一是成功將Mamba與Transformer架構(gòu)相結(jié)合，二是將新形態(tài)的模型（SSM-Transformer）成功提升到了生產(chǎn)級(jí)的規(guī)模和質(zhì)量。

當(dāng)前性能最強(qiáng)的大模型全是基于Transformer的，盡管大家也都認(rèn)識(shí)到了Transformer架構(gòu)存在的兩個(gè)主要缺點(diǎn)：

內(nèi)存占用量大：Transformer的內(nèi)存占用量隨上下文長(zhǎng)度而擴(kuò)展。想要運(yùn)行長(zhǎng)上下文窗口，或大量并行批處理就需要大量硬件資源，這限制了大規(guī)模的實(shí)驗(yàn)和部署。

隨著上下文的增長(zhǎng)，推理速度會(huì)變慢：Transformer的注意力機(jī)制導(dǎo)致推理時(shí)間相對(duì)于序列長(zhǎng)度呈平方增長(zhǎng)，吞吐會(huì)越來(lái)越慢。因?yàn)槊總€(gè)token都依賴(lài)于它之前的整個(gè)序列，所以要做到超長(zhǎng)上下文就變得相當(dāng)困難。

年前，來(lái)自卡內(nèi)基梅隆和普林斯頓的兩位大佬提出了Mamba，一下子就點(diǎn)燃了人們的希望。

Mamba以SSM為基礎(chǔ)，增加了選擇性提取信息的能力、以及硬件上高效的算法，一舉解決了Transformer存在的問(wèn)題。

這個(gè)新領(lǐng)域馬上就吸引了大量的研究者，arXiv上一時(shí)涌現(xiàn)了大量關(guān)于Mamba的應(yīng)用和改進(jìn)，比如將Mamba用于視覺(jué)的Vision Mamba。

不得不說(shuō)，現(xiàn)在的科研領(lǐng)域?qū)嵲谑翘砹耍裈ransformer引入視覺(jué)（ViT）用了三年，但Mamba到Vision Mamba只用了一個(gè)月。

不過(guò)原始Mamba的上下文長(zhǎng)度較短，加上模型本身也沒(méi)有做大，所以很難打過(guò)SOTA的Transformer模型，尤其是在與召回相關(guān)的任務(wù)上。

Jamba于是更進(jìn)一步，通過(guò)Joint Attention and Mamba架構(gòu)，整合了Transformer、Mamba、以及專(zhuān)家混合（MoE）的優(yōu)勢(shì)，同時(shí)優(yōu)化了內(nèi)存、吞吐量和性能。

Jamba是第一個(gè)達(dá)到生產(chǎn)級(jí)規(guī)模（52B參數(shù)）的混合架構(gòu)。

如下圖所示，AI21的Jamba架構(gòu)采用blocks-and-layers的方法，使Jamba能夠成功集成這兩種架構(gòu)。

每個(gè)Jamba塊都包含一個(gè)注意力層或一個(gè)Mamba層，然后是一個(gè)多層感知器（MLP）。

Jamba的第二個(gè)特點(diǎn)，是利用MoE來(lái)增加模型參數(shù)的總數(shù)，同時(shí)簡(jiǎn)化推理中使用的活動(dòng)參數(shù)的數(shù)量，從而在不增加計(jì)算要求的情況下提高模型容量。

為了在單個(gè)80GB GPU上最大限度地提高模型的質(zhì)量和吞吐量，研究人員優(yōu)化了使用的MoE層和專(zhuān)家的數(shù)量，為常見(jiàn)的推理工作負(fù)載留出足夠的內(nèi)存。

對(duì)比Mixtral 8x7B等類(lèi)似大小的基于Transformer的模型，Jamba在長(zhǎng)上下文上做到了3倍的加速。

Jamba將在不久之后加入NVIDIA API目錄。

長(zhǎng)上下文又出新選手

最近，各大公司都在卷長(zhǎng)上下文。

具有較小上下文窗口的模型，往往會(huì)忘記最近對(duì)話的內(nèi)容，而具有較大上下文的模型則避免了這種陷阱，可以更好地掌握所接收的數(shù)據(jù)流。

不過(guò)，具有長(zhǎng)上下文窗口的模型，往往是計(jì)算密集的。

初創(chuàng)公司AI21 Labs的生成式模型就證明，事實(shí)并非如此。

Jamba在具有至少80GB顯存的單個(gè)GPU（如A100）上運(yùn)行時(shí)，可以處理多達(dá)140,000個(gè)token。

這相當(dāng)于大約105,000字，或210頁(yè)，是一本長(zhǎng)度適中的長(zhǎng)篇小說(shuō)的篇幅。

相比之下，Meta Llama 2的上下文窗口，只有32,000個(gè)token，需要12GB的GPU顯存。

按今天的標(biāo)準(zhǔn)來(lái)看，這種上下文窗口顯然是偏小的。

對(duì)此，有網(wǎng)友也第一時(shí)間表示，性能什么的都不重要，關(guān)鍵的是Jamba有256K的上下文，除了Gemini，其他人都沒(méi)有這么長(zhǎng)，——而Jamba可是開(kāi)源的。

Jamba真正的獨(dú)特之處

從表面上看，Jamba似乎并不起眼。

無(wú)論是昨天風(fēng)頭正盛的DBRX，還是Llama 2，現(xiàn)在都已經(jīng)有大量免費(fèi)提供、可下載的生成式AI模型。

而Jamba的獨(dú)特之處，是藏在模型之下的：它同時(shí)結(jié)合了兩種模型架構(gòu)——Transformer和狀態(tài)空間模型SSM。

一方面，Transformer是復(fù)雜推理任務(wù)的首選架構(gòu)。它最核心的定義特征，就是「注意力機(jī)制」。對(duì)于每條輸入數(shù)據(jù)，Transformer會(huì)權(quán)衡所有其他輸入的相關(guān)性，并從中提取以生成輸出。

另一方面，SSM結(jié)合了早前AI模型的多個(gè)優(yōu)點(diǎn)，例如遞歸神經(jīng)網(wǎng)絡(luò)和卷積神經(jīng)網(wǎng)絡(luò)，因此能夠?qū)崿F(xiàn)長(zhǎng)序列數(shù)據(jù)的處理，且計(jì)算效率更高。

雖然SSM有自己的局限性。但一些早期的代表，比如由普林斯頓和CMU提出的Mamba，就可以處理比Transformer模型更大的輸出，在語(yǔ)言生成任務(wù)上也更優(yōu)。

對(duì)此，AI21 Labs產(chǎn)品負(fù)責(zé)人Dagan表示——

雖然也有一些SSM模型的初步樣例，但Jamba是第一個(gè)生產(chǎn)規(guī)模的商業(yè)級(jí)模型。

在他看來(lái)，Jamba除了創(chuàng)新性和趣味性可供社區(qū)進(jìn)一步研究，還提供了巨大的效率，和吞吐量的可能性。

目前，Jamba是基于Apache 2.0許可發(fā)布的，使用限制較少但不能商用。后續(xù)的微調(diào)版本，預(yù)計(jì)會(huì)在幾周內(nèi)推出。

即便還處在研究的早期階段，但Dagan斷言，Jamba無(wú)疑展示了SSM架構(gòu)的巨大前景。

「這種模型的附加價(jià)值——無(wú)論是因?yàn)槌叽邕€是架構(gòu)的創(chuàng)新——都可以很容易地安裝到單個(gè)GPU上?！?/span>