标签： AGI

引言

自2017年被提出以来，Transformer已成为AI大模型的主流架构，在语言建模领域长期占据主导地位。然而，随着模型规模不断扩大、需要处理的序列越来越长，Transformer的局限性也日益凸显。其中一个明显的缺陷是：Transformer模型中自注意力机制的计算量会随着上下文长度的增加呈平方级增长。

几个月前，一种名为Mamba的新型架构的出现打破了这一局面。Mamba能够随上下文长度的增加实现线性扩展，在中小型规模上已经实现了与Transformers匹敌甚至超越的性能。

仅仅六个月后，Mamba的原作者团队再接再厉，推出了更强大的Mamba 2。本文将详细介绍Mamba 2的核心创新点及其相比Transformer的优势。

Mamba 2的核心创新

1. SSD框架：连接SSM和Transformer

Mamba 2的核心贡献是提出了SSD（State Space Duality）框架。这一框架揭示了状态空间模型(SSM)与一类称为半可分矩阵的结构化矩阵族之间的等价性。通过SSD框架，研究者建立了SSM和Transformer之间的理论联系，为理解和改进序列模型开辟了新的方向。

SSD框架的主要内容包括：

• 展示了状态空间模型与半可分矩阵之间的等价性
• 改进了线性注意力理论，推广出新的结构化掩码注意力(SMA)族
• 证明了SSM和SMA有很大的交集，它们是对偶的
• 证明了任何具有快速循环形式的核注意方法都是SSM

2. 高效的SSD算法

基于SSD框架，研究者提出了一种新的高效算法来计算SSM。这种基于半可分离矩阵块分解的SSD算法，利用了SSM的线性递推和二次对偶形式，在各个效率维度上都取得了最优权衡。

与Mamba的实现相比，SSD算法的速度提高了2到8倍。同时，它还允许使用更大的循环状态大小(是Mamba的8倍甚至更高)，而几乎不影响速度。在长序列处理上，SSD算法的优势更加明显 – 在16K长度的序列上，它比优化过的softmax注意力实现(FlashAttention-2)快6倍。

3. 改进的架构设计

Mamba 2在网络架构上也做了一些改进：

• 从顺序生成变为并行生成SSM参数
• 引入分组值注意力(GVA)头结构
• 更适合张量并行等扩展方法

这些改进使得Mamba 2在保持模型表达能力的同时，显著提高了训练效率，特别是能够更好地利用现代加速器上的矩阵乘法单元。

Mamba 2的性能优势

1. 语言建模任务

研究者在Pile数据集上训练了一系列Mamba 2模型，结果显示Mamba 2在标准下游评估中匹配或超过了Mamba和开源的Transformers。

例如，在Pile上训练了3000亿token的2.7B参数Mamba 2模型，其性能超过了:

• 在同一数据集上训练的2.8B参数Mamba模型
• 2.8B参数的Pythia模型
• 6.9B参数的Pythia模型

这一结果表明，Mamba 2不仅能够与同等规模的Transformer模型相匹敌，甚至能够在更小的参数量下超越更大的Transformer模型。

2. 复杂关联回忆任务

研究团队在MQAR(multi-query associative recall)任务上对比了Mamba 2和Mamba 1的性能。MQAR是一种比文献中通常报告的版本更难的任务，要求模型具有更强的长程依赖建模能力。

实验结果显示，Mamba 2明显优于Mamba 1。研究者认为，性能提升的一个重要原因是Mamba 2使用了更大的状态大小(比Mamba 1大约16倍)。这说明Mamba 2在处理需要更大状态容量的任务上有显著优势。

3. 训练效率

Mamba 2在训练效率方面也有明显提升。研究者在与Mamba相同的设置中研究了Mamba 2的Chinchilla扩展法则，发现它在困惑度和实际运行时间方面均优于Mamba和Transformer++。

这意味着，在相同的计算资源下，Mamba 2能够更快地收敛到更好的性能，从而大幅提高模型训练的成本效益比。

Mamba 2的理论意义

Mamba 2不仅在实际性能上有显著提升，其背后的理论创新也具有重要意义：

1. SSD框架提供了状态空间模型、注意力机制和结构化矩阵之间丰富的联系，为未来的序列模型研究开辟了新的方向。
2. 通过建立SSM和Transformer之间的理论联系，Mamba 2为两种不同范式的模型架构搭建了桥梁，有助于研究者更好地理解和改进这两类模型。
3. SSD算法的提出为高效计算SSM提供了新的思路，这一算法不仅适用于Mamba系列模型，也可能被应用到其他基于SSM的模型中。

结论

Mamba 2的出现为序列建模领域带来了新的可能性。它不仅在性能上挑战了长期占据主导地位的Transformer，更重要的是提供了一种新的思路来构建和理解序列模型。

虽然目前Mamba 2主要在中小规模模型上展现出优势，但其线性扩展的特性使它在处理超长序列时具有巨大潜力。随着进一步的研究和优化，Mamba系列模型很可能在更多任务和更大规模上挑战Transformer的地位。

然而，需要注意的是，Transformer模型经过多年发展已经形成了成熟的生态系统。Mamba要真正取代Transformer还需要时间和更多的实践验证。未来，我们可能会看到Transformer和Mamba各自在不同场景下发挥优势，或者两者结合形成新的混合架构。

无论如何，Mamba 2的出现无疑为AI领域注入了新的活力，推动了序列模型的进一步发展。它的成功再次证明，在人工智能领域，创新永不止步。

参考文献

1. Gu, A., Dao, T. et al. (2024). Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality. arXiv preprint.
2. Zhou, W. et al. (2024). 再战Transformer！原作者带队的Mamba 2来了，新架构训练效率大幅提升. 腾讯云开发者社区.
3. Vaswani, A. et al. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems.
4. Gu, A., Dao, T. et al. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv preprint.

引言

人工智能领域正在探索通向通用人工智能(AGI)的道路,其中一个重要方向是开发”语言智能体”。语言智能体是复杂的大型语言模型(LLM)管道,涉及提示技术和工具使用方法。虽然语言智能体在许多现实世界任务中展现了令人印象深刻的能力,但当前语言智能体研究的一个根本局限性在于它们是以模型为中心或以工程为中心的。也就是说,语言智能体的提示、工具和管道的进展需要人类专家而非自动从数据中学习进行大量的手动工程努力。

我们认为,语言智能体从以模型为中心或以工程为中心向以数据为中心的转变,即语言智能体在环境中自主学习和进化的能力,是它们可能实现AGI的关键。

智能体符号学习

为了实现这一目标,本文介绍了智能体符号学习,这是一个系统框架,能够使语言智能体以数据为中心的方式使用符号优化器自我优化。具体而言,我们将智能体视为符号网络,其中可学习的权重由提示、工具以及它们的堆叠方式定义。智能体符号学习旨在通过模仿连接主义学习中的两个基本算法来优化语言智能体中的符号网络:反向传播和梯度下降。与处理数值权重不同,智能体符号学习使用自然语言模拟权重、损失和梯度。

智能体符号学习框架的主要组成部分包括:

1. 智能体管道(Agent Pipeline):类似于神经网络中的计算图,表示智能体处理输入数据的节点(或步骤)序列。
2. 节点(Node):智能体管道中的单个步骤,类似于神经网络中的层。每个节点接收输入,使用LLM、提示和工具进行处理,然后输出结果。
3. 轨迹(Trajectory):存储前向传播过程中的所有信息,包括每个节点的输入、输出、提示和工具使用情况。
4. 语言损失(Language Loss):通过精心设计的提示实现的自然语言损失函数,用于评估智能体的输出结果。
5. 语言梯度(Language Gradient):用于更新智能体中每个组件的文本分析和反思。

智能体符号学习的工作流程如下:

1. 前向传播:执行标准智能体操作,同时将信息存储在轨迹中。
2. 语言损失计算:使用精心设计的提示模板计算语言损失。
3. 语言梯度反向传播:从最后一个节点到第一个节点迭代,计算每个节点的梯度。
4. 基于语言梯度的更新:使用符号优化器更新每个节点中的提示和工具,以及整体智能体管道。

实验结果

为了验证智能体符号学习框架的有效性,我们在标准LLM基准测试和复杂的智能体任务上进行了一系列实验。

标准LLM基准测试结果

在HotpotQA、MATH和HumanEval等标准LLM基准测试上,智能体符号学习框架始终优于所有比较方法。特别是在MATH(一个竞赛级别的基准测试)上,性能提升尤为显著。相比之下,传统的基于LLM的提示优化方法和基于搜索的提示优化方法并不稳定:在某些情况下会带来良好的性能提升,但在其他情况下会导致显著的性能下降。这表明智能体符号学习框架更加稳健,能够更有效地优化语言智能体的整体性能。

复杂任务结果

在软件开发和创意写作这两个复杂任务上,我们的方法显著优于所有比较基线,性能差距比在常规LLM基准测试上更大。有趣的是,我们的方法在创意写作任务上甚至超过了精心设计的思维树(Tree-of-Thought)提示工程和推理算法。我们发现,我们的方法成功地找到了计划、写作和修订的管道,并且每个步骤中的提示都得到了很好的优化。我们还发现,智能体符号学习框架恢复了类似于MetaGPT(一个专门为软件开发设计的智能体框架)中开发的标准操作程序。

案例研究与分析

我们展示了智能体符号学习框架在创意写作任务上的优化动态案例研究。结果表明,我们的方法能够有效地进行提示工程并设计智能体管道,方式类似于人类专家开发语言智能体的方式。

此外,我们发现智能体系统的初始化对最终性能有不可忽视的影响,就像神经网络的初始化对训练很重要一样。总的来说,我们发现以最简单的方式初始化智能体,然后让符号优化器进行优化通常是有帮助的。相比之下,如果初始智能体系统过度工程化,性能往往会变得不稳定。

结论

本文介绍了智能体符号学习,这是一个能够联合优化智能体系统内所有符号组件的智能体学习框架。智能体符号学习框架借鉴了标准连接主义学习程序的灵感来进行符号学习。它使用基于语言的损失、梯度和优化器来优化提示、工具和智能体管道,以提高智能体系统的整体性能。

我们的框架是首次尝试优化能够使用复杂管道解决复杂现实世界任务的智能体。我们的框架使语言智能体能够”从数据中学习”并在创建和部署后在野外进行”自我进化”。我们进行了几个概念验证实验,表明智能体符号学习框架可以有效地优化不同任务复杂度的智能体。

我们相信,从以模型为中心到以数据为中心的智能体研究转变是朝着通用人工智能迈进的一个有意义的步骤。我们开源了智能体符号学习框架的代码和提示,以加速这一转变。

参考文献

1. Zhou, W. et al. (2023). Symbolic Learning Enables Self-Evolving Agents. arXiv preprint arXiv:2406.18532v1.
2. Brown, T. et al. (2020). Language models are few-shot learners. Advances in Neural Information Processing Systems, 33, 1877-1901.
3. Park, J. S. et al. (2023). Generative agents: Interactive simulacra of human behavior. arXiv preprint.
4. Khattab, O. et al. (2023). Dspy: Compiling declarative language model calls into self-improving pipelines. arXiv preprint arXiv:2310.03714.
5. Zhang, S. et al. (2024). Offline training of language model agents with functions as learnable weights. arXiv preprint.