DYNAMAX: Dynamic computing for Transformers and Mamba based architectures

早期退出（EEs）作为一种承诺通过在数据样本达到满意的预测置信度后动态终止推断，从而减少计算成本和延迟的有前景的方法。尽管许多工作将EEs集成到仅编码器的Transformer中，但其应用于解码器架构以及更重要的是，Mamba模型——这是一种在大语言模型（LLM）领域的新型状态空间架构——仍不充分探索。本工作引入DYNAMAX，第一个利用Mamba架构独特性质以实现早期退出机制的框架。我们不仅将EEs集成到Mamba中，还将Mamba用作基于Mamba和基于Transformer的LLMs的高效EE分类器，展示了其灵活性。我们的实验采用TruthfulQA、CoQA和TriviaQA等数据集，对Mistral 7B transformer与Codestral 7B Mamba模型进行计算节省、准确率和一致性评估。结果突出了Mamba作为强大EE分类器的适应性以及其在NLP任务中在计算成本和性能质量之间取得平衡的效率。通过利用Mamba的动态处理内在设计，我们为嵌入式应用程序和资源受限环境中的可扩展高效推断开辟了道路。本研究凸显了Mamba在重新定义LLM动态计算范式方面的变革性潜力。