RsViT – A Scalable Vision Transformer Backbone for Diffusion Model

Denoising diffusion probabilistic models and vision transformers have shown tremendous progress in the image-generative field in recent years. This research proposes a more scalable vision transformer-based backbone for a diffusion model called RsViT. While RsViT has a similar architecture to U-Net, a hierarchical symmetrical encoder-decoder architecture. RsViT’s primary mechanism is more similar to a transformer than CNNs. Instead of the ResNet block, RsViT uses Restormer’s transformer block, which eliminates the weaknesses of vision transformer: the computational overhead of the attention map and the lack of localization, making the model more scalable and superior in generative tasks. From our experiments, on average, RsViT achieves more PSNR and SSIM than the UNet-based backbone by 0.12 and 0.0032 on 256 \(\,\times \,\) 256 ImageNet-100 dataset in restoration tasks.