定位和编辑 Mamba 中的事实关联
本文通过信息流的视角研究了基于 Transformer 的语言模型如何在推理过程中检索参数中的实际知识,并通过对模型中注意力机制的介入实现了对预测中信息流的详细分析和理解。通过此研究,我们阐述了在语言模型中实现知识局部化和编辑的方法。
Apr, 2023
这项研究提供了实证证据,证明了一种新提出的选择性结构化状态空间模型 Mamba 具有类似于 transformers 的上下文学习(ICL)能力。我们在涉及简单函数逼近和更复杂的自然语言处理问题的任务上评估了 Mamba。我们的结果表明,在两类任务中,Mamba 与 transformer 模型在 ICL 方面的性能相当。进一步的分析表明,像 transformer 一样,Mamba 似乎通过逐步优化其内部表示来解决 ICL 问题。总体而言,我们的研究表明,对于涉及更长输入序列的 ICL 任务,Mamba 可以作为 transformers 的高效替代品。
Feb, 2024
本文使用因果干预技术研究了自回归转换语言模型中实际关联的存储和检索,并发现这些关联对应于本地化的可直接编辑的计算。研究表明中间层前馈模块在存储实际关联方面具有重要作用,并且为模型编辑提供了直接操作计算机制的方法。
Feb, 2022
基于 Transformer 架构的基础模型凭借其核心注意力模块,驱动着深度学习中大部分令人兴奋的应用。我们发现这种模型的一个关键弱点是其无法进行内容导向的推理,并对此进行了改进,通过让结构状态空间模型(SSMs)参数成为输入的函数来解决离散模态的弱点,该模型在长度可选的维度上选择性地传播或遗忘信息,并且通过在递归模式下设计一种硬件感知并行算法,将这些选择性 SSMs 集成到简化的端到端神经网络架构中。该模型(Mamba)具有快速推断速度(比 Transformers 快 5 倍)和序列长度的线性扩展,并在实际数据上对长达百万长度的序列显示出改进。作为一种基于通用序列模型的支持,Mamba 在语言、音频和基因组等多个模态上实现了最先进的性能。在语言建模中,我们的 Mamba-3B 模型在预训练和下游评估中均优于同样大小的 Transformers,与其两倍大小的模型性能相当。
Dec, 2023
探索了 Mamba 模型与线性注意力 Transformer 之间的相似性和差异性,发现忘记门和块设计是 Mamba 模型成功的核心因素,并通过融合这两个关键设计提出了 Mamba-Like Linear Attention (MLLA) 模型,该模型在图像分类和高分辨率密集预测任务中优于多种视觉任务中的 Mamba 模型,同时具有可并行计算和快速推理速度。
May, 2024
Mamba4KT 是第一个探索知识追踪中增强效率和资源利用的新模型,通过对顺序级别和题目级别的解释性进行研究,实验证明其在三个公共数据集上达到可比较的预测准确性,同时显著提高了训练和推理效率以及资源利用效率。
May, 2024
选择性状态空间模型(SSMs)如 Mamba 克服了 Transformer 的一些缺点,例如与序列长度呈二次增长的计算复杂度和从键值缓存中获取大量的推理时间内存需求。此外,最近的研究显示,SSMs 可以达到或超越 Transformer 的语言建模能力,使其成为一种有吸引力的替代选择。然而,迄今为止的研究只在相同数据的受控环境中进行了小规模实验,比较了 SSMs 和 Transformers。为了了解这些体系结构在更大规模上的优缺点,我们在相同数据集上直接比较了 8B 参数的 Mamba、Mamba-2 和 Transformer 模型,数据集涵盖了多达 3.5T 个标记。我们还将这些模型与由 43% 的 Mamba-2、7% 的注意力和 50% 的 MLP 层(Mamba-2-Hybrid)组成的混合体系结构进行了比较。通过使用多种任务,我们回答了 Mamba 模型是否能在较大的训练预算下与 Transformers 相匹配的问题。我们的结果表明,纯 SSMs 在许多任务上达到或超越了 Transformers,但在需要强大的复制或上下文学习能力(例如,5-shot MMLU、电话簿)或长期推理的任务上,它们落后于 Transformers。相反,我们发现 8B 的 Mamba-2-Hybrid 在我们评估的所有 12 个标准任务上超过了 8B Transformer(平均增加 2.65 个点),并且在生成推理标记时预计速度最多快 8 倍。为了验证长期上下文能力,我们进行了其他实验,评估了支持 16K、32K 和 128K 序列的 Mamba-2-Hybrid 和 Transformer 的变体。在 23 个额外的长期上下文任务中,混合模型在平均水平上继续紧密匹配或超越了 Transformer。为了进一步的研究,我们将检查点以及用于训练我们的模型的代码作为 NVIDIA 的 Megatron-LM 项目的一部分发布。
Jun, 2024
使用混合模型能够增强在语言模型中的上下文学习能力,这项研究对几种模型进行了比较,发现混合模型可以在某些任务中克服各个模型独立处理时的局限性。
Feb, 2024
通过在经典的 IR 任务 - 文档排序中考察 Mamba 的效果,我们发现 Mamba 模型在与具有相同训练方法的基于 Transformer 的模型相比具有有竞争力的性能,但与 flash attention 等高效的 Transformer 实现相比,训练吞吐量较低。我们希望这项研究能够成为探索 Mamba 模型在其他经典 IR 任务中的起点。
Mar, 2024
借助相关性传播的层次和 Mamba 架构的分析,我们提出了 MambaLRP 算法,以增强深度学习模型的可解释性,并取得了多个模型和数据集上的最先进解释性能。
Jun, 2024