提出共享注意力机制,让视觉Transformer在边缘设备上更快更省内存
UniForm: A Reuse Attention Mechanism Optimized for Efficient Vision Transformers on Edge Devices
- 用共享注意力矩阵替代多头独立计算,减少内存和算力开销
- 在Jetson AGX Orin上实现21.8ms推理时间,比基准快5倍,精度达76.7%
- 适合部署在资源受限的边缘设备,兼顾性能与效率
基于Transformer的架构在多个领域表现卓越,但在边缘设备上部署仍面临高内存和计算需求的挑战。本文提出一种新型的重用注意力(Reuse Attention)机制,专为高效内存访问和计算优化设计,可在不损失性能的前提下实现在资源受限平台上的无缝运行。与传统多头注意力(MHA)冗余计算各头注意力矩阵不同,重用注意力将这些计算合并为一个共享矩阵,显著降低内存占用和计算复杂度。在ImageNet-1K及下游任务上的全面实验表明,采用重用注意力的UniForm模型在图像分类精度上达到顶尖水平,同时在推理速度和内存可扩展性方面优于线性注意力和Flash Attention等现有机制。值得注意的是,UniForm-l在边缘设备Jetson AGX Orin上实现76.7%的Top-1准确率,推理时间仅21.8ms,相比竞品方法最高提速5倍。这些结果证明了重用注意力在高性能GPU与边缘平台间的普适性,为更广泛实时应用铺平道路。
原文摘要 · Abstract (English)
Transformer-based architectures have demonstrated remarkable success across various domains, but their deployment on edge devices remains challenging due to high memory and computational demands. In this paper, we introduce a novel Reuse Attention mechanism, tailored for efficient memory access and computational optimization, enabling seamless operation on resource-constrained platforms without compromising performance. Unlike traditional multi-head attention (MHA), which redundantly computes separate attention matrices for each head, Reuse Attention consolidates these computations into a shared attention matrix, significantly reducing memory overhead and computational complexity. Comprehensive experiments on ImageNet-1K and downstream tasks show that the proposed UniForm models leveraging Reuse Attention achieve state-of-the-art imagenet classification accuracy while outperforming existing attention mechanisms, such as Linear Attention and Flash Attention, in inference speed and memory scalability. Notably, UniForm-l achieves a 76.7% Top-1 accuracy on ImageNet-1K with 21.8ms inference time on edge devices like the Jetson AGX Orin, representing up to a 5x speedup over competing benchmark methods. These results demonstrate the versatility of Reuse Attention across high-performance GPUs and edge platforms, paving the way for broader real-time applications
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。