自2017年开创性研究论文《Attention is All You Need》发布以来,变换器(Transformers)一直是生成式人工智能(AI)领域的核心。然而,变换器并不是生成式AI的唯一有效方法。AI21 Labs推出了一种创新框架“Jamba”,旨在超越传统变换器。
Jamba结合了基于结构状态空间模型(SSM)的Mamba模型与变换器架构,以创建一个优化的生成式AI解决方案。“Jamba”代表联合注意力和Mamba架构,旨在发挥SSM和变换器的优势。该模型作为开源项目在Apache 2.0许可证下发布。
尽管Jamba并不打算取代现有的变换器基础大语言模型(LLMs),但预计将在特定应用中发挥重要补充作用。AI21 Labs表示,Jamba在生成推理任务中优于传统变换器模型,如HellaSwag基准测试所示。然而,在解决能力的关键基准测试大规模多任务语言理解(MMLU)中,Jamba尚未超过变换器模型。
AI21 Labs专注于企业应用的生成式AI,最近在2023年8月筹集了1.55亿美元以推进其项目。其中一项企业产品是Wordtune,它帮助组织生成符合其语气和品牌的内容。该公司在2023年声称,已成功与生成式AI巨头OpenAI竞争,获得企业客户。
在此之前,AI21 Labs的LLM技术主要基于变换器架构,包括其Jurassic-2 LLM家族,作为AI21 Studio自然语言处理(NLP)平台的一部分,通过API供企业集成使用。然而,Jamba标志着向结合SSM与变换器模型的混合方式转变。
尽管变换器在生成式AI中占据重要地位,但它们存在一定局限。例如,随着上下文窗口的扩大,推理速度往往减慢。AI21 Labs的研究人员解释,变换器的注意力机制会随着序列长度的增加而扩展,导致吞吐量减少,因为每个token依赖于整个前序序列。这使得长上下文应用效率低下。
另一个挑战是扩大变换器所需的内存占用。随着上下文长度的增长,其内存要求也随之增加,这使得在没有大量硬件资源的情况下处理长上下文或多个并行任务变得复杂。SSM方法旨在解决这些上下文和内存问题。
Mamba SSM架构最初由卡内基梅隆大学和普林斯顿大学的研究人员开发,旨在减少内存使用并利用不同的注意力机制来管理大上下文窗口。然而,它在输出质量方面无法与变换器模型相媲美。Jamba的混合方法结合了SSM的资源和上下文优化与变换器的输出能力。
AI21 Labs声称,Jamba模型提供256K的上下文窗口,并在长上下文处理上吞吐量是Mixtral 8x7B的三倍。值得注意的是,Jamba被定位为同类模型中唯一能在单个GPU上适应最多140K上下文的模型。
与Mixtral类似,Jamba也采用了混合专家模型(MoE)。但Jamba在其混合SSM变换器框架中利用MoE,实现了更高的优化水平。具体而言,Jamba的MoE层在推理过程中仅激活其可用的520亿参数中的120亿,使其效率高于同等规模的变换器模型。
目前,Jamba仍处于早期阶段,并尚未成为AI21 Labs的企业产品,不过该公司计划很快在AI21平台上推出一个测试版的教学版。