[论文解读] Aggregating Nested Transformers.

Zizhao Zhang, Han Zhang+3 more|arXiv (Cornell University)|May 26, 2021

Advanced Neural Network Applications参考文献 64被引用 56

一句话总结

NesT 提出了一种简化的分层视觉 Transformer 架构，通过在非重叠图像块上嵌套局部自注意力机制，并利用关键的块聚合函数实现跨块通信，从而实现更快速的收敛和最先进性能。该设计在参数量减少高达 57% 的情况下实现了最先进准确率，且收敛速度更快，包括一个 600 万参数的模型仅使用单张 GPU 从零开始训练即可在 CIFAR-10 上达到 96% 的准确率。

ABSTRACT

Although hierarchical structures are popular in recent vision transformers, they require sophisticated designs and massive datasets to work well. In this work, we explore the idea of nesting basic local transformers on non-overlapping image blocks and aggregating them in a hierarchical manner. We find that the block aggregation function plays a critical role in enabling cross-block non-local information communication. This observation leads us to design a simplified architecture with minor code changes upon the original vision transformer and obtains improved performance compared to existing methods. Our empirical results show that the proposed method NesT converges faster and requires much less training data to achieve good generalization. For example, a NesT with 68M parameters trained on ImageNet for 100/300 epochs achieves $82.3\%/83.8\%$ accuracy evaluated on $224 imes 224$ image size, outperforming previous methods with up to $57\%$ parameter reduction. Training a NesT with 6M parameters from scratch on CIFAR10 achieves $96\%$ accuracy using a single GPU, setting a new state of the art for vision transformers. Beyond image classification, we extend the key idea to image generation and show NesT leads to a strong decoder that is 8$ imes$ faster than previous transformer based generators. Furthermore, we also propose a novel method for visually interpreting the learned model. Source code is available this https URL.

研究动机与目标

• 通过最小化架构复杂度来简化分层视觉 Transformer 架构，同时保持或提升性能。
• 探究块聚合在实现嵌套 Transformer 中有效跨块非局部信息流动中的作用。
• 降低视觉 Transformer 模型实现强泛化能力所需的数据量和参数量。
• 将 NesT 框架扩展至图像生成任务，提升自回归 Transformer 的推理速度。
• 提出一种新颖的方法，用于可视化模型所学习注意力机制的可解释性。

提出的方法

• 在非重叠图像块上嵌套基本的局部自注意力模块，构建分层结构。
• 引入块聚合函数，实现非相邻图像块之间的长距离依赖关系。
• 对原始视觉 Transformer 的架构改动最小，重点在于聚合机制而非注意力头设计。
• 采用多尺度特征聚合策略，结合分层结构中不同层级的特征。
• 对编码器和解码器头使用相同架构，使模型可应用于图像生成任务。
• 提出一种新的可视化方法，用于解释模型所学特征中注意力模式与行为。

实验结果

研究问题

• RQ1 块聚合函数如何影响嵌套 Transformer 中的跨块信息流动？
• RQ2 在参数量更少、数据量更少的情况下，通过最小化架构修改的简化分层 Transformer 是否能超越现有视觉 Transformer？
• RQ3 当在小样本数据集（如 CIFAR-10）上从零开始训练时，NesT 能在多大程度上实现强性能？
• RQ4 NesT 架构能否有效扩展至图像生成任务？与先前基于 Transformer 的生成器相比，其在速度和质量方面表现如何？
• RQ5 所提出的新型可视化方法能否有效解释 NesT 模型所学习注意力模式？

主要发现

• 一个 6800 万参数的 NesT 模型在 ImageNet 上训练 100 个周期后达到 82.3% 的 top-1 准确率，300 个周期后达到 83.8%，相比之前方法参数量减少高达 57%。
• 一个 600 万参数的 NesT 模型在 CIFAR-10 上从零开始训练，仅使用单张 GPU 即达到 96% 的准确率，创下该基准上视觉 Transformer 的新 SOTA 记录。
• 与先前方法相比，NesT 模型收敛更快，泛化能力更强，尤其在训练数据有限时表现更优。
• 在扩展至图像生成任务时，基于 NesT 的解码器速度比先前基于 Transformer 的生成器快 8 倍，同时保持了强大的生成质量。
• 所提出的块聚合机制对于实现有效的跨块通信至关重要，显著优于简单堆叠局部注意力模块的基线方法。
• 新型可视化方法成功揭示了有意义的注意力模式，显著提升了模型内部推理过程的可解释性。