多头注意力的自动通道剪枝

尽管Transformer性能强劲，但其二次计算复杂度给应用于视觉任务带来了挑战。自动剪枝是无需启发式方法即可降低计算复杂度的有效方法之一。然而，由于通道不对齐，直接将其应用于多头注意力并不直接。本文提出了一种自动通道剪枝方法，以考虑多头注意力机制。首先，我们将基于通道相似性的权重纳入剪枝指标，以在每个头中保留更多信息丰富的通道。然后，我们调整剪枝指标，强制在所有头中按等比例移除通道，防止通道不对齐。我们还添加了一个重新加权模块以补偿通道移除导致的信息损失，以及一个基于原始结构与每个通道之间注意力差异的剪枝指标有效初始化步骤。我们提出的方法不仅可用于原始注意力，还可用于线性注意力，后者在令牌数量上具有线性复杂度，效率更高。在ImageNet-1K上，将我们的剪枝方法应用于包含两种注意力机制的FLattenTransformer，在多个MACs下相比先前最先进的 efficient 模型和剪枝方法显示出更优的准确率。代码即将发布。