Low-latency vision transformers via large-scale multi-head attention

作者:Ronit D. Gross et.al.

论文链接:http://arxiv.org/abs/2506.23832

发布日期:2025-06-30

解读时间:2025-07-19 19:11:05

论文摘要

The emergence of spontaneous symmetry breaking among a few heads of multi-head attention (MHA) across transformer blocks in classification tasks was recently demonstrated through the quantification of single-nodal performance (SNP). This finding indicates that each head focuses its attention on a subset of labels through cooperation among its SNPs. This underlying learning mechanism is generalized to large-scale MHA (LS-MHA) using a single matrix value representing single-head performance (SHP), analogous to single-filter performance in convolutional neural networks (CNNs). The results indicate that each SHP matrix comprises multiple unit clusters such that each label being explicitly recognized by a few heads with negligible noise. This leads to an increased signal-to-noise ratio (SNR) along the transformer blocks, thereby improving classification accuracy. These features give rise to several distinct vision transformer (ViT) architectures that achieve the same accuracy but differ in their LS-MHA structures. As a result, their soft committee yields superior accuracy, an outcome not typically observed in CNNs which rely on hundreds of filters. In addition, a significant reduction in latency is achieved without affecting the accuracy by replacing the initial transformer blocks with convolutional layers. This substitution accelerates early-stage learning, which is then improved by subsequent transformer layers. The extension of this learning mechanism to natural language processing tasks, based on quantitative differences between CNNs and ViT architectures, has the potential to yield new insights in deep learning. The findings are demonstrated using compact convolutional transformer architectures trained on the CIFAR-100 dataset.

AI解读

好的,我将按照您的要求对这篇论文进行详细分析。

1. 核心方法与创新点

这篇论文的核心方法是研究在视觉Transformer (ViT) 中使用大规模多头注意力 (LS-MHA) 的底层学习机制,并利用这种机制来构建低延迟的ViT架构。作者并没有提出全新的模型结构,而是专注于分析现有ViT模型中MHA的工作方式,并通过实验验证其理论,最终找到了一种降低ViT延迟的方法。

主要创新点:

* Single-Head Performance (SHP) 指标: 引入SHP指标来量化LS-MHA中每个头的性能,类似于CNN中的单滤波器性能 (SFP)。这种指标提供了一种更细粒度的分析MHA的方法,允许研究人员了解每个头关注的标签子集和信号-噪声比 (SNR)。这与之前基于Single-Nodal Performance(SNP)的研究相比,在高头数的情况下简化了分析。

* LS-MHA的学习机制分析: 揭示了LS-MHA的学习机制,发现每个头明确识别一组标签,簇大小接近于1,噪声可忽略。这与CNN中的滤波器形成较大簇的情况不同。强调了各头之间的部分重叠,形成微弱但显著的信号,从而实现高SNR。

* 低延迟ViT架构: 通过将ViT的初始Transformer块替换为卷积层 (CLs),在不影响准确性的前提下显著降低了延迟。这种混合架构利用了CNN在早期学习阶段的优势,以及Transformer块在后续阶段的优势,实现了加速早期学习的目的。

* 软委员会决策 (Soft Committee Decision) 提升准确率: 发现具有相似准确率的不同ViT架构的SHP矩阵存在差异,因此可以通过软委员会决策方法来提升准确率。实验结果表明,结合多个ViT架构可以显著提升性能,优于传统的CNN集成方法。

2. 算法细节与流程

核心算法主要围绕SHP的计算和分析展开。算法流程如下:

1. 预训练ViT模型: 首先,使用CIFAR-100数据集预训练一个ViT模型,如CCT-7/3x1。

2. 提取中间层特征: 选择ViT模型中的一个Transformer块,并提取该块的QKV Attention的输出特征。

3. 训练分类器头: 将提取的特征连接到一个分类器头,分类器头包含序列池化 (SP) 和全连接 (FC) 层,在CIFAR-100数据集上训练该分类器头。

4. 计算SHP矩阵:
* 静默所有头: 对于每个头,静默(设置为零)分类器头的所有输入节点,除了属于被评估的头的那些节点。
* 输入验证集: 将验证数据集输入到预训练的ViT模型,并经过部分静默的分类器头。
* 生成矩阵: 对于每个验证集标签,计算100个输出单元的平均激活值,得到一个100x100的矩阵。该矩阵的每个元素(i, j)表示标签为i的输入在输出单元j上的平均激活值,并进行归一化。

5. SHP矩阵分析:
* 阈值处理: 对SHP矩阵应用阈值,得到一个布尔矩阵。
* 聚类: 对布尔矩阵进行排列,形成对角线簇。
* 噪声分类: 将阈值以上但位于簇外的元素分类为噪声。
* 计算统计量: 计算各种统计量,如簇的数量、簇的大小、对角线元素数量和噪声量。

6. 混合架构构建与训练: 将初始Transformer块替换为CLs,构建混合架构。然后,在CIFAR-100数据集上训练这些混合架构。

7. 软委员会决策: 创建多个具有不同架构的ViT模型,并使用软委员会决策方法组合它们的预测,以提高准确率。

算法的技术优势和创新之处:

* SHP简化分析: SHP指标简化了LS-MHA的分析,特别是在头数较多的情况下,它提供了一种量化每个头性能的方法。

* 混合架构提升效率: 混合架构能够以较低的延迟获得与纯ViT模型相当的准确率。

* 软委员会决策提升准确率: 软委员会决策方法能够利用不同ViT架构的优势,显著提高准确率。

3. 详细解读论文第三部分

论文的第三部分(Advanced accuracy with low latency)主要介绍了通过构建混合模型(卷积层 + Transformer)以及使用软委员会决策来实现低延迟和高精度的ViT模型。

* 混合模型(CCT-2/3x5)

CCT-2/3x5模型由5个卷积层和2个Transformer块组成,模型总层数为14层(5 + 2 * 4 + 1 = 14),而CCT-7/3x1模型层数为30层。这里一个Transformer块包含4层(QKV Attention, FC projection, FF sub-block中两层FC),再加上一个分类器层,所以CCT-7/3x1模型层数为1 + 7 * 4 + 1 = 30。作者通过实验证明,CCT-2/3x5模型可以达到与CCT-7/3x1相似的精度(~0.81),但延迟更低,可训练参数更少。

作者发现,如果将CLs和Transformer块的顺序颠倒,精度会明显下降。这表明直接从输入数据中学习短程相关性比在通过Transformer的FC层进行空间混合后执行此操作更有效。

* 软委员会决策

多个CCT模型结构(CLs数量,Transformer块数量,head数量,MHA维度)相似的架构,其SHP矩阵性质不同,软委员会决策能提高精度。每个验证输入通过训练好的CCT架构,生成一个100维度的输出向量。$N$ 个CCT架构的软委员会决策通过对这些输出向量进行求和,以和向量作为最终决策的依据。

在只有$N=4$个CCT架构的委员会中,每个架构由七个组件(CL和Transformer块)组成,平均独立精度为0.813,软委员会将精度提高到0.851。注意:这里的软委员会决策没有经过Softmax,激活,或者归一化等操作。

* 数学公式(没有显式的数学公式推导)

在论文第三部分,并没有显式的数学公式推导。但是,软委员会决策可以简单地用以下公式表示:

$$
\hat{y} = \arg \max_c \sum_{i=1}^{N} p_i(c|x)
$$

其中:
- $\hat{y}$ 是软委员会预测的最终标签。
- $c$ 是一个可能的类别标签(在CIFAR-100数据集中,c ∈ {1, ..., 100})。
- $N$ 是委员会中模型的数量。
- $p_i(c|x)$ 是第 i 个模型预测输入 x 属于类别 c 的概率。

这个公式表示,对于每个可能的类别标签,软委员会将所有模型的预测概率相加,并将具有最高总和概率的类别作为最终预测。

* CCT-2/3x2架构

CCT-2/3x2架构包含两个CLs,然后是两个Transformer块。此架构的延迟为 $2 + 4 × 2 + 1 = 11$,远远低于CCT-7/3x1的延迟,为30。

通过使用具有$dim$(MHA维度)为512和1024的LS-MHA,并改变每个Transformer块的head数量在8到64之间,发现六个CCT-2/3x2架构的平均独立精度约为0.79。然后,这些架构的软委员会分别提供了$N=5$个架构的~0.827精度,$N=6$个架构的0.83精度,甚至超过了训练5000多个epoch的CCT-7/3x1的精度。

* CCT-1/3x1和CCT-1/3x2架构

CCT-1/3x1由一个CL和一个Transformer块组成。CCT-1/3x2由两个CLs和一个Transformer块组成。

对于CCT-1/3x1,总延迟为$1 + 4 + 1 = 6$,而对于CCT-1/3x2,延迟为$2 + 4 + 1 = 7$,均小于CCT-7/3x1延迟的四分之一。包含五个CCT-1/3x2和五个CCT-1/3x1架构的软委员会达到了~0.811精度,与CCT-7/3x1经过1000个epoch获得的0.809精度相当。

虽然CCT-1/3x2的平均独立精度比CCT-1/3x1高0.017,但单独的$N=10$个CCT-1/3x2架构的软委员会只能达到0.802精度。作者将CCT-1/3x1和CCT-1/3x2架构的组合提高的软委员会精度归因于它们在验证集上的较低相关性。

总而言之,论文的第三部分提供了以下结论:
1. 通过用卷积层替换初始的Transformer块,可以构建低延迟且高精度的ViT架构。
2. 软委员会决策方法可以通过组合多个不同架构的ViT模型来进一步提高准确率。

4. 实现细节与注意事项

关键实现细节:

* SHP计算: SHP计算的关键在于正确的头静默和平均激活值的计算。务必确保在计算SHP矩阵时,只有当前被评估的头的节点处于激活状态。
* 混合架构构建: 在构建混合架构时,需要注意卷积层和Transformer块之间的维度匹配。可以使用填充 (padding) 和池化 (pooling) 来调整维度。
* 软委员会决策: 软委员会决策的实现非常简单,只需将不同模型的输出向量相加即可。但是,选择具有较低相关性的模型对于最大化软委员会的性能至关重要。

实现难点与解决方案:

* 阈值选择: SHP矩阵的阈值选择会影响聚类和噪声分类的结果。需要仔细调整阈值,以获得最佳的性能。可以通过实验来选择最佳阈值。
* 头静默: 在头静默过程中,务必确保所有不属于被评估的头的节点都被完全静默。可以使用掩码 (mask) 操作来实现头静默。
* 模型选择: 选择具有较低相关性的模型对于最大化软委员会的性能至关重要。可以使用相关性分析来评估模型之间的相关性。

优化建议与最佳实践:

* 自动化超参数搜索: 使用自动化超参数搜索技术来优化模型的超参数,例如学习率、权重衰减和批次大小。
* 使用预训练模型: 使用预训练模型可以显著提高模型的性能和收敛速度。
* 数据增强: 使用数据增强技术可以提高模型的泛化能力。
* 模型压缩: 可以使用模型压缩技术来进一步降低模型的延迟。

参数设置与调优:

* 学习率 (η): 学习率是一个非常重要的超参数。可以尝试不同的学习率,并选择能够获得最佳性能的学习率。
* 权重衰减 (α): 权重衰减可以防止模型过拟合。可以尝试不同的权重衰减值,并选择能够获得最佳性能的权重衰减值。
* 批次大小: 批次大小会影响模型的训练速度和性能。可以尝试不同的批次大小,并选择能够获得最佳性能的批次大小。
* 阈值: SHP矩阵的阈值需要仔细调整,以获得最佳的性能。
* 模型选择: 在构建软委员会时,需要选择具有较低相关性的模型。可以使用相关性分析来评估模型之间的相关性。

希望这些分析对您有所帮助!
返回论文列表