Multi scale transformers improve image generation quality significantly

Scale as a Modality: Multi-Scale Pixel-Space Diffusion Transformers

Computer Vision and Pattern Recognition

Summary

Generating high-quality images with AI often involves simplifying or compressing images in ways that limit detail. The authors designed a new approach that looks at images at many scales or sizes simultaneously, treating these scales as different information types in a single model. This method improves how well the AI can generate detailed images, better than just making the model deeper or bigger. Their model works well at different image sizes and also adapts naturally for creating images from text.

What this means in practice

  • For image generation teams: Produce higher quality images by integrating multiple spatial scales simultaneously in diffusion transformer models.
  • For text-to-image system developers: Enhance text-to-image generation pipelines by utilizing multi-scale transformers that better capture fine details and spatial context.

Authors

Ruibo Chen, Kaishen Wang, Tianyi Xiong, Chenxi Liu, Yanshuo Chen, Heng Huang, Jian Du

Abstract

Latent diffusion models rely on pretrained autoencoders that introduce lossy compression and constrain spatial granularity. Pixel-space diffusion avoids this bottleneck and enables flexible tokenization at multiple spatial scales, yet existing approaches largely adopt two-level coarse-to-fine architectures with limited cross-scale interaction. We introduce MS-pDiT, a multi-scale pixel-space diffusion transformer that treats spatial scale as a modality. MS-pDiT tokenizes the same noisy image at multiple patch sizes, processes each scale in a dedicated transformer stream, and enables bidirectional information exchange through compacted MM-DiT-style joint attention, with attention cost independent of the finest patch size. This design establishes the number of scales as a new architectural scaling dimension. On ImageNet $256\times256$, increasing the number of scales from two to five consistently improves FID from 4.39 to 2.31, 2.26, and 2.19, while adding an intermediate scale substantially outperforms increasing depth under comparable parameter budgets. Further analysis highlights the importance of bidirectional cross-scale interaction and late introduction of fine-grained streams. MS-pDiT achieves FID scores of 1.43 and 1.44 on ImageNet $256\times256$ and $512\times512$, respectively, and naturally extends to text-to-image generation, achieving 0.86 on GenEval and 83.9 on DPG-Bench.