Jamba 的架構以大約 1:7 的比例交錯 Transformer 區塊(標準注意力)和 Mamba 區塊(選擇性狀態空間)——每七個 Mamba 層配一個注意力層。這兼取兩者之長:Mamba 層高效處理大部分序列(在序列長度上是線性的),而注意力層提供純 SSM 有時缺乏的全域 token 交互。結果:一個在 256K 上下文下可放入單個 80GB GPU 的模型,品質上與純 Transformer 模型匹配。
Jamba 還使用了混合專家(MoE),總共 520 億參數但每個 token 只有約 120 億活躍。這種 SSM + 注意力 + MoE 的組合是生產中最複雜的混合架構,展示了這些技術可以良好地組合。與同等品質的純 Transformer 相比,KV 快取記憶體減少 3 倍,對於長上下文工作負載的服務具有實際意義。