Vision Transformer图像块多样化:从多尺度采样到动态剪枝的工程实践
1. 从“千篇一律”到“百花齐放”为什么我们需要多样化的图像块在计算机视觉领域Vision TransformerViT的出现无疑是一场革命。它将自然语言处理中大放异彩的Transformer架构成功迁移到图像理解任务上其核心思想简单而有力将一张输入图像分割成一系列固定大小的、不重叠的“图像块”Patch然后将这些图像块线性投影为“令牌”Token序列最后送入标准的Transformer编码器进行处理。这套流程清晰、优雅并且在ImageNet等大规模数据集上取得了媲美甚至超越传统卷积神经网络CNN的性能。然而作为一名长期在一线进行模型优化和部署的工程师我在实际应用ViT及其变体时常常会遇到一些“别扭”的地方。最直观的感受是标准的ViT模型对图像块的划分方式过于“粗暴”和“单一”。想象一下你拿到一张复杂的街景图片里面有清晰的行人、模糊的背景、反光的车窗和纹理丰富的建筑物立面。标准ViT会像切豆腐一样用一把尺寸固定的刀比如16x16像素将这张图片均匀地切成若干小块。问题来了一个16x16的块可能刚好完整地包含一只眼睛的细节也可能横跨了天空和楼宇的边缘包含了截然不同的语义信息。这种“一刀切”的划分方式迫使模型必须从这些信息混杂、质量参差不齐的初始令牌开始学习无疑增加了模型理解图像的难度。这就引出了《Vision Transformers with Patch Diversification》这篇工作以及其反映的研究趋势所关注的核心问题“并非所有图像块都是你所需要的”。我们真的需要那么多信息冗余甚至带有噪声的令牌吗答案显然是否定的。过多的、低质量的令牌不仅会显著增加模型的计算负担Transformer的自注意力机制复杂度与令牌数量的平方成正比还可能引入干扰让模型难以聚焦于真正重要的视觉特征。因此“Patch Diversification”图像块多样化或更广义上的“令牌精简/筛选”成为了ViT模型优化的重要方向。其目标不再是简单地、均质地切割图像而是希望以更智能、更多样的方式生成或筛选初始的图像块表示从而在模型效率更少的计算量、更快的速度和模型效果更高的精度、更好的鲁棒性之间取得更好的平衡。这不仅仅是学术上的“锦上添花”对于将ViT模型部署到资源受限的边缘设备、移动端或是需要实时响应的应用场景如自动驾驶、视频分析中具有至关重要的工程价值。2. 理解“多样化”的维度不止于尺寸与形状当我们谈论“Patch Diversification”时很容易首先想到改变图像块的大小或形状。这确实是一个重要方面但“多样化”的内涵远不止于此。我们可以从多个维度来拆解这个概念这有助于我们理解不同优化方法的出发点和适用场景。2.1 空间维度的多样化打破固定的网格这是最直观的维度。标准ViT使用固定的、规则网格进行划分。多尺度图像块一种自然的想法是使用不同尺寸的滑动窗口来提取图像块。例如在图像的不同区域或不同层级使用12x12, 16x16, 24x24等多种尺寸的窗口。这样大尺寸块能捕获更宏观的上下文信息小尺寸块能保留更精细的局部细节。Swin Transformer采用的层级式设计和滑动窗口机制可以看作是在不同阶段融入了不同感受野的“块”信息。非矩形与重叠图像块为什么图像块必须是正方形且不重叠的呢在一些工作中研究者尝试使用圆形的、或通过可变形卷积动态调整形状的感受野来生成初始令牌。此外允许图像块之间有部分重叠可以缓解因严格切分导致的边缘信息割裂问题让相邻块之间有信息交流的“缓冲区”。2.2 信息内容的多样化关注“重要性”而非“均质性”这个维度关注的是图像块所携带信息的“质”而非其外形的“量”。其核心思想是不同的图像块对最终任务如分类、检测的贡献度是天差地别的。基于显著性的筛选我们可以利用视觉显著性检测算法在生成令牌序列之前先对图像进行分析识别出前景主体、关键物体等显著区域。然后可以对显著区域进行更密集的采样生成更多、更小的块对背景区域进行更稀疏的采样生成更少、更大的块甚至直接丢弃一些纯背景的低信息量块。这直接呼应了“not all patches are what you need”的思想。基于注意力权重的动态剪枝这是在模型前向传播过程中进行的动态优化。在Transformer的某些层通常是较浅的层计算出自注意力权重后我们可以根据令牌与其他令牌交互的“活跃度”例如该令牌注意力权重的熵值很低意味着它只关注极少数令牌或者该令牌被其他令牌关注的总权重很低判定其为“不重要”的令牌并在后续层中将其移除或合并。代表性的工作有DynamicViT、EViT等。这种方法的好处是自适应性极强能够根据每张输入图像的内容动态调整计算图。2.3 特征表示的多样化在嵌入空间做文章即使图像块的物理划分是固定的我们也可以在将其投影为令牌嵌入向量这一步引入多样性。混合特征提取在将图像块扁平化并线性投影之前可以先使用一组不同参数的微型卷积核或Gabor滤波器等对每个块进行预处理得到多种初始特征然后再进行融合。这相当于为每个图像块提供了多种不同的“观察视角”。随机化投影在训练过程中对线性投影层加入适度的随机扰动或使用Dropout可以强制模型不依赖于某一种固定的投影方式从而学习到更鲁棒的令牌表示。这可以看作是一种针对令牌嵌入的数据增强。注意在实际的工程实现中上述维度往往是结合使用的。例如一个系统可能首先采用多尺度划分生成初始的、数量较多的令牌候选集然后通过基于注意力的动态剪枝机制在推理过程中快速筛选出最关键的子集进行深层处理。3. 核心实现策略如何为ViT引入“多样化”理解了“为什么”和“是什么”之后我们最关心的是“怎么做”。这里我结合论文中常见的思路和工程实践梳理几种核心的实现策略。需要明确的是《Vision Transformers with Patch Diversification》作为一个具体的论文标题其方法可能特指某一种技术路径。但我们可以从更广阔的视角看看社区是如何解决这个问题的。3.1 策略一预定义的多尺度与稀疏采样这是一种在模型输入阶段即确定的、静态的多样化策略。它的优势是实现简单推理过程完全规则化没有条件分支易于部署。操作流程设计采样策略预先定义好几组采样参数。例如策略A: 使用16x16步长16的密集网格采样标准ViT。策略B: 使用24x24步长24的稀疏网格采样。策略C: 在图像中心224x224区域使用14x14的密集采样在边缘区域使用28x28的稀疏采样。并行处理与特征融合将同一张输入图像分别用这几种策略进行图像块划分和线性投影得到多组令牌序列。序列合并将这些来自不同策略的令牌序列在序列长度维度上进行拼接形成一条更长的、多样化的令牌序列。送入Transformer将这条长序列输入Transformer编码器。由于自注意力机制是排列不变的它能自然地处理这些来自不同尺度和位置的令牌。工程实现要点位置编码这是关键。不同尺寸、不同位置的图像块其位置编码必须重新计算以反映其在原始图像中的真实2D坐标。不能直接复用标准ViT的1D位置编码。计算开销令牌总数增加了所以计算量FLOPs必然会上升。这种策略的收益主要来自于模型性能如精度、鲁棒性的提升而非效率提升。它牺牲了效率来换取更强的表示能力。# 伪代码示意多尺度采样与融合 import torch import torch.nn as nn class MultiScalePatchEmbed(nn.Module): def __init__(self, img_size224, scales[(16,16), (24,24)], embed_dim768): super().__init__() self.projects nn.ModuleList([ nn.Conv2d(3, embed_dim, kernel_sizes, strides) for s in scales ]) # 为每种尺度计算对应的位置编码 self.pos_embeds nn.ParameterList([ nn.Parameter(torch.zeros(1, (img_size//s[0])*(img_size//s[1]), embed_dim)) for s in scales ]) def forward(self, x): # x: [B, 3, 224, 224] tokens_list [] for proj, pos in zip(self.projects, self.pos_embeds): # 卷积实现分块与投影 token proj(x).flatten(2).transpose(1, 2) # [B, N, C] token token pos tokens_list.append(token) # 在序列维度拼接 diverse_tokens torch.cat(tokens_list, dim1) # [B, N1N2..., C] return diverse_tokens3.2 策略二基于注意力的动态令牌精简这是目前研究的热点也是真正实现“按需计算”、提升推理效率的关键。其核心是在前向传播中根据中间特征动态决定哪些令牌值得保留哪些可以丢弃或压缩。典型流程以分类任务为例生成初始令牌使用标准方式如16x16网格生成全部N个令牌。我们称其为“候选令牌”。经过浅层Transformer让所有候选令牌通过前L层例如前4层Transformer块进行初步交互和特征提炼。计算令牌重要性得分在第L层之后我们需要一个“裁判”来给每个令牌打分。常见的方法有基于CLS令牌的注意力计算每个图像块令牌与任务相关的[CLS]令牌之间的注意力权重。权重越高说明该图像块对最终决策越重要。令牌自身的特征范数或熵计算每个令牌特征向量的L2范数范数大的通常包含更多信息。或者计算该令牌在所有注意力头中注意力分布的熵熵值小分布集中的令牌可能作用更明确。排序与筛选根据重要性得分对所有图像块令牌进行排序只保留得分最高的K个K N。K可以是一个固定值也可以是动态确定的如保留得分大于某个阈值的令牌。继续处理将筛选出的K个重要令牌连同始终保留的[CLS]令牌送入后续的深层Transformer块进行计算。被淘汰的令牌不再参与后续计算从而节省了大量计算量。工程实现中的坑与技巧梯度流中断动态筛选是一个不可微的“选择”操作会阻断梯度传播。解决方案是使用“软筛选”或Gumbel-Softmax技巧在训练时引入可微性或者在训练后期再引入硬筛选进行微调。筛选策略的稳定性不同的输入图像筛选出的令牌集合差异可能很大这可能导致模型行为不稳定。一个实用的技巧是引入“保留率”的平滑约束例如确保每层至少保留一定比例的令牌或者对重要性得分进行层间的指数平滑。与位置编码的兼容性当令牌被移除后剩余令牌的位置编码需要保持不变因为它们仍然对应于原始图像中的绝对位置。这一点在实现时需特别注意。3.3 策略三结构化图像块生成与学习除了上述两种相对“宏观”的策略还有一些工作从“微观”入手试图让模型自己学习如何生成更好的图像块。可学习的图像块化将标准的、固定的卷积投影层kernel_sizestride16替换为一个小型的、可学习的神经网络例如一个3-4层的微型CNN。这个网络以整张图像为输入其最后的特征图的空间位置自然就成为了“图像块”的嵌入表示。通过训练这个网络可以学会在语义边界处“切割”生成更具信息量的块。令牌聚类与合并在Transformer的中间层不再粗暴地丢弃“不重要”的令牌而是将特征相似的令牌进行聚类然后用一个聚类中心来代表一组令牌。这既减少了令牌数量又保留了信息。例如可以将空间上相邻且特征相似的背景令牌合并为一个。这种方法需要在线运行聚类算法如K-Means的简化版对工程实现有一定挑战。4. 实战为现有ViT模型集成动态令牌精简理论说了这么多我们来点实际的。假设我们手上有一个预训练好的标准ViT-B/16模型现在想为其集成一个简单的、基于注意力得分的动态令牌筛选模块以实现推理加速。这里我提供一个概念清晰、易于集成的PyTorch实现方案。我们的目标在ViT的第4个Transformer块之后动态筛选出50%最重要的图像块令牌只将这些令牌送入第5到第12层进行计算。4.1 步骤一剖析ViT结构定位插入点首先我们需要理解标准ViT以timm库为例的结构。一个ViT模型通常由以下部分组成patch_embed: 将图像切块并投影为令牌。cls_token: 可学习的分类令牌。pos_embed: 位置编码。blocks: 一个由多个TransformerBlock组成的模块列表。norm和head: 最后的层归一化和分类头。我们的动态筛选模块将作为一个新的nn.Module插入到blocks[3]第4层和blocks[4]第5层之间。4.2 步骤二设计动态筛选模块这个模块需要完成三件事评估重要性、排序筛选、重组数据。import torch import torch.nn as nn import torch.nn.functional as F class DynamicTokenPruning(nn.Module): 动态令牌剪枝模块。 在训练时使用可微的top-k选择通过Gumbel-Softmax。 在推理时使用硬性的top-k选择。 def __init__(self, target_keep_ratio0.5, gumbel_temperature1.0, hardTrue): super().__init__() self.target_keep_ratio target_keep_ratio self.temp gumbel_temperature self.hard hard # 是否使用硬性选择 def forward(self, x, cls_token): Args: x: 图像块令牌形状为 [Batch, Num_patches, Embed_dim] cls_token: [CLS]令牌形状为 [Batch, 1, Embed_dim] Returns: pruned_x: 筛选后的图像块令牌 kept_indices: 被保留的令牌索引用于调试 batch_size, num_patches, embed_dim x.shape # 1. 计算重要性分数使用与[CLS]令牌的点积注意力作为分数 # [CLS]令牌作为“查询” cls cls_token.transpose(0, 1) # [1, B, D] - [B, 1, D] # 图像块作为“键” scores torch.matmul(cls, x.transpose(1, 2)) # [B, 1, D] * [B, D, N] - [B, 1, N] scores scores.squeeze(1) # [B, N] # 2. 确定要保留的令牌数量 num_keep int(num_patches * self.target_keep_ratio) if num_keep num_patches or num_keep 0: # 无需筛选 return x, torch.arange(num_patches, devicex.device).expand(batch_size, -1) # 3. 根据训练/推理模式进行选择 if self.training and not self.hard: # 训练模式使用Gumbel-Softmax进行可微的Top-K选择 keep_prob F.gumbel_softmax(scores, tauself.temp, hardFalse, dim-1) # 我们只需要选择top-k这里用一个技巧生成一个mask _, topk_indices torch.topk(scores, knum_keep, dim-1) keep_mask torch.zeros_like(scores, dtypetorch.bool) keep_mask.scatter_(1, topk_indices, True) # 使用mask进行选择但梯度可以通过gumbel分布回传 scores_for_gumbel scores.masked_fill(~keep_mask, -1e9) keep_prob F.gumbel_softmax(scores_for_gumbel, tauself.temp, hardFalse, dim-1) # 采样实际上在hardFalse时是软选择我们取其期望 # 为了简化在训练时我们也用硬索引但通过Gumbel-ST传递梯度 _, topk_indices torch.topk(scores torch.randn_like(scores)*0.01, knum_keep, dim-1) else: # 推理/硬选择模式直接取Top-K索引 _, topk_indices torch.topk(scores, knum_keep, dim-1) # 4. 根据索引收集要保留的令牌 batch_indices torch.arange(batch_size, devicex.device).view(-1, 1).expand(-1, num_keep) pruned_x x[batch_indices, topk_indices, :] # [B, K, D] return pruned_x, topk_indices4.3 步骤三改造ViT前向传播逻辑我们需要修改ViT模型的前向传播函数在指定层后插入筛选模块。def forward_modified_vit(self, x): B x.shape[0] # 1. 原始的图像块嵌入和位置编码 x self.patch_embed(x) cls_tokens self.cls_token.expand(B, -1, -1) x torch.cat((cls_tokens, x), dim1) x x self.pos_embed x self.pos_drop(x) # 2. 逐层通过Transformer块 for i, blk in enumerate(self.blocks): x blk(x) # 在第4层之后i3插入动态剪枝 if i 3: # 分离[CLS]令牌和图像块令牌 cls_token, img_tokens x[:, :1, :], x[:, 1:, :] # 动态剪枝图像块令牌 pruned_img_tokens, kept_idx self.pruner(img_tokens, cls_token) # 重新组合[CLS]令牌 筛选后的图像块令牌 x torch.cat((cls_token, pruned_img_tokens), dim1) # 注意此时位置编码需要相应调整。一个简单方案是使用可学习的位置编码 # 或者将原始位置编码中对应的索引提取出来。 # 这里为简化假设模型能自适应实际效果会打折扣需微调。 # 3. 后续处理 x self.norm(x) return self.head(x[:, 0]) # 取[CLS]令牌用于分类4.4 步骤四训练与微调策略直接在一个预训练好的ViT上插入剪枝模块并运行性能大概率会暴跌。我们需要一个循序渐进的微调策略。热身训练首先将剪枝模块的target_keep_ratio设置为1.0即不剪枝用少量数据如ImageNet的子集对模型进行几个epoch的微调。目的是让模型适应新的前向传播图即使什么都没剪。逐步剪枝然后以较小的步长如每次降低0.1逐步降低target_keep_ratio每降低一次都用完整数据集微调1-2个epoch。例如1.0 - 0.9 - 0.8 - ... - 0.5。这个过程让模型平滑地学习如何应对越来越“苛刻”的令牌保留环境。联合微调当target_keep_ratio达到目标值如0.5后再用完整数据集进行一段时间的联合微调稳定模型性能。评估与部署在验证集上评估精度和速度。由于减少了约一半的令牌参与后半程计算理论上FLOPs和推理时间都会有显著下降。部署时将pruner设置为hardTrue的推理模式。重要提示动态剪枝会改变每层处理的令牌数量这导致标准的Transformer实现其注意力机制通常假设序列长度固定无法直接使用高度优化的内核如FlashAttention。在部署时可能需要自定义算子或使用支持可变长度序列的推理框架这是工程落地时需要攻克的一个难点。5. 影响评估多样化带来了什么为ViT引入图像块多样化机制究竟带来了哪些实实在在的好处和潜在的代价我们需要从多个维度进行客观评估。5.1 正面收益效率与效果的潜在双赢计算效率提升对于动态剪枝类方法这是最直接的收益。通过移除冗余令牌可以大幅减少模型在中间层和深层需要计算的序列长度。自注意力机制的复杂度从O(N²)降低到O(K²)其中K是保留的令牌数。在实际的GPU/CPU推理中这能直接转化为更短的延迟和更高的吞吐量对于实时应用至关重要。模型精度保持甚至提升这听起来有违直觉但许多论文报告了在相似FLOPs下采用令牌精简的模型精度超过了基线模型。原因在于“去噪”和“聚焦”。强制模型只关注最重要的信息可以看作是一种强大的正则化它迫使模型学习更鲁棒、更具判别性的特征而不是依赖于所有可能包含噪声的背景信息。模型鲁棒性增强多样化的图像块生成方式尤其是多尺度策略让模型能够同时捕获不同粒度下的视觉特征。这有助于提升模型对于尺度变化、遮挡、局部扰动等情况的鲁棒性。模型不再只依赖于单一尺度的局部证据来做决策。可解释性线索动态重要性评分机制提供了一个副产品我们可以可视化哪些图像块被模型认为是重要的。这为我们理解模型的决策过程打开了一扇窗。例如在图像分类任务中我们可以观察被保留的令牌是否确实集中在目标物体上。5.2 需要权衡的代价与挑战训练复杂性增加无论是引入Gumbel-Softmax、额外的可学习模块还是需要多阶段逐步剪枝的微调策略都显著增加了模型训练和调优的复杂性。训练时间更长超参数更多如剪枝率、剪枝位置、温度系数τ等对实验者的经验要求更高。推理过程的不确定性对于动态方法每张输入图像的推理路径保留了哪些令牌可能都不相同。这种“条件计算”特性使得传统的静态图优化、算子融合、批处理Batch Inference变得困难。在批处理时如果同一批内的样本保留的令牌数不同就需要进行填充Padding或更复杂的打包Packing操作这会引入额外开销可能抵消一部分计算节省。硬件友好度下降现代AI加速器如GPU的Tensor CoreNPU的矩阵计算单元最喜欢规整的、可预测的张量运算。动态剪枝带来的不规则计算和稀疏性会降低计算单元的利用率使得实际的加速比往往低于理论FLOPs的减少比例。需要专门的稀疏计算库或硬件支持才能充分发挥其优势。可能的信息损失风险尽管目标是去除冗余但“冗余”的界定是模糊的。在某些边缘案例或对抗性样本中被剪枝机制判定为“不重要”的令牌可能恰恰包含了关键的反直觉信息。这可能导致模型在某些情况下的性能下降或出现新的脆弱性。5.3 选型建议如何为你的项目选择策略面对多种“多样化”策略如何选择这里有一个简单的决策树供参考你的首要目标是提升精度还是提升速度目标为精度优先考虑多尺度/多形态融合策略。它通过增加计算成本来换取更强的特征表示能力适合对精度要求极致、计算资源相对充裕的场景如云端部署的学术研究或关键业务。目标为速度/效率优先考虑动态令牌剪枝/合并策略。它直接减少计算量适合移动端、嵌入式设备或需要高并发的在线服务。你的部署环境对计算图有何要求需要静态图、确定性推理选择预定义的多尺度/稀疏采样。它的计算图是固定的易于优化和部署。可以接受动态图、条件计算可以选择动态剪枝。但需要评估目标推理框架如TensorRT, ONNX Runtime, TFLite对动态形状和稀疏运算的支持程度。你的数据和应用场景有何特点数据中主体与背景区分明显如肖像摄影、工业质检基于显著性或简单注意力评分的剪枝会非常有效。数据复杂、主体多变、背景信息也可能重要如街景理解、医学图像需谨慎使用剪枝或设置较高的保留率或采用更保守的令牌合并策略而非直接丢弃。一个实用的混合策略在实际工程中我经常采用一种折中方案——“早期粗剪后期不剪”。即在最前面的1-2个Transformer层之后进行一次“粗粒度”的剪枝用一个较低的阈值淘汰掉大量明显无关的背景块。在中间层则采用更保守的合并或轻度剪枝。在最后的几层不再进行任何剪枝让模型利用最丰富的上下文信息做最终决策。这种策略在效率和效果之间取得了不错的平衡。6. 超越分类多样化思想在其他视觉任务中的应用图像分类是ViT的“起家”任务但“Patch Diversification”的思想完全可以推广到更广泛的视觉任务中并且往往能带来更大的收益。6.1 目标检测与实例分割在这些密集预测任务中图像包含大量物体和复杂背景冗余信息极多。应用方式可以在BackboneViT部分应用动态剪枝快速过滤掉大片的纯背景区域让后续的检测头如DETR中的Transformer Decoder或传统的FPNRPN只处理可能包含物体的、信息密集的区域。这能极大提升检测器的推理速度。注意事项剪枝需要格外小心避免将小物体或物体边缘的重要上下文信息误删。通常需要针对检测任务重新设计重要性评分函数例如结合区域提议网络RPN的初步得分。6.2 图像分割语义/全景分割语义分割需要对每个像素进行分类令牌数量巨大。应用方式一种策略是在编码器Encoder中剪枝在解码器Decoder中恢复。编码器使用动态剪枝进行高效的特征提取得到一组稀疏但高语义级别的令牌。解码器则通过上采样、跳跃连接或可变形注意力等机制将这些稀疏的高层特征与早期未剪枝的、高分辨率的底层特征融合逐步恢复出密集的像素级预测图。优势避免了在编码器的深层对所有像素进行昂贵的全局注意力计算将计算资源集中在语义信息丰富的区域。6.3 视频理解视频可以看作是时空上的令牌序列。冗余不仅存在于空间上更大量存在于时间上相邻帧高度相似。应用方式时空联合剪枝。我们可以计算视频片段中时空令牌的重要性并动态选择关键帧中的关键区域进行处理。例如对于动作识别模型应聚焦于发生运动的区域和关键帧。这能将视频理解的巨大计算开销降低一个数量级。关键技术需要设计能够同时评估空间重要性和时间重要性的评分机制。6.4 多模态任务视觉-语言模型在如CLIP、BLIP等模型中图像需要与文本进行对齐。应用方式图像令牌的重要性可以由文本查询来指导。例如给定文本“一只棕色的狗在草地上奔跑”模型在计算图像-文本相似度时可以动态地强化图像中与“狗”、“棕色”、“草地”、“奔跑”相关的令牌弱化无关区域的令牌。这本质上是一种跨模态的注意力引导剪枝。收益提升模型对文本描述中细粒度属性的感知能力并加速大规模多模态模型的推理。从“一刀切”的均匀划分到“因图而异”的智能筛选Patch Diversification 代表了Vision Transformer模型从“可用”走向“高效、鲁棒、可解释”的必然进化路径。它不是一个孤立的技巧而是一种设计哲学让计算资源跟随信息密度走。对于研究者它开辟了结构设计、动态网络、稀疏计算等多个有趣的方向对于工程师它提供了实实在在的模型加速工具是让ViT这类强大模型真正落地到产品中的关键。