自2017年突破性研究論文《注意力是你所需要的一切》發表以來,transformers在生成式AI領域成為焦點。然而,transformers並非生成式AI的唯一可行方法。AI21 Labs推出了一種名為“Jamba”的新框架,旨在超越傳統的transformers。
Jamba結合了基於結構狀態空間模型(SSM)的Mamba模型和transformer架構,以創造出優化的生成式AI解決方案。“Jamba”代表聯合注意力和Mamba架構,旨在充分利用SSM和transformers的優勢。此模型以Apache 2.0許可證作為開源發布。
雖然Jamba尚未取代現有的基於transformer的大型語言模型(LLMs),但它預計將在某些應用中作為有價值的補充。AI21 Labs表示,Jamba在生成推理任務上可能表現優於傳統的transformer模型,這在HellaSwag等基準測試中得到了證明。然而,在評估問題解決能力的關鍵基準如大規模多任務語言理解(MMLU)中,Jamba仍未超越transformer模型。
AI21 Labs專注於企業應用的生成式AI,於2023年8月成功籌集1.55億美元以推進其計劃。其中一項企業產品是Wordtune,旨在幫助組織生成與其語調和品牌一致的內容。該公司在2023年報告中提到,它已成功與生成式AI巨頭OpenAI競爭,獲得企業客戶。
傳統上,AI21 Labs的LLM技術利用transformer架構,包括其Jurassic-2 LLM系列,這是AI21 Studio自然語言處理(NLP)平台的一部分,並可通過API進行企業整合。然而,Jamba代表了向混合SSM和transformer模型的轉變。
儘管transformers在生成式AI中扮演了重要角色,但也存在一些限制。最大問題之一就是隨著上下文窗口的擴大,推理速度會減慢。AI21 Labs的研究人員解釋道,transformer的注意力機制會隨著序列長度的增加而擴展,這會導致吞吐量降低,因為每個token都依賴於前面整個序列的資訊,這使得在長上下文中的應用效率低下。
另一個挑戰是擴展transformers所需的巨大內存。隨著上下文長度的增加,它們的內存需求也隨之增長,這使得在沒有大量硬體資源的情況下,處理長上下文或多個平行任務變得更加複雜。SSM方法旨在解決這些上下文和內存問題。
Mamba SSM架構最初由卡內基梅隆大學和普林斯頓大學的研究人員開發,設計上需要更少的內存並使用不同的注意力機制來管理大型上下文窗口。然而,它在輸出質量上未能達到transformer模型的水準。Jamba的混合方法結合了SSM在資源和上下文上的優化以及transformer的輸出能力。
AI21 Labs聲稱,Jamba模型具有256K的上下文窗口,並在長上下文中提供三倍於Mixtral 8x7B的吞吐量。值得注意的是,Jamba被定位為其大小類別中唯一一個能在單個GPU上支持高達140K上下文的模型。
與Mixtral類似,Jamba也整合了專家混合模型(MoE)。不過,Jamba在其混合SSM transformer框架內使用MoE,使優化水平更高。具體來說,據AI21 Labs稱,Jamba的MoE層在推理過程中僅激活其可用的520億參數中的120億,這使其比等效大小的僅基於transformer的模型更高效。
目前,Jamba仍處於早期階段,尚未成為AI21 Labs的企業產品之一,不過該公司計劃在AI21平台上很快推出一個教學版本的測試版。