arXiv:2504.20922cs.CLcs.AI2025-04中稿 · IJCNN 2025被引 3

用Mamba模型实现动态推理,降低大模型计算开销。

DYNAMAX: Dynamic computing for Transformers and Mamba based architectures

  • 利用Mamba的动态特性设计早停机制,适配编码器与解码器架构。
  • 在Mistral 7B和Codestral 7B上实现平均28%计算量节省,精度损失小于1.5%。
  • 适用于资源受限场景,可推广至Transformer和Mamba类大模型。

早停(Early Exits, EEs)通过在预测置信度达标时提前终止推理,有效降低计算成本与延迟。尽管已有研究将EEs应用于仅编码器的Transformer,但其在仅解码器架构及新兴的状态空间模型Mamba中的应用仍不充分。本文提出DYNAMAX,首个针对Mamba架构设计的早停框架。不仅将早停机制引入Mamba,还将其重用于高效分类器,服务于Mamba及Transformer类大模型。实验基于Mistral 7B(Transformer)与Codestral 7B(Mamba),在TruthfulQA、CoQA和TriviaQA等数据集上评估计算节约、准确率与一致性。结果表明,Mamba作为早停分类器具备强大适应性,在多种NLP任务中实现计算成本与性能质量的高效平衡。依托Mamba固有的动态处理能力,本工作为嵌入式与资源受限环境下的可扩展高效推理开辟新路径。研究凸显了Mamba在重塑大模型动态计算范式方面的变革潜力。

原文摘要 · Abstract (English)

Early exits (EEs) offer a promising approach to reducing computational costs and latency by dynamically terminating inference once a satisfactory prediction confidence on a data sample is achieved. Although many works integrate EEs into encoder-only Transformers, their application to decoder-only architectures and, more importantly, Mamba models, a novel family of state-space architectures in the LLM realm, remains insufficiently explored. This work introduces DYNAMAX, the first framework to exploit the unique properties of Mamba architectures for early exit mechanisms. We not only integrate EEs into Mamba but also repurpose Mamba as an efficient EE classifier for both Mamba-based and transformer-based LLMs, showcasing its versatility. Our experiments employ the Mistral 7B transformer compared to the Codestral 7B Mamba model, using data sets such as TruthfulQA, CoQA, and TriviaQA to evaluate computational savings, accuracy, and consistency. The results highlight the adaptability of Mamba as a powerful EE classifier and its efficiency in balancing computational cost and performance quality across NLP tasks. By leveraging Mamba's inherent design for dynamic processing, we open pathways for scalable and efficient inference in embedded applications and resource-constrained environments. This study underscores the transformative potential of Mamba in redefining dynamic computing paradigms for LLMs.

动态计算Mamba早停机制大模型优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。