Stylizing ViT: Anatomy-Preserving Instance Style Transfer for Domain Generalization

TL;DR

Stylizing ViT introduces a single-encoder Vision Transformer with weight-sharing cross-attention blocks to perform anatomy-preserving style transfer for domain generalization in medical imaging. By fusing anatomy and style via cross-attention and reconstructing through a compact MLP with a dot-product stage, the method yields realistic stylizations without artifacts and enhances downstream robustness during training and inference. Across three medical datasets, it achieves state-of-the-art or near-state-of-the-art performance for training-time augmentation and demonstrates meaningful gains with test-time augmentation, indicating practical impact for improving model generalization under domain shifts. The approach balances style diversity with anatomical fidelity, offering a scalable augmentation strategy suitable for real-time pipelines in histopathology and dermatology applications.

Abstract

Deep learning models in medical image analysis often struggle with generalizability across domains and demographic groups due to data heterogeneity and scarcity. Traditional augmentation improves robustness, but fails under substantial domain shifts. Recent advances in stylistic augmentation enhance domain generalization by varying image styles but fall short in terms of style diversity or by introducing artifacts into the generated images. To address these limitations, we propose Stylizing ViT, a novel Vision Transformer encoder that utilizes weight-shared attention blocks for both self- and cross-attention. This design allows the same attention block to maintain anatomical consistency through self-attention while performing style transfer via cross-attention. We assess the effectiveness of our method for domain generalization by employing it for data augmentation on three distinct image classification tasks in the context of histopathology and dermatology. Results demonstrate an improved robustness (up to +13% accuracy) over the state of the art while generating perceptually convincing images without artifacts. Additionally, we show that Stylizing ViT is effective beyond training, achieving a 17% performance improvement during inference when used for test-time augmentation. The source code is available at https://github.com/sdoerrich97/stylizing-vit .

Figures (11)

• Figure 1: Overview of how Stylizing ViT is used to improve domain generalization during training and inference. During training, it generates stylistically diverse but anatomically consistent images, encouraging downstream classifiers to learn structure-aware representations. At test time, it is used to align unseen input styles with the training distribution, thereby mitigating style-induced domain shifts.
• Figure 2: Overview of our proposed Stylizing ViT. The input image pair ($I$, $S$) is processed by the encoder that fuses the anatomical structure of $I$ with the style characteristics of $S$ using cross-attention. Subsequently the stylized image $T$ is reconstructed through a two-layer MLP, a dot-product operation, and a convolutional layer. A frozen VGG19 encoder is used during training to compute perceptual losses.
• Figure 3: Qualitative comparison of style transfer quality on training image pairs. Our method consistently outperforms all reference methods in applying diverse styles without compromising anatomy.
• Figure 4: Style transfer results for Camelyon17-WILDS (left) and Fitzpatrick17k (right), showing high quality images across different dataset splits and domains such as staining variations and skin tones.
• Figure 5: Demonstration of our method’s ability to transfer style while preserving anatomical structure. Given an input image pair $(I, S)$, we first apply a shared color jitter augmentation $f(\cdot)$ to produce $\Tilde{I}$ and $\Tilde{S}$. Our model generates $T$, which closely matches $\Tilde{I}$, indicating successful transfer of style while maintaining anatomical fidelity.