无需微调！微软亚研院：高效加速大规模视觉 Transformer 密集预测任务的方法

49 无需微调加速大规模视觉 Transformer 密集预测任务的方法
(来自微软亚洲研究院)
49.1 论文解读
49.1.1 本文背景和动机
49.1.2 本文方法总体思路
49.1.3 Token Clustering 层
49.1.4 Token Reconstruction 层
49.1.5 本文方法如何拓展到 Swin Transformer？
49.1.6 为什么本文方法可以避免微调？
49.1.7 实验结果

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制，不采用 RNN 的顺序结构，使得模型可以并行化训练，而且能够拥有全局信息。

视觉 Transformer 在诸多计算机视觉任务中取得了可观的表现，但是在处理一张高分辨率图片 (图片的 token (Patch) 数比较多) 时需要较大的计算代价。因此之前的一些工作倾向于减少输入给 Transformer 模型的 token 数，以减小大规模视觉 Transformer 模型的计算代价。具体而言，可以从所有的图片 token 中选择出一小部分比较重要的 token，然后微调整个模型的权重。但是这种微调的范式其实对于下游的密集预测任务是不太友好的，因为下游任务输入的分辨率普遍要比 ImageNet-1K 图像分类任务大很多，这就带来了更大的 GPU memory cost 和计算代价。因此，本文的目的就是想去探索如何在不微调权重的情况下加速大规模视觉 Transformer 密集预测任务的方法。

论文名称：Expediting Large-Scale Vision Transformer for Dense Prediction without Fine-tuning

论文地址：

https://arxiv.org/pdf/2210.01035.pdf

本文背景和动机

视觉 Transformer 在诸多计算机视觉任务中取得了可观的表现, 定义输入图片为 , Transformer 模型先将其展平成 。其中, ( 代表每个 Patch 的大小, 代表图片 token 或者 Patch 的数量。

但是，大规模的视觉 Transformer 模型在处理一张高分辨率图片 (图片的 token (Patch) 数比较多) 时需要较大的计算代价，尤其是是在处理高分辨率输入图片的时候，这也限制了它们在更多资源受限设备上的广泛应用。因此之前的一些工作倾向于减少输入给 Transformer 模型的 token 数，以减小大规模视觉 Transformer 模型的计算代价，以及加速视觉 Transformer 模型。具体而言，可以从所有的图片 token 中选择出一小部分比较重要的 token，然后微调整个模型的权重。这种范式的例子是 Dynamic ViT[1] 以及 EViT[2]，提出两种不同的动态标记稀疏化框架，以逐步减少多余的 token，根据额外训练的预测模块所预测的分数或与[class] token 的相关性，选择信息量最大的 token。

但是这类方法都需要额外在下游任务中微调模型的权重，但这种微调的范式其实对于下游的密集预测任务是不太友好的，因为下游任务输入的分辨率普遍要比 ImageNet-1K 图像分类任务大很多 (普遍是 1024×1024)，这就带来了更大的 GPU memory cost 和计算代价。因此，本文的目的就是想去探索如何在不微调权重的情况下加速大规模视觉 Transformer 密集预测任务的方法。

作者认为：

• 训练好的视觉 Transformer 模型包含了大量的局部空间信息的冗余。
• 密集预测任务 (如检测，分割等) 往往需要高分辨率的输入图片。

因此，作者希望在使用大规模视觉 Transformer 模型时，把输入图片的分辨率从大变小再变大，借此减小高分辨率的输入图片带来的计算代价过高的问题。而且还能够加速视觉 Transformer 模型以及在不进行任何微调操作的情况下在下游任务中使用。

本文方法总体思路

视觉 Transformer 模型由一系列的 multi-head self-attention (MHSA) 和 feed-forward network，搭配 Layer Normalization 和残差链接组成，如下式所示：

式中, 代表着层的索引, 为第 层的特征。视觉 Transformer 模型在处理高分辨率表征时的计算成本明显增加。而高分辨率表征又对密集预测任务至关重要。如前文所述, 本文希望提出一种不需要微调任何权重即可显著减小 "大规模视觉 Transformer 模型 高分辨率输入图片" 前向传播计算代价的方法。

如上图1所示，本文方法包含两种无参数的层：Token Clustering 层以及 Token Reconstruction 层。Token Clustering 层通过聚合局部语义相似的 token 来降低特征的分辨率，然后在低分辨率表征中通过 个 Transformer 层，低分辨率表征可以大大加快推理速度并节省计算资源。经过数层之后，再通过 Token Reconstruction 层恢复特征的分辨率。

一张图片的高分辨率特征首先通过 个 Transformer 层，再通过 Token Clustering 层转化成低分辨率特征通过 个 Transformer 层， 再通过 Token Reconstruction 层恢复成高分辨率特征通过 个 Transformer 层。

Token Clustering层

Token Clustering 层的基本思路是: 模型中的高分辨率特征有 个向量 ，我们 希望把它简化成 个 , 所以这个问题可以视为一个基本的聚类 问题, 即: 如何把 个向量给聚类成 个? 本文方法基于 , 也基于 EM 算法。

首先对第 层 (低分辨率的第1层) 高分辨率的特征应用 adaptive average pooling (AAP) 获得 个聚类中心:

其中, 。

E 步：

代表 个聚类中心，那么第 个 pixel 与第 个聚类中心 (虽然一共有 hw 个聚类中心, 但是只计算该 pixel与周围 个的相似度) 的相似度 可以计算为:

M 步：

得到了所有的相似度之后，再根据它们计算新的聚类中心：

以上 步和 步一共迭代 次, 是温度系数, 。

Token Clustering 层的 PyTorch 伪代码是：

def token_clustering_layer(features, cluster_features_shape, num_iters, teu):
    # args :
        # features : shape [B, C, H, W]
        # cluster_features_shape : [h, w]
        # num_iters : num of iterations of updating cluster features
        # teu: the temperture of distance matrix
    # output : 
        # cluster_features : shape [B, hw, C]   

    B, C, H, W = features.shape

    # initialize the cluster features
    cluster_features = interpolate(features, cluster_features_shape)

    # construct mask to constrain the interactions within local range
    mask = calculate_mask(features.shape, cluster_features_shape)
    mask = (~ mask) * 1e16

    features = features.reshape(B, C, -1).permute(0, 2, 1) # (B, HW, C)
    cluster_features = cluster_features.reshape(B, C, -1).permute(0, 2, 1) # (B, hw, C)

    for _ in range (num_iters):
    # calculate L2 distance of features and cluster features, the shape distance_matrix is (B, hw, HW)
        distance_matrix = L2_distance(features, cluster_features)
        # mask remote distance through softmax
        distance_matrix += mask
        weights = (-distance_matrix / teu).softmax(dim=1)
        # let the sum of weight of each cluster feature be 1
        weights = weights / weights.sum(dim=2, keepdim=True).clamp_(min=1e-16)
        cluster_features = matrix_product(weights, features)

    return cluster_features

Token Reconstruction层

Token Reconstruction 层的基本思路是: 现在有 个聚类中心 , 那么要如何根据这些聚类中心恢复出高分辨率的 个向量

基本的思路就是通过聚类中心 和 Token 之间的相似性 来计算和完成低分辨率特征到高分辨率特征的重建过程:

具体写出来就是：

式中, 代表距离 最近的 个聚类中心。

最后将剩余的 个 Transformer 层应用于重建的高分辨率特征上面, 以预测目标结果, 比如语义分割的结果或者深度估计的结果。

Token Reconstruction 层的 PyTorch 伪代码是 (部分张量的维度并非来自原论文，是我补充上的)：

def token_reconstruction_layer(cluster_features, features_before_clustering, features_after_clustering, k, teu):
    # args :
        # cluster_features: shape [B, hw , C]
        # features_before_clustering: features of alpha -th layer before clustering, shape [B, hw , C]
        # features_after_clustering: features of alpha -th layer before clustering, shape [B, HW , C]
        # k: topk parameter
        # teu: the temperture of weight matrix
    # output :
        # features : reconstruction features, shape [B, HW , C]

    # calculate L2 distance between features and cluster_features
    distance = L2_distance(features_before_clustering, features_after_clustering) # (B, HW, hw)
    weight = exp(-teu * distance)
    # only remain the k weight of the most simliar features, calculating mask
    topk, indices = topk(weight, k=k, dim=2) # (B, HW, k)
    mink = min(topk, dim=-1).values # (B, HW)
    mink = mink.unsqueeze(-1).repeat(1, 1, weight.shape[-1])  # (B, HW, hw)
    mask = greater_or_equal(weight, mink) # (B, HW, hw)
    weight = weight * mask # (B, HW, hw)

    weight = weight / weight.sum(dim=2, keepdim=True).clamp_(min=1e-16) # (B, HW, hw)
    features = matrix_product(weight, cluster_features)

    return features

本文方法如何拓展到Swin Transformer

下图3展示了本文方法如何拓展到 Swin Transformer。Swin Transformer 是在一个 window 的内部进行 Self-attention 的操作，所以要在 Swin Transformer 里面使用本文的方法，就需要在一个 window 的内部进行 Token Clustering 的操作和 Token Reconstruction 的操作。

首先把一个 window 内部的 个 token 进行 49.1.3 节所述的聚类操作, 得到 个 token, 再进行 节所述的重建操作, 将其恢复成 个 token。

对于 Swin Transformer 里面使用到的相对位置编码, 作者通过双线性的揷值策略, 把维度为 的相对位置编码, 在处理 clustered representation 时变为 的维度。

以上策略使得本文方法也能够提升 Swin Transformer 模型的效率，且无需微调其权重。

为什么本文方法可以避免微调？

1. 本文方法引入的 Token Clustering 层以及 Token Reconstruction 层都是不包含任何参数的，所以把它们放到大规模的视觉 Transformer 模型以后，无需重新训练任何参数。
2. 重建之后的高分辨率特征保持着与低分辨率特征很高的语义相似性，以 Segmenter + ViT-L/16 (on ADE20K, 为例，作者计算了重建之后得到的高分辨率特征和原始模型的高分辨率特征之间的余弦相似度，结果如下图4所示。可以看到，在不同的 Transformer 层中，重建的高分辨率特征和原始的高分辨率特征之间的余弦相似度一直很高。换言之，本文方法很好地保持了原始高分辨率特征图中的语义信息，因此能够避免微调。

实验结果

作者在一系列较为复杂的下有任务上面验证本文方法的有效性，任务包括：object detection, semantic segmentation, instance segmentation, panoptic segmentation, 以及 monocular depth estimation。在实验过程中作者使用了 ViT 和 Swin 提供的预训练权重，没有任何微调操作。

数据集：

COCO [object detection, instance segmentation, panoptic segmentation]： 训练集 118K images，验证集 5K images。

ADE20K [semantic segmentation]： 训练集 20210 images，验证集 2000 images。

PASCAL-Context [semantic segmentation]： 训练集 4996 images，验证集 5104 images。

Cityscapes [semantic segmentation]： 训练集 2975 images，验证集 500 images。

KITTI [monocular depth estimation]： 训练集 26K images，验证集 698 images。

NYUv2 [monocular depth estimation]： 训练集 795 images，验证集 654 images。

评价指标：

AP (average precision) [object detection]。

mask AP (mask average precision) [instance segmentation]。

PQ (panoptic quality) [panoptic segmentation]。

mIoU (mean intersection-over-union) [semantic segmentation]。

RMSE (root mean squared error) [monocular depth estimation]。

消融实验结果

消融实验基于 ADE20K 语义分割任务，模型为 Segmenter+ViT-L/16。

首先探索的是 Token Clustering 层以及 Token Reconstruction 层，式3和4中的一些超参数的影响，结果如下图5所示。结果显示本文方法对超参数而言并不敏感，最终选择的是：。

然后探索聚类数 的影响。实验中使用的高分辨率特征是 , 图片块的大小是 , 所以 , 结果如下图6所示。我们发现当 聚类数量过小时会严重影响密集预测任务的性能, 把聚类中心数量设置为 可以 获得性能-计算量的 trade-off。

作者又将本文提出的 token clustering/reconstruction 范式与 adaptive average pooling/bilinear upsample 范式进行了对比。对比实验设置了两种聚类数量，分别是 20×20 和 28×28，在实验中使用 adaptive average pooling 以及 bilinear upsample 替换 token clustering/reconstruction，结果如下图7所示。可以看到，本文提出的提出的 token clustering/reconstruction 范式的性能在不同的聚类数量设置中一直优于 adaptive average pooling/bilinear upsample 范式。

高分辨率层数 的影响如下图8所示。结果显示, 当 大于10时取得了不错的性能, 综从性能和 计算量的角度综合考虑之后选择了 。

高分辨率层数+低分辨率层数 的影响如下图9所示 (实验固定了 )。结果显示, 当 时取得了最佳性能, 因此最终选择了 。

目标检测实验结果

目标检测实验所基于的框架是 SwinV2-L + HTC++, 原始的 SwinV2-L 模型的4个 stage 的层数分别为 , 作者只将本文方法应用于第3个 stage 的18层里面, 因为这部分占据了整个模型较大的计算代价。具体而言, 将 window token clustering/reconstruction 层分别揷在了第3个 stage 的8/18层之后, 也就是 。默认窗口大小为 , 将聚类窗口大小设置为 。与 SwinV2-L 相比, 本文方法提高了 的 FPS，节省了近 的 GFLOPs, 同时保持了大约 的目标检测和实例分割性能。

分割任务实验结果

分割任务实验所基于的框架有两种，分别是 Segmenter + ViT-L/16 和 Mask2Former + Swin-L。

原始的 ViT-L 首先将图像分割成大小为 16×16 的图像 Patch，并使用 Patch Embedding 将通道维度增加到1024，然后应用24个连续的 Transformer 层进行学习图片的表征。作者把 token clustering/reconstruction 层插入 ViT-L 的 backbone 中，实验结果如下图11所示。

Swin- 的默认窗口大小为 , 将聚类窗口大小设置为 , 实验 结果如下图12所示。可以看到, 本文方法在 COCO 全景分割任务上显著提高了 FPS 35% 以上, 而性能略有下降。

单目深度估计实验结果

目标检测实验所基于的框架是 DPT (Dense Prediction Transformer) + R50+ViT-B/16 (ResNet50 后面连接一个 ViT-B/16)，原始的 ViT-B/16 架构由12个 Transformer 层组成。实验结果如下图13所示，本文方法在KITTI/NYUv2 上分别将 DPT 加速了近 30%/37%。

ImageNet-1K 图像分类实验结果

这部分实验基于的基线模型是 EViT，它的做法是在图片通过模型的过程中，识别出并且只保留与 [class] token 最相关的 top-k 个 tokens。训练策略使用 SWAG[4]，它使用了36亿张弱标记的数据来预训练大规模视觉 Transformer 模型，SWAG + ViT-H/14 和 SWAG + ViT-L/16 分别是 88.55% 和 88.07%。本文方法与 EViT 在 SWAG 训练的加持下的性能，FPS，和 GFLOPs 的比较结果如下图所示。本文方法取得了与 EViT 相当的结果，同时效率更高，这进一步验证了本文策略可以通用于图像分类任务，而无需进行微调。

本文提出一种提高大规模视觉 Transformer 模型进行密集预测任务效率的简单实用的策略。下游任务输入的分辨率普遍要比 ImageNet-1K 图像分类任务大很多，这就带来了更大的 GPU memory cost 和计算代价。因此，本文的目的就是想去探索如何在不微调权重的情况下加速大规模视觉 Transformer 密集预测任务的方法。本文提出2种无参数的层：Token Clustering 层和 Token Reconstruction 层，分别用来降低和恢复中间特征的分辨率。作者在多种下游密集预测任务中验证了本文方法的有效性。

参考

1. ^Dynamicvit: Efficient vision transformers with dynamic token sparsification
2. ^Evit: Expediting vision transformers via token reorganizations
3. ^Superpixel sampling networks
4. ^Revisiting weakly supervised pre-training of visual perception models

关于我们