Zubnet AI學習Wiki › AI21 Labs
公司

AI21 Labs

別名:Jamba
一家以色列 AI 公司,以 Jamba 聞名——第一個結合 Transformer 注意力層與 Mamba SSM 層的生產級混合架構。AI21 由 AI 研究者(包括 Yoav Shoham)創立,自 2017 年以來一直在建造語言模型,早於 ChatGPT。其模型透過 API 和雲端供應商提供。

為什麼重要

AI21 Labs 之所以重要,是因為 Jamba 證明了混合 Transformer-SSM 架構在實踐中可行,而不僅僅存在於研究論文中。透過交錯注意力和 Mamba 層,Jamba 以比同等品質的純 Transformer 模型更低的記憶體使用量實現了 256K 上下文視窗。這種混合方法可能是 LLM 架構的未來。

深度解析

Jamba 的架構以大約 1:7 的比例交錯 Transformer 區塊(標準注意力)和 Mamba 區塊(選擇性狀態空間)——每七個 Mamba 層配一個注意力層。這兼取兩者之長:Mamba 層高效處理大部分序列(在序列長度上是線性的),而注意力層提供純 SSM 有時缺乏的全域 token 交互。結果:一個在 256K 上下文下可放入單個 80GB GPU 的模型,品質上與純 Transformer 模型匹配。

MoE 組件

Jamba 還使用了混合專家(MoE),總共 520 億參數但每個 token 只有約 120 億活躍。這種 SSM + 注意力 + MoE 的組合是生產中最複雜的混合架構,展示了這些技術可以良好地組合。與同等品質的純 Transformer 相比,KV 快取記憶體減少 3 倍,對於長上下文工作負載的服務具有實際意義。

相關概念

← 所有術語
← AI 隱私 AlexNet →
ESC