MambaLiteSR: Image Super-Resolution with Low-Rank Mamba using Knowledge Distillation

2025-02-19

AI Summary

Deploying image super-resolution (SR) on resource-constrained edge devices faces significant challenges in computational efficiency and memory footprint. To address this, we propose MambaLiteSR—a lightweight SR model tailored for edge deployment. Our method introduces three key innovations: (1) the first low-rank approximation of the Mamba architecture, synergistically integrating Vision Mamba’s global receptive field with state-space modeling’s linear complexity; (2) a mathematically grounded, PSNR-oriented hyperparameter optimization strategy; and (3) knowledge distillation to enhance feature learning efficiency. Evaluated under stringent edge constraints, MambaLiteSR achieves a 15% reduction in parameter count and up to 58% lower power consumption compared to state-of-the-art edge SR methods, while maintaining competitive PSNR and SSIM performance. Notably, it enables real-time, energy-efficient inference on the Jetson Orin Nano platform.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation

Abstract

Generative Artificial Intelligence (AI) has gained significant attention in recent years, revolutionizing various applications across industries. Among these, advanced vision models for image super-resolution are in high demand, particularly for deployment on edge devices where real-time processing is crucial. However, deploying such models on edge devices is challenging due to limited computing power and memory. In this paper, we present MambaLiteSR, a novel lightweight image Super-Resolution (SR) model that utilizes the architecture of Vision Mamba. It integrates State Space Blocks and a reconstruction module for efficient feature extraction. To optimize efficiency without affecting performance, MambaLiteSR employs knowledge distillation to transfer key insights from a larger Mamba-based teacher model to a smaller student model via hyperparameter tuning. Through mathematical analysis of model parameters and their impact on PSNR, we identify key factors and adjust them accordingly. Our comprehensive evaluation shows that MambaLiteSR outperforms state-of-the-art edge SR methods by reducing power consumption while maintaining competitive PSNR and SSIM scores across benchmark datasets. It also reduces power usage during training via low-rank approximation. Moreover, MambaLiteSR reduces parameters with minimal performance loss, enabling efficient deployment of generative AI models on resource-constrained devices. Deployment on the embedded NVIDIA Jetson Orin Nano confirms the superior balance of MambaLiteSR size, latency, and efficiency. Experiments show that MambaLiteSR achieves performance comparable to both the baseline and other edge models while using 15% fewer parameters. It also improves power consumption by up to 58% compared to state-of-the-art SR edge models, all while maintaining low energy use during training.

Problem

• Enhance image super-resolution on edge devices
• Reduce model size and power consumption
• Maintain performance with fewer parameters

Innovation

• Lightweight Vision Mamba architecture
• Knowledge distillation technique
• Low-rank approximation optimization