PaperlayerView exact revision on arXiv

Source-linked AI summary

Towards Robust Vision Transformer

Xiaofeng Mao, Gege Qi, Yuefeng Chen, Xiaodan Li, Ranjie Duan, Shaokai Ye, Yuan He, Hui Xue

arXiv:2105.07926v4cs.CV

TL;DR

Existing ViT work largely evaluates standard accuracy and computation cost rather than robustness and generalization. The paper systematically studies ViT components, constructs RVT from robust components, and adds PAAS and patch-wise augmentation; experiments show strong performance across ImageNet and six robustness benchmarks, including Top-1 ranks for RVT-S* on several benchmarks.

  • Problem

    Existing ViT studies focus on standard accuracy and computation cost, leaving their intrinsic effects on robustness and generalization insufficiently investigated.

  • Method

    The paper evaluates ViT components, combines robust choices into RVT, and augments it with position-aware attention scaling and patch-wise augmentation.

  • Results

    RVT shows strong performance on ImageNet and six robustness benchmarks, while RVT-S* achieves Top-1 rank on ImageNet-C, ImageNet-Sketch, and ImageNet-R.

  • Takeaways & Limitations

    The proposed robust components and plug-and-play techniques provide a ViT design with strong robustness and generalization across evaluated benchmarks.

Abstract

from arXiv · show

Recent advances on Vision Transformer (ViT) and its improved variants have shown that self-attention-based networks surpass traditional Convolutional Neural Networks (CNNs) in most vision tasks. However, existing ViTs focus on the standard accuracy and computation cost, lacking the investigation of the intrinsic influence on model robustness and generalization. In this work, we conduct systematic evaluation on components of ViTs in terms of their impact on robustness to adversarial examples, common corruptions and distribution shifts. We find some components can be harmful to robustness. By using and combining robust components as building blocks of ViTs, we propose Robust Vision Transformer (RVT), which is a new vision transformer and has superior performance with strong robustness. We further propose two new plug-and-play techniques called position-aware attention scaling and patch-wise augmentation to augment our RVT, which we abbreviate as RVT*. The experimental results on ImageNet and six robustness benchmarks show the advanced robustness and generalization ability of RVT compared with previous ViTs and state-of-the-art CNNs. Furthermore, RVT-S* also achieves Top-1 rank on multiple robustness leaderboards including ImageNet-C and ImageNet-Sketch. The code will be available at \url{https://github.com/alibaba/easyrobust}.

1. Introduction

Existing ViT research emphasizes standard accuracy and computation cost, leaving robustness and generalization insufficiently examined. The paper analyzes ViT components, builds RVT from robust choices, and augments it with PAAS and patch-wise augmentation.

  • Existing ViT variants primarily optimize standard accuracy and computation cost, with limited study of robustness and generalization.
  • 26.6% robust accuracy: PVT-Small trails original DeiT-S by 14.1% under the paper’s comparison, despite high standard accuracy.
  • The paper systematically analyzes patch embedding, position embedding, transformer blocks, and classification heads for robustness effects.
  • PAAS and patch-wise augmentation further boost RVT by refining attention correlations and diversifying patch-sequence training data.
  • RVT combines robust component choices into a transformer that improves robustness while exceeding most other transformers in accuracy.
  • RVT-S* achieves Top-1 rank on ImageNet-C, ImageNet-Sketch, and ImageNet-R.

2. Related Work

Prior robustness studies span adversarial examples, common corruptions, and distribution shifts, with substantial attention historically devoted to CNNs and increasing study of ViTs.

  • ImageNet-C evaluates corruption robustness, while ImageNet-R and ImageNet-Sketch represent naturally occurring distribution changes.
  • CNN robustness research explores augmentation, architecture, training, weight compression, pooling, and activation-function strategies.
  • Early ViT studies report greater adversarial robustness than CNNs and low adversarial-example transferability between the architectures.
  • Later ViT studies extend evaluation to common image corruption and distribution shifts, reporting that ViTs are more robust learners.

3. Robustness Analysis of Designed Components

The study evaluates four main ViT components across robustness benchmarks and finds that their design choices can substantially affect robustness. It identifies robust building blocks for RVT, including convolutional patch embedding and FFNs, suitable stage distributions, attention-head choices, and token feature pooling.

  • Robustness analysis covers patch embedding, position embedding, transformer blocks, and classification heads using DeiT-Ti across all robustness benchmarks.FGSM robust accuracy is reported in the main text, with other benchmark results in the appendix.
  • Patch Embedding: Low-level patch embedding improves both robustness and standard accuracy; tokens-to-tokens embedding performs best, while convolutional embedding offers lower computation cost.Tokens-to-tokens embedding has quadratic complexity as image size expands, motivating the convolutional choice.
  • Position Embedding: Position encoding is critical for shape-biased features, while existing encoding methods have limited robustness impact and input-conditioned encoding can harm robustness.Without position encoding, ViT shows an 8% accuracy drop on ImageNet-Sketch and fails to recognize shape-bias objects.
  • Transformer Blocks: Reducing spatial resolution in later transformer blocks modestly improves robustness, whereas more blocks at large spatial resolutions worsen robustness and increase FLOPs and memory.Among the tested stage variants, V2 is identified as the best speed-performance trade-off.
  • Transformer Blocks: Convolutional FFNs improve standard accuracy and robustness without the long-term dependency damage associated with local self-attention.
  • Classification Head: Replacing the CLS token with global average pooling improves adversarial robustness and slightly helps standard accuracy through translation invariance.
  • Combination of Robust Components: RVT combines convolutional patch embedding, multi-stage blocks, suitable head counts, convolutional FFNs, and token feature pooling; PAAS and patch-wise augmentation further boost performance.Both plug-and-play techniques are also applicable to other ViT models.

4. Position-Aware Attention Scaling

PAAS makes dot-product attention position-aware by learning pairwise position importance, replacing traditional position embeddings and separating content- from position-based attention. This design is intended to suppress noisy positional correlations and improve robustness.

  • Scaled Dot-Product Attention: Scaled dot-product attention produces value-weighted outputs from query-key matching scores, with Q, K, and V representing queries, keys, and values.The standard operation uses Softmax(QK^T/√d)V.
  • Position-Aware Attention Scaling: PAAS introduces a learnable position importance matrix Wp that rescales each query-key pair before the softmax operation.Its attention is Softmax(QK^T ⊙ (Wp/√d))V.
  • Position-Aware Attention Scaling: Because Wp depends only on sequence positions, PAAS serves as a position representation and replaces traditional position embedding in RVT.The resulting attention separates content-based QK^T from position-based Wp/√d terms.
  • Robustness of PAAS: PAAS acts as a soft attention mask that suppresses unrelated regions activated by adversarial inputs and retains stronger positional correlations.The paper visualizes self-attention before and after scaling and the learned scaling factor.

5. Patch-Wise Augmentation

Patch-wise augmentation increases training diversity by applying transformations independently to image patches after conventional image-level augmentation. The strategy targets inter-patch diversity while retaining intra-patch augmentation and improves robustness across ViTs.

  • Method: Patch-wise augmentation applies diverse transformations independently to each image patch after the original image-level augmentation.The method retains whole-image augmentation while adding patch-level transformations during training.
  • Motivation: The strategy addresses ViTs’ inter-patch diversity because whole-image transformations give every patch the same transformation at one time.The motivation follows ViTs’ modeling of both intra-patch features and inter-patch relations.
  • Robustness: Patch-wise augmentation improves standard and robust accuracy by more than 1% and 5%, respectively, across different ViT models.The paper describes it as broadly applicable across ViT architectures.

6. Experiments

Experiments evaluate RVT on ImageNet and six robustness benchmarks, covering standard accuracy, adversarial attacks, common corruptions, and distribution shifts. RVT and its augmented RVT* variants achieve strong accuracy-robustness trade-offs, with gains from PAAS and patch-wise augmentation across architectures.

  • Standard Performance Evaluation: RVT-Ti* achieves 79.2% Top-1 accuracy on ImageNet-1K with 1.3G FLOPs and 10.9M parameters.The model is described as competitive with popular ResNet and RegNet series while using around 60% fewer resources than compared CNNs.
  • Standard Performance Evaluation: At the same computation cost, RVT-S* reaches 81.9% test accuracy, 2.9% higher than ResNet-50.The result is close to EfficientNet-B4, which uses 380×380 inputs and has lower throughput.
  • Standard Performance Evaluation: RVT-Ti reaches 78.4% Top-1 accuracy, while RVT-Ti* adds 0.8% with little additional computation cost.RVT-S* and RVT-B* also improve over DeiT-S and DeiT-B, although gains become smaller with increasing model capacity.
  • Adversarial Robustness: RVT-Ti and RVT-S improve FGSM robust accuracy by over 10% versus previous ViT variants, with further gains from PAAS and patch-wise augmentation.The evaluation uses FGSM and PGD attacks with t = 5, α = 0.5, and ϵ = 1.
  • Common Corruption Robustness: RVT-S* obtains 49.4 mCE on ImageNet-C, improving by 4.2 over leaderboard method DeepAugment and establishing a new state-of-the-art.Lower mCE indicates greater corruption robustness.
  • Out-of-Distribution Robustness: RVT-S achieves 35.0% and 46.9% test accuracy on ImageNet-Sketch and ImageNet-R, respectively, surpassing most CNN and ViT models.RVT and RVT* are reported to outperform other ViTs on out-of-distribution generalization.
  • Ablation Studies: Using PAAS on more transformer blocks increases standard and robust accuracy until gains saturate after five blocks.The ablation compares single- and multiple-block PAAS settings against learned absolute position embeddings.
  • Ablation Studies: Applying PAAS and patch-wise augmentation to DeiT-Ti, ConViT-Ti, and PiT-Ti yields average gains exceeding 1% in standard accuracy and 5% in robust accuracy.This evaluates the techniques as plug-and-play additions beyond RVT.

7. Conclusion

The paper systematically studies how ViT components affect robustness and builds RVT by replacing harmful modifications with robust components. PAAS and patch-wise augmentation further improve RVT, whose experiments show strong performance across ImageNet and six robustness benchmarks.

  • Conclusion: RVT is constructed by alternating ViT modifications that damage robustness and using robust components as its building blocks.The design is based on a systematic study of key ViT components.
  • Conclusion: Patch-wise augmentation adds affinity and diversity to training data, while PAAS addresses missing spatial-information correlation in scaled dot-product attention.Both techniques are presented as additions that further boost RVT.
  • Conclusion: Experiments report outstanding RVT performance consistently on ImageNet and six robustness benchmarks across FLOPs, standard accuracy, and robust accuracy trade-offs.The conclusion specifically validates RVT-Ti and RVT-S.

A. Additional Results of Robustness Analysis on Designed Components

Additional analyses largely follow the overall robustness trend but reveal benchmark-specific exceptions, showing that robustness cannot be assessed comprehensively with only a few benchmarks.

  • Benchmark-specific findings: Some stage-distribution variants trade robustness across benchmarks: V6 performs poorly adversarially but best on IN-A and IN-R.This indicates stronger generalization on some distribution-shift datasets does not guarantee adversarial robustness.
  • Benchmark-specific findings: The token-to-token embedder improves IN-C, IN-A, IN-R, and IN-SK over the original linear embedder but reaches only 4.7% robust accuracy under PGD.Its performance therefore depends strongly on the robustness axis being measured.
  • Assessment implication: Because individual benchmarks can disagree, relying on only several robustness benchmarks may produce a biased assessment.The authors advocate broader evaluation across robustness settings.
  • Analysis coverage: Additional robustness analyses cover head-number choices and other designed components through separate tables.Table 9 specifically reports results for different head numbers.

B. Feature Visualization

Feature visualizations compare RVT with CNN and transformer baselines using t-SNE features and intermediate feature maps, showing more compact, separable, diverse, and less redundant representations for RVT.

  • t-SNE visualization: RVT features are closest to the desired combination of intra-class compactness and inter-class separability in t-SNE visualizations.Features were extracted on ImageNet and ImageNet-C validation data, with 10 classes selected for visualization.
  • Interpretation: The visualized representation patterns are presented as evidence of RVT's stronger robustness, classification performance, and generalization ability.The t-SNE and feature-map analyses support complementary claims about separability and feature diversity.
  • Feature-map visualization: RVT-S reduces redundant feature maps and preserves feature diversity compared with ResNet50 and DeiT-S.The comparison uses features extracted from the models' fifth layers; redundant regions are highlighted with red boxes.

C. Loss Landscape Visualization

The paper relates loss-surface geometry to model behavior and presents RVT-S as having a flatter loss surface, while additional tables document component and training analyses.

  • Loss-surface comparison: RVT-S has a flatter loss surface than ResNet50 in the visualization.The paper associates flatter loss surfaces with stability under input changes.
  • Interpretation: The loss-landscape discussion frames surface geometry as relevant to model generalization and trainability.The stated interpretation connects the visualization to behavior under input changes.
  • Additional analyses: Related additional analyses cover position encoding, patch embeddings, attention locality, convolutional FFN, CLS-token replacement, stage distribution, and attention scaling.These topics are organized across Tables 10–13 and Table 15.
Loading 2105.07926v4…