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

文章插图
这个过程可以通过对输入序列进行卷积来实现：

文章插图

文章插图
在文章所提出的 DenseSSM 中，可以获得隐藏状态加强的 SSM 的输出：

文章插图
DenseSSM 方法的并行实现示例图：

文章插图
【DenseMamba：大模型的DenseNet时刻，Mamba和RetNet精度显著提升】Selective Transition Module （选择性转换模块）

文章插图
选择性转换模块 φ(·) 的目的是将输入投影到目标子空间，并同时选择隐藏信息的有用部分。通过投影层和门控选择机制实现了选择性转换模块，如上图所示。首先，前 m 个 SSM 块中的隐藏状态会被投影到相同的空间：

文章插图
然后，根据输入生成门控权重，并使用它们来选择有用的隐藏状态：

文章插图

文章插图
在实践中作者保持了简单且高效的实现。投影层使用线性变换实现，而门控模块则使用参数高效的带有激活函数的两层 MLP 。
Hidden Fusion Module (隐藏层融合模块)
选择性转换模块后从浅层获得了选择的隐藏状态，即后，DenseSSM 方法利用一个隐藏融合模块将这些精选的浅层隐藏状态与当前层的隐藏状态结合起来。由于这些精选状态已经被投影到相同的空间，因此可以简单地将它们累加到当前层的隐藏状态上：

文章插图

文章插图
为了保持模型的高效性，其他可能的实现方式，例如拼接和交叉注意力机制没有被使用。
扩展到 RetNet
RetNet 可以被视为一种状态空间模型，它利用线性注意力来简化自注意力的计算复杂度。与标准 Transformer 相比具有快速推理和并行化训练兼得的优势。

文章插图
其中，是循环状态，RetNet 的密集 KV 连接执行方式如下。首先，浅层的 K 和 V 被连接起来：

文章插图

文章插图
然后，这些 K 和 V 被注入到当前层的原始键（或值）中：

文章插图
配备了使用所提出 DenseSSM 方法的密集键值（KV）连接的 RetNet 被称为 DenseRetNet，如下图所示。

文章插图
此外，DenseRetNet 也可以在并行模式下实现，也就是说，可以在 GPU 或 NPU 上并行训练。DenseRetNet 的并行模式公式如下：

文章插图
实验
文章进行了全面的实验，以验证所提出的 DenseSSM 的有效性。这些实验在不同的架构上进行，包括 RetNet 和 Mamba 。
预训练数据
在实验中，选择了 The Pile 数据集的一个子集，并从头开始训练所有模型。为了确保训练集包含 150 亿（15B）个 tokens，对数据集进行了随机抽样。在所有实验中，统一使用了 LLaMA 分词器来处理这些数据。
评估数据集
在评估模型性能时，特别关注了模型在多种下游任务上的零样本和少样本学习能力。这些任务包括了一系列测试常识推理和问答的数据集，例如 HellaSwag、BoolQ、COPA、PIQA、Winograd、Winogrande、StoryCloze、OpenBookQA、SciQ、ARC-easy 和 ARC-challenge 。此外，文章还报告了 WikiText 和 LAMBADA 的词困惑度指标。所有评估都通过使用 LM evaluation harness 标准化的评估工具进行，以确保评估模型能力的一致性。
实验设置
为了验证提出的 DenseSSM 机制的有效性，选择了 350M 和 1.3B 两种模型规格进行实验。所有模型都是从头开始训练的，并进行了一个 Epoch 的训练，共使用了 1.5B tokens 。训练时，设置训练的 batch size 为 0.5M，序列长度为 2048 个 token 。训练过程中使用了 AdamW 优化器，并采用了多项式学习率衰减， warm-up 比例设置为总训练步数的 1.5% 。权重衰减设置为 0.01 ，梯度裁剪设置为 1 。