Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DYNAMAX: Dynamic computing for Transformers and Mamba based architectures

Miguel Nogales, Matteo Gambella, Manuel Roveri

TL;DR

DYNAMAX tackles the high computational cost of decoder-only LLMs by introducing early exit mechanisms (EEs) into both Transformer and Mamba architectures. It trains EE classifiers via knowledge distillation and employs a token-forwarding inference strategy that gracefully handles missing-state updates, enabling per-token adaptive computation. The framework demonstrates that Mamba-based EE classifiers can deliver strong cost-accuracy trade-offs across tasks such as TriviaQA, CoQA, and TruthfulQA, while Transformer EEs also show substantial gains, particularly in generation tasks. These results highlight the potential of combining state-space models with adaptive computation to enable efficient, scalable inference on resource-constrained devices.

Abstract

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.

DYNAMAX: Dynamic computing for Transformers and Mamba based architectures

TL;DR

DYNAMAX tackles the high computational cost of decoder-only LLMs by introducing early exit mechanisms (EEs) into both Transformer and Mamba architectures. It trains EE classifiers via knowledge distillation and employs a token-forwarding inference strategy that gracefully handles missing-state updates, enabling per-token adaptive computation. The framework demonstrates that Mamba-based EE classifiers can deliver strong cost-accuracy trade-offs across tasks such as TriviaQA, CoQA, and TruthfulQA, while Transformer EEs also show substantial gains, particularly in generation tasks. These results highlight the potential of combining state-space models with adaptive computation to enable efficient, scalable inference on resource-constrained devices.

Abstract

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.
Paper Structure (17 sections, 7 equations, 5 figures, 1 algorithm)

This paper contains 17 sections, 7 equations, 5 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related literature
  3. Technical background
  4. Transformer and Mamba Architectures
  5. Early Exit Mechanisms
  6. Proposed solution and implementation
  7. Models
  8. Early Exits for Computational Efficiency in Autoregressive Generation
  9. Training Early Exit Classifiers
  10. Inference with Early Exits
  11. Results
  12. Datasets for training and benchmarking
  13. Performance in the Tasks
  14. TriviaQA
  15. CoQA
  16. ...and 2 more sections

Figures (5)

  • Figure 1: Comparison of EE and layer pruning performances in the TriviaQA set for Mistral and Mamba 7B models.
  • Figure 2: Performance of EE and layer pruning in the CoQA set for Mistral 7B.
  • Figure 3: Performance of EE and layer pruning in the CoQA set for Mamba.
  • Figure 4: Performance of EE and layer pruning in the Truthful QA generation (acc) set for Mistral 7B.
  • Figure 5: Performance of EE and layer pruning in the Truthful QA generation (acc) set for Mamba.