DenseMamba：大模型的DenseNet时刻，Mamba和RetNet精度显著提升 _大模型

近期，来自华为诺亚方舟实验室的研究者提出了 DenseSSM，用于增强 SSM 中各层间隐藏信息的流动。通过将浅层隐藏状态有选择地整合到深层中，DenseSSM 保留了对最终输出至关重要的精细信息。DenseSSM 在保持训练并行性和推理效率的同时，通过密集连接实现了性能提升。该方法可广泛应用于各种 SSM 类型，如 Mamba 和 Re.NET 。

随着 ChatGPT 的突破性进展，大型语言模型（LLMs）迎来了一个崭新的里程碑。这些模型在语言理解、对话交互和逻辑推理方面展现了卓越的性能。过去一年，人们目睹了 LLaMA、ChatGLM 等模型的诞生，它们基于 Transformer 架构，采用多头自注意力（MHSA）机制来捕捉词汇间的复杂关系，尽管 MHSA 模块在模型中扮演着核心角色，但其在推理过程中对计算和内存资源的需求却极为庞大。具体来说，对于长度为 N 的输入句子，自注意力的计算复杂度高达 O (N^2) ，而内存占用则达到了 O (N^2D)，其中 D 是模型的维度。
为了应对这一挑战，最新的研究致力于简化 Transformer 架构，以降低其在计算和空间上的复杂度。研究者们探索了多种创新方法，包括卷积语言模型、循环单元、长上下文模型，以及状态空间模型（SSMs）。这些新兴技术为构建高效能的 LLMs 提供了强有力的替代方案。SSMs 通过引入高效的隐藏状态机制，有效处理长距离依赖问题，同时保持了训练的并行性和推理的高效率。隐藏状态能够在时间维度上传递信息，减少了在每一步中访问历史词汇的计算负担。通过状态转移参数 A，隐藏状态能够将前一时间步的信息传递至当前时间步，实现对下一个词汇的自回归预测。
尽管隐藏状态在 SSMs 中起着至关重要的作用，但其在以往的研究中并未得到充分研究。不同层的权重和隐藏特征包含了从细粒度到粗粒度的多层次信息。然而，在早期的 SSMs 版本中，隐藏状态仅在当前层内流动，限制了其传递更深层信息的能力，从而影响了模型捕获丰富层次信息的能力。
为了解决这个挑战，华为诺亚方舟实验室的科研团队发表了新工作《DenseMamba: State Space Models with Dense Hidden Connection for Efficient Large Language Models》, 提出一个适用于各类 SSM 模型例如 Mamba 和 RetNet 的 DenseSSM 方法，该方法有选择地将浅层隐藏状态整合到深层，保留了对最终输出至关重要的浅层细粒度信息，以增强深层感知原始文本信息的能力。

DenseMamba：大模型的DenseNet时刻，Mamba和RetNet精度显著提升

文章插图

论文链接：https://arxiv.org/abs/2403.00818
项目主页：https://Github.com/WAIlordHe/DenseSSM

文章首先分析了状态空间模型（SSMs）中的隐藏状态退化问题，

文章插图
上标 “l” 表示第 l 个块。其中，Θ(·) 是从 SSM 模块的最后一个输出到输入的转换，例如卷积和前馈网络（FFN）。从公式 (7) 可以看出，从第 (l-m) 层到第 l 层的隐藏信息传递需要经过 m 个变换块和 m 次 BC 矩阵乘法。这样复杂的计算过程可能导致显著的信息丢失，这意味着在第 l 层尝试检索浅层的某些信息变得非常困难和不清晰。
方法
密集（Dense）隐藏层连接
在上述分析中发现随着层深度的增加，SSM 中重要隐藏状态的衰减。因此，DenseSSM 提出了一种密集连接的隐藏状态方法，以更好地保留来自浅层的细粒度信息，增强深层感知原始文本信息的能力。对于第 l 个块，DenseSSM 在其前 m 个块中密集连接隐藏状态。