ALGM: Adaptive Local-then-Global Token Merging for Efficient Semantic Segmentation with Plain Vision Transformers
研究背景
随着计算机视觉技术的发展，语义分割作为图像理解的重要任务，在自动驾驶、医疗影像分析等领域得到了广泛应用。然而，传统的卷积神经网络在处理高分辨率图像时面临计算效率低下的问题。近年来，Vision Transformers（ViTs）因其强大的特征提取能力，在图像分割任务中取得了显著成效。然而，ViTs中的多头自注意力机制的计算复杂度与输入像素数的平方成正比，这极大地限制了其在高分辨率图像分割任务中的应用。因此，如何提高ViT基语义分割网络的计算效率，同时保持甚至提高其分割质量，成为了当前研究的热点。
研究意义
本研究提出了一种自适应的局部到全局的令牌合并方法（ALGM），旨在解决ViT基语义分割网络中的效率问题。通过合并冗余的令牌，ALGM能够显著减少ViT需要处理的令牌数量，从而提高计算效率。同时，该方法还能够保持甚至提高分割质量，为ViT在高分辨率图像分割任务中的应用提供了新的解决方案。这一研究对于推动ViT在实际应用中的落地具有重要意义。
文献综述
令牌减少方法
令牌减少是提高ViT效率的主要手段之一。现有方法主要分为令牌修剪、令牌暂停、令牌共享和令牌合并等几类。令牌修剪方法通过丢弃不重要的令牌来减少计算量，但这种方法不适用于需要为每个令牌进行预测的语义分割任务。令牌暂停方法保留被暂停的令牌并在后续层中重建它们，但这可能导致有用信息的丢失。令牌共享和令牌合并方法通过合并冗余令牌来减少计算量，但现有的方法要么需要引入额外的计算开销，要么仅在一次合并中限制了效率提升。
语义分割中的令牌减少
针对语义分割任务，现有的令牌减少方法通常面临分割质量下降的问题。一些方法通过重建合并后的令牌来恢复原始位置以进行预测，但这往往会导致分割质量的损失。因此，如何在减少令牌数量的同时保持甚至提高分割质量，是当前研究的难点。
具体方法
ALGM方法概述
ALGM方法通过两个阶段来合并令牌：第一阶段在第一网络层中合并局部窗口内的相似令牌；第二阶段在中间层中合并整个图像中的相似令牌。这种方法基于一个分析发现，即在早期层中，局部窗口内的同类令牌之间的相似度远高于异类令牌，而在中间层中，全局相似度逐渐成为合并令牌的有效指标。
具体步骤
局部合并：在第一层中，使用条件局部平均池化（CLAP）模块来合并局部窗口内的相似令牌。该模块首先计算窗口内令牌之间的相似度，并合并相似度高于阈值的令牌。
全局合并：在中间层中，使用全局二分图匹配（GBM）模块来合并整个图像中的相似令牌。该模块首先构建令牌之间的二分图，并保留每个令牌与其最相似令牌之间的边，然后合并这些令牌。
令牌恢复：在解码器之前，使用令牌恢复模块来恢复原始令牌分辨率，以便进行分割预测。
自适应合并策略
为了处理不同复杂度的图像，ALGM方法还引入了一种自适应合并策略。通过在训练集上计算令牌相似度的统计信息，动态确定合并阈值，以确保在不同图像中合并适当数量的令牌。
图1展示了ALGM（Adaptive Local-then-Global Merging）方法在提升语义分割效率和质量方面的效果。该方法应用于Segmenter、SegViT和SETR模型，并在ADE20K数据集上进行评估。
关键点：
- 吞吐量提升：应用ALGM后，模型的吞吐量（以每秒处理图像数衡量）显著提高，最高可达100%的提升。
- 分割质量：在提升效率的同时，ALGM还提高了分割质量，平均mIoU（均交并比）提升约0.7%。
图2比较了同类（intra-class）和不同类（inter-class）令牌之间的余弦相似度。
关键点：
- 局部相似度：在第一层网络中，小局部窗口内的同类令牌相似度远高于不同类令牌相似度。
- 全局相似度：随着网络层数的增加，同类令牌的全局相似度逐渐高于不同类令牌相似度。
图3展示了ALGM方法的两个主要模块：条件局部平均池化（CLAP）用于局部合并，全局二部匹配（GBM）用于全局合并。
关键点：
- CLAP模块：在第一层网络中，通过计算小窗口内令牌的余弦相似度并合并高相似度令牌来减少令牌数量。
- GBM模块：在中间层，通过构建令牌之间的全连接二部图并合并最相似令牌来进一步减少全局冗余
- 表1展示了ALGM方法在不同ViT骨干网络（ViT-T/S/B/L）和分割模型（Segmenter, SegViT, SETR）上的性能表现。
- 指标：使用mIoU（平均交并比）评估分割质量，使用im/sec（每秒处理图像数）和GFLOPs（十亿次浮点运算数）评估计算效率。
- 结果：ALGM在所有设置下均显著提高了吞吐量和减少了GFLOPs，同时mIoU也有所提升。例如，在Segmenter+ViT-S模型上，ALGM将吞吐量从134 im/sec提升至192 im/sec，GFLOPs从38.6减少到26.3，同时mIoU从45.3%提升至46.4%。ALGM*版本进一步优化了效率，吞吐量最高可提升100%（如Segmenter+ViT-L），同时保持或略微提升mIoU。
- 表2展示了ALGM在COCO-Stuff、Cityscapes和Pascal-Context数据集上的性能。
- 结果：ALGM在所有这些数据集上均显著提高了吞吐量和分割质量。例如，在COCO-Stuff数据集上，使用Segmenter+ViT-S模型时，ALGM将吞吐量从41 im/sec提升至65 im/sec，同时mIoU从76.5%提升至76.9%。ALGM*版本进一步优化了效率，如在Cityscapes数据集上，Segmenter+ViT-S模型的吞吐量提升了90%。
- 表3比较了ALGM与DToP和DoViT方法在SETR和Segmenter模型上的性能。
- 结果：ALGM在保持或略微提升mIoU的同时，显著减少了GFLOPs。例如，在SETR-B模型上，ALGM*在保持mIoU为46.4%的同时，将GFLOPs减少了30%。相比之下，DToP仅减少了25%的GFLOPs，而DoViT甚至导致mIoU下降。
- 表4展示了ALGM在EVA+ViT-Adapter+Mask2Former模型上的应用效果。
- 结果：无需训练，ALGM即可将吞吐量提升26%，同时保持mIoU不变。经过训练后，ALGM进一步将mIoU提升了0.2%。
- 表5通过消融实验评估了ALGM中不同组件和参数的影响。
- 结果：CLAP模块窗口大小的影响：较小窗口尺寸提高mIoU，较大窗口尺寸提高效率。合并模块的影响：结合CLAP和GBM模块可以获得最佳性能。GBM模块位置的影响：在ViT-S模型中，第5层是最佳位置。
- 表6分析了ALGM提升mIoU的原因。
- 结果：注意力平衡：合并同类别令牌可以减少大类别在自注意力操作中的主导地位，从而提高稀有类别的分割质量。令牌去噪：合并过程中通过取平均值对令牌进行去噪，有助于模型学习。实验表明，禁用注意力平衡或令牌去噪会导致mIoU显著下降。
图4展示了使用不同相似度阈值进行令牌合并时，模型在效率和分割质量之间的权衡。
关键点：
- 阈值影响：较低的阈值会导致更多的令牌合并，从而提高效率但可能降低分割质量；较高的阈值则相反。
- 自动阈值：ALGM方法使用自动生成的阈值，能够在效率和分割质量之间找到最佳平衡点。
图5通过可视化展示了CLAP和GBM模块在不同层合并的令牌。
关键点：
- 局部合并示例：在第一层，CLAP模块合并了局部窗口内的高相似度令牌。
- 全局合并示例：在中间层，GBM模块进一步合并了全局范围内的高相似度令牌。
- 合并效果：合并后的令牌数量显著减少，同时保留了重要的语义信息，有助于提升分割效率和质量。