PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding The Robustness in Vision Transformers

Daquan Zhou, Zhiding Yu, Enze Xie, Chaowei Xiao, Anima Anandkumar, Jiashi Feng, Jose M. Alvarez

arXiv:2204.12451v4cs.CV

TL;DR

The paper addresses the limited systematic understanding of self-attention’s role in Vision Transformer robustness. It analyzes visual grouping and information bottlenecks, then introduces fully attentional networks with attentional channel processing; FANs achieve improved robustness across classification and downstream tasks.

  • Problem

    The actual role of self-attention in Vision Transformers’ robust generalization remains unclear, particularly after ConvNeXt performs favorably without self-attention.

  • Method

    The paper studies token grouping and interprets self-attention through information bottleneck theory, then proposes FANs using self-attention for both token mixing and channel processing.

  • Results

    35.8% mCE is achieved by the 76.8M-parameter FAN on ImageNet-C, with robustness gains also transferring to semantic segmentation and object detection.

  • Takeaways & Limitations

    Self-attention promotes naturally formed token clusters, and fully attentional designs substantially improve robustness over CNN and ViT counterparts.

  • Takeaways & Limitations

    Direct channel attention has quadratic computational complexity in the channel dimension and can impose large overhead as channel widths grow in pyramid models.

Abstract

from arXiv · show

Recent studies show that Vision Transformers(ViTs) exhibit strong robustness against various corruptions. Although this property is partly attributed to the self-attention mechanism, there is still a lack of systematic understanding. In this paper, we examine the role of self-attention in learning robust representations. Our study is motivated by the intriguing properties of the emerging visual grouping in Vision Transformers, which indicates that self-attention may promote robustness through improved mid-level representations. We further propose a family of fully attentional networks (FANs) that strengthen this capability by incorporating an attentional channel processing design. We validate the design comprehensively on various hierarchical backbones. Our model achieves a state-of-the-art 87.1% accuracy and 35.8% mCE on ImageNet-1k and ImageNet-C with 76.8M parameters. We also demonstrate state-of-the-art accuracy and robustness in two downstream tasks: semantic segmentation and object detection. Code is available at: https://github.com/NVlabs/FAN.

1. Introduction

The paper examines why self-attention is associated with robustness in Vision Transformers and proposes fully attentional networks that strengthen robust representation learning through attentional channel processing.

  • Motivation: ViTs model nonlocal relations with self-attention and consistently outperform ConvNets on corruption robustness, but ConvNeXt challenges self-attention as the sole explanation.This motivates investigating the actual role of self-attention in robust generalization.
  • Approach: The study links self-attention to naturally emerging visual grouping and investigates whether grouping improves mid-level representations and robustness.It analyzes output tokens from each ViT layer using spectral clustering and relates significant eigenvalues to perturbations from the input.
  • Approach: The paper interprets self-attention through information bottleneck theory, showing under mild assumptions that it can be written as an iterative optimization step of the information bottleneck objective.The objective compresses unimportant information while preserving information about inputs and target labels.
  • Method: FANs replace conventional MLP channel processing with attentional channel selection, making the network fully attentional and producing dynamic, content-dependent feature transformations.The design aims to strengthen the interaction between visual grouping and robustness while avoiding static convolution operations in the MLP block.
  • Results: 47.7% mCE is achieved on ImageNet-C with 28M parameters, while the 76.8M-parameter FAN reaches 35.8% mCE, reported as state-of-the-art robustness among supervised models.At comparable model size, the 28M model improves over ResNet-50, Swin-T, and ConvNeXt-T by 29.0%, 11.9%, and 5.5%, respectively.
  • Results: The proposed design’s robustness gains transfer to semantic segmentation and object detection, while the paper frames grouping, information bottleneck, and robust generalization as a unified explanatory framework.The authors also argue that attention representations provide a non-trivial benefit for robust generalization beyond clean-image recognition.

2. Fully Attentional Networks

The paper interprets self-attention as promoting visual grouping and noise filtering, then proposes fully attentional networks with efficient channel attention to strengthen robust representations.

  • Visual grouping and robustness: Spectral clustering reveals meaningful token groups whose emergence across ViT blocks coincides with decreasing perturbation norms.The analysis measures token affinity through feature correlations and tracks significant eigenvalues and normalized Gaussian-noise norms.
  • Visual grouping and robustness: Compared with ResNet-50, ViT and FAN noise decay continues through self-attention blocks, whereas ResNet-50 shows weaker final decay and plateaus over regular convolution blocks.The comparison links deeper token grouping and perturbation decay in ViT-family models, while the reported ResNet behavior differs.
  • Information bottleneck perspective: Self-attention can be interpreted as iteratively aggregating similar inputs into soft cluster-center representations under an information bottleneck principle.The key matrix stores temporary cluster centers, while softmax associations assign input features to them.
  • Fully Attentional Networks: Multi-head self-attention balances robustness and clean accuracy, with more heads improving expressivity and robustness but fewer channels per head reducing clean accuracy.The reported best trade-off uses 32 channels per head.
  • Fully Attentional Networks: Fully Attentional Networks extend transformer blocks with attentional channel processing that reweights channels across information bottlenecks.The design targets more compositional representations and spatially holistic channel aggregation.
  • Efficient channel self-attention: Efficient channel attention replaces quadratic channel self-attention costs with a token-prototype-based mechanism and is reported to improve performance while consuming less computation.The conventional channel-attention complexity is quadratically proportional to D^2; the proposed mechanism uses a prototype formed by averaging across channels.

3. Experiment Results & Analysis

Experiments evaluate corruption robustness and out-of-distribution generalization across classification, segmentation, and detection. FAN improves robustness across architectures and tasks, with efficiency and clean-accuracy benefits in several comparisons.

  • Experiment details: Robustness is evaluated on ImageNet-C, Cityscapes-C, COCO-C, ImageNet-A, and ImageNet-R without corruption-specific fine-tuning.Clean ImageNet-1K accuracy and corruption robustness are both reported; mCE quantifies corruption resilience.
  • Analysis: Advanced training tricks improve absolute accuracy, but CutMix and RandAugmentation provide the main robustness gains, unlike knowledge distillation and large-dataset pretraining.The comparison uses ViT-S* and applies the advanced techniques to ResNet-50 for fairness.
  • Analysis: Even with matched training recipes and model sizes, ViT-S* outperforms ResNet-50* in robustness, supporting an architectural contribution from self-attention.ResNet-50* uses the DeiT recipe, SE, and strided convolution; ViT-S* uses the DeiT recipe with 12 blocks and 384-dimensional embeddings.
  • Fully Attentional Networks: FAN channel attention improves robustness with modest resource costs: CSA achieves the best robustness while using memory comparable to SE attention.SE consumes less memory and maintains comparable clean accuracy and robustness relative to conventional self-attention.
  • Fully Attentional Networks: Adding FAN blocks improves Swin robustness by 5%, while FAN-Hybrid matches FAN-ViT and FAN-SWIN robustness and provides higher clean and corrupted-dataset accuracy.FAN-Hybrid is selected for downstream tasks because it also remains efficient for high-resolution inputs and dense prediction.
  • Comparison to SOTAs on various tasks: FAN models surpass competing models across corruption types and tasks, including 6.8% mIoU over SegFormer and 6.2% mAP over Swin at comparable model sizes.Object detection also reaches 42.0% mAP with 76.8M encoder parameters.
  • Comparison to SOTAs on various tasks: FANs outperform ConvNeXt, Swin, DeiT, and other baselines in generalization, indicating that fully attentional processing benefits learned representations.The cited comparison relates this benefit to processing irrelevant features effectively.

4. Related Works

Related work contrasts ViTs’ global self-attention with convolutional inductive biases and distinguishes natural-corruption robustness from adversarial robustness. Prior studies report stronger corruption robustness for ViTs than ConvNets.

  • Vision Transformers: ViTs perform global interactions among visual tokens through self-attention, whereas CNNs rely on inductive biases such as locality and translation invariance.The related work positions self-attention as a central architectural difference across vision tasks.
  • Vision Transformers: Several studies report that ViTs consistently outperform ConvNets on robustness to observable corruptions.The passage distinguishes this robustness advantage from the broader discussion of self-attention versus convolutions.
  • Robustness scope: The paper focuses on zero-shot robustness to natural corruptions and architectural approaches rather than adversarial-robustness training algorithms.The distinction defines the scope of the work.
  • Robustness scope: The broader evaluation context includes out-of-distribution generalization results, with FAN models reported to improve generalization across datasets.The cited table caption identifies IN-C as measured by mCE and notes that some results use 384 × 384 fine-tuning.

5. Conclusion

The paper concludes that self-attention contributes to robustness through token grouping and information extraction, and introduces FANs that extend attention into channel processing. FANs improve robustness over CNN and ViT counterparts across evaluated applications.

  • Conclusion: Self-attention promotes naturally formed token clusters, which the paper relates to visual grouping and robust representation learning.The conclusion frames this as an observed property of vision transformers.
  • Conclusion: An information-bottleneck framework is established to explain properties of self-attention.The conclusion presents this as an explanatory framework for the observed grouping and robustness behavior.
  • Conclusion: FAN architectures use self-attention for both token mixing and channel processing, producing improved robustness over CNN and ViT counterparts.The conclusion describes fully attentional design as the mechanism used to push robust representation learning further.
  • Conclusion: The work suggests that vision-transformer inductive biases can extend beyond convolutions and may benefit safety-critical applications such as autonomous driving.This is presented as a potential real-world implication within the conclusion.

A.1. Proof on the relationship between the Information Bottleneck and Self-Attention

The proof connects iterative information-bottleneck clustering with self-attention under shared-covariance and normalized-mean assumptions. The resulting update has row-wise softmax normalization, establishing Proposition 2.1.

  • Information bottleneck formulation: IB seeks a representation that compresses information while retaining information relevant to predicting the clean code.The mapping satisfies a Markov constraint, and β trades off compression against relevant-information preservation.
  • Iterative clustering updates: The IB clustering formulation assigns data points to clusters and iteratively updates cluster assignments using a normalizing factor and cluster-member sets.The data-point marginal is assumed uniform, and the cluster likelihood is approximated with a Gaussian.
  • Cluster-mean update: The cluster means are updated by minimizing the KL-divergence between the Gaussian approximation and the cluster-conditional distribution.The derivation then assumes a covariance shared across clusters and means normalized with respect to its inverse.
  • Connection to self-attention: Under these assumptions, the cluster update can be rearranged into a compact matrix form that clarifies its relationship to self-attention.The construction defines Q, K, and V from the data, cluster means, covariance, and value projection.
  • Proof conclusion: The resulting expression uses row-wise softmax normalization, completing the proof of Proposition 2.1.The proof follows the stated derivation from the iterative IB clustering update to the self-attention form.

A.2. Implementation details

The implementation uses established training recipes and task-specific finetuning for classification, segmentation, and detection. Robustness is evaluated on corrupted datasets using retention and mCE, with head-number effects shown separately.

  • Architecture implementation: FAN comparisons add depthwise convolution in the MLP block and scale the residual path with a learnable parameter for training stability.These choices follow practices used in prior methods.
  • ImageNet classification: FAN classification models are trained for 300 epochs with AdamW, a 2e-3 learning rate, five warmup epochs, cosine decay, and batch size 2048.The same recipe is used for the baseline and FAN model family unless otherwise specified.
  • Downstream tasks: Segmentation and detection use established decoders and finetuning settings, including SETR-style decoding, SegFormer training, Faster R-CNN, and multi-scale resolutions from 640×640 to 896×896.Testing uses a deterministic 896×896 image resolution for detection.
  • Downstream tasks: FAN-Swin and FAN-Hybrid are finetuned with Mask R-CNN on COCO using multi-scale training, AdamW, and a 3x schedule.The implementations use MMSegmentation and MMDetection toolboxes.
  • Corruption datasets: ImageNet-C is downloaded from a provided mirror, while Cityscapes-C and COCO-C use 16 generated corruptions spanning noise, blur, weather, and digital categories.The corruption datasets are prepared separately for the evaluated tasks.
  • Evaluation metrics: Retention measures the ratio of ImageNet-C accuracy to clean ImageNet accuracy, while mCE is also reported for ImageNet-C comparisons.Downstream robustness is evaluated on corrupted Cityscapes and COCO after clean-dataset finetuning.

A.4. Detailed benchmark results on corrupted images on classification, segmentation and detection

FAN’s stronger robustness extends beyond ImageNet-C classification to corrupted-image semantic segmentation and object detection across downstream tasks.

  • Cross-task robustness: FAN’s strong robustness transfers to all evaluated downstream tasks and corruption categories shown in the detailed benchmark tables.The reported comparisons cover semantic segmentation and object detection in addition to classification.

A.5. Architecture details of FAN-Swin and FAN-Hybrid

FAN-Swin preserves the hierarchical Swin design while replacing conventional self-attention with shifted-window attention, and FAN-Hybrid incorporates ConvNeXt convolutional blocks. Feature clustering visualizations compare FAN with ViT and CNN models.

  • FAN-Swin: FAN-Swin follows Swin’s macro-architecture and replaces conventional self-attention with efficient shifted-window self-attention.The replacement is made within the same macro-architecture design.
  • FAN-Hybrid: FAN-Hybrid uses three convolutional building blocks per stage within the ConvNeXt architecture.This defines the convolutional component of the hybrid hierarchy.
  • Feature clustering: Token features are normalized with SoftMax, converted into a self-correlation affinity matrix, and clustered spectrally.The features come from the second-last block’s output.
  • Visualization: Figure 9 compares clustering visualizations for FAN, ViT, and CNN token features, while Figure 8 visualizes results on Cityscapes.The figure caption states that FAN provides clearer clusters highlighting important foreground-object regions.
  • Robustness comparisons: ImageNet-C, Cityscapes-C, and COCO-C tables compare FAN robustness with other models across classification, segmentation, and detection settings.The table captions state stronger robustness for FAN across the listed corruption settings.
Loading 2204.12451v4…