NAISS
SUPR
NAISS Projects
SUPR
GPU-Accelerated Training of Mamba Models
Dnr:

NAISS 2026/4-1236

Type:

NAISS Small

Principal Investigator:

Diletta Romano

Affiliation:

Uppsala universitet

Start Date:

2026-07-01

End Date:

2027-07-01

Primary Classification:

10210: Artificial Intelligence

Allocation

Abstract

Advances in machine learning have enabled increasingly powerful models for sequence modelling, perception, and decision-making. However, most modern neural architectures are developed and trained with cloud-scale resources in mind, while many real-world applications require intelligence directly on resource-constrained embedded devices such as microcontrollers. This creates an important gap between state-of-the-art machine learning and embedded AI, where memory, computation, latency, and energy consumption are severely limited. In my PhD research, I study embedded AI for microcontrollers and other constrained devices. In this project, I focus on Mamba/state space models, a recent class of sequence models with strong potential for efficient long-sequence modelling. Mamba models are particularly interesting because they offer an alternative to Transformer-based architectures, whose attention mechanisms can be expensive for long sequences. For embedded AI applications involving sensor signals, time-series data, and sequential system behaviour, models that can capture long-range dependencies efficiently are especially relevant. The aim of this project is to train and evaluate Mamba models in order to understand their suitability for resource-constrained inference. I will study how model size, sequence length, training configuration, and hyperparameters affect accuracy, memory use, computational cost, and deployment potential. The project will also compare Mamba-based models with relevant baseline architectures, with the goal of identifying when Mamba models are a practical choice for embedded AI. Although the long-term motivation is efficient inference on microcontrollers, the training phase requires GPU acceleration. Efficient Mamba implementations are designed around hardware-aware selective scan operations, also referred to as parallel scan, implemented through CUDA/GPU kernels. This makes GPUs essential for practical training. Without GPU access, only very small toy examples could be trained, while meaningful experiments with realistic sequence lengths, model sizes, and repeated comparisons would be prohibitively slow. The expected outcome of the project is a better understanding of how Mamba models can be trained, evaluated, and potentially adapted for embedded AI. This will contribute to the broader goal of bringing modern sequence modelling methods closer to deployment on low-power, resource-constrained devices. I am a doctoral student at Uppsala University. My main supervisor is Professor Thiemo Voigt, Department of Electrical Engineering, Uppsala University.