面向边缘设备的高效视觉 Transformer 重用注意力机制——UniForm
计算机视觉与模式识别
2024-12-04 v1
摘要
Transformer-based 架构在各个领域展现出巨大成功,但由于高内存和计算需求,其在边缘设备上的部署仍具有挑战性。本文引入一种新颖的重用注意力机制,专为高效内存访问和计算优化而设计,使其能在资源受限的平台上无需牺牲性能地顺利运行。不同于传统的多头注意力(MHA),后者为每个头冗余计算独立的注意力矩阵,重用注意力将这些计算整合到一个共享注意力矩阵中,显著降低内存开销和计算复杂度。在 ImageNet-1K 和下游任务上的全面实验表明,利用重用注意力实现的 UniForm 模型在实现最先进的 ImageNet 分类准确率的同时,优于线性注意力和 Flash 注意力等现有注意力机制,在推理速度和内存可扩展性方面表现更佳。值得注意的是,UniForm-l 在 Jetson AGX Orin 等边缘设备上实现了 21.8ms 的推理时间,达到了 76.7% 的 Top-1 准确率,速度可提高最高 5 倍。这些结果表明了重用注意力在高性能 GPU 和边缘平台上的通用性,为更广泛的实时应用铺平了道路。
引用
@article{arxiv.2412.02344,
title = {UniForm: A Reuse Attention Mechanism Optimized for Efficient Vision Transformers on Edge Devices},
author = {Seul-Ki Yeom and Tae-Ho Kim},
journal= {arXiv preprint arXiv:2412.02344},
year = {2024}
}
备注
13 Pages, 8 Tables, 7 Figures