[KongchangAI]
· 1 min read· 889 words

Rethinking the Design Space of Pixel Diffusion Transformers: MFlow and ConvRefiner

Rethinking the Design Space of Pixel Diffusion Transformers: MFlow and ConvRefiner

MFlow lets pixel diffusion models adaptively choose x0 or u prediction and uses ConvRefiner to repair high-frequency deficiencies.

This article re-examines the parameterization design of pixel-space diffusion Transformers. Prior JiT work showed that replacing velocity prediction with clean-image prediction alone enables pixel-space generation. The new study proposes MFlow, which learns adaptive interpolation between x0 and u prediction and is jointly optimized end-to-end with diffusion training. It finds the optimal parameterization shifts toward velocity prediction as model capacity grows, and MFlow achieves competitive or superior results on ImageNet, UCF-101 video, and text-to-image tasks. The authors also reveal a persistent high-frequency deficiency independent of the parameterization choice, and introduce ConvRefiner, a lightweight conv module with a frozen backbone trained post hoc with adversarial loss to restore high-frequency statistics while keeping FID from degrading.

Pixel-space diffusion models have recently re-emerged as a strong alternative to latent diffusion, avoiding the lossy latent bottleneck and separately trained autoencoders. Prior work, JiT, showed that a plain diffusion Transformer can generate images directly in pixel space simply by replacing the commonly used velocity prediction (u-prediction) with clean-image prediction (x0-prediction). This new study revisits that choice and proposes MFlow—an adaptive parameterization that learns to interpolate between x0-prediction and u-prediction along the principal directions of the data, and is learned end-to-end jointly with diffusion training.

Rethinking the design space of Vanilla Pixel Diffusion Transformers

The Parameterization Choice Is Not Set in Stone

The research team found that the optimal parameterization shifts systematically with data representation and model capacity. A key observation is that as model capacity increases, the preference shifts toward velocity prediction. This means that x0-prediction, which performs better on small models, is not necessarily optimal for larger models. The value of MFlow lies in not requiring this choice to be manually preset—instead, it lets the model adaptively learn the interpolation weights during training.

Cross-Scenario Validation

The authors evaluated MFlow on ImageNet, UCF-101 video, and large-scale text-to-image generation tasks, showing competitive or superior generation performance across a variety of settings. This indicates that adaptive parameterization is not a tuning trick for a single benchmark, but a design principle with generalization.

The difference between x0-prediction and u-prediction: In diffusion models, the network needs to predict some target to guide denoising. x0-prediction means the network directly estimates the clean image (i.e., the denoising endpoint), while u-prediction (velocity prediction) estimates the "velocity" vector from the current noisy image to the clean image; the two are mathematically convertible. The x0-prediction loss is more sensitive to high-noise timesteps, since predicting the clean image itself is difficult there; u-prediction provides a more balanced gradient signal across noise levels. Early work found x0-prediction to be more stable on small models, but as model capacity grows, u-prediction can more fully exploit the network's expressive power. The core of MFlow is precisely to let the model learn an interpolation coefficient that smoothly transitions between the two parameterizations, rather than artificially choosing one.

High-Frequency Deficiencies: The Hidden Weakness of Pixel Diffusion

The study also revealed a problem independent of the parameterization choice: regardless of which prediction method is used, images generated by plain pixel diffusion Transformers consistently lack high-frequency information relative to real images. Visually, this deficiency may manifest as blurred details or textures that are not sharp enough, and the traditional FID metric may not fully capture this.

ConvRefiner: Post-Processing Repair with a Frozen Backbone

To address the high-frequency deficiency, the authors introduce ConvRefiner—a lightweight convolutional module. Its training approach is quite clever: with the diffusion backbone completely frozen, ConvRefiner is trained post hoc with an adversarial loss. With only a small amount of additional training, ConvRefiner can restore the high-frequency statistical properties of generated images while keeping FID from degrading. This "frozen backbone + lightweight post-processing" idea provides a low-cost paradigm for post-processing optimization of diffusion models.

Limitations of the FID metric: FID (Fréchet Inception Distance) compares generated and real distributions using features extracted by a pretrained Inception network. It is more sensitive to semantics and overall structure, but insensitive to differences in high-frequency details (such as sharp edges and fine textures). As a result, generated images may look blurry or "overly smooth" visually while still achieving good FID scores. High-frequency information deficiency typically manifests as insufficient energy in high-wavenumber components of the image spectrum, which is related to the regression-to-the-mean effect on local details when diffusion models directly model in pixel space. ConvRefiner's adversarial training is precisely aimed at closing the high-frequency distribution gap at the level of feature statistics, rather than merely optimizing FID.

Implications for the Pixel Diffusion Route

This work conveys a message on two levels. At the parameterization level, there is no universally optimal solution between x0-prediction and u-prediction; model capacity and data characteristics jointly determine the direction of preference, and adaptive interpolation is more robust than a fixed choice. At the architecture level, the high-frequency deficiency of pixel diffusion Transformers suggests that simply scaling up the model or adjusting the prediction target is not enough to close the statistical gap with real images, and targeted post-processing modules still have independent value.

The project webpage is now public, containing more implementation details and visualization results for researchers to further reproduce and compare.

One-line Summary

MFlow lets pixel diffusion models adaptively choose between x0 and u prediction, and uses ConvRefiner to repair high-frequency deficiencies.

Paragraph Summary

This article re-examines the parameterization design of pixel-space diffusion Transformers. Prior work, JiT, showed that simply replacing velocity prediction with clean-image prediction enables image generation in pixel space. The new study proposes MFlow, which learns to adaptively interpolate between x0-prediction and u-prediction and is jointly optimized end-to-end with diffusion training. The research finds that the optimal parameterization shifts toward velocity prediction as model capacity increases, and that MFlow achieves competitive or superior performance on ImageNet, UCF-101 video, and text-to-image tasks. In addition, the authors reveal that plain pixel diffusion Transformers suffer from a persistent high-frequency information deficiency independent of the parameterization choice, and introduce ConvRefiner, a lightweight convolutional module with a frozen backbone, trained post hoc with an adversarial loss to restore high-frequency statistics while keeping FID from degrading.

#pixel diffusion#diffusion Transformer#MFlow#ConvRefiner#x0 prediction#u prediction#ImageNet#text-to-image
Share:

Related articles