TransMamba：一种序列级混合 Transformer-Mamba 语言模型

TransMamba: A Sequence-Level Hybrid Transformer-Mamba Language Model

孙兴武 Xingwu Sun · Tencent Hunyuan · 2025-03-31 · arXiv:2503.24067 ↗ · 被引 3

摘要 · Abstract

Transformer 是现代大型语言模型的基石，但其二次计算复杂度限制了长序列处理的效率。最近 Mamba（一种具有线性复杂度的状态空间模型）的进展提供了有前景的效率提升，但存在上下文学习不稳定和多任务泛化能力差的问题。一些工作采用层级混合结构，结合 Transformer 和 Mamba 层，旨在充分利用两者的优势。本文提出 TransMamba，一种新颖的序列级混合框架，通过共享参数矩阵（QKV 和 CBx）统一 Transformer 和 Mamba，从而能在不同 token 长度和层上动态切换注意力与 SSM 机制。我们设计了 Memory Converter，通过将注意力输出转换为 SSM 兼容状态来桥接 Transformer 和 Mamba，确保在转换发生的 TransPoint 处信息流无缝衔接。还深入探索了 TransPoint 调度以平衡效果与效率。大量实验表明，与单一和混合基线相比，TransMamba 实现了更优的训练效率和性能，并在序列级验证了 Transformer 与 Mamba 范式之间更深层的一致性，为下一代语言建模提供了可扩展的解决方案。代码和数据见 https://github.com/Yixing-Li/TransMamba。

Transformers are the cornerstone of modern large language models, but their quadratic computational complexity limits efficiency in long-sequence processing. Recent advancements in Mamba, a state space model (SSM) with linear complexity, offer promising efficiency gains but suffer from unstable contextual learning and multitask generalization. Some works conduct layer-level hybrid structures that combine Transformer and Mamba layers, aiming to make full use of both advantages. This paper proposes TransMamba, a novel sequence-level hybrid framework that unifies Transformer and Mamba through shared parameter matrices (QKV and CBx), and thus could dynamically switch between attention and SSM mechanisms at different token lengths and layers. We design the Memory Converter to bridge Transformer and Mamba by converting attention outputs into SSM-compatible states, ensuring seamless information flow at TransPoints where the transformation happens. The TransPoint scheduling is also thoroughly explored for balancing effectiveness and efficiency. We conducted extensive experiments demonstrating that TransMamba achieves superior training efficiency and performance compared to single and hybrid baselines, and validated the deeper consistency between Transformer and Mamba paradigms at sequence level, offering a scalable solution for next-generation language modeling. Code and data are available at https://github.com/Yixing-Li/TransMamba

核心贡献 · Key contributions

提出 TransMamba，一种通过共享 QKV 和 CBx 参数在序列级别统一 Transformer 和 Mamba 的混合框架。
Proposes TransMamba, a sequence-level hybrid framework unifying Transformer and Mamba via shared QKV and CBx parameters.
设计记忆转换器，在转换点实现从注意力机制到 SSM 的无损信息传递。
Designs Memory Converter for lossless information transfer from attention to SSM at TransPoints.
探索转换点调度策略，以平衡不同层和词元长度下的效果与效率。
Explores TransPoint scheduling to balance effectiveness and efficiency across layers and token lengths.
在训练效率和性能上优于 Transformer、Mamba 及混合基线模型。
Demonstrates superior training efficiency and performance over Transformer, Mamba, and hybrid baselines.
验证了 Transformer 与 Mamba 范式在序列级别更深层次的一致性。
Validates deeper consistency between Transformer and Mamba paradigms at the sequence level.
为下一代语言建模提供可扩展的解决方案，支持灵活的结构切换。
Provides a scalable solution for next-generation language modeling with flexible structure switching.

局限 · Limitations

TransMamba 的效率提升受工程优化限制，理论加速未完全实现。
TransMamba's efficiency gain is limited by engineering optimizations; theoretical speedup not fully realized.
转换点调度需要仔细调优，次优调度会降低性能。
TransPoint scheduling requires careful tuning; suboptimal schedules degrade performance.
评估仅限于 8 个英文任务，未验证对其他语言和领域的泛化能力。
Evaluation limited to 8 English tasks; generalization to other languages and domains not verified.
模型规模和训练数据量适中，未探索向更大模型和数据集的扩展性。
Model size and training data scale are moderate; scalability to larger models and datasets not explored.
记忆转换器假设理论无损，实际中可能存在数值精度问题。
Memory Converter assumes theoretical losslessness; practical numerical precision issues may arise.

论文章节 · Sections（共 14）

摘要 Abstract
1 引言 1 Introduction
2 方法 2 Method
2.1 预备知识 2.1 Preliminary
2.2 TransMamba 总体框架 2.2 Overall Framework of TransMamba
2.3 无损记忆转换器 2.3 Lossless Memory Converter
2.4 灵活 TransPoint 调度 2.4 Flexible TransPoint Scheduling
3.1 实验设置 3.1 Experiment Setup
3.2 主要结果 3.2 Main Results
3.3 不同 TransPoint 调度的深入分析 3.3 In-depth Analyses on Different TransPoint Schedule
3.4 训练/推理不一致 TransPoint 调度的探索 3.4 Explorations on Inconsistent Training/Inference TransPoint Scheduling
3.5 其他模型组件的消融研究 3.5 Ablation Study on Other Model Components
4 相关工作 4 Related Works
5 结论 5 Conclusion

阅读逐段中英对照全文 →