专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

性能远超Mamba！北大林宙辰团队提出全新混合序列建模架构MixCon

PaperWeekly · 公众号 · 科研 · 2024-10-17 17:56

正文

©作者 | 林宙辰、徐鑫

单位 | 北京大学

来源 | 机器之心

在自然语言处理、语音识别和时间序列分析等众多领域中，序列建模是一项至关重要的任务。然而，现有的模型在捕捉长程依赖关系和高效建模序列方面仍面临诸多挑战。

因此，北京大学林宙辰、徐鑫提出了一种全新混合序列建模架构 ——MixCon，它为解决这些难题带来了创新性的方案。经实验验证，其性能远超 Mixtral、Mamba 和 Jamba。论文已在 European Conference on Artificial Intelligence (ECAI) 2024 上发表。

论文标题：

MixCon: A Hybrid Architecture for Efficient and Adaptive Sequence Modeling

论文链接：

https://zhouchenlin.github.io/Publications/2024-ECAI-MixCon.pdf

现有序列建模模型的困境

线性注意力Transformer

线性注意力 Transformer 旨在通过近似注意力机制来提高原始 Transformer 模型的效率，将计算复杂度从

降低到

或

，但在处理长序列时可能会面临性能下降和计算开销增加的问题。

例如，早期利用局部敏感哈希方案虽降低复杂度，但引入大常数因子；近期通过改变计算顺序等方法近似 Softmax 函数，但仍存在性能不如 Softmax 注意力且可能增加额外开销的情况。

线性RNN模型

线性 RNN 模型如 Mamba 等通过将序列表示为状态空间并利用扫描操作，以线性时间复杂度提供了序列建模的新解决方案。

然而，它们可能缺乏复杂序列建模任务所需的适应性和动态特性，并且像传统序列模型一样，缺少反馈机制和自适应控制。

MoE模型

MoE 模型通过结合专家模块，能有效处理长序列并保持计算效率，根据输入数据自适应选择专家模块。

但 MoE 模型的专家模块稀疏激活可能导致训练稳定性问题，部分参数不常使用降低参数效率，在处理长序列时可能在计算效率和训练稳定性方面面临挑战，且对动态变化适应性不足。

MixCon的核心架构与技术

Conba 模型架构

1. 状态空间方程

Conba 将序列建模任务表示为状态空间系统，状态空间定义为

和

，其中

、

和

分别为时间步

的状态、输入和输出，

和

是非线性函数，可由神经网络近似。

，其中

和

是可学习参数矩阵。

，

是可学习参数矩阵。

为处理长序列，Conba 采用选择性状态空间机制

，以及引入延迟状态

和动态状态缩放机制

。

最后状态空间系统表示为

来捕捉长程依赖和适应序列动态变化。

2. 自适应控制机制

设计目标是最小化实际输出

和期望输出

之间的跟踪误差

。

控制增益矩阵

通过

更新，其中

是跟踪误差向量

的 2 范数，

是学习率。

3. 实施细节

神经网络近似：利用 MLP 近似状态转移函数

的非线性部分和观察函数

，通过端到端训练确定最佳参数。

选择性状态空间函数

用 SwiGLU（基于 Swish 和 GLU 的混合激活函数）近似，其公式为

。

进一步扩展和增强稳健性：引入过程噪声

和观察噪声

，

，增强模型对扰动的适应性。

4. 模型架构图如下所示：

MixCon模型架构

MixCon 是结合注意力机制的 Transformer 层、Conba 层和 MoE 组件的创新混合解码器架构。

在内存使用方面，通过平衡注意力和 Conba 层，相比 Mamba 可将 KV 缓存减少 32 倍。例如，在 256K 令牌上下文环境中，MixCon 仍能保持较小的 KV 缓存优势（如表 1 所示）。

在吞吐量方面，处理长序列时，Conba 层计算效率更高，增加其比例可提高整体吞吐量。

基本配置单位是 MixCon 块，由 Conba 或注意力层组合而成，每个层包含注意力模块或 Conba 模块，后接 MLP 或 MoE 层。MixCon 中的 MLP 层被 MoE 层替换，以增加模型容量同时保持较低计算负载。

对于 Conba 层实施，采用 RMSNorm 等技术，模型词汇量为 256K，使用 BPE 进行训练，每个数字为单独令牌。

模型架构图如下所示：

MixCon的实验与评估

实施细节

选择特定配置适应单块 80GB A800 NVIDIA GPU 的计算能力，实现质量和吞吐量的优化。

序列由 4 个 MixCon 块组成，每个 MixCon 块含 8 层 L = 8，注意力层和 Conba 层比例为 2:6 (a:c = 2:6)，每隔一层 (e = 2) 用 MoE 替换 MLP 模块，模型有 16 个专家 (n = 16)，每个令牌使用 2 个顶级专家 (K = 2)。

上下文长度分析

MixCon 在单块 80GB A800 GPU 上的最大上下文长度是 Jamba 的两倍、Mixtral 的四倍、Llama - 2 - 70B 的十四倍（如图 3 所示）。

吞吐量分析

1. 配置一：考虑不同批大小，在单块 A800 80GB GPU（int8 量化）、8K 上下文长度下生成 512 个输出令牌，MixCon 吞吐量是 Mixtral 的三倍、Jamba 的两倍（如图 4 所示）。

2. 配置二：单批次（批大小 = 1）、四块 A800 GPUs（无量化）、不同上下文长度下生成 512 个输出令牌，处理 128K 令牌时，MixCon 吞吐量是 Jamba 的 1.5 倍、Mixtral 的 4.5 倍（如图 5 所示）。

数据集评估

本文在一系列标准学术基准测试中评估 Conba 性能，包括常识推理任务（如 HellaSwag、WinoGrande、ARC - E、ARC - Challenge）、阅读理解任务（如 BoolQ、QuAC）、聚合基准测试（如 MMLU、BBH），采用不同的学习策略。

MixCon 性能与类似或更大规模的先进公开模型相当或更优，尽管总参数比 Llama - 2 少，但作为稀疏模型，其活跃参数仅 5B，处理长序列时 KV 缓存仅需 2GB，而 Mixtral 需 32GB（如表 2 所示）。

消融实验

展示注意力和 Conba 层结合的优势及最佳比例和交织技术。纯 Conba 模型在上下文学习有困难，Attention - Conba 混合模型有类似纯 Transformer 模型的上下文学习能力。

以 HellaSwag（10 - shot）、WinoGrande（5 - shot）、Natural Questions（NQ，5 - shot）为指标，MixCon 表现稳健（如表 3 所示），MixCon（无 MoE）训练过程损失更低（如图 6 所示）。

长上下文评估

利用问答基准测试评估 MixCon 处理长上下文能力，使用 L - Eval 中最长上下文数据集的五个数据集，以少样本格式（每个实验用三个例子）进行实验。

在 NarrativeQA、LongFQA、Natural Questions（NQ）、CUAD 等数据集上评估，MixCon 在多数数据集上优于 Mixtral 和 Jamba，平均性能优越，且在长上下文任务中具有更好的吞吐量（如表 4 所示）。

结合注意力和Conba的优势及混合专家的影响

1. 注意力和 Conba 比例研究

用 13 亿参数模型在 2500 亿令牌上训练，MixCon 性能优于纯注意力或纯 Mamba，注意力和 Conba 层比例为 2:6 或 1:7 时性能差异小（如表 5 所示）。

2. 混合专家的影响

当在 MixCon 架构的大规模情境（5B 参数，在 50B 令牌上训练）中应用 MoE 技术时，性能有显著提升（如表 6 所示）。

MixCon的优势与展望

MixCon 作为创新的混合序列建模架构，通过整合多种技术，在处理复杂动态序列时具有高效的计算效率，在各项任务中展现出显著优势，能高效处理长序列、内存使用低且吞吐量高，具有高可扩展性和实用性。然而，它仍有改进空间，如进一步优化状态空间表示、长序列的自适应控制、特定领域微调以及训练算法等。

总体而言，MixCon 为序列建模提供了新的解决方案，在复杂序列处理方面表现出色，为 NLP 及其他领域的应用开辟了新道路。未来，我们期待它在更多领域发挥更大的作用，为技术发展带来更多的突破和创新。

更多阅读

#投稿通道#

让你的文字被更多人看到

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。

📝 稿件基本要求：

• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注

• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题

• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算

📬 投稿通道：

• 投稿邮箱：[email protected]

• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者

• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿

△长按添加PaperWeekly小编

🔍

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧