Pixel diffusion models improve image quality with adaptive prediction
Rethinking the design space of Vanilla Pixel Diffusion Transformers
Computer Vision and Pattern RecognitionArtificial Intelligence
Summary
Some AI models create pictures by gradually filling in pixels, but they often struggle with small details and can use methods that lose image quality. The authors revisited how these models predict pixels and developed MFlow, which mixes two ways of predicting images based on the data and the model's size. They show that this adaptive approach works well across different kinds of images and videos. To fix missing fine details in generated images, they also added a small extra module called ConvRefiner that sharpens images while keeping the main model unchanged.
What this means in practice
- •For image generation engineers: Improve the quality of AI-generated images by using an adaptive prediction method that adjusts based on data type and model size.
- •For video synthesis developers: Enhance video frame generation by applying MFlow’s learned mix of prediction parameters for better detail and accuracy.
Authors
Soumik Mukhopadhyay, Long Mai, Abhinav Shrivastava, Aniruddha Mahapatra
Abstract
Pixel-space diffusion models have recently re-emerged as a compelling alternative to latent diffusion, avoiding the lossy latent bottleneck and separately trained autoencoder. JiT shows that a vanilla diffusion transformer can generate directly in pixel space by replacing the commonly used velocity prediction ($\boldsymbol{u}$-prediction) with clean image prediction ($\boldsymbol{x_0}$-prediction). In this work, we revisit this choice and introduce MFlow, an adaptive parameterization that learns to interpolate between $\boldsymbol{x_0}$- and $\boldsymbol{u}$-prediction across principal directions of the data, learned end-to-end with diffusion training. We find that the preferred parameterization varies systematically with the data representation and model capacity -- notably, increasing model capacity shifts preference toward velocity prediction. We evaluate MFlow across ImageNet, UCF-101 video, and large-scale text-to-image generation, demonstrating competitive or improved generative performance across diverse settings. In addition, we find that vanilla pixel diffusion transformers, irrespective of parameterization, retain a persistent high-frequency deficit relative to real images. To address this, we introduce ConvRefiner, a lightweight convolutional module trained post-hoc with an adversarial loss while keeping the diffusion backbone frozen. With only a small amount of additional training, ConvRefiner restores the high-frequency statistics of generated images while preserving FID. Project webpage: https://soumik-kanad.github.io/mflow-diffusion/.