[论文解读] DenseMamba: State Space Models with Dense Hidden Connection for Efficient Large Language Models

Wei He, Kai Han+5 more|arXiv (Cornell University)|Feb 26, 2024

Topic Modeling被引用 5

一句话总结

DenseMamba 引入 DenseSSM，一种用于状态空间模型的密集隐藏连接机制，以改善跨层信息流，提升 RetNet/Mamba 风格架构在极小参数和计算开销下的性能。

ABSTRACT

Large language models (LLMs) face a daunting challenge due to the excessive computational and memory requirements of the commonly used Transformer architecture. While state space model (SSM) is a new type of foundational network architecture offering lower computational complexity, their performance has yet to fully rival that of Transformers. This paper introduces DenseSSM, a novel approach to enhance the flow of hidden information between layers in SSMs. By selectively integrating shallowlayer hidden states into deeper layers, DenseSSM retains fine-grained information crucial for the final output. Dense connections enhanced DenseSSM still maintains the training parallelizability and inference efficiency. The proposed method can be widely applicable to various SSM types like RetNet and Mamba. With similar model size, DenseSSM achieves significant improvements, exemplified by DenseRetNet outperforming the original RetNet with up to 5% accuracy improvement on public benchmarks. code is avalaible at https://github.com/WailordHe/DenseSSM

研究动机与目标

• 激励使用高效序列模型（SSM）作为大型语言模型的 Transformer 替代方案。
• 识别 SSM 层之间隐藏状态的退化及其对信息流的影响。
• 提出 DenseSSM，通过将浅层隐藏状态密集连接到更深层同时保持训练/推理效率。
• 证明 DenseSSM 在公开基准上提升 RetNet 和 Mamba 变体的性能。
• 提供切实可行、对硬件友好的实现策略，用于选择性过渡模块和隐藏融合模块。

提出的方法

• 通过选择性过渡模块 φ 和隐藏融合模块 Fuse，将前 m 层的隐藏状态密集连接到当前层，提出 DenseSSM。
• 在 φ 中使用投影+门控机制，将浅层隐藏状态投影到目标子空间并选择有用的分量。
• 通过加性融合（h_t^L = h_t^L + sum φ(h_t^{l-i})）将选定的浅层状态与当前层的隐藏状态融合。
• 将 DenseSSM 扩展到递归（DenseRetNet）和卷积/GAU -like 块变体，以保持并行化训练和高效自回归推理。
• 将 DenseSSM 应用于 RetNet 和 Mamba 架构，创建 DenseRetNet 与 DenseMamba，模型规模相似但性能提升。
• 使用 56GB The Pile-derived 数据、15B tokens、LLaMA 分词器，在多样化的问答/常识任务上评估零-shot 和 4-shot 能力。

实验结果

研究问题

• RQ1 引入跨 SSM 层的密集隐藏连接是否能保留或提升信息流，相较于传统的仅限于层的隐藏状态传播？
• RQ2 DenseSSM 在不牺牲训练/推理效率的前提下，是否能提高基于 RetNet 与 Mamba 的模型的准确性和困惑度？
• RQ3 选择性过渡模块和隐藏融合模块如何影响性能与参数效率？
• RQ4 在零-shot 和少数-shot（4-shot）评估的 commonsense QA 与推理基准上，改进是否具有一致性？
• RQ5 DenseSSM 是否能在不同的 SSM 变体（如 RetNet、Mamba）中，在最小的架构改动下实现适用？

主要发现

• DenseSSM 通过密集连接浅层隐藏状态到深层，并结合选择性投影与门控机制，增强隐藏信息流。
• DenseRetNet 和 DenseMamba 在公开基准上，在零-shot 和少数-shot 设置下，表现优于基线 RetNet、Mamba 和基于 Transformer 的模型。
• DenseRetNet 在公开基准上相较于原始 RetNet 实现最高达 5% 的准确性提升，如摘要所述。
• DenseSSM 保持训练并行性与自回归推理效率，只有参数数量略有增加。
• 消融研究表明，使用 Identity 投影配合 MLP 的选择性过渡，以及加性融合，能在性能与参数数量之间取得有利的折衷。