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

TL;DR

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 by transforms token features by mixing the normalization statistics of images from different domains.

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.

Figures (8)

• Figure 1: Overview of the proposed architecture for Token-level Feature Stylization (TFS-ViT).
• Figure 2: Synthesized features using our proposed method. Different colors denote different styles. By randomly selecting a subset of tokens to stylize at each layer, our method generates diverse samples while preserving the underlying structure of the tokens. This leads to forcing the network to only focus on the structure-related information which eventually results in improving the generalization performance. It is worth mentioning that we perform our stylization method on multiple layers of the ViT network.
• Figure 3: Effects of varying hyperparameters on the PACS dataset. The figure shows the influence of $n$, the number of layers where stylization is performed, with results averaged over different $d$ values, alongside the impact of $d$, the fraction of tokens to be replaced with their stylized counterparts, averaged over various $n$ values.
• Figure 4: Performance comparison of stylization applied to a fixed initial set of layers versus random layer selection on the PACS dataset. Results are averaged over the different $d$ values.
• Figure 5: Comparison of ERM-ViT and TFS-ViT performance in Single-Source Domain Generalization setting on the PACS dataset.