NuWa: Deriving Lightweight Task-Specific Vision Transformers for Edge Devices

2025-04-04

AI Summary

To address the mismatch between resource-constrained edge devices and over-parameterized general-purpose Vision Transformers (ViTs), which leads to suboptimal accuracy–latency trade-offs for specialized edge vision tasks, this paper proposes a task-oriented knowledge distillation framework for ViTs. Our method integrates task-aware attention reweighting and inter-layer feature alignment within a knowledge distillation pipeline, coupled with structured pruning, to jointly optimize model size, accuracy, and inference latency. To the best of our knowledge, this is the first lightweight, task-specialized ViT distillation framework tailored for edge deployment. Evaluated on three benchmark datasets, it achieves up to 11.83% higher accuracy and 1.29×–2.79× faster inference compared to state-of-the-art methods, significantly enhancing both practicality and deployment efficiency for edge vision applications.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Computer Vision: Large Vision Models
• Search and Optimization: Distributed Search

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Semantics and Knowledge: Scalable techniques for the creation, curation, publication, maintenance, and consumption of large, Web-based, structured, reusable, knowledge graphs and ontologies
• Systems and Infrastructure for Web, Mobile and WoT: Cloud, edge and content delivery systems for the Web

Abstract

Vision Transformers (ViTs) excel in computer vision tasks but lack flexibility for edge devices' diverse needs. A vital issue is that ViTs pre-trained to cover a broad range of tasks are extit{over-qualified} for edge devices that usually demand only part of a ViT's knowledge for specific tasks. Their task-specific accuracy on these edge devices is suboptimal. We discovered that small ViTs that focus on device-specific tasks can improve model accuracy and in the meantime, accelerate model inference. This paper presents NuWa, an approach that derives small ViTs from the base ViT for edge devices with specific task requirements. NuWa can transfer task-specific knowledge extracted from the base ViT into small ViTs that fully leverage constrained resources on edge devices to maximize model accuracy with inference latency assurance. Experiments with three base ViTs on three public datasets demonstrate that compared with state-of-the-art solutions, NuWa improves model accuracy by up to $ ext{11.83}%$ and accelerates model inference by 1.29$ imes$ - 2.79$ imes$. Code for reproduction is available at https://anonymous.4open.science/r/Task_Specific-3A5E.

Problem

• Optimizing Vision Transformers for edge device tasks
• Reducing model size while maintaining task accuracy
• Accelerating inference speed on resource-limited devices

Innovation

• Derives small ViTs from base ViT
• Transfers task-specific knowledge efficiently
• Ensures inference latency and accuracy