TFS-ViT: Token-Level Feature Stylization for Domain Generalization

TL;DR

TFS-ViT introduces a token-level feature stylization method for Vision Transformers, enhancing their ability to generalize to unseen domains by synthesizing new domain styles through normalization statistic mixing, with attention-aware enhancements.

Contribution

This work is the first to propose token-level feature stylization for ViTs, improving domain generalization by synthesizing diverse domain styles using attention-guided normalization mixing.

Findings

Achieves state-of-the-art results on five domain generalization benchmarks.

Effectively handles various types of domain shifts.

Flexible and compatible with any ViT architecture.

Abstract

Standard deep learning models such as convolutional neural networks (CNNs) lack the ability of generalizing to domains which have not been seen during training. This problem is mainly due to the common but often wrong assumption of such models that the source and target data come from the same i.i.d. distribution. Recently, Vision Transformers (ViTs) have shown outstanding performance for a broad range of computer vision tasks. However, very few studies have investigated their ability to generalize to new domains. This paper presents a first Token-level Feature Stylization (TFS-ViT) approach for domain generalization, which improves the performance of ViTs to unseen data by synthesizing new domains. Our approach transforms token features by mixing the normalization statistics of images from different domains. We further improve this approach with a novel strategy for attention-aware stylization, which uses the attention maps of class (CLS) tokens to compute and mix normalization statistics of tokens corresponding to different image regions. The proposed method is flexible to the choice of backbone model and can be easily applied to any ViT-based architecture with a negligible increase in computational complexity. Comprehensive experiments show that our approach is able to achieve state-of-the-art performance on five challenging benchmarks for domain generalization, and demonstrate its ability to deal with different types of domain shifts. The implementation is available at: https://github.com/Mehrdad-Noori/TFS-ViT_Token-level_Feature_Stylization.