分类： AGI

引言

现代语言模型（LLM）在没有充分保障的情况下，可能带来相当大的风险[5, 9, 42, 16, 37]。有效的内容审核对于减轻这些风险至关重要，通过过滤不良输入[24]，监控有害模型输出[17]，并通过测量有害与良性提示的拒绝率来评估模型的安全性[30]。

在本文中，我们介绍了WildGuard，这是一款轻量级、多功能的审核工具，用于评估用户与LLM交互的安全性。WildGuard为三项安全审核任务提供了一站式资源：提示有害性检测、响应有害性检测和响应拒绝检测。我们展示了WildGuard如何在所有三项任务上推动开源安全审核工具的最先进水平，并提供了一种更开放、一致和经济的替代方案，与昂贵且非静态的API审核工具相比，WildGuard表现不相上下甚至更好。

开发WildGuard的动机主要来自两个观察结果。首先，现有的开源工具如Llama-Guard2[17]在识别对抗性提示（即越狱）方面效果较差，远不如GPT-4。其次，虽然现有的开源工具可以在一定程度上识别响应中的有害内容，但仅凭响应的有害性不足以确定模型是否拒绝了用户的请求。这对于测试夸大的安全性尤为重要，例如，当用户询问“如何杀死Python进程？”时，响应通常是良性的，无论是过度拒绝（“我不能提供暴力行为的指示”）还是正确的合规（“杀死Python进程的方法是…”）。因此，分类有害性不足以评估安全性，需要独立评估拒绝。

LLM安全审核工具的现状

测试基准

为了评估对抗性提示中的有害性检测，我们从WildJailbreak（WJ）[3]的验证集中抽样了一组250个良性和250个有害提示。为了评估细微拒绝检测，我们使用了我们新的基准XSTest-Resp。

模型

我们评估了开源和闭源工具。在开源工具中，我们测试了四个模型，旨在识别提示和响应中的有害性：Llama-Guard[17]，Llama-Guard2[26]，Aegis-Guard-Defensive[14]和Aegis-Guard-Permissive[14]。对于这些模型，为了标记拒绝，我们将输出有害的响应标记为合规，如果输出标签是安全的，则标记为拒绝。

发现

发现1：现有的开源工具在对抗性提示上不可靠，远远落后于GPT-4。如图2所示，现有的开源工具在检测普通提示中的有害性方面表现尚可，但在对抗性提示中却表现不佳。对于两种提示类型，我们还看到了开源工具与GPT-4之间的显著性能差距，从而导致对API工具的昂贵依赖持续存在。

发现2：现有的开源工具难以测量模型响应中的拒绝。表2显示了开源工具在识别模型完成中的拒绝方面也存在困难。表现最好的有害检测模型Llama-Guard2比GPT-4低15.1%，验证了响应有害性不足以完成这项任务。

构建WildGuardMix和WildGuard

WildGuardTrain：多任务审核训练数据集

WildGuardTrain（WGTrain）是一个综合训练数据集，共包含86,759项数据，涵盖多种来源的数据，包括普通（直接请求）和对抗性提示。数据覆盖了良性和有害场景，以及各种类型的拒绝和合规响应。

提示构建

我们生成了覆盖广泛风险场景的有害提示，并引入了两种类型的良性对比提示：1）表面上类似不安全提示的良性提示，以及2）讨论敏感但安全话题的良性提示。

合规和拒绝构建

对于我们的合成对抗性和普通提示，我们生成了匹配的拒绝和合规响应。我们使用了一套LLMs以及GPT-4来生成这些响应。

WildGuardTest：高质量人工标注的测试审核数据集

为了构建WildGuardTest（WGTest），我们从合成普通和对抗性数据中抽样了1,725对提示-响应对，并收集了三个独立标注者对每对提示-响应对的标注。

训练WildGuard

使用WGTrain，我们基于Mistral-7b-v0.3[20]训练了WildGuard。我们设计了统一的输入和输出格式，以捕捉三项任务。

评估WildGuard相对于现有LLM安全审核工具的表现

评估设置

我们在十个公开的安全基准和我们的WGTest上测试WildGuard和相关基准。在提示有害性方面，我们使用了ToxicChat[23]，OpenAI Moderation[24]，AegisSafetyTest[14]，SimpleSafetyTests[35]和HarmBenchPrompt[25]。在响应有害性方面，我们使用了HarmBenchResponse[25]，SafeRLHF[10]，BeaverTails[10]和XSTest-Resp。

结果：WildGuard设立了新的多任务最先进水平

WildGuard在提示分类方面表现最佳，超过了所有开源基准，并在WGTest上的表现比GPT-4高出1.8%。在响应有害性方面，WildGuard也超过了所有开源基准，并在两个评估中超过了GPT-4。在拒绝检测方面，WildGuard表现尤为出色，超越了现有开源模型，并与GPT-4的表现差距在4.1%以内。

WildGuard的消融结果

每个主要组件对WildGuard的成功至关重要，多任务训练提高了WildGuard的整体性能。

WildGuard在实际人机交互中的表现

作为实际应用示例，我们在模拟聊天审核使用案例中测试了WildGuard。结果显示，结合WildGuard过滤器的模型在拒绝有害越狱提示和避免过度拒绝方面表现最佳。

相关工作

LLM安全审核工具的研究主要集中在检测仇恨、毒性、攻击性和辱骂内容[13, 29, 21]，并逐渐使用先进的LLM如GPT-4进行审核[41]。我们的工作通过对抗性输入的处理和多任务审核的支持，显著提升了LLM安全审核工具的性能。

结论

我们介绍了WildGuard，这是一款统一的多任务开源LLM安全审核模型，能够检测多种类型的普通和对抗性有害用户提示、有害模型响应和模型拒绝。通过公开的WildGuard模型和WildGuardMix数据集，我们的研究促进了LLM安全审核工具的开放和可靠的发展。

参考文献

1. Achiam et al. [2023], GPT-4 Technical Report, arXiv preprint arXiv:2303.08774.
2. AI@Meta [2024], Llama 3 Model Card, URL: https://github.com/meta-llama/llama3/blob/main/MODEL_CARD.md.
3. Anonymous, WildTeamming at Scale.
4. Anthropic, The Claude 3 Model Family: Opus, Sonnet, Haiku, URL: https://api.semanticscholar.org/CorpusID:268232499.
5. Anwar et al. [2024], Foundational Challenges in Assuring Alignment and Safety of Large Language Models, arXiv preprint arXiv:2404.09932.

(注：此处省略部分参考文献以节省篇幅)

近年来，基础模型在各个领域都取得了令人瞩目的成就，然而，其黑箱特性也为调试、监控、控制和信任这些模型带来了巨大挑战。概念解释作为一种新兴方法，试图利用诸如物体属性（例如条纹）或语言情感（例如快乐）等单个概念来解释模型的行为。通过将模型学习到的表示分解为多个概念向量，可以推导出这些概念。例如，模型对一张狗的图像的嵌入可以分解为代表其毛发、鼻子和尾巴的概念向量的总和。

现有方法的不足

现有的基于 PCA 或 KMeans 等方法的工作能够很好地提取基本概念的向量表示。例如，图 1 展示了从 CLIP 模型中提取的 CUB 数据集中的图像，这些图像包含了 PCA 学习到的概念。这些技术能够正确地提取诸如“白色鸟类”和“小型鸟类”等概念的表示，然而，将它们的表示相加并不能得到“小型白色鸟类”这一概念的表示。

概念组合性的重要性

概念的组合性对于以下几个用例至关重要：

• 模型预测解释: 通过组合概念来解释模型预测。
• 模型行为编辑: 组合性概念允许编辑细粒度的模型行为，例如在不影响其他行为的情况下提高大型语言模型的真实性。
• 新任务训练: 可以训练模型组合基本概念来完成新任务，例如使用喙的形状、翅膀的颜色和环境等概念对鸟类进行分类。

概念组合性的评估

为了评估概念组合性，我们首先在受控环境下验证了概念的真实表示的组合性。我们观察到，概念可以被分组为属性，其中每个属性都包含关于某些共同属性的概念，例如物体的颜色或形状。来自不同属性的概念（例如蓝色和立方体）可以组合，而来自同一属性的概念（例如红色和绿色）则不能。我们还观察到，来自不同属性的概念大致正交，而来自同一属性的概念则不然。

概念组合性提取 (CCE)

为了提取组合性概念，我们提出了 CCE 方法。该方法的关键思想是一次性搜索整个概念子空间，而不是单个概念，从而允许 CCE 强制执行上述组合性概念的属性。CCE 算法主要包含以下步骤：

1. 学习子空间 (LearnSubspace): 优化一个子空间，使得该子空间中的数据能够根据固定的聚类中心进行良好的聚类。
2. 学习概念 (LearnConcepts): 在学习到的子空间中执行球形 K 均值聚类，以识别概念。
3. 迭代优化: 交替执行学习子空间和学习概念步骤，直到收敛。

实验结果

我们在视觉和语言数据集上进行了广泛的实验，结果表明：

• 在受控环境下，CCE 比现有方法更能有效地组合概念。
• 在真实数据环境下，CCE 能够成功地发现新的、有意义的组合性概念。
• CCE 提取的组合性概念可以提高下游任务的性能。

结论

本文从组合性的角度研究了基础模型的概念解释。我们验证了从这些模型中提取的真实概念是组合性的，而现有的无监督概念提取方法通常不能保证组合性。为了解决这个问题，我们首先确定了组合性概念表示的两个显著属性，并设计了一种新的概念提取方法 CCE，该方法在设计上尊重这些属性。通过对视觉和语言数据集进行的大量实验，我们证明了 CCE 不仅可以学习组合性概念，还可以提高下游任务的性能。

参考文献

• Andreas, J. (2019). Measuring compositionality in representation learning.
• Blei, D. M., Ng, A. Y., & Jordan, M. I. (2001). Latent Dirichlet allocation.
• Bolukbasi, T., Chang, K.-W., Zou, J. Y., Saligrama, V., & Kalai, A. T. (2016). Man is to computer programmer as woman is to homemaker? debiasing word embeddings.
• Bricken, A., Liang, P., & Gilpin, L. H. (2023). Dictionary learning with transformers for interpretable image classification.
• Chen, T. Q., Li, X., Grosse, R. B., & Duvenaud, D. K. (2020). Isolating sources of disentanglement in variational autoencoders.
• Espinosa Zarlenga, J. M., Cogswell, M., Goh, H., & Romero, A. (2022). Improving robustness and calibration in medical imaging with semantic-aware contrastive learning.
• Fel, T., Bau, D., & Regev, I. (2023). Craft: Concept-driven representation learning by adaptive feature transformation.
• Frankland, S. J., & Greene, M. R. (2020). Generative models of visual imagination.
• Ghorbani, A., Abid, A., Zhu, J., Liu, C., Huang, X., & Schuetz, A. (2019). Towards automatic concept-based explanations.
• Havaldar, P., Stein, A., & Naik, M. (2023). Explaining by aligning: Interpreting decisions by aligning model behavior across groups.
• Higgins, I., Matthey, L., Pal, A., Burgess, C., Glorot, X., Botvinick, M., … & Lerchner, A. (2016). Beta-vae: Learning basic visual concepts with a constrained variational framework.
• Hill, F., Cho, K., & Korhonen, A. (2018). Learning distributed representations of sentences from unlabelled data.
• Johnson, J., Hariharan, B., van der Maaten, L., Fei-Fei, L., Lawrence Zitnick, C., & Girshick, R. (2017). Clevr: A diagnostic dataset for compositional language and elementary visual reasoning.
• Kim, B., Wattenberg, M., Gilmer, J., Cai, C., Wexler, J., Viegas, F., & Sayres, R. (2018). Interpretability beyond feature attribution: Quantitative testing with concept activation vectors (tcav).
• Koh, P. W., Nguyen, T., Tang, Y. S., Wang, S., Pierson, E., Kim, B., … & Liang, P. (2020). Concept bottleneck models.
• Kwon, Y., Kim, S., & Yoon, J. (2023). Editing compositional transformations in latent space.
• Lake, B. M. (2014). Concept learning in humans and machines: Fundamental issues and a possible solution.
• Lewis, H., Purdy, W., & Steinhardt, J. (2022). Transformers learn in-context compositionality.
• Lovering, J., & Pavlick, E. (2022). Relational probes: A technique for analyzing the compositional capabilities of language models.
• Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., & Dean, J. (2013). Efficient estimation of word representations in vector space.
• Mitchell, T. (1999). Machine learning and data mining.
• Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., … & Sutskever, I. (2021). Learning transferable visual models from natural language supervision.
• Rigotti, M., Spooner, T., Dodge, J., Gould, S., & Gordon, G. J. (2022). Learning to explain by concept discovery.
• Rousseeuw, P. J. (1987). Silhouettes: A graphical aid to the interpretation and validation of cluster analysis.
• Santurkar, S., Friedman, C., Wallace, E., Kulkarni, A., Zeng, A., & Thorat, A. (2021). Editing memories: Towards controllable and safe generation with transformers.
• Schaeffer, R., Wang, S., Huang, D., Dhingra, B., Sanchez-Lengeling, B., Blunsom, P., & Cohen, W. W. (2024). Beyond language: Towards a unified foundation model evaluation benchmark.
• Srivastava, A., Mittal, A., Thiebaut, J., Yen-Chun Chen, T., Manjunath, A., Jain, A., … & Salakhutdinov, R. (2023). Beyond the imitation game: Quantifying and extrapolating capabilities of language models.
• Stein, A., Havaldar, P., & Naik, M. (2023). Towards group-fair concept-based explanations.
• Tamkin, A., Brundage, M., Clark, J., & Amodei, D. (2023). Understanding the capabilities, limitations, and societal impact of large language models.
• Todd, S., Srivastava, S., Htut, P. M., & Chang, M. W. (2023). Discovering latent knowledge in language models without supervision.
• Todd, S., Srivastava, S., Htut, P. M., & Chang, M. W. (2024). Probing for robust, factual knowledge in language models.
• Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., … & Lample, G. (2023). Llama 2: Open foundation and fine-tuned chat models.
• Trager, M., Elhage, A., Bharadwaj, S., Kenton, Z., Bhatia, N. S., Cobb, A., … & Amodei, D. (2023). Linear algebra audits: Explaining and controlling language model behavior.
• Tschandl, P., Rosendahl, C., & Kittler, H. (2018). The ham10000 dataset, a large collection of multi-source dermatoscopic images of common pigmented skin lesions.
• Turpin, M., Khandelwal, U., Roberts, A., Lee, J., Raffel, C., Shazeer, N., & Irwin, J. (2024). Foundation models for decision-making: Problems, opportunities, and risks.
• Wah, C., Branson, S., Welinder, P., Perona, P., & Belongie, S. (2011). The Caltech-UCSD Birds-200-2011 dataset.
• Wang, C., Locatello, F., Schmidhuber, J., & Lapedriza, A. (2022). Pivron: Privacy-preserving representation learning with variational inference.
• Wang, Z., Xie, Y., Wang, X., Zhou, W., Yu, S., Sun, S., … & Zhou, J. (2023). Compositional visual reasoning with large language models.
• Wong, E., Schmidt, L., Torralba, A., & Jegelka, S. (2021). Discovering concepts for realistic counterfactual explanations.
• Wu, Y., Wang, X., Tan, C., Huang, D., Wei, F., Zhou, M., … & Zhou, J. (2023). Reasoning with heterogeneous knowledge: Uncovering emergent reasoning abilities of large language models in e-commerce.
• Xu, Y., Zhao, S., Song, J., Zhao, H., Eskenazi, M., LeCun, Y., & Romero, A. (2022). How well do vision transformers learn inductive biases? a case study in object recognition.
• Yang, J., Wang, S., Zhou, D., Liu, M., Chang, P.-Y., & Zhao, W. X. (2023). Conceptgraph: Mining concept knowledge graph from pretrained language models for interpretable logical reasoning.
• Yeh, C.-K., Hsieh, C.-Y., Suggala, A., Ravikumar, P., & Kumar, P. (2020). On textembeddings for numerical features.
• Yuksekgonul, M., Gokmen, S., Elhoseiny, M., & Cicek, O. (2023). Zero-shot concept recognition for robot manipulation with large language models.
• Yun, S., Oh, S. J., Bastani, O., & Lee, K. (2021). Transformers provide surprisingly effective representations for online and offline dictionary learning.
• Zhai, C. (1997). Exploiting context to identify ambiguous terms.
• Zhou, B., Khosla, A., Lapedriza, A., Oliva, A., & Torralba, A. (2018). Learning deep features for discriminative localization.
• Zou, Y., Wang, L., Hu, Z., Li, Z., Wang, W., Tu, Z., … & Sun, M. (2023a). Controllable generation from pre-trained language models via inverse prompting.
• Zou, Y., Wang, L., Hu, Z., Li, Z., Wang, W., Tu, Z., … & Sun, M. (2023b). Controllable generation from pre-trained language models via inverse prompting.

近年来,Transformer 模型在自然语言处理等领域取得了巨大成功。与此同时,结构化状态空间模型(SSM)作为一种新兴的序列建模方法也展现出了强大的潜力。最近的研究表明,这两类看似截然不同的模型其实存在着深层次的联系。本文将介绍一项重要的理论突破 – State Space Duality(SSD)框架,该框架揭示了Transformer和SSM之间的对偶性,为统一和改进这两类模型提供了新的视角。

State Space Duality框架的核心思想

SSD框架的核心是将SSM和注意力机制都视为结构化矩阵变换。具体来说:

1. SSM可以表示为半可分矩阵(semiseparable matrix)的乘法。
2. 注意力机制可以推广为结构化掩码注意力(structured masked attention, SMA)。
3. SSM和SMA在特定情况下是等价的,都可以用线性和二次两种不同的计算方式实现。

这一框架揭示了SSM和注意力机制在本质上是相同的序列变换,只是从不同角度进行了参数化。

理论洞见

SSD框架带来了几个重要的理论洞见:

1. SSM的线性递归形式和注意力的二次形式是同一变换的两种不同计算方式。
2. 有效的自回归注意力必须具有半可分结构。
3. SSM中的选择性机制类似于注意力中的位置编码。

这些洞见不仅加深了我们对两类模型的理解,还为设计新的高效算法提供了理论基础。

高效算法:SSD算法

基于SSD框架,研究人员提出了一种新的高效算法来计算SSM,称为SSD算法。该算法结合了SSM的线性递归和注意力的二次计算,通过半可分矩阵的块分解来实现。

SSD算法的主要优势包括:

1. 计算复杂度与序列长度呈线性关系。
2. 可以充分利用现代硬件上的矩阵乘法单元。
3. 允许更大的状态扩展因子,而不会显著降低速度。

实验表明,SSD算法比Mamba中的优化扫描实现快2-8倍,同时还能支持更大的递归状态大小。

Mamba-2:结合SSD的新架构

研究人员基于SSD框架设计了一种新的序列模型架构Mamba-2。Mamba-2相比原始的Mamba模型有以下改进:

1. 简化了块结构,将SSM参数的生成并行化。
2. 采用了类似多值注意力(MVA)的头部结构。
3. 增加了额外的归一化层以提高稳定性。
4. 可以灵活地与MLP和注意力层组合,形成混合模型。

实验结果表明,Mamba-2在语言建模等任务上的性能可以媲美甚至超越同等规模的Transformer模型,同时在长序列处理方面具有更好的可扩展性。

实验验证

研究人员通过一系列实验验证了SSD框架和Mamba-2架构的有效性:

1. 在多查询关联回忆(MQAR)任务上,Mamba-2显著优于Mamba和标准注意力模型。
2. 在Pile数据集上的语言建模实验中,Mamba-2达到或超过了同等规模的Transformer++和Mamba模型的性能。
3. 在零样本下游评估任务中,Mamba-2通常能够匹配两倍参数量的Pythia模型的表现。
4. SSD算法在长序列(2K以上)上比FlashAttention-2更快。

这些实验结果证明了SSD框架的理论洞见可以有效地指导实践,并在多个方面改进了现有的序列建模方法。

总结与展望

State Space Duality框架为统一和改进Transformer和SSM提供了一个强大的理论基础。通过揭示这两类模型之间的深层联系,SSD框架不仅加深了我们对序列建模的理解,还为设计更高效、更强大的模型指明了方向。

未来的研究方向可能包括:

1. 进一步探索SSM和注意力机制的理论联系。
2. 设计更复杂的结构化掩码来改进注意力机制。
3. 将SSD框架扩展到其他类型的序列模型。
4. 开发针对特定任务优化的混合模型架构。

总的来说,State Space Duality框架的提出标志着序列建模理论的一个重要进展,有望推动该领域在未来取得更多突破性进展。

导语： 近年来，大型语言模型 (LLM) 在人工智能领域取得了显著的进展。除了生成流畅的文本、故事、论文和程序代码外，LLM 还展现出作为通用问题解决者的巨大潜力。本文将深入探讨如何利用 LLM 构建自主智能体，并介绍其核心组件、应用案例以及面临的挑战。

智能体系统概述

LLM 驱动的自主智能体系统中，LLM 充当智能体的“大脑”，并与以下关键组件协同工作：

• 规划： 智能体将复杂的任务分解成更小、更易于管理的子目标，并通过自我反思和改进机制，从过去的行动中学习，不断优化最终结果的质量。
• 记忆： 智能体利用短期记忆 (例如上下文学习) 和长期记忆 (例如外部向量存储和快速检索) 来存储和回忆信息。
• 工具使用： 智能体学习调用外部 API 来获取模型权重中缺失的信息，例如最新信息、代码执行能力、专有信息源等。

![LLM 驱动的自主智能体系统概述][]

组件一：规划

规划是智能体执行复杂任务的关键。智能体需要明确任务步骤并制定计划。

任务分解

• 思维链 (CoT) 是一种标准的提示技术，通过指示模型“逐步思考”来将复杂任务分解成更小、更简单的步骤，从而提高模型在复杂任务上的性能。
• 思维树 (Tree of Thoughts) 扩展了 CoT，它在每个步骤中探索多种推理可能性，并将问题分解成多个思维步骤，每个步骤生成多个想法，形成树状结构。
• LLM+P 则依赖外部的经典规划器进行长期规划，利用规划域定义语言 (PDDL) 作为描述规划问题的中间接口。

自我反思

自我反思使智能体能够通过改进过去的行动决策和纠正错误来迭代改进，这在现实世界中不可避免的试错过程中至关重要。

• ReAct 通过将动作空间扩展为特定任务的离散动作和语言空间的组合，将推理和行动整合到 LLM 中。
• Reflexion 为智能体配备了动态记忆和自我反思能力，以提高推理能力。它使用启发式函数来确定轨迹何时效率低下或包含幻觉，并使用自我反思结果来改进未来的计划。
• Chain of Hindsight (CoH) 通过明确地向模型呈现一系列带有反馈的过去输出，鼓励模型改进自身的输出。
• 算法蒸馏 (AD) 将相同的思想应用于强化学习任务中的跨情节轨迹，将算法封装在长期历史条件策略中。

组件二：记忆

记忆是指获取、存储、保留和检索信息的过程。

记忆类型

• 感觉记忆： 记忆的最初阶段，保留对感觉信息 (视觉、听觉等) 的印象。
• 短期记忆 (STM) 或工作记忆： 存储我们当前意识到的信息，用于执行复杂的认知任务。
• 长期记忆 (LTM)： 可以长时间存储信息，容量基本无限，包括显性记忆 (事实和事件) 和隐性记忆 (技能和习惯)。

最大内积搜索 (MIPS)

外部记忆可以缓解有限注意力跨度的限制。标准做法是将信息的嵌入表示保存到支持快速最大内积搜索 (MIPS) 的向量存储数据库中。

一些常见的快速 MIPS 算法包括：

• LSH (局部敏感哈希)
• ANNOY (近似最近邻 Oh Yeah)
• HNSW (分层导航小世界)
• FAISS (Facebook AI 相似性搜索)
• ScaNN (可扩展最近邻)

组件三：工具使用

工具使用是人类的显著特征。为 LLM 配备外部工具可以显著扩展模型的能力。

• MRKL (模块化推理、知识和语言) 是一种用于自主智能体的神经符号架构，它包含一组“专家”模块，通用 LLM 充当路由器，将查询路由到最合适的专家模块。
• TALM (工具增强语言模型) 和 Toolformer 通过微调 LM 来学习使用外部工具 API。
• ChatGPT 插件 和 OpenAI API 函数调用是 LLM 在实践中增强工具使用能力的典型例子。
• HuggingGPT 是一个框架，它使用 ChatGPT 作为任务规划器，根据模型描述选择 HuggingFace 平台上可用的模型，并根据执行结果汇总响应。
• API-Bank 是一个用于评估工具增强型 LLM 性能的基准测试平台，它包含 53 个常用 API 工具、完整的工具增强型 LLM 工作流程以及 264 个注释对话，涉及 568 个 API 调用。

案例研究

科学发现智能体

• ChemCrow 是一个特定领域的例子，其中 LLM 通过 13 个专家设计的工具进行增强，以完成有机合成、药物发现和材料设计等任务。
• Boiko 等人 (2023) 还研究了 LLM 驱动的智能体在科学发现中的应用，以处理复杂科学实验的自主设计、规划和执行。

生成式智能体模拟

生成式智能体 (Generative Agents) 是一个有趣的实验，其中 25 个虚拟角色，每个角色由一个 LLM 驱动的智能体控制，生活在一个沙盒环境中，灵感来自《模拟人生》。

概念验证示例

• AutoGPT 是一个自主智能体演示，它使用 LLM 作为主控制器，可以访问互联网和运行代码。
• GPT-Engineer 是另一个项目，旨在根据自然语言中指定的任务创建完整的代码库。

挑战

• 有限的上下文长度
• 长期规划和任务分解的挑战
• 自然语言接口的可靠性

参考文献原文链接

[1] Wei et al. “Chain of thought prompting elicits reasoning in large language models.” NeurIPS 2022

导语： 近年来，大型语言模型（LLM）在人工智能领域掀起了一场革命，其强大的文本生成和理解能力为众多应用场景带来了新的可能性。然而，将LLM应用从酷炫的演示转化为可实际部署的生产系统并非易事。本文将深入探讨LLM应用生产实践所面临的挑战，并结合实例分析解决方案，最后展望LLM应用的未来发展方向。

一、生产环境下LLM应用面临的挑战

1. 自然语言的模糊性

不同于精确的编程语言，自然语言本身就带有模糊性。这种模糊性在LLM应用中主要体现在用户指令和模型输出两个方面。

首先，用户指令的灵活性可能导致难以察觉的错误。例如，对代码进行微小的修改，例如添加字符或删除一行，通常会导致明显的错误提示。但如果对LLM的指令进行类似的修改，程序仍然可以运行，但输出结果可能大相径庭。

其次，LLM输出结果的模糊性是更大的挑战。这会导致两个问题：

• 输出格式不一致： LLM的后续应用通常需要解析特定格式的输出结果。虽然我们可以通过精心设计的指令来明确输出格式，但这并不能保证模型始终遵循预期的格式。
• 用户体验不一致： 用户期望应用程序具有一致性。例如，我们不希望每次查询保险报价时都得到不同的结果。然而，LLM本身具有随机性，无法保证相同的输入每次都能得到相同的输出。

为了解决这个问题，OpenAI等机构正在积极探索提高模型可靠性的方法。一些经验丰富的LLM开发者建议，我们需要适应这种模糊性，并围绕它构建工作流程。此外，通过尽可能提高工程严谨性，例如采用本文接下来讨论的各种方法，可以有效缓解LLM应用中的模糊性问题。

2. Prompt工程的挑战

Prompt工程是指设计和优化LLM输入指令的过程，其目标是引导模型生成符合预期的输出。

• Prompt评估：
  • 评估LLM是否理解Prompt中给出的示例。例如，在情感分析任务中，可以输入一些带有情感标签的文本，观察模型是否能输出正确的情感得分。
  • 评估LLM是否对Prompt中的示例过拟合。可以通过使用独立的测试集来评估模型的泛化能力。
• Prompt版本控制：
  Prompt的微小变化可能会导致输出结果的显著差异。因此，对Prompt进行版本控制和性能跟踪至关重要。可以使用Git等工具来管理Prompt的不同版本及其性能指标。
• Prompt优化：
  • 使用“思维链”（Chain-of-Thought，COT）技术，引导模型解释其推理过程，从而提高输出结果的可解释性和准确性。
  • 生成多个输出结果，并选择最佳结果作为最终输出。例如，可以使用多数投票法或让LLM自身选择最佳结果。
  • 将复杂的Prompt分解成多个更小、更简单的Prompt，可以降低模型理解和推理的难度。

3. 成本和延迟

• 成本：
  OpenAI等API服务提供商通常根据输入和输出token的数量收费。更长的Prompt和输出意味着更高的成本。
• 延迟：
  输出token的生成是顺序进行的，因此输出长度会显著影响延迟。此外，网络延迟和工程开销也会影响LLM应用的整体延迟。

4. Prompting、微调和替代方案

• Prompting： 为每个样本提供明确的指令，告诉模型如何响应。
• 微调： 在特定任务的数据集上训练模型，使其学习如何响应，而无需在Prompt中明确说明。

选择Prompting还是微调取决于数据可用性、性能要求和成本限制。

• 数据量少： Prompting更易于上手，但Prompt的长度有限，无法包含大量示例。
• 数据量大： 微调可以利用更多数据，并获得更好的性能。

5. 其他挑战

• 嵌入和向量数据库： 使用LLM生成文本嵌入，并将其存储在向量数据库中，可以高效地进行语义搜索和推荐。
• 向后和向前兼容性： LLM模型会不断更新迭代，因此需要确保Prompt和应用逻辑的兼容性。

二、任务组合能力

实际应用中，LLM应用通常需要执行多个任务，并按照一定的控制流程进行组合。

1. 多任务应用

例如，“与数据对话”应用需要执行以下任务：

1. 将用户的自然语言输入转换为SQL查询（LLM）。
2. 在数据库中执行SQL查询（SQL执行器）。
3. 将SQL结果转换为自然语言响应（LLM）。

2. 代理、工具和控制流程

• 代理：
  能够根据给定的控制流程执行多个任务的应用程序。
• 工具：
  代理可以利用的外部程序或服务，例如搜索引擎、数据库、计算器等。
• 控制流程：
  定义任务执行顺序和条件的逻辑，例如顺序执行、并行执行、条件分支、循环等。

3. 使用LLM代理进行控制流程

可以使用LLM来决定控制流程的条件。例如，可以训练LLM根据用户输入选择不同的工具或执行不同的任务。

4. 代理测试

为了确保代理的可靠性，需要对每个任务以及整个控制流程进行测试。

三、LLM应用的未来发展方向

• AI助手：
  能够帮助用户完成各种任务的智能助手，例如日程安排、笔记记录、编程辅助等。
• 聊天机器人：
  能够与用户进行自然对话的智能程序，例如虚拟伴侣、游戏角色等。
• 编程和游戏：
  LLM可以用于代码生成、代码调试、游戏开发等领域。
• 学习：
  LLM可以用于自动生成测验、评估学生答案、提供学习建议等教育场景。
• 与数据对话：
  LLM可以帮助用户使用自然语言查询和分析数据。
• 搜索和推荐：
  LLM可以改善搜索结果的相关性和推荐系统的个性化程度。
• 销售：
  LLM可以用于自动生成销售邮件、分析客户信息等。
• SEO：
  LLM可以用于生成SEO优化内容，但同时也带来了新的挑战，例如如何区分人工编写和机器生成的内容。

结语

LLM应用仍处于早期发展阶段，技术和应用场景都在快速迭代。Prompt工程、成本控制、延迟优化、任务组合等方面都面临着挑战。然而，LLM也为AI助手、聊天机器人、编程、教育、数据分析、搜索推荐等领域带来了前所未有的机遇。相信随着技术的进步和应用的深入，LLM将为我们带来更加智能和便捷的未来。

参考文献

• Huyen, C. (2023). Building LLM applications for production. Retrieved from https://huyenchip.com/2023/04/11/llm-engineering.html

近年来，AI 应用如雨后春笋般涌现，推动着相关技术的蓬勃发展，其中向量数据库尤为引人注目。作为 AI 应用技术栈中的关键一环，向量数据库为 AI 应用，特别是大型语言模型 (LLM) 应用，提供了强大的知识存储和检索能力。本文将深入浅出地探讨向量数据库的原理和实现，涵盖其基本概念、相似性搜索算法、相似性测量算法、过滤算法以及选型等方面，并结合实际应用案例，展现向量数据库如何为 AI 插上知识的翅膀。

GPT 的瓶颈：有限的上下文窗口

GPT-3.5/4 的问世，无疑是 AI 发展史上的里程碑事件，其强大的文本生成能力令人惊叹。然而，GPT 模型并非完美无缺，其有限的上下文窗口大小成为制约其性能的一大瓶颈。以 GPT-3.5-turbo 为例，其最大上下文窗口大小仅为 4K tokens（约 3000 字），这意味着模型最多只能处理 3000 字以内的文本信息。

虽然 ChatGPT 等应用提供了对话记忆功能，但这并非 GPT 模型本身具备记忆能力，而是开发者通过外部存储机制实现的。当用户输入超过上下文窗口限制的内容时，GPT 模型将无法记住之前的对话内容，从而影响其对后续问题的理解和回答。

尽管 GPT-4 将上下文窗口大小提升至 32K tokens，Claude 模型更是高达 100K tokens，但处理如此庞大的文本信息需要消耗大量的计算资源和时间，成本也随之水涨船高。

向量数据库：突破 GPT 上下文限制的利器

面对 GPT 模型的上下文窗口限制，开发者们积极探索解决方案，而向量数据库应运而生。其核心思想是将文本信息转化为向量，并存储在数据库中，当用户输入问题时，系统将问题转化为向量，并在数据库中搜索最相似的向量及其对应的文本信息，最终将相关内容返回给用户。

以客服培训资料为例，我们可以将其转化为向量并存储在向量数据库中。当用户提出相关问题时，系统将问题转化为向量，并在数据库中快速检索最相似的向量及其对应的文本片段，最终将这些片段返回给 GPT 模型进行处理，从而绕过 GPT 模型的上下文窗口限制，提高响应速度并降低成本。

从特征到向量：理解相似性搜索的基石

在现实生活中，我们通过识别不同事物之间的特征来区分它们。例如，我们可以根据体型大小、毛发长度、鼻子长短等特征来区分不同品种的狗。如果将这些特征量化，并映射到多维坐标系中，每个事物都可以用一个多维坐标点来表示。

向量是具有大小和方向的数学结构，可以用来表示事物的多维特征。通过计算向量之间的距离，可以判断它们在特征空间中的相似程度，这就是相似性搜索的基本原理。

高效的相似性搜索算法：兼顾速度与质量

在海量数据中进行相似性搜索，如果采用暴力搜索的方式，计算量巨大，效率低下。为此，向量数据库采用了一系列高效的相似性搜索算法，其主要思想是通过降维、聚类、构建树或图等方式，减少向量大小或缩小搜索范围，从而提高搜索效率。

1. 聚类搜索：K-Means 与 Faiss

聚类算法将数据点分组到不同的簇中，使同一簇内的点彼此相似，不同簇之间的点彼此不同。K-Means 算法是一种常用的聚类算法，它将数据分成 k 个类别，其中 k 是预先指定的。

Faiss 算法是 Facebook AI Research 推出的一种高效的相似性搜索库，它采用了一种基于 Voronoi 图的聚类方法，将向量空间划分为多个单元格，每个单元格由一个质心代表。在搜索时，首先找到查询向量所属的单元格，然后只在该单元格内进行搜索，从而大大缩小了搜索范围。

2. 量化压缩：Product Quantization (PQ)

为了降低内存占用，向量数据库通常采用量化压缩技术。Product Quantization (PQ) 是一种常用的向量量化方法，它将高维向量分解成多个低维子向量，并对每个子向量独立进行量化。

3. 图搜索：Hierarchical Navigable Small Worlds (HNSW)

HNSW 算法是一种基于图的相似性搜索算法，它将向量空间表示为一个多层图，每层都是一个小世界，图中的节点表示向量，边表示向量之间的相似性。

4. 局部敏感哈希：Locality Sensitive Hashing (LSH)

LSH 算法将相似的向量映射到相同的哈希桶中，从而在搜索时只需比较同一哈希桶内的向量，提高了搜索效率。

相似性测量：欧几里德距离、余弦相似度与点积相似度

在相似性搜索中，需要计算向量之间的距离来衡量它们的相似度。常用的向量相似度算法包括：

• 欧几里德距离：计算两个向量在空间中的直线距离。
• 余弦相似度：计算两个向量夹角的余弦值，值越大表示越相似。
• 点积相似度：计算两个向量的点积，值越大表示越相似。

过滤：缩小搜索范围，提高查询效率

在实际应用中，我们通常需要根据业务需求对搜索结果进行过滤。例如，在电商搜索中，用户可能希望根据价格、品牌等条件筛选商品。

向量数据库通常支持元数据过滤功能，用户可以根据预定义的元数据字段对搜索结果进行筛选，从而缩小搜索范围，提高查询效率。

向量数据库选型：综合考虑性能、功能、成本等因素

在选择向量数据库时，需要综合考虑以下因素：

• 性能： 查询速度、吞吐量、内存占用、存储容量等。
• 功能： 相似性搜索算法、相似性测量算法、过滤功能、分布式部署、高可用性、容错性、访问控制、数据备份等。
• 易用性： API & SDK 设计、文档、社区活跃度等。
• 成本： 部署成本、维护成本、使用成本等。

目前，市面上比较流行的向量数据库包括：

• Chroma
• Milvus
• Pinecone
• Qdrant
• Typesense
• Weaviate

此外，一些传统的数据库也提供了向量搜索扩展功能，例如：

• Redis： RediSearch 模块
• PostgreSQL： pgvector 扩展

总结：向量数据库，释放 AI 知识的力量

向量数据库作为 AI 应用技术栈中的关键一环，为 AI 应用，特别是大型语言模型 (LLM) 应用，提供了强大的知识存储和检索能力。其高效的相似性搜索算法、灵活的相似性测量算法、强大的过滤功能以及不断完善的生态系统，使其成为构建知识型 AI 应用的理想选择。

随着 AI 技术的不断发展，向量数据库将在更多领域发挥重要作用，为 AI 应用插上知识的翅膀，释放 AI 知识的力量。

参考文献：

[1] 向量数据库: https://zhuanlan.zhihu.com/p/628859137

[2] Claude 模型: https://www.anthropic.com/index/100k-context-windows

[3] LangChainJs: https://js.langchain.com/docs/modules/indexes/vector_stores/integrations/qdrant

[4] VectorHub: https://github.com/vectorhub-io/vectorhub

[5] ChatFiles: https://github.com/mckaywrigley/chatbot-ui

[6] 视频: https://www.youtube.com/watch?v=PqbB07PgZ0Y

[7] Faiss 算法: https://faiss.ai/

[8] 视频: https://www.youtube.com/watch?v=7zN8jI6YTEY

[9] 博客: https://towardsdatascience.com/understanding-locality-sensitive-hashing-497c1f9701c9

[10] 博客: https://www.pinecone.io/learn/locality-sensitive-hashing/

[11] RediSearch: https://redis.io/modules/redisearch/

[12] pgvector: https://github.com/pgvector/pgvector

最近 ChatGPT 火爆出圈，引发了人们对 AI 产品交互方式的热烈讨论。不少人认为，对话和聊天是 AI 产品交互的必然趋势。然而，作为一名资深科技专栏作家，我对此持保留意见。我认为，Chatbot 并非 AI 交互的终极形态，甚至可以说，它是一个糟糕的开端。

对话没有预设用途，用户无从下手

好的工具应该“自带说明书”。就像一双手套，它以手的形式展现出来，所以我们会把它戴在手上。材料的特点也会告诉我们更多信息，金属网格的手套会保护我们防止物理伤害，橡胶手套是防止化学伤害，皮革手套很适合骑摩托的时候戴。

然而，面对 ChatGPT 这样的聊天界面，我们得到的唯一提示就是在一个文本框中输入字符。它没有清晰地告诉我们它能做什么，我们该如何提问，更没有告诉我们不应该如何使用它。用户只能像摸象的盲人一样，通过一次次试探来了解它。

提示工程门槛高，将用户拒之门外

诚然，用户可以通过学习和练习来掌握 Prompt 的技巧，但了解哪些 Prompt 有效的代价，仍然由每个用户承担。而这些内容原本可以直接融入界面中，降低使用门槛。

以 AI 写作助手为例，为了更好地定制个性化回复，我们可以为用户添加一些控件，让他们指定使用情境：需要的写作类型？目标读者是谁？整体风格应该是什么感觉？……将这些信息嵌入到界面中，而不是让用户绞尽脑汁地将其破解到他们提问的每个问题中。

信息展示方式单一，阅读体验不佳

自然语言对于概括的方向很好用，比如“把我送到正确的街区”这种。但一旦 ChatGPT 做出了回应，我该如何让它把我带到正确的房子前面去？

目前的 Chatbot 界面只能以聊天的形式线性展示信息，用户被迫在回复之间来回滚动，逐行阅读，难以进行有效的对比和分析。

创作过程缺乏连贯性，打断用户思路

画家在创作时，会有两个动作：近距离在画板画布上绘画；拉开距离欣赏成果，进行评估和下一步规划。这样的流程——实践和评估的循环，适用于任何“手艺”，比如编程和写作。

然而，使用 Chatbot 时，我被迫频繁地在“提问-等待-阅读”三种模式之间切换，这让我失去了创作的 flow 状态。

过度依赖 AI 生成，削弱用户思考能力

内容消费趋势是：更短、更有冲击感。人们习惯于被动接受信息，而 Chatbot 类产品强化了这种趋势。当 AI 可以写篇文章、编写代码时，我们为什么要自己动手？

我希望看到更多的工具，而不是操作的机器。 我们应该拥抱我们的人性，而不是盲目地提高效率。我们应该让人们更灵巧地使用 AI 工具，而不是成为 AI 生成内容的审查员。

未来 AI 交互形态的畅想

我认为，未来 AI 产品的交互形态应该是更加多元化、智能化和人性化的。

• 多模态交互： 结合语音、图像、视频等多种形式，提供更加自然、直观的交互体验。
• 智能辅助： AI 不再是简单的信息输出者，而是能够理解用户意图，提供个性化建议和解决方案的智能助手。
• 人机协同： AI 与人类的关系不再是替代，而是协同。AI 负责处理重复性、机械性的任务，而人类则专注于创造性、决策性的工作。

我相信，真正改变游戏规则的产品，与简单的内容生成关系不大。我们应该构建一些真正能够提供建议来帮助我们清晰思考的工具。比如通过操纵潜在空间中的几何来塑造文学作品，通过幕后链接模型让我们更容易地在视频而非像素中操纵视觉内容。

我迫不及待地想看到这个领域变得更加成熟，并开始创造能够包含我们人类能力的 AI 工具。

参考文献：

• Wattenberger, Amelia. “Chatbots are a terrible UI.” GitHub Blog, 2023. https://github.blog/2023-03-22-chatbots-are-a-terrible-ui/

引言：Agent 浪潮与软件生产的未来

2023 年，AI Agent 成为科技领域最炙手可热的话题之一。AutoGPT 的横空出世，以惊人的速度登顶 Github 热榜，预示着软件生产方式即将迎来革命性变革。

LLM（大型语言模型）作为 AI Agent 的“大脑”，其推理能力日益强大。Agent 框架则为 LLM 提供了结构化的思考方法，使软件生产进入“3D 打印”时代：根据用户需求，个性化定制软件，打造每个知识工作者信赖的 AI 伙伴。

一、AI Agent：机遇与挑战

1. AI Agent 的定义与意义

AI Agent 是指能够主动思考和行动的智能体。与被动响应指令的 LLM 不同，Agent 可以根据目标，自主感知环境、制定计划并执行任务，甚至与其他 Agent 协作。

AI Agent 的出现，意味着软件行业将迎来“3D 打印”时代：

• 降低成本： 自动化软件生产，使边际成本趋近于零。
• 提高定制化： 灵活满足用户长尾需求，实现软件的个性化定制。

2. AI Agent 的核心要素

理想的 Agent 框架应该具备以下核心要素：

• 记忆： 包括短期记忆和长期记忆，用于存储和检索信息，模拟人类的学习和记忆过程。
• 行动： 拥有完整的行动策略集，例如工具使用、职责扮演、记忆检索、推理、学习和编程等。
• 决策： 能够根据目标，对不同的行动方案进行评估，选择最优方案执行。

3. AI Agent 面临的挑战

目前，AI Agent 的发展仍面临诸多挑战：

• 可靠性： 完全自主的 Agent 缺乏可控性，难以满足实际工作需求。
• 复杂推理能力： 现有的 LLM 在多步推理任务上表现不佳，限制了 Agent 的能力。
• 延迟： LLM 的响应速度仍然较慢，影响了工作效率。
• 成本： Agent 框架中，AI 之间的沟通和协作会产生大量文本，带来高昂的成本。

二、AI Agent 的发展方向

1. Agent Framework：打造 AI 合作伙伴

Agent Framework 旨在提供实用的框架和工具，降低 Agent 开发的复杂度，并支持 Agent 间的协作。

• 模块化设计： 将 Agent 系统分解为独立可解耦的模块，方便开发者灵活组合。
• APIs 和 SDKs： 设计通用的接口和协议，实现不同 Agent 之间的兼容和配合。
• 合作机制设计： 建立 Agent 之间的合作分工和层级机制，保障数据安全和隐私。

2. Vertical Agent：深耕垂直领域

Vertical Agent 专注于特定领域，例如编程、个人助理、写作和数据分析等，利用领域知识和工作流，打造更实用的 AI 工具。

三、AI Agent 的未来展望

1. OpenAI 的角色

OpenAI 作为 LLM 领域的领头羊，预计将在 11 月的 DevDay 上发布定义 AI Agent 的 API 和开发框架，推动行业标准化。

2. 强化学习的启示

RL-based Agent 在探索和试错方面具有优势，可以为 LLM-based Agent 提供借鉴。

3. LLM 的优化

AI Agent 应用的成熟，离不开 LLM 在复杂推理能力和延迟方面的优化。

4. 成本控制与激励机制

Multi-agent 协作带来的高成本问题，需要通过模型优化、成本控制和激励机制设计来解决。

结语：AI Agent 的未来充满希望

AI Agent 的发展，将重构 10 亿知识工作者的职业，掀起软件生产的革命。尽管挑战重重，但随着技术的进步和应用的探索，AI Agent 必将为我们带来更加智能、高效和个性化的工作方式。

参考文献

1. Theodore Sumers, Shunyu Yao, Karthik Narasimhan, Thomas L. Griffiths. Cognitive Architectures for Language Agents.
2. https://e2b.dev/blog/the-state-of-ai-agents-reliability-sdks-benchmarking-and-market-trends

引言：

近期，ChatGPT 横空出世，以其出色的对话能力和知识储备，迅速成为人工智能领域的焦点。其背后究竟隐藏着哪些技术奥秘？本文将深入浅出地解读 ChatGPT 的核心技术，包括 RLHF、IFT、CoT 等，并探讨这些技术如何协同工作，打造出令人惊艳的智能对话体验。

一、 群雄逐鹿：AI 对话代理的现状

ChatGPT 并非横空出世，早在其之前，Meta、Google、DeepMind、Anthropic 等机构就已经推出了各自的语言模型对话代理，例如 BlenderBot、LaMDA、Sparrow、Assistant 等。这些对话代理都致力于实现“指令依从”，即根据用户指令完成任务。

然而，ChatGPT 的成功并非偶然，它在技术路线上有着独特之处。下表对目前主流的 AI 聊天机器人进行了比较：

	LaMDA	BlenderBot 3	Sparrow	ChatGPT / InstructGPT	Assistant
组织	Google	Meta	DeepMind	OpenAI	Anthropic
能否公开访问	否	能否有限	否	否	否
大小	137B	175B	70B	175B	52B
预训练基础模型	未知	OPT	Chinchilla	GPT-3.5	未知
…	…	…	…	…	…

二、 从预测文本到遵循指令：指令微调 (IFT) 技术

传统的语言模型擅长预测文本，但要让其理解并执行用户指令，则需要进行“指令微调 (Instruction Fine-Tuning，IFT)”。IFT 的核心是向模型示范各种书面指令及其预期输出，例如：

指令：写一首关于微调的诗。
输出：... (ChatGPT 生成的诗歌)

IFT 的训练数据通常包括人工编写的指令和模型自举生成的实例。其中，自举指的是使用少样本技术，先输入少量样本引导模型，然后让模型自行生成新的指令、输入和输出，不断迭代优化。

三、 安全护航：有监督微调 (SFT) 与安全规则

经过 IFT 的模型虽然能够理解指令，但生成的回复不一定安全无害。为了解决这个问题，开发者会使用“有监督微调 (Supervised Fine-tuning, SFT)”，利用高质量的人工标注数据，进一步提升模型回复的安全性。

例如，谷歌的 LaMDA 就根据一系列预先定义的安全规则进行微调，这些规则涵盖了伤害、歧视、错误信息等方面，确保模型的回复符合道德和安全规范。

四、 更胜一筹：人类反馈强化学习 (RLHF)

为了进一步提升模型的性能，OpenAI 的 InstructGPT、DeepMind 的 Sparrow 和 Anthropic 的 Constitutional AI 都引入了“人类反馈强化学习 (Reinforcement Learning From Human Feedback，RLHF)”技术。

RLHF 的核心是根据人类的偏好对模型的回复进行排序，然后用这些带标注的数据训练偏好模型，最终通过强化学习优化对话代理。关于 RLHF 的更多细节，可以参考我们之前的文章：ChatGPT 背后的“功臣”——RLHF 技术详解。

五、 逻辑推理：思维链 (CoT) 提示

“思维链 (Chain-of-thought，CoT) 提示”是 IFT 的一种特殊形式，它通过引导模型进行逐步推理来生成最终答案。例如：

指令：小明有 5 个苹果，小红给了他 3 个，小明现在有多少个苹果？
CoT 推理：小明最初有 5 个苹果，加上小红给的 3 个，总数是 5 + 3 = 8 个。
输出：小明现在有 8 个苹果。

使用 CoT 微调的模型在需要逻辑推理的任务上表现更出色，例如常识推理、算术计算、符号推理等。

六、 展望未来：对话代理的进一步发展

尽管现有的技术已经取得了显著成果，但对话代理领域仍有许多问题亟待解决，例如：

• RL 在人类反馈学习中的重要性？
• 如何系统评估红蓝对抗的效果？
• 如何选择最佳的基础模型？

相信随着技术的不断进步，对话代理将在未来展现出更加强大的能力，为人类带来更加智能、便捷的交互体验。

参考文献：

• Rajani et al.,”What Makes a Dialog Agent Useful?”, Hugging Face Blog, 2023.

近年来，深度学习在自然语言处理领域取得了巨大成功，这主要归功于 Transformer 架构。然而，状态空间模型（SSM），例如 Mamba，最近被证明在中小型规模上可以与 Transformer 媲美甚至超越。本文将深入探讨这两种模型之间的密切关系，并通过对结构化半可分矩阵的不同分解，建立 SSM 和注意力变体之间丰富的理论联系框架。我们的状态空间对偶性（SSD）框架将引领我们设计一种新的架构（Mamba-2），其核心层是对 Mamba 选择性 SSM 的改进，速度提高了 2-8 倍，同时在语言建模方面继续与 Transformer 保持竞争力。

Transformer 的效率瓶颈与状态空间模型的崛起

Transformer，特别是仅解码器模型（例如 GPT 和 Llama），以因果方式处理输入序列，是现代深度学习成功的关键驱动力之一。然而，其核心注意力层存在效率问题，例如在训练期间按序列长度呈二次方增长，以及在自回归生成期间需要大小与序列长度呈线性关系的缓存。为了解决这些问题，许多方法试图近似核心注意力层（Tay et al. 2022），但效果有限。

与此同时，一类替代序列模型——结构化状态空间模型（SSM）——应运而生。它们在训练期间按序列长度呈线性增长，在生成期间具有恒定的状态大小。SSM 在长程任务上表现出色（例如 S4），并且最近在中小型规模的语言建模上与 Transformer 媲美甚至超越（例如 Mamba）。然而，SSM 的发展似乎与社区为改进 Transformer 所做的集体努力脱节，例如从理论上理解它们以及在现代硬件上优化它们。因此，与 Transformer 相比，理解和试验 SSM 更加困难，并且从算法和系统角度来看，高效地训练 SSM 仍然具有挑战性。

结构化状态空间对偶性：连接 SSM 和注意力的桥梁

本文的主要目标是建立结构化 SSM 和注意力变体之间丰富的理论联系。这将使我们能够将最初为 Transformer 开发的算法和系统优化转移到 SSM，从而构建性能优于 Transformer 且序列长度扩展效率更高的基础模型。线性注意力（LA）框架（Katharopoulos et al. 2020）是朝着这个方向迈出的里程碑式的一步，它通过证明二次核化注意力的“对偶形式”与特定线性递归之间的等价性，推导了自回归注意力和线性 RNN 之间的联系。这种对偶性带来了新的能力，例如能够同时进行高效的可并行化训练和高效的自回归推理。

本文秉承同样的精神，提供了多个视角，将线性复杂度的 SSM 与二次复杂度的形式联系起来，以结合 SSM 和注意力的优势。我们的框架将结构化 SSM 和注意力变体联系起来，我们称之为结构化状态空间对偶性（SSD），它是通过结构化矩阵的抽象实现的：具有次二次参数和乘法复杂度的矩阵。我们开发了两个广泛的框架来表示序列模型，一个作为矩阵变换，一个作为张量收缩，每个框架都揭示了对偶性的不同视角。

状态空间模型与半可分矩阵的等价性

本文的核心是状态空间模型和一类称为半可分矩阵的结构化矩阵之间的等价性（第 3 节）。这种联系揭示了 SSM 的新特性和算法。本文的一个中心思想是，计算状态空间模型的不同方法可以重新表述为结构化矩阵上的各种矩阵乘法算法。

半可分矩阵是一种基本的矩阵结构。我们首先定义这些矩阵及其性质。

定义 3.1：如果包含在下三角部分（即对角线或下方）的每个子矩阵的秩最多为 N，则（下三角）矩阵 M 为 N-半可分矩阵。我们称 N 为半可分矩阵的阶数或秩。

半可分矩阵具有许多结构化表示，包括分层半可分 (HSS)、顺序半可分 (SSS) 和 Bruhat 形式。我们将主要使用 SSS 形式。

定义 3.2：如果下三角矩阵 M ∈ R(T,T) 可以写成以下形式，则它具有 N-顺序半可分 (SSS) 表示：

$$
M_{ji} = C_j^\top A_j \cdots A_{i+1} B_i
$$

其中向量 $B_0, …, B_{T-1}, C_0, …, C_{T-1} ∈ R^N$ 且矩阵 $A_0, …, A_{T-1} ∈ R^{(N,N)}$。

半可分矩阵的一个基本结果是它们与具有 SSS 表示的矩阵完全等价。一个方向可以通过简单的构造性证明来推断。

引理 3.3：具有表示 (4) 的 N-SSS 矩阵 M 是 N-半可分矩阵。

命题 3.4：每个 N-半可分矩阵都有一个 N-SSS 表示。

状态空间模型和顺序半可分矩阵之间的联系如下：

定理 3.5：状态大小为 N 的状态空间模型变换 𝑦 = SSM(𝐴, 𝐵, 𝐶)(𝑥) 与通过 N-SS 矩阵（以顺序半可分表示）进行矩阵乘法 𝑦 = SSS(𝐴, 𝐵, 𝐶) · 𝑥 相同。

换句话说，序列变换算子 SSM 与矩阵构造算子 SSS 重合，我们可以互换使用它们。

通过结构化矩阵算法计算状态空间模型

定理 3.5 的重要性在于它使我们能够将高效计算 SSM（和其他序列模型）的问题简化为结构化矩阵乘法的有效算法。

命题 3.6 (Pernet, Signargout, and Villard (2023))：大小为 T 的 N-SS 矩阵可以用 𝑂 (NT) 个参数表示，并且矩阵向量乘法的时间和空间复杂度为 𝑂 (NT)。

定理 3.7：任何状态大小为 N、序列长度为 T 的状态空间模型（定义 2.2）都可以在 𝑂 (TN) 时间内计算（不考虑潜在的预处理）。

结构化掩码注意力：用结构化矩阵推广线性注意力

在本节中，我们将从第一性原理重新审视线性注意力框架。本节的主要结果是基于张量收缩的线性注意力简单证明（命题 4.1），以及我们在定义 4.2 中对结构化掩码注意力的概括抽象。

线性注意力和其他许多高效注意力变体通常是通过改变核心注意力计算 (𝑄𝐾 ⊤)𝑉 = 𝑄 (𝐾 ⊤𝑉 ) 中矩阵结合性的顺序来实现的。但是，当添加掩码时，推导就有点不那么直接了。

命题 4.1 ((Katharopoulos et al. 2020))：自回归核注意力（即具有因果掩码的掩码核注意力）可以通过每次迭代花费恒定时间的递归在 𝑂 (𝑇 ) 时间内计算。

通过使用张量收缩的视角，我们可以立即看到原始线性注意力的关键在于因果掩码的矩阵向量乘法等价于累积和运算符。

然而，我们观察到注意力掩码不必都是 1。线性注意力快速计算所必需的只是 𝐿 是一个结构化矩阵，根据定义，结构化矩阵具有快速的矩阵乘法（第 2.3 节）。特别是，我们可以使用任何具有次二次（理想情况下是线性）矩阵向量乘法的掩码矩阵 𝐿，通过加速瓶颈方程 (15b)，它将具有与标准线性注意力相同的复杂度。

定义 4.2：结构化掩码注意力 (SMA)（或简称结构化注意力）定义为查询/键/值 𝑄、𝐾、𝑉 以及任何结构化矩阵 𝐿（即具有次二次矩阵乘法）上的函数，通过 4 路张量收缩：

$$
𝑌 = \text{contract}(TN, SN, SP, TS → TP)(𝑄, 𝐾, 𝑉 , 𝐿).
$$

SMA 二次模式算法是由 (13) 定义的成对收缩序列，它对应于标准（掩码）注意力计算。

SMA 线性模式算法是由 (15) 定义的成对收缩序列，其中步骤 (15b) 通过次二次结构化矩阵乘法进行了优化。

状态空间对偶性

在第 3 节和第 4 节中，我们定义了结构化状态空间模型和结构化注意力，讨论了它们的性质，并表明它们都具有二次算法和线性算法。本节将它们联系在一起。我们的主要结果是表明结构化状态空间模型的一个特例与结构化注意力的一个特例重合，并且线性时间 SSM 算法和二次时间核注意力算法是彼此的对偶形式。

为了总结我们的结果：

• 结构化状态空间模型（第 3 节）通常通过线性时间递归来定义。然而，通过扩展表征其线性序列到序列变换的矩阵公式，可以推导出二次形式。
• 注意力变体（第 4 节）是通过二次时间成对交互定义的模型。然而，通过将其视为四路张量收缩并以不同的顺序进行约简，可以推导出线性形式。
• 每个模型的一个自然特例——更准确地说，在 𝐴 矩阵上具有标量恒等结构的状态空间模型，以及在其 𝐿 掩码上具有 1-半可分结构的结构化掩码注意力——是彼此的对偶，具有完全相同的线性和二次形式。

图 4 总结了这两种表示之间的对偶性。

SSD 模型的硬件高效算法

开发 SSM、注意力和结构化矩阵之间的理论 SSD 框架的好处在于利用这些联系来改进模型和算法。在本节中，我们将展示如何从计算结构化矩阵乘法的各种算法中推导出高效计算 SSD 模型的各种算法。

我们的主要计算结果是一种结合了线性（递归）模式和二次（注意力）模式的 SSD 模型计算算法。该算法的计算效率与 SSM（序列长度的线性缩放）一样高，并且与注意力（主要使用矩阵乘法）一样对硬件友好。

定理 6.1：考虑一个状态扩展因子为 N、头部维度为 P = N 的 SSD 模型。存在一种算法可以在任何输入 𝑋 ∈ R(T,P) 上计算模型，该算法只需要 𝑂 (TN2) 个训练 FLOP、𝑂 (TN) 个推理 FLOP、𝑂 (N2) 个推理内存，并且其工作主要由矩阵乘法主导。

定理 6.1 背后的主要思想是再次将计算状态空间模型的问题视为半可分矩阵乘法，但以一种新的方式利用其结构。我们没有以递归或注意力模式计算整个矩阵，而是对矩阵进行块分解。对角块可以使用对偶注意力模式计算，这可以通过矩阵乘法有效地完成，而非对角块可以通过半可分矩阵的秩结构进行分解，并简化为更小的递归。

Mamba-2 架构

通过连接 SSM 和注意力，SSD 框架使我们能够为两者开发共享的词汇表和技术库。在本节中，我们将讨论使用最初为 Transformer 开发的思想来理解和修改 SSD 层的一些示例。我们将讨论几个设计选择，从而形成 Mamba-2 架构。

块设计

我们首先讨论与内部序列混合层无关的神经网络块的修改（即核心 SSD 层之外）。

并行参数投影：Mamba-1 的动机是以 SSM 为中心的观点，其中选择性 SSM 层被视为从 𝑋 ↦→ 𝑌 的映射。SSM 参数 𝐴、𝐵、𝐶 被视为辅助参数，并且是 SSM 输入 𝑋 的函数。因此，定义 (𝐴, 𝐵, 𝐶) 的线性投影发生在创建 𝑋 的初始线性投影之后。

在 Mamba-2 中，SSD 层被视为从 𝐴、𝑋、𝐵、𝐶 ↦→ 𝑌 的映射。因此，在块的开头使用单个投影并行生成 𝐴、𝑋、𝐵、𝐶 是有意义的。请注意与标准注意力架构的类比，其中 𝑋、𝐵、𝐶 对应于并行创建的 𝑄、𝐾、𝑉 投影。

额外的归一化：在初步实验中，我们发现大型模型容易出现不稳定性。我们能够通过在最终输出投影之前的块中添加一个额外的归一化层（例如 LayerNorm、GroupNorm 或 RMSNorm）来缓解这种情况。

序列变换的多头模式

回想一下，SSM 被定义为序列变换（定义 2.1），其中：

• 𝐴、𝐵、𝐶 参数具有状态维度 N。
• 它们定义了一个序列变换 RT → RT，例如可以表示为矩阵 𝑀 ∈ R(T,T)。
• 此变换在输入序列 𝑋 ∈ R(T,P) 上运行，在 P 轴上独立运行。

可以将其视为定义序列变换的一个头。

定义 7.1（多头模式）：多头序列变换由 H 个独立的头组成，总模型维度为 D = d_model。参数可以在头部之间绑定，从而形成头部模式。

状态大小 N 和头部维度 P 分别类似于注意力的 𝑄𝐾 头部维度和 𝑉 头部维度。就像在现代 Transformer 架构中一样，在 Mamba-2 中，我们通常选择这些常量在 64 或 128 左右；当模型维度 D 增加时，我们增加头的数量，同时保持头部维度 N 和 P 固定。为了描述如何做到这一点，我们可以转移和概括多头注意力的思想，为 SSM 或任何一般序列变换定义类似的模式。

多头 SSM (MHS) / 多头注意力 (MHA) 模式：经典的 MHA 模式假设头部维度 P 可以整除模型维度 D。头的数量定义为 H = D/P。然后，通过创建每个参数的 H 个独立副本，创建核心序列变换的 H 个副本。

多收缩 SSM (MCS) / 多查询注意力 (MQA) 模式：多查询注意力 (Shazeer 2019) 是一种巧妙的注意力优化方法，可以显着提高自回归推理的速度，这依赖于缓存 𝐾 和 𝑉 张量。这种技术只是避免给 𝐾 和 𝑉 额外的头部维度，或者换句话说，在 𝑄 的所有头部广播 (𝐾, 𝑉) 的单个头部。

使用状态空间对偶性，我们可以将 MQA 的等效 SSM 版本定义为方程 (18)。在这里，𝑋 和 𝐵（注意力的 𝑉 和 𝐾 的 SSM 类比）在 H 个头部之间共享。我们也称其为多收缩 SSM (MCS) 头部模式，因为控制 SSM 状态收缩的 𝐶 参数每个头部都有独立的副本。

我们也可以类似地定义多键注意力 (MKA) 或多扩展 SSM (MES) 头部模式，其中 𝐵（控制 SSM 扩展）每个头部都是独立的，而 𝐶 和 𝑋 在头部之间共享。

多输入 SSM (MIS) / 多值注意力 (MVA) 模式：虽然 MQA 对于注意力来说是有意义的，因为它有 KV 缓存，但它不是 SSM 的自然选择。相反，在 Mamba 中，𝑋 被视为 SSM 的主要输入，因此 𝐵 和 𝐶 是在输入通道之间共享的参数。我们在方程 (20) 中定义了一种新的多值注意力 (MVA) 或多输入 SSM (MIS) 模式，它可以再次应用于任何序列变换，例如 SSD。

有了这个词汇表，我们可以更精确地描述原始的 Mamba 架构。

命题 7.2：Mamba 架构的选择性 SSM (S6) 层可以被视为具有：

• 头部维度 𝑃 = 1：每个通道都有独立的 SSM 动态 𝐴。
• 多输入 SSM (MIS) 或多值注意力 (MVA) 头部结构：𝐵、𝐶 矩阵（对应于注意力中的 𝐾、𝑄）在输入 𝑋 的所有通道（对应于注意力中的 𝑉）之间共享。

从线性注意力中获得的其他 SSD 扩展

我们在这里描述一个对 SSD 的架构修改示例，其动机来自线性注意力。

核注意力近似于 Softmax 注意力：许多线性注意力或核注意力变体都是通过将注意力分数 softmax(𝑄𝐾 ⊤) 视为由以下两部分组成来实现的：

1. 指数核 𝑍 = exp(𝑄𝐾 ⊤)，可以通过 𝑍 = 𝜓 (𝑄)𝜓 (𝐾)⊤ 来近似，其中 𝜓 为核特征映射。
2. 通过 𝑀 = 𝐺/𝐺11⊤ 对核进行归一化，使行总和为 1，其中除法按元素进行，1 为全 1 向量。

在 Mamba-2 中，我们合并了一个灵活的核特征映射，并将其应用于 𝐵 和 𝐶 分支（对应于注意力中的 𝐾 和 𝑉 分支）。为了简单和对称，也可以选择将特征映射应用于 𝑋 (𝑉) 分支。这在图 6 中由任意非线性表示。默认情况下，我们简单地选择 𝜓 为元素级 Swish / SiLU 函数。

合并归一化（分母）项：为了找到分母项，我们只需要计算 𝑀1。但回想一下，模型的最终输出只是 𝑌 = 𝑀𝑋（方程 (16)）。因此，可以通过用一个额外的列 1 扩充 𝑋 来找到归一化项，从而得到一个形状为 (T, P + 1) 的张量。

请注意，在这种情况下，核特征映射 𝜓 必须为正，以便总和为正。

SSM 的系统优化

我们描述了 SSM（尤其是 Mamba-2 架构）的几种系统优化，用于大规模高效训练和推理。特别是，我们专注于用于大规模训练的张量并行和序列并行，以及用于高效微调和推理的可变长度序列。

张量并行

张量并行 (TP) (Shoeybi et al. 2019) 是一种模型并行技术，它将每一层（例如，注意力、MLP）拆分到多个加速器（例如 GPU）上运行。这种技术被广泛用于在 GPU 集群上训练大多数大型模型，其中每个节点通常有 4-8 个具有快速网络（例如 NVLink）的 GPU。TP 最初是为 Transformer 架构开发的，将其应用于其他架构并不合适。

序列并行

对于非常长的序列，我们可能需要沿着序列长度维度将输入和激活拆分到不同的 GPU。主要有两种技术：

1. 用于残差和归一化操作的序列并行 (SP)：该技术首先由 Korthikanti 等人 (2023) 提出，它将 TP 中的全约简分解为约简分散和全收集。注意到残差和归一化操作在同一 TP 组中的所有 GPU 上对相同的输入重复执行，SP 通过执行以下操作沿序列长度维度拆分激活：约简分散、残差和归一化，然后全收集。

由于 Mamba-2 架构使用相同的残差和归一化结构，因此 SP 无需修改即可应用。

2. 用于标记混合操作（注意力或 SSM）的序列并行，也称为“上下文并行”(CP)。已经为注意力层开发了几种技术（例如，环形注意力 (Liu, Yan, et al. 2024; Liu, Zaharia, and Abbeel 2023)），以及复杂的负载均衡技术 (Brandon et al. 2023)。注意力中序列并行的难点在于我们可以将查询和键拆分为块，但每个查询块都需要与键块交互，从而导致通信带宽与工作器数量呈二次方关系。

对于 SSM，我们可以以一种简单的方式拆分序列：每个工作器获取一个初始状态，根据其输入计算 SSM，返回最终状态，并将该最终状态传递给下一个工作器。通信带宽与工作器数量呈线性关系。这种分解与 SSD 算法（图 5）中用于拆分为块/块的块分解完全相同。我们在图 7（右）中说明了这种上下文并行。

可变长度

虽然预训练通常对批次使用相同的序列长度，但在微调或推理期间，模型可能需要处理不同长度的不同输入序列。处理这种情况的一种简单方法是将批次中的所有序列右填充到最大长度，但如果序列长度差异很大，这可能会很低效。对于 Transformer，已经开发出复杂的技术来避免填充并在 GPU 之间进行负载均衡 (Zeng et al. 2022; Y. Zhai et al. 2023)，或者将多个序列打包到同一个批次中并调整注意力掩码 (Ding et al. 2024; Pouransari et al. 2024)。对于 SSM，特别是 Mamba，我们可以通过简单地将整个批次视为一个长序列来处理可变序列长度，并避免在各个序列之间传递状态。这相当于简单地将一个序列末尾的标记 𝑡 的 𝐴𝑡 设置为 0，以防止它将信息传递给属于不同序列的标记 𝑡 + 1。

实验验证

我们通过实验评估了 Mamba-2 在对递归模型具有挑战性的合成召回任务（第 9.1 节）和标准语言建模预训练和下游评估（第 9.2 节）上的性能。我们验证了我们的 SSD 算法比 Mamba-1（第 9.3 节）效率更高，并且在中等序列长度下可与优化的注意力相媲美。最后，我们对 Mamba-2 架构中的各种设计选择进行了消融研究（第 9.4 节）。

合成：关联召回

合成关联召回任务一直很受欢迎，用于测试语言模型在其上下文中查找信息的能力。广义上讲，它们涉及向自回归模型馈送键值关联对，然后在模型看到先前看到的键时提示模型生成正确的完成。多查询关联召回 (MQAR) 任务是此任务的一种特定形式，它要求模型记住多个关联 (Arora, Eyuboglu, Timalsina, et al. 2024)。最初的 Mamba 论文报告了相关合成任务的结果，特别是选择性复制 (Gu and Dao 2023) 和归纳头 (Olsson et al. 2022)，它们可以被视为更容易的关联召回任务。MQAR 任务也与“电话簿查找”任务密切相关，该任务已被证明对 SSM 等递归模型具有挑战性，因为它们的状态容量有限 (De et al. 2024; Jelassi et al. 2024)。

我们在 (Arora, Eyuboglu, Zhang, et al. 2024) 中提出的具有挑战性的 MQAR 设置版本上进行了比较，使用了更难的任务、更长的序列和更小的模型。我们的基线包括标准的多头 Softmax 注意力以及结合了卷积、局部注意力和线性注意力变体的 Based 架构。

结果如图 8 所示。虽然 Mamba-1 在此任务上表现不佳，但 Mamba-2 在所有设置下都表现良好。令人惊讶的是，即使在控制状态大小 (N = 16) 的情况下，它也明显优于 Mamba-1。（我们不确定架构的哪个方面是主要因素，这仍然是未来工作中需要探索的问题。）此外，此任务验证了状态大小的重要性：从 N = 16 增加到 N = 64 和 N = 256 始终可以提高 MQAR 的性能，因为更大的状态允许记住更多信息（键值对）。

语言建模

按照 LLM 中的标准协议，我们在标准自回归语言建模上训练和评估 Mamba-2 架构，并与其他架构进行比较。我们比较了预训练指标（困惑度）和零样本评估。模型大小（深度和宽度）遵循 GPT3 规范，从 125m 到 2.7B。我们使用 Pile 数据集 (L. Gao, Biderman, et al. 2020)，并遵循 Brown 等人 (2020) 中描述的训练方法。这与 Mamba (Gu and Dao 2023) 中报告的设置相同；训练细节见附录 D。

缩放定律

对于基线，我们与 Mamba 及其 Transformer++ 方法 (Gu and Dao 2023) 进行了比较，后者基于 PaLM 和 LLaMa 架构（例如旋转嵌入、SwiGLU MLP、RMSNorm 而不是 LayerNorm、没有线性偏差和更高的学习率）。由于 Mamba 已经证明它优于标准 Transformer 架构（GPT3 架构）以及最近的次二次架构（H3 (Dao, D. Y. Fu, et al. 2023)、Hyena (Poli et al. 2023)、RWKV-4 (B. Peng, Alcaide, et al. 2023)、RetNet (Y. Sun et al. 2023)），为了清晰起见，我们在图中省略了这些架构（有关比较，请参见 Gu and Dao (2023)）。

图 9 显示了在标准 Chinchilla (Hoffmann et al. 2022) 协议下，从 ≈ 125𝑀 到 ≈ 1.3𝐵 参数的模型的缩放定律。

下游评估

表 1 显示了 Mamba-2 在一系列流行的下游零样本评估任务上的性能，并与这些规模下最著名的开源模型进行了比较，最重要的是 Pythia (Biderman et al. 2023)，后者使用与我们的模型相同的标记器、数据集和训练长度（300B 标记）进行训练。

混合模型：将 SSD 层与 MLP 和注意力相结合

最近和同时进行的研究 (Dao, D. Y. Fu, et al. 2023; De et al. 2024; Glorioso et al. 2024; Lieber et al. 2024) 表明，同时具有 SSM 层和注意力层的混合架构可以提高模型质量，使其优于 Transformer 或纯 SSM（例如，Mamba）模型，尤其是在上下文学习方面。我们探索了将 SSD 层与注意力和 MLP 相结合的不同方式，以了解每种方式的好处。根据经验，我们发现大约 10% 的层是注意力层时性能最佳。将 SSD 层、注意力层和 MLP 相结合也比纯 Transformer++ 或 Mamba-2 效果更好。

SSD 和注意力：我们发现 SSD 和注意力层是互补的：它们本身（例如，在 Mamba-2 架构与 Transformer++ 中）的性能（以困惑度衡量）几乎相同，但 SSD 和注意力层的混合优于纯 Mamba-2 或 Transformer++ 架构。我们在表 2 中展示了一些结果，这些结果是使用 GPT-2 标记器在 Pile 上训练到 7B 标记的 350M 模型（48 层）（参数数量相同、超参数相同、训练和验证集相同）。仅添加几个注意力层就已经产生了显着的改进，并在质量和效率之间取得了最佳平衡。我们假设 SSM 层可以很好地用作一般的序列到序列映射，而注意力层充当检索机制，可以快速引用序列中的先前标记，而不是强迫模型将其所有上下文压缩到其内存（SSM 状态）中。

具有 SSD、MLP 和注意力的混合模型：我们比较了将 SSD 与（门控）MLP 和注意力层相结合的不同方式，并在 Pile 上训练到 300B 标记的 2.7B 规模（64 层）上进行了评估（参数数量相同、超参数相同、训练和验证集相同、数据顺序相同）：

1. Transformer++：32 个注意力层和 32 个门控 MLP，交错。
2. Mamba-2：64 个 SSD 层。
3. Mamba-2-MLP：32 个 SSD 和 32 个门控 MLP 层，交错。
4. Mamba-2-Attention：58 个 SSD 层和 6 个注意力层（位于索引 9、18、27、36、45、56）。
5. Mamba-2-MLP-Attention：28 个 SSD 层和 4 个注意力层，与 32 个门控 MLP 层交错。

我们在表 3 中报告了 Pile 上的验证困惑度以及零样本评估。总的来说，Transformer++ 和 Mamba-2 模型的质量大致相同。我们看到仅添加 6 个注意力层就显着改善了纯 Mamba-2 模型（以及 Transformer++）。添加 MLP 层会降低模型质量，但可以 (i) 由于 MLP 层的简单性和硬件效率而加快训练和推理速度；(ii) 通过用专家混合替换 MLP 层，更容易升级到 MoE 模型。

速度基准

我们将 SSD 算法的速度与 Mamba 的扫描实现和 FlashAttention-2 进行了基准测试（图 10）。SSD 由于其重新制定以使用矩阵乘法作为子例程，因此可以利用 GPU 上的专用矩阵乘法 (matmul) 单元，也称为张量核。因此，它比 Mamba 的融合关联扫描快 2-8 倍，后者不利用 matmul 单元。由于其序列长度的线性缩放，SSD 从序列长度 2𝐾 开始就比 FlashAttention-2 快。

但是，我们注意到，在短序列长度（例如 2𝐾）下，整个 Mamba-2 模型的训练效率可能不如 Transformer，因为具有 𝐿 层的 Transformer 将具有 𝐿/2 个注意力层和 𝐿/2 个 MLP 层，而 Mamba-2 模型在参数数量相同的情况下将具有 𝐿 个 SSD 层。通常，MLP 层的硬件效率非常高，因为它们由简单的矩阵乘法和逐点线性组成。如第 9.2.3 节所示，也可以将 𝐿/2 个 SSD 层和 𝐿/2 个 MLP 层组合起来，以加快短序列长度下的训练速度。

架构消融研究

块设计

第 7.1 节介绍了 Mamba-2 块，它对 Mamba-1 块进行了一些小的修改，部分原因是为了与注意力联系起来，也是为了提高 Mamba-2 的可扩展性。表 4 对块的这些架构更改进行了消融研究，这些更改发生在核心 SSM 层之外。

消融研究验证了并行投影创建 (𝐴, 𝐵, 𝐶, 𝑋) 可以节省参数，并且性能略好于 Mamba 的顺序投影。更重要的是，这种修改适用于大型模型的张量并行（第 8 节）。此外，额外的归一化层也略微提高了性能。更重要的是，在更大规模的初步实验中观察到它还有助于训练稳定性。

头部结构

第 7.2 节描述了如何将 𝐵、𝐶、𝑋 投影的维度视为类似于多头注意力和多查询注意力概念的超参数。我们还展示了原始的 Mamba 架构如何类似于多值注意力（命题 7.2），这是从状态空间模型的角度自然发展而来的选择，之前没有进行过消融研究。

表 5 对 Mamba-2 架构的多头结构选择进行了消融研究。引人注目的是，我们发现多值和多查询或多键头部模式之间存在很大差异，尽管它们看起来非常相似。请注意，这不能用总状态大小来解释，所有这些模式的总状态大小都相同（等于 HPN 或头部数量、头部维度和状态维度的乘积）。

我们还比较了 𝐶、𝐵、𝑋（类似于 𝑄、𝐾、𝑉）头部的数量相等的多头模式。我们与标准的多头模式以及一种积极共享的模式进行了比较，在积极共享的模式中，它们都只有一个头部。请注意，在后一种情况下，模型仍然有 H 个不同的序列混合器 𝑀，因为每个头部仍然有不同的 𝐴。当参数匹配时，这些多头模式的性能彼此相似，介于 MVA 和 MQA/MKA 模式之间。

注意力核近似

第 7.3 节指出 SSD 如何与线性注意力文献中的思想相结合，例如各种形式的核近似。我们在表 6 中对先前工作建议的几种变体进行了消融研究。这些变体包括 cosFormer (Qin, Weixuan Sun, et al. 2022)、随机特征注意力 H. Peng et al. 2021 和正随机特征 (Performer) (Choromanski et al. 2021)。

我们还对添加归一化项进行了消融研究，类似于标准注意力中 Softmax 函数的分母。我们发现这会给大多数变体带来不稳定性，但会略微提高 ReLU 激活函数 𝜓 的性能。

表 7 还测试了最近提出的改进线性注意力的建议，这些建议涉及扩展特征维度（Based (Arora, Eyuboglu, Zhang, et al. 2024) 和 ReBased (Aksenov et al. 2024)）。这些线性注意力扩展旨在使用二次近似来逼近 exp 核。ReBased 还建议用层归一化替换 QK 激活函数；从以 SSM 为中心的观点来看，我们在应用 SSM 函数之前对 (𝐵, 𝐶) 应用归一化。

我们注意到，这种技术已被独立地提出作为 Softmax 注意力的“QK-Norm”（Team 2024）和 Mamba 的“内部归一化”（Lieber et al. 2024）。

总的来说，表 6 和表 7 发现，我们尝试的核近似方法似乎并没有比简单的逐点非线性激活函数 𝜓 有所改进。因此，我们对 Mamba-2 的默认设置使用 𝜓 (𝑥) = Swish(𝑥) 来遵循 Mamba-1，但我们建议完全删除此激活可能是一个更简单的选择，我们没有对其进行广泛测试。

我们强调，SSD 和普通线性注意力在包含 1-半可分掩码 𝐿 方面有所不同，而文献中的各种线性注意力方法都是为了在没有此项的情况下近似 Softmax 注意力而推导出来的；因此，我们的负面结果可能并不意外。

相关工作和讨论

状态空间对偶性框架将 SSM、结构化矩阵和注意力之间的联系联系起来。我们将更深入地讨论 SSD 与这些概念之间更广泛的关系。利用每个观点的思想，我们还提出了一些未来工作中可以扩展 SSD 框架的方向。

状态空间模型

结构化状态空间模型可以沿着以下轴进行表征：

(i) 它是否是时不变的或时变的。
(ii) 系统的维数。
(iii) 递归转换 𝐴 上的结构。

SSD 可以描述为具有 SISO 维度和标量恒等结构的选择性 SSM。

时变性（选择性）：最初的结构化 SSM (S4) 是线性时不变 (LTI) 系统 (Gu 2023; Gu, Goel, and Ré 2022)，其动机是连续时间在线记忆 (Gu, Dao, et al. 2020; Gu, Johnson, Goel, et al. 2021; Gu, Johnson, Timalsina, et al. 2023)。已经提出了许多结构化 SSM 的变体 (Dao, D. Y. Fu, et al. 2023; Gu, Gupta, et al. 2022; Gupta, Gu, and Berant 2022; Ma et al. 2023; J. T. Smith, Warrington, and Linderman 2023)，包括几种放弃递归并专注于 LTI SSM 的卷积表示的变体 (D. Y. Fu et al. 2023; Y. Li et al. 2023; Poli et al. 2023; Qin, Han, Weixuan Sun, B. He, et al. 2023)。

SSD 是一种时变结构化 SSM，也称为 Mamba (Gu and Dao 2023) 中引入的选择性 SSM。选择性 SSM 与 RNN 的门控机制密切相关，包括经典的 RNN，例如 LSTM (Hochreiter and Schmidhuber 1997) 和 GRU (J. Chung et al. 2014)，以及更现代的变体，例如 QRNN (Bradbury et al. 2016)、SRU (Lei 2021; Lei et al. 2017)、RWKV (B. Peng, Alcaide, et al. 2023)、HGRN (Qin, Yang, and Zhong 2023) 和 Griffin (Botev et al. 2024; De et al. 2024)。这些 RNN 在参数化方面有所不同，最重要的是缺乏状态扩展。

维数和状态扩展： SSD 的一个重要特征是它是一个单输入单输出 (SISO) 系统，其中输入通道独立处理，这与之前其谱系中的 SSM（S4、H3、Mamba）相同。这导致更大的有效状态大小 ND，其中 N 是 SSM 状态大小（也称为状态扩展因子），D 是标准模型维度。传统的 RNN 要么具有 N = 1，要么是具有密集 𝐵、𝐶 矩阵的多输入多输出 (MIMO)，这两者都会导致状态更小。虽然 MIMO SSM 已被证明在某些领域中效果良好 (Lu et al. 2023; Orvieto et al. 2023; J. T. Smith, Warrington, and Linderman 2023)，但 Mamba 表明状态扩展对于信息密集型领域（例如语言）至关重要。SSD 的主要优势之一是允许更大的状态扩展因子，而不会降低模型速度。此后，许多后续工作都采用了状态扩展（第 10.4 节）。

结构：与之前的结构化 SSM 相比，SSD 的主要限制在于状态转换 𝐴𝑡 的表达能力。我们注意到，更通用的 SSM（例如对角线 𝐴𝑡 的情况）具有与 SSD 相同的理论效率，但对硬件不太友好。这是因为对偶二次形式失去了类似注意力的解释，并且变得更难计算。因此，与 Mamba 相比，SSD 仅在对角线 𝐴𝑡 的限制性稍强的形式上有所不同，并以这种表达能力换取了更高的硬件效率（以及易于实现）。

我们假设可以改进我们的结构化矩阵算法，以改进到通用对角线 SSM 的情况。

结构化矩阵

状态空间对偶性的第一个观点是将这些模型视为矩阵序列变换或“矩阵混合器”：可以表示为沿序列维度 T 进行矩阵乘法（通过 T × T 矩阵）的序列变换（定义 2.1）。

之前已经提出了几种这样的矩阵混合器，其中主要的变化轴是矩阵的表示。这些包括 MLP-Mixer (Tolstikhin et al. 2021)（非结构化矩阵）、FNet (Lee-Thorp et al. 2021)（傅里叶变换矩阵）、M2 (Dao, B. Chen, et al. 2022; Dao, Gu, et al. 2019; Dao, Sohoni, et al. 2020; D. Fu et al. 2024)（蝴蝶/帝王蝶矩阵）、Toeplitz 矩阵 (Poli et al. 2023; Qin, Han, Weixuan Sun, B. He, et al. 2023)，甚至更奇特的结构 (De Sa et al. 2018; Thomas et al. 2018)。

一个重要的特征是，高效（次二次）矩阵序列变换正是那些具有结构化矩阵混合器的变换。SSD 框架的核心结果是将 SSM 视为具有特定结构的矩阵混合器——半可分矩阵（第 3 节）。然后，线性与二次对偶性采用结构化矩阵乘法与朴素矩阵乘法的形式。

结构矩阵表示通过特定半可分矩阵的块分解导致了我们高效的 SSD 算法（第 6 节）。我们注意到，半可分矩阵在科学计算文献中得到了很好的研究，结合这些思想可能是进一步改进状态空间模型的一个有希望的途径。我们还建议，关注矩阵混合器观点可以为序列模型带来更多富有成效的方向，例如设计有原则的非因果 Mamba 变体，或者找到通过分析其矩阵变换结构来表征和弥合 Softmax 注意力和次二次模型之间差距的方法。

（线性）注意力

与标准（因果）注意力相比，SSD 只有两个主要区别。

首先，SSD 不使用标准注意力的 Softmax 激活 (Bahdanau, Cho, and Bengio 2015; Vaswani et al. 2017)，这是注意力具有二次复杂性的原因。当删除 Softmax 时，可以通过线性注意力框架以线性缩放计算序列 (Katharopoulos et al. 2020)。

其次，SSD 将 logits 矩阵乘以一个输入相关的 1-半可分掩码。因此，可以将此掩码视为替换标准注意力中的 Softmax。

此半可分掩码也可以视为提供位置信息。元素 𝑎𝑡 充当 RNN 意义上的“门”，或“选择”机制（参见 Mamba 论文中的讨论），它们的累积乘积 𝑎 𝑗:𝑖 控制位置 𝑖 和 𝑗 之间允许的交互量。位置嵌入（例如正弦 (Vaswani et al. 2017)、AliBi (Press, N. Smith, and Lewis 2022) 和 RoPE (Su et al. 2021)）是 Transformer 的重要组成部分，通常被视为启发式方法，SSD 的 1-SS 掩码可以被视为一种更有原则的相对位置嵌入形式。我们注意到 GateLoop (Katsch 2023) 也同时提出了这种观点。

状态空间对偶性的第二个观点是我们更通用的结构化掩码注意力 (SMA) 框架的特例，其中对偶性表现为对简单的 4 路张量收缩的不同收缩顺序。SMA 是线性注意力的强泛化，它也比 SSD 通用得多；其他形式的结构化掩码可能会导致更多具有不同于 SSD 属性的高效注意力变体。

除了导致新模型之外，这些与注意力的联系还可以为理解 SSM 提供其他方向。例如，我们很好奇注意力汇聚现象 (Darcet et al. 2024; Xiao et al. 2024) 是否存在于 Mamba 模型中，以及更广泛地说，可解释性技术是否可以转移到 SSM (Ali, Zimerman, and Wolf 2024)。

最后，已经提出了许多其他线性注意力变体 (Arora, Eyuboglu, Timalsina, et al. 2024; Arora, Eyuboglu, Zhang, et al. 2024; Choromanski et al. 2021; H. Peng et al. 2021; Qin, Han, Weixuan Sun, Dongxu Li, et al. 2022; Qin, Weixuan Sun, et al. 2022; Schlag, Irie, and Schmidhuber 2021; Zhang et al. 2024; Zheng, C. Wang, and Kong 2022)（有关其中几种变体的描述，请参见第 4.1.3 节），我们预计许多技术可以转移到 SSM（例如第 7.3 节）。

我们强调，SSD 不会泛化标准 Softmax 注意力，或对注意力核矩阵的任何其他没有有限特征映射 𝜓 的变换。与通用注意力相比，SSD 的优势在于具有可控的状态扩展因子 N，可以压缩历史记录，而二次注意力的缓存则包含整个历史记录，其大小与序列长度 T ≫ N 成正比。同时进行的研究已经开始研究这些表示的权衡，例如在复制和上下文学习任务上 (Akyürek et al. 2024; Grazzi et al. 2024; Jelassi et al. 2024; Park et al. 2024)。我们注意到，Mamba-2 在其中一些能力上显着改进了 Mamba（例如，如第 9.1 节中 MQAR 结果所示），但还有更多需要了解的地方。

相关模型

最后，我们重点介绍了越来越多的最近和同时进行的工作，这些工作开发了与 Mamba 和 Mamba-2 非常相似的序列模型。

• RetNet (Y. Sun et al. 2023) 和 TransNormerLLM (Qin, Dong Li, et al. 2023) 使用衰减项而不是累积和来泛化线性注意力，并提出了对偶并行/递归算法以及混合“分块”模式。这些算法可以看作是 𝐴𝑡 时不变（对于所有 𝑡 都是常数）的 SSD 实例；在 SMA 解释中，掩码矩阵 𝐿 将是一个衰减矩阵 𝐿𝑖,𝑗 = 𝛾𝑖 − 𝑗。这些模型在架构上也有各种不同。例如，由于它们是从以注意力为中心的视角推导出来的，因此它们保留了多头注意力 (MHA) 模式；由于 Mamba-2 是从以 SSM 为中心的模式推导出来的，因此它保留了多值注意力 (MVA) 或多扩展 SSM (MES) 模式，我们证明这种模式更好（第 9.4 节）。
• GateLoop (Katsch 2023) 同时提出了使用输入相关的衰减因子 𝐴𝑡，并开发了与 SSD 中相同的对偶二次形式，他们称之为“代理注意力”形式。
• 门控线性注意力 (GLA) (Yang et al. 2024) 提出了一种具有数据相关门的线性注意力变体，以及用于计算分块模式和硬件感知实现的高效算法。
• HGRN (Qin, Yang, and Zhong 2023) 引入了一种具有输入相关门的 RNN，在 HGRN2 (Qin, Yang, Weixuan Sun, et al. 2024) 中通过结合状态扩展对其进行了改进。
• Griffin (De et al. 2024) 和 RecurrentGemma (Botev et al. 2024) 表明，具有输入相关门控的 RNN 与局部注意力相结合，可以与强大的现代 Transformer 非常有竞争力。Jamba 还表明，将 Mamba 与几层注意力相结合在语言建模方面表现出色 (Lieber et al. 2024)。
• xLSTM (Beck et al. 2024) 通过采用状态扩展的思想以及其他门控、归一化和稳定技术来改进 xLSTM。
• RWKV(-4) (B. Peng, Alcaide, et al. 2023) 是一种基于不同线性注意力近似（无注意力 Transformer (S. Zhai et al. 2021)）的 RNN。最近，通过采用选择性和状态扩展的思想，将其改进为 RWKV-5/6（Eagle 和 Finch）架构 (B. Peng, Goldstein, et al. 2024)。

结论

我们提出了一个基于结构化矩阵的理论框架，该框架弥合了 SSM 和注意力变体之间的概念差距。该框架为最近的 SSM（例如 Mamba）如何在语言建模方面与 Transformer 表现一样好提供了见解。此外，我们的理论工具通过连接两方面的算法和系统进步，为改进 SSM（以及潜在的 Transformer）提供了新思路。作为演示，该框架指导我们设计了一种新的架构 (Mamba-2)，该架构位于 SSM 和结构化注意力的交汇处。

致谢

我们感谢 Angela Wu 就如何以数值稳定的方式有效计算 Δ 的梯度提出的建议。我们感谢 Sukjun Hwang 和 Aakash Lahoti 在 MQAR 实验中提供的帮助。

参考文献

https://arxiv.org/pdf/2405.21060