【论文阅读】TWDP: A Vision Transformer Accelerator with Token-Weight Dual-Pruning Strategy for Edge Device-CSDN博客

本文链接：https://blog.csdn.net/qq_52050257/article/details/146610467

【论文阅读】

Title： TWDP: A Vision Transformer Accelerator with Token-Weight Dual-Pruning Strategy for Edge Device Deployment
Authors：Guang Yang, Xinming Yan, Hui Kou, Zihan Zou, Qingwen Wei, Hao Cai, Bo Liu
Organizations：Southeast University, Nanjing, Jiangsu, China
Conference：ASPDAC ’25

摘要

视觉变换器（ViTs）因其在各种计算机视觉任务中相较于卷积神经网络（CNNs）具有更高的准确性而备受关注。然而，其庞大的计算负载和显著的内存占用导致了过长的延迟和巨大的数据存储开销，给资源受限的边缘设备部署带来了挑战。为了解决这些问题，我们提出了TWDP，这是一种采用令牌-权重双剪枝策略的视觉变换器加速器，旨在提高推理过程的效率。首先，我们提出了一种无参数的自适应令牌剪枝方法，以图像依赖的方式跳过冗余计算。其次，我们应用了一种基于Hessian的逐层N:M权重剪枝方法，以最小化存储开销、内存访问和计算功耗。此外，为了管理ViTs中复杂的计算模式，我们采用了重叠数据流，以进一步减少临时存储和推理延迟。在工业28nm技术下实现和评估，所提出的TWDP框架将权重存储需求减少了66.1%，并实现了2070.9 FPS/W的能量效率。与最先进的架构相比，TWDP在能量效率上提高了1.6倍，且准确率损失可忽略不计，证明了TWDP在边缘设备部署场景中的优越性。

1.引言

在这里插入图片描述

受Transformer模型在自然语言处理（NLP）领域取得的重大成功启发，Vision Transformer（ViTs）被提出用于处理计算机视觉任务，并取得了SOTA性能。在推理过程中，图像数据首先被嵌入为令牌，然后通过堆叠的L层编码器，并最终通过全连接层生成分类结果。编码器层是关键组件，占据了大部分的计算和内存访问工作，如图1左上角所示。多头自注意力（MSA）模块负责捕捉图像的全局特征，而前馈网络（FFN）模块则增强了模型的表达能力。

然而，根据目前的ViT研究，仍有三大挑战阻碍了ViTs在边缘设备上的部署，如图1所示。首先，ViTs的计算复杂度与令牌数量N密切相关，而每张图像中包含的信息性令牌数量各不相同。相关研究[8, 14]通过实时移除冗余令牌，降低了MSA和FFN模块的计算复杂度。然而，这些方法要么无法为不同图像自适应地实现令牌剪枝，要么引入了可训练参数，从而需要额外的再训练工作。其次，像DeiT-S这样的典型边缘模型，其权重参数在8位量化后仍高达22 MB，远远超出了边缘设备的片上存储能力，导致频繁且代价高昂的数据传输。此前的研究[22]和[9]在不同维度上探索了权重稀疏性，但采用了统一的剪枝率，而未深入探索细粒度的逐层稀疏性。此外，很少有研究将令牌剪枝和权重剪枝结合到一个完整的框架中，这一空白有待填补。第三，注意力过程中动态生成的矩阵需要额外的临时存储空间，而非线性函数（如Softmax和LayerNorm）由于除法运算和高数据依赖性而耗时较长。这些特性需要一种细粒度优化的数据流来解决。

为应对这些挑战，我们开发了TWDP，这是一种采用令牌-权重双剪枝策略的可用于边缘设备部署的高效视觉Transformer加速器。本文的贡献如下：

我们提出了一个名为TWDP的框架，将自适应令牌剪枝方法与基于Hessian的逐层N:M权重剪枝方法相结合，有效减少了计算和存储开销，并采用逐步微调策略以最小化准确率损失；
利用带有预移位除法器的重叠数据流，最小化MSA模块的临时存储需求，并高效降低编码器层的延迟；
在工业28nm工艺下实现并评估，我们提出的TWDP原型加速器实现了2070.9 FPS/W的能量效率，比现有的SOTA架构提高了1.6倍。据我们所知，TWDP是第一个将自适应令牌剪枝与逐层权重剪枝相结合的框架，用于ViTs的高效推理。

2.背景与相关工作

2.1 ViTs的剪枝策略

剪枝是将ViTs部署到边缘设备的有效方法。相关研究[14]在ViT骨干网络中引入额外的全连接层和非线性函数以生成令牌掩码，并据此丢弃冗余token。基于这种方法，[8]进一步将丢弃的token压缩成一个单一令牌，以便后续处理，从而提高精度。[11]通过引入一个可解释性感知单元来动态减少冗余token，从而跳过冗余token的计算。在减少权重参数方面，[22]沿通道维度实现权重剪枝，而[20]通过计算重要性分数实现通道和块剪枝。[1]提出了基于衰减的训练策略，并采用N:M权重稀疏方案。然而，这些方法中很少有考虑到高效的硬件部署，因此在将ViTs部署到硬件加速器时仍有很大的优化空间。

2.2 采用剪枝方法的ViT加速器

已有数项工作在FPGA和ASIC平台上部署了采用各种剪枝策略的ViTs，包括token剪枝、权重剪枝和注意力剪枝。对于token剪枝，[3]实现了一个token选择器，该选择器大量重用硬件资源以最小化令牌剪枝开销。[7]采用了类似的token剪枝逻辑，并结合乒乓流水线调度和数字存内计算（CIM），进一步提高了能效。对于权重剪枝，[9]提出了一个特定的硬件加速器，支持按头部、块和行的方式实现权重稀疏性。[13]结合了token剪枝和块权重剪枝，这与我们的框架类似，但未能实现自适应的令牌剪枝和逐层的权重剪枝，导致更高的精度损失。此外，尽管注意力剪枝在自然语言处理任务中得到了广泛探索[21]，但由于ViT中令牌数量较少，其效果有限，因此带来的改进也很有限。

3.令牌和权重剪枝的软件方案

3.1 自适应令牌剪枝

Vision Transformers（ViTs）的最终分类结果来源于类别令牌（class token）。在下文中，我们将 $\frac{QK^T} {\sqrt{d}}$ 矩阵称为注意力分数，而应用Softmax后的结果称为注意力概率。类别令牌通过注意力概率矩阵 $ap_{cls}$ 的第一行和V矩阵 $V = [v1, v2, ..., vn]^T$ 的线性组合进行更新，其中（v1, v2, …, vn）表示令牌的高维信息。因此， $ap_{cls}$ 向量中的条目可以被视为令牌的重要性分数。需要注意的是，第一个条目对应于类别令牌的重要性，因此将被忽略。整个自适应令牌剪枝过程在算法1中进行了描述。
在这里插入图片描述

具体来说，如果当前层处于剪枝位置（在我们的工作中为第3、6、9层，与[14]相同），则所有头部的 $ap_{cls}$ 向量将被累加为token_score，然后token_score中的条目按降序排列以获得 $ap_{clsr}$ 。然后根据预定义的比例r计算阈值T，并再次对 $ap_{clsr}$ 向量中的条目进行累加，直到结果超过T。参与此累加过程的令牌将被保留，而其他令牌将被丢弃。此过程不会引入额外的可训练参数，从而有助于训练过程的收敛。此外，如果所有令牌的重要性分数相似，表明当前层尚未区分出重要的令牌，则会保留更多的令牌。相反，如果重要性分数集中在少数几个令牌上，累加将迅速超过T，从而允许丢弃更多的令牌。图2中的一个例子展示了在不同复杂度的图像上应用所提出的剪枝方法。
在这里插入图片描述

3.2 基于Hessian的N:M权重剪枝

Vision Transformers（ViTs）中的不同编码器层负责提取不同层次的抽象信息，因此对参数剪枝具有不同的敏感度。因此，剪枝比例应以层为单位进行设置。我们将Average Hessian Trace （AHT）与N:M剪枝相结合，以实现逐层权重剪枝，最大化剪枝比例的同时最小化准确度损失。

首先，AHT测量损失函数相对于模型参数的平均曲率，指示模型性能对权重扰动的敏感性，并已被证明是非常有效的。其次，N:M剪枝是一种结构化的剪枝策略，它在M个元素中保留N个值，并将其他值置零。这种剪枝方法不仅对硬件友好，而且有效减少冗余参数。为了实现细粒度的权重剪枝，我们使用AHT指导逐层的N:M稀疏比率，以充分探索权重数据的冗余性。应用基于Hessian的逐层N:M稀疏性的第一个问题是决定N和M的值。在这项工作中，我们为了便于硬件部署，将所有层中的N值固定为1。根据建议，Transformer模型中的稀疏性范围应为50%到87.5%。我们遵循这一建议，并根据相应块的AHT从2、4、8中选择M值。

第二个问题是确定区分不同稀疏度水平的阈值AHT值。我们测量了DeiT-S模型中四种类型的块的权重参数的AHT：MSA块中的QKV生成和注意力投影，以及FFN块中的FC1和FC2线性投影。结果如图3所示。T1和T2将AHT值分为三部分，每部分都有特定的稀疏比率。我们测试了几种可能的AHT阈值组合，结果如表1所示。选择T₁(MSA)=T₁(FFN)=0.1 × 10^-3 和 T₂(MSA)=T₂(FFN)=0.3 × 10^-3 作为阈值配置，以在权重参数减少和模型准确率之间实现平衡性能。
在这里插入图片描述

3.3 分步训练策略

由于我们提出的TWDP结合了令牌剪枝和权重剪枝，直接训练整个模型难以收敛，并可能导致不可接受的准确率损失。为了解决这个问题，我们将训练过程分为两个不同的步骤。第一步专注于权重剪枝，我们应用了MDGF-Exponential训练策略来逐层稀疏化权重矩阵，该策略由AHT指导。这种方法确保只保留最关键的权重，从而在保持关键信息的同时降低模型的复杂性。

然后我们进入第二步，目标是令牌剪枝。我们停止对置零的权重参数进行梯度计算，以防止任何进一步更新，然后使用训练策略来微调非零权重。第一步中的权重剪枝过程已经确保剩余参数足以提取关键信息，第二步旨在帮助所提出的自适应令牌剪枝方法识别这些信息。此外，自适应令牌剪枝方法不引入额外参数，简化了微调过程，使训练过程更容易达到收敛。

4.架构

4.1 架构概览

在这里插入图片描述

图4展示了提出的TWDP加速器整体架构。该加速器支持N:M稀疏场景下的动态自适应令牌剪枝（on-the-fly self-adaptive token pruning）和稀疏-稠密矩阵乘法（SDMM），以及稠密-稠密矩阵乘法（DDMM）。整个推理过程量化为8位，采用训练策略实现，同时支持非线性函数（如Softmax、LayerNorm、GELU）。TWDP由向量乘法单元阵列（VMU array）、整数非线性函数单元、片上存储体（memory banks）、带宽匹配器（bandwidth matcher）和自适应令牌剪枝引擎（SATPE）组成。系统还包含一个Cortex-M3处理器，用于控制计算流程并从外部DDR内存中搬运数据。带宽匹配器和SATPE的细节将在4.2节详细说明。76.5KB的层存储体和特征存储体分别存储LayerNorm前后的数据，64KB的权重存储体和8KB的位图存储体用于保存稠密权重参数和稀疏位置索引。此外，另有一个64KB的片上缓冲区用于Softmax和LayerNorm过程，以协调数据流重叠。

4.2 剪枝的硬件支持

带宽匹配器：N:M稀疏化的权重数据以稠密形式存储，显著减少存储开销，但需选择对应的激活值以正确执行矩阵乘法。本文通过带宽匹配器（由缓冲器和多路复用器构成）确保特征存储体与权重存储体的带宽一致性，如图4蓝色部分所示。假设稀疏比为1:2，每周期从特征存储体和权重存储体各读取4个数据，VMU执行4次乘累加（MAC）操作。在周期0，激活矩阵（0:7,0:3）从特征存储体读取，但仅根据位图（bitmap）选择其中2个有效数据存入缓冲器。周期1时，另一组2个数据被选中，与缓冲器内数据合并为4个数据块，同时权重数据被读取至VMU执行MAC操作。在TWDP加速器中，每个VMU每周期执行16次8位乘累加，并将32位部分和存储在本地寄存器中。

自适应令牌剪枝引擎（SATPE）：如3.1节所述，自适应令牌剪枝包含三步：累加注意力概率向量、排序和令牌过滤。如图4橙色部分所示，步骤1通过寄存器和加法器直接实现。步骤2采用插入排序（而非资源密集的Bitonic排序）以节省硬件资源。排序阶段由n个比较单元组成，每个单元包含寄存器、比较器、多路复用器和编码器。经过n周期后生成排序后的注意力概率向量。由于仅需注意力概率矩阵的首行，排序延迟可被矩阵其余部分的计算完全覆盖，避免额外推理延迟。步骤3中，累加排序后的向量直至超过阈值T，确定保留的令牌索引，后续层中跳过剪枝令牌的计算。

4.3 重叠数据流

在这里插入图片描述

如前所述，MSA块中动态生成的矩阵需要额外临时存储，而非线性函数因硬件资源限制成为推理速度瓶颈。为此，我们提出重叠数据流以解决这些问题。如图5(a)所示，首先优化MSA块的计算顺序：将单个注意力头的计算分解为n个重复步骤，VMU阵列在此期间持续运行不中断。单个头的完整K矩阵在这些重复步骤前预计算并存入K缓冲区。首先生成8行Q矩阵存入Q缓冲区，随后与K矩阵相乘并将结果发送至QKT缓冲区。接着计算8行V矩阵以填补矩阵乘法与Softmax之间的延迟间隙。在下一重复步骤中，新生成的8行Q矩阵覆盖旧数据，避免存储完整Q矩阵。类似优化减少MSA块42.7%的临时存储需求。

为降低非线性函数的延迟，我们优化了Softmax、GELU和LayerNorm中的除法操作。传统逐位递归除法器（每周期生成1位商）虽节省资源但延迟较高。我们发现，仅当被除数与除数的最高有效位（MSB）对齐时才能生成有效商值，此前商值始终为零。因此，利用前导一检测器（Leading One Detector）预移位被除数与除数，使除法器首周期即可输出有效商值，Softmax、GELU和LayerNorm的除法延迟分别降低42%、33%和37%。

结合优化后的计算顺序与预移位除法器，实现图5所示的重叠数据流。图5(a)中，Softmax操作包含最大值查找、指数计算、求和及除法，其延迟因预移位除法器和计算顺序优化得以与矩阵乘法重叠。图5(b)展示了完整编码器层的数据流：LayerNorm与FC2数据回写并行执行，GELU与FC1计算重叠。SATPE仅在剪枝位置（本文中为第3、6、9层）激活，其余时间保持空闲。通过重叠数据流，MSA块和编码器整层的延迟分别降低28.2%和21.7%。

5.实验

5.1 实验设置

为评估提出的TWDP框架有效性，我们在ImageNet1k数据集上对DeiT-S模型进行重新训练与精度测试。权重剪枝的微调采用MDGF-Exponential方法，结合逐层稀疏配置；令牌剪枝则基于自适应令牌剪枝方法改进其评分计算流程。INT8量化在微调阶段完成，整个训练过程使用Python 3.9环境在8块NVIDIA-V100 GPU上运行300个周期。TWDP原型加速器基于Verilog HDL实现，并通过Synopsys Design Compiler在0.8V电压、400MHz频率、28nm工艺下完成综合，功耗由PrimeTime PX测量。

5.2 性能分析与对比

在这里插入图片描述
TWDP原型加速器版图如图6所示，芯片面积1.69mm²，片上存储容量288KB。在0.8V、400MHz工况下，其能效达2070.9 FPS/W，验证了边缘部署的高效性。图7展示了硬件资源分布，自适应令牌剪枝引擎（SATPE）与带宽匹配器的硬件开销可忽略。表2显示令牌剪枝阈值比率可手动调节以适配不同场景，实验表明最大可节省37.6%能耗与37.1%延迟，精度损失仅1.4%。

在这里插入图片描述

图8分解了本文三项核心技术的贡献：令牌剪枝与权重剪枝分别减少1.59倍和2.84倍乘累加（MAC）操作，权重剪枝降低2.95倍参数量；令牌剪枝与重叠数据流分别减少1.59倍与1.28倍计算延迟。三者协同后，MAC操作、权重参数及延迟分别降至基线的22.2%、33.9%和49.5%。表3对比了TWDP与现有方案：其他工作或因仅优化注意力稀疏性（如[19]）、采用固定剪枝率（如[7]）导致性能受限，而TWDP凭借令牌-权重双剪枝框架实现1.6倍能效提升及66.1%权重参数压缩，精度损失可忽略。
在这里插入图片描述

6 结论

本文提出TWDP——一种面向边缘设备部署的Vision Transformers加速器，采用令牌-权重双剪枝策略。双剪枝显著提升推理效率，而重叠数据流进一步减少临时存储需求与计算延迟。基于工业级28nm工艺实现的TWDP原型加速器能效达到2070.9 FPS/W，相比现有ASIC方案提升1.6倍，在ImageNet1k数据集上DeiT-S模型的精度损失仅1.4%，验证了其在边缘设备部署中的优势。