[论文解读] Discrete Representations Strengthen Vision Transformer Robustness

Chengzhi Mao, Lu Jiang+4 more|arXiv (Cornell University)|Nov 20, 2021

Advanced Neural Network Applications参考文献 41被引用 10

一句话总结

本文提出 Dr. ViT，一种 Vision Transformer 变体，通过将标准的连续像素标记替换为来自向量量化 VQ-GAN 编码器的离散标记，提升了模型的鲁棒性。通过将离散标记与残差像素嵌入相结合，模型学习到更具不变性、更关注形状的表征，从而在七个 ImageNet 鲁棒性基准上实现最高 12% 的性能提升，同时不损失 ImageNet 准确率。

ABSTRACT

Vision Transformer (ViT) is emerging as the state-of-the-art architecture for image recognition. While recent studies suggest that ViTs are more robust than their convolutional counterparts, our experiments find that ViTs trained on ImageNet are overly reliant on local textures and fail to make adequate use of shape information. ViTs thus have difficulties generalizing to out-of-distribution, real-world data. To address this deficiency, we present a simple and effective architecture modification to ViT's input layer by adding discrete tokens produced by a vector-quantized encoder. Different from the standard continuous pixel tokens, discrete tokens are invariant under small perturbations and contain less information individually, which promote ViTs to learn global information that is invariant. Experimental results demonstrate that adding discrete representation on four architecture variants strengthens ViT robustness by up to 12% across seven ImageNet robustness benchmarks while maintaining the performance on ImageNet.

研究动机与目标

• 解决 Vision Transformer 对局部纹理的过度依赖及其在分布偏移下的泛化能力差的问题。
• 在不损害 ImageNet 上域内性能的前提下提升鲁棒性。
• 探究离散标记化是否能促进 ViT 学习全局、形状不变的特征。
• 开发一种即插即用的输入层改进方法，以增强各类 ViT 变体的鲁棒性。
• 证明先前用于生成任务的离散表征可提升分类任务的泛化能力。

提出的方法

• 用来自 VQ-GAN 编码器的离散标记替代标准的连续像素嵌入，该编码器将图像块量化为码本索引。
• 使用可学习的码本将图像块映射为离散视觉词，同时保留结构和形状信息。
• 将离散标记嵌入与原始像素嵌入拼接，以保留局部细节，特别是对小物体的表征。
• 使用标准优化方法（Adam，余弦衰减）和数据增强（RandAug，Mixup）联合训练 ViT，同时使用离散和连续表征。
• 为小型模型使用轻量级 VQ-编码器，为 ViT-B 使用更大模型，码本大小分别为 1,024（ImageNet）和 8,192（ImageNet-21K）。
• 对 ViT-B 的最终特征应用平均池化，并在更高分辨率输入（384×384，512×512）上进行微调，以提升泛化能力。

实验结果

研究问题

• RQ1 离散标记化能否提升 Vision Transformer 对分布偏移的鲁棒性？
• RQ2 用离散表征替代连续像素标记是否能促使 ViT 学习到更多形状不变的特征？
• RQ3 离散与连续标记的结合对域内与域外基准性能有何影响？
• RQ4 该方法是否可无需架构修改地应用于不同 ViT 模型？
• RQ5 使用离散标记是否能带来反映全局上下文而非局部纹理的注意力模式？

主要发现

• Dr. ViT 在 Stylized-ImageNet 上的鲁棒性最高提升 12%，在 ImageNet-Sketch 和 ImageNet-C 上最高提升 10%，优于标准 ViT。
• 该方法在七个 ImageNet 鲁棒性基准中的四个上达到最先进性能，且未使用数据集特定的数据增强。
• 添加离散标记可改善注意力模式，使模型对全局结构更敏感，对局部纹理更不敏感，可视化对比结果已证实。
• 仅使用离散标记的变体在小型模型（如 ViT-Ti）上表现较差，但与像素嵌入结合后性能显著提升。
• 将 ViT-B 训练 800 个周期而非 300 个，可在 ImageNet 上额外获得 0.1–1% 的准确率增益，表明使用离散标记进行更长训练具有优势。
• 该方法与现有鲁棒性技术正交，可仅通过极少修改集成到任何基于 ViT 的架构中。