PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking the design space of Vanilla Pixel Diffusion Transformers

Soumik Mukhopadhyay, Long Mai, Abhinav Shrivastava, Aniruddha Mahapatra

arXiv:2610.23854v1cs.CVcs.AI

TL;DR

Pixel-space diffusion avoids the lossy bottleneck and separately trained autoencoder of latent diffusion, but it leaves questions about prediction targets and fine-scale detail. The paper introduces MFLOW, which learns direction-wise interpolation between clean-image and velocity prediction, and ConvRefiner, a post-hoc refinement module. MFLOW adapts across representations and model capacities, while ConvRefiner restores high-frequency statistics without changing FID; the paper also notes that FID alone does not capture the refiner’s spectral changes.

  • Problem

    Latent diffusion relies on a separately trained autoencoder whose lossy bottleneck can make reconstruction unfaithful, while large-patch pixel-space transformers raise questions about prediction parameterization and fine-detail modeling.

    Source: §1 · p. 2

  • Method

    MFLOW learns continuous prediction gains across principal data directions, while ConvRefiner is a post-hoc convolutional module trained with the diffusion backbone frozen.

    Source: §1 · p. 2

  • Results

    MFLOW’s preferred parameterization varies with data representation and model capacity, and ConvRefiner brings generated high-frequency statistics closer to real images while preserving FID.

    Source: §1 · p. 2, §6 · p. 11, §7 · p. 12

  • Takeaways & Limitations

    Adaptive prediction targets and explicit treatment of fine-scale signal structure are complementary ingredients for scaling simple pixel-space diffusion transformers.

    Source: §7 · p. 12

    ConvRefiner’s spectral improvement produces little or no change in FID, so spectral fidelity and FID are complementary measures and texture-sensitive metrics or human evaluation are needed to assess perceived quality.

    Source: §A.8 · p. 25

Abstract

from arXiv · show

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/.

1 INTRODUCTION

Pixel-space diffusion avoids the lossy bottleneck and separately trained autoencoder used by latent diffusion, but large-patch transformers raise questions about prediction targets and fine detail. The paper introduces direction-adaptive MFLOW and post-hoc ConvRefiner to address these challenges.

1 INTRODUCTION

  • Latent diffusion uses a separately trained compression autoencoder, and its lossy bottleneck can make reconstruction unfaithful after compression.The paper frames direct pixel-space generation as avoiding these limitations.Source: §1 · p. 2
  • Vanilla pixel-space transformers face two open design questions: which prediction parameterization suits high-dimensional patches, and whether they capture fine-scale image details.Aggressive patchification keeps sequence length manageable but makes each token model a high-dimensional block of raw pixels.Source: §1 · p. 2
  • MFLOW learns a continuous gain for each principal data direction, adapting the mix of clean-image and velocity prediction to the representation, objective, and model capacity.The learned direction-wise gains can shift toward velocity-like prediction as model capacity increases.Source: §1 · p. 2
  • MFLOW extends beyond pixel-space images to pixel-space video, text-to-image generation at multiple capacities, and high-dimensional latents.The paper also reports that vanilla JiT-family models retain a high-frequency deficit regardless of prediction parameterization.Source: §1 · p. 2
  • ConvRefiner is a lightweight post-hoc convolutional module that brings generated high-frequency spectra closer to real images while preserving FID.It is trained with perceptual and adversarial supervision while the diffusion backbone remains frozen; the reported post-hoc training uses 5%.Source: §1 · p. 2

2 RELATED WORK

Prior work has advanced pixel-space generation through both architectural and formulation changes, while recent approaches also revisit vanilla transformers and prediction parameterizations. Fine-detail methods use perceptual or frequency-aware training, specialized architectures, or improved information pathways; ConvRefiner instead operates post hoc with the backbone frozen.

2 RELATED WORK

  • Pixel-space generation developed from early diffusion models to later architectural and formulation advances, followed by renewed interest in vanilla and specialized transformers.The cited examples include early models (Ho et al. (2020); Nichol & Dhariwal (2021); Dhariwal & Nichol (2021); Karras et al. (2022)) and later advances (Jabri et al. (2023); Hoogeboom et al. (2023; 2025); Crowson et al. (2024); Chen et al. (2025b); Mukhopadhyay et al. (2026)).Source: §2 · p. 4
  • JiT advocates clean-image prediction for high-dimensional pixel modeling, while k-Diff and AsymFlow explore adaptive combinations of clean-data and velocity prediction.MFLOW lets the model learn a fit between clean-image and velocity prediction rather than prescribing a fixed formulation.Source: §2 · p. 4
  • Fine-detail approaches include perceptual supervision, frequency-aware objectives or flows, specialized architectures, and improved information pathways; ConvRefiner instead trains post hoc with the backbone frozen.The cited approaches are described as generally modifying the objective or adding specialized components, and some still show a high-frequency gap relative to real images; ConvRefiner is reported to recover missing high-frequency statistics while preserving FID.Source: §2 · p. 4

3 MODEL PREDICTION PARAMETERIZATION

The section motivates adaptive prediction targets because fixed choices and hard subspace cutoffs can be poorly suited to different data directions. MFLOW replaces the cutoff with learnable continuous gains, recovers velocity for standard flow-matching training, and remains effective in a high-dimensional toy experiment.

3.1 MOTIVATION

  • The motivation contrasts clean-image, velocity, and noise prediction as algebraically equivalent targets with different conditioning, and notes that a globally shared interpolation cannot capture within-representation directional differences.JiT favors clean-image prediction for high-dimensional pixel patches, while k-Diff uses a shared coefficient.Source: §3.1 · p. 4
  • AsymFlow assigns velocity prediction to a selected low-rank subspace and clean-image prediction to its complement, but its hard rank cutoff requires setting-specific search.The cited ImageNet example reports FID 3.54 for PCA[9 : 16] versus 3.59 for PCA[1 : 8], illustrating that the top-variance directions need not be the best selected subspace.Source: §3.1 · pp. 4–5

3.2 MFLOW: DIRECTION-ADAPTIVE LEARNABLE REPARAMETERIZATION

  • MFLOW replaces the binary projector with a symmetric graded operator, giving each principal direction its own continuous velocity–data mixing weight.A gain of 0 corresponds to x0-prediction and 1 to u-prediction; binary gains recover the AsymFlow rule.Source: §3.2 · pp. 5–6
  • MFLOW maps its learned target to the full velocity and trains with a velocity-space flow-matching loss, without changing the forward process or transport target.At sampling, the recovered velocity is used with standard Euler or Heun integration.Source: §3.2 · p. 6

3.3 TOY EXAMPLE

  • At D=512, full u-prediction degrades while x0-prediction, rank-matched AsymFlow, and fixed or learned MFLOW recover the spiral.Across D ∈ {2, 8, 16, 512}, the continuous-gain MFLOW variants remain tractable, and learned gains can be optimized end-to-end.Source: §3.3 · p. 6, Figure 4 · p. 6

4 SIGNAL MODELING

JiT-family pixel diffusion models lack high radial frequencies, a deficit not seen in the compared latent and other pixel models. ConvRefiner adds a post-hoc convolutional residual pathway to recover fine-scale detail while keeping the backbone fixed.

4.1 MOTIVATION

  • JiT-family generations lack radial frequencies in the 0.5–0.7 range, while the compared VAE and PixelDiT generations do not.The authors suggest the deficit may be specific to JiT’s large-patch design without intra-patch local operations.Source: §4.1 · p. 7

4.2 CONVREFINER: A SIMPLE POST-HOC FIX

  • On ImageNet-1K 256×256, MFLOW plus ConvRefiner is reported to solve the high-frequency deficit while improving FID and avoiding expensive parameterization sweeps.The table caption defines the deficit measure as mean radial power deficit relative to real images.Source: Table 1 · p. 7
  • ConvRefiner adds lightweight convolutional blocks, trains them post hoc with GAN loss, and keeps the JiT backbone frozen.It predicts a residual over the clean-image prediction; other parameterizations are first converted to a clean image.Source: §4.2 · p. 7

5 EXPERIMENTS

Experiments evaluate MFLOW and ConvRefiner on ImageNet, UCF-101 video, and text-to-image generation. MFLOW is competitive or best in the reported comparisons, while ConvRefiner recovers high-frequency detail with little reported change to generative performance.

ImageNet

  • On ImageNet, adding LPIPS makes MFLOW the best of the compared models at the B-level configuration, while MFLOW remains competitive at H-scale.The authors hypothesize AsymFlow’s very slight H-scale lead reflects heavier ImageNet tuning.Source: §5 · p. 8, Table 2 · p. 8, Table 3 · p. 9

ConvRefiner

  • Post-training ConvRefiner fixes the high radial-frequency deficit across JiT and MFLOW at B- and H-levels while maintaining FID on ImageNet.For text-to-image MFLOW-XXL, ConvRefiner closes the 0.5–0.7 frequency-band gap without a noticeable benchmark-performance effect.Source: §5 · pp. 8–9, Figure 8 · p. 9, Figure 9 · p. 10

Videos

  • On UCF-101, MFLOW outperforms WAN2.1, JiT, and AsymFlow at B-level and beats JiT and AsymFlow at L-level; post-hoc Framewise-LPIPS yields the best performance.Training with Framewise-LPIPS from the start hurts video generation.Source: §5 · p. 8, Table 4 · p. 9, Table 5 · p. 9

Text-to-Image

  • At 13B total parameters and 2.9B active parameters, MFLOW achieves the best performance among the text-to-image models in the comparison.At the 1.1B-parameter scale, MFLOW is best on every reported benchmark, including versus previous pixel-space models.Source: §5 · p. 9, Table 6 · p. 10

MFLOW Ablations

  • For MFLOW ablations, the uG loss worsens FID from 3.66 to 4.88, while norm-var initialization and a 10× gain-parameter learning-rate multiplier are selected for the final setting.The authors attribute the uG-loss result to an incentive to reduce gains toward x0-prediction without improving the recovered transport field.Source: §5 · p. 10

ConvRefiner Ablations

  • For ConvRefiner, more blocks and higher GAN-loss weight move the frequency spectrum closer to real, and the selected setting is b = 4 and w = 0.02.Finetuning the unpatchify layer or using a residual-attention refiner has negligible effect; DiP patch-detailer improves high frequency but degrades JiT FID from 3.42 to 3.90 and 4.46 with GAN loss.Source: §5 · pp. 10–11, Figure 11 · p. 10

6 DISCUSSION

Learned MFLOW gains vary with objective, domain, representation, and capacity. The results support data-dependent allocation of velocity prediction rather than a manually chosen fixed subspace.

M-curve Analysis

  • MFLOW’s learned gains shift across objectives, domains, representations, and model capacities, with more directions moving toward u-prediction as text-to-image model capacity increases.With RAEv2 latents, gains collapse nearly to zero; adding u-prediction to a small leading subspace changes FID only marginally.Source: §6 · p. 11, Figure 12 · p. 11
  • On ImageNet, adding LPIPS moves all but the leading few directions toward x0-prediction, while video gains are non-negligible in fewer directions.Without LPIPS, gains decrease approximately with PCA variance; with LPIPS, the leading gains are non-monotonic, including a dip at PC2–PC3 and a rise at PC4.Source: §6 · p. 11

Discussion

  • For a fixed number of u-predicted directions, allocation should remain data-dependent, with leading high-variance PCA directions a natural choice.MFLOW can adapt to increased capacity or domain change without a manually chosen rank.Source: §6 · p. 11

7 CONCLUSION

The paper revisits prediction parameterization and fine-scale signal modeling in vanilla pixel diffusion transformers. It proposes MFLOW’s learned interpolation and ConvRefiner’s post-hoc convolutional refinement as complementary responses.

Conclusion

  • MFLOW learns a continuous interpolation between clean-data and velocity prediction across principal data directions, avoiding manual subspace selection.Its learned M-curves vary systematically with data statistics and model capacity.Source: §7 · pp. 11–12
  • ConvRefiner addresses a persistent high-frequency deficit in JiT-style models with a lightweight post-hoc convolutional module that restores high-frequency statistics while preserving generative quality.The paper concludes that adaptive prediction targets and explicit treatment of fine-scale signal structure are important ingredients for scaling simple pixel-space diffusion transformers.Source: §7 · p. 12
  • The paper reports that AI tools were used only for supporting roles such as grammar correction and visualization tooling, not for research ideation, experiment design, or scientific claims.The authors state that they reviewed the AI-assisted text and figures for correctness.Source: AI USE STATEMENT · p. 12

A.1 HYPERPARAMETERS

The appendix specifies training and sampling configurations for ImageNet, UCF-101, and text-to-image experiments. Evaluation settings are tuned per model in most experiments, while text-to-image uses one shared sampler.

A.1 HYPERPARAMETERS

  • ImageNet B-level experiments use Heun with 50 steps, CFG 2.9, CFG Interval [0.10, 1.0], and σmin = 0.04 to isolate training-objective effects.H-scale evaluation uses a coarse CFG scan followed by a CFG×σmin grid for the strongest LPIPS variant.Source: §A.1 · p. 16
  • UCF-101 uses 8 × 8 spatial patches and a two-stage CFG×σmin search, with Heun/100 and full-scale FVD@10K evaluation around the coarse optimum.The cited Skorokhodov et al. (2024) results are omitted because their resolution and temporal subsampling differ and no public code or checkpoint is available.Source: §A.1 · p. 16
  • Text-to-image uses a shared DDIM sampler with 50 steps and CFG 3.0, so its numbers reflect out-of-the-box behavior rather than per-model CFG×σmin tuning.The evaluation hyperparameters are summarized in Table 8.Source: §A.1 · p. 17, Table 8 · p. 18
  • ImageNet and UCF-101 retain the base architecture unless otherwise specified, with training configurations listed in Table 7 and evaluation settings in Table 8.Table 1 uses a shared base JiT-B evaluation configuration; other reported tables use per-model tuning.Source: §A.1 · p. 16, Table 7 · p. 17, Table 8 · p. 18

A.2 MFLOW IMPLEMENTATION DETAILS

MFLOW performs its operations in a fixed patch-PCA basis, using direction-wise gains to form a graded target and recover clean-image and velocity predictions. Training optimizes the standard velocity-space flow-matching loss, and sampling integrates recovered velocity.

A.2 MFLOW IMPLEMENTATION DETAILS

  • MFLOW changes basis to the fixed patch-PCA coordinates, applies per-direction scalar operations, then changes back to image space.The implementation uses x as a batch of patch-token tensors, t ∈ [0,1] as the flow-matching timestep, and precomputed PCA basis U and variances lam.Source: §A.2 · p. 19
  • The default learnable gains are direction-wise sigmoid values, initialized from normalized PCA variances; a fixed alternative sets m_fixed = lam_n / (lam_n + c).The learnable variant uses D logits phi, with m_i = sigmoid(phi_i).Source: §A.2 · p. 19
  • The graded training target is u_G = x0 − Mϵ, formed by applying the gains to noise in PCA coordinates before transforming back.MFLOW converts the network output to a clean-image estimate and then recovers velocity as (x0_hat − x_t) / clamp_min(1 − t, t_eps).Source: §A.2 · p. 19
  • Training interpolates noise and clean data, predicts recovered velocity, and minimizes weighted MSE against x0 − eps; sampling starts from noise and integrates velocity with Euler/Heun.The sampling schedule runs from t = 0 to 1.Source: §A.2 · pp. 19–20

A.3 WHY MFLOW USES u-PREDICTION

MFLOW assigns each PCA direction a continuous gain that mixes velocity-like and clean-image-like prediction. The gains are trained with velocity-space supervision because direct graded-target loss can improve through a parameterization-scale shortcut rather than better transport prediction.

A.3 WHY MFLOW USES u-PREDICTION

  • MFLOW uses a graded target whose direction-wise gains determine the mix between clean-image and velocity prediction.The target in direction i is u_G,i = (1 − m_i)x_t,i + ∆_i(t)u_i, with ∆_i(t) = (1 − t) + t m_i.Source: §A.3 · p. 20
  • Direct graded-target supervision weights recovered-velocity error differently by gain and timestep, allowing loss changes unrelated to transport accuracy.The objective admits a parameterization-scale shortcut even though gains receive gradients.Source: §A.3 · p. 20
  • For a fixed recovered predictor, decreasing m_i when t > 0 and velocity error is nonzero lowers direct graded-target loss without changing the recovered transport field.Thus the direct objective can favor a more x0-like target instead of a more accurate transport field.Source: §A.3 · pp. 20–21
  • In the clean-image-prediction subspace, direct u_G supervision suppresses recovered-velocity errors by (1 − t)^2, nearly removing their penalty near t → 1.Those errors are amplified when predictions are converted back to velocity.Source: §A.3 · p. 21
  • Velocity-space supervision evaluates parameterizations with the same transport objective, while gains still affect the network’s conditioning and finite-capacity approximation problem.It prevents gains from lowering the objective merely by shrinking regression-target scale.Source: §A.3 · p. 21

A.4 SURPRISING BEHAVIOUR LPIPS ON THE LEARNED M-CURVE

LPIPS changes the learned M-curve so gains align with patch-global PCA modes rather than variance ordering alone. The capacity-constrained view frames parameterization as balancing velocity-recovery conditioning against the difficulty of predicting noise, with the balance varying by representation and model capacity.

A.4 SURPRISING BEHAVIOUR LPIPS ON THE LEARNED M-CURVE

  • With LPIPS, gains are non-monotonic across leading PCA directions: PC0–PC1 and PC4 are high, while gradient modes PC2–PC3 dip.PC4 has lower variance than PC2–PC3 (λ4 = 3.74 versus λ2 = 9.41 and λ3 = 7.93), yet receives a larger gain; its patch-global chromatic structure distinguishes it from spatially varying modes.Source: §A.4 · pp. 21–22

A.5 A CAPACITY-CONSTRAINED VIEW OF x0- AND u-PREDICTION

  • MFlow balances the improved near-clean velocity recovery of velocity-like prediction against the greater noise-prediction burden it imposes on a finite-capacity network.Pure x0-prediction avoids that noise burden but retains clean-endpoint amplification; MFlow's soft velocity rank summarizes how much velocity prediction the model can support.Source: §A.5 · p. 23
  • The learned velocity budget depends on representation and capacity: RAEv2 favors x0-prediction, while text-to-image models shift more directions toward velocity prediction as capacity increases.On RAEv2, pure x0-prediction obtains FID 1.05, versus 1.08 for AsymFlow (r = 8) and 1.09 for MFlow after 80 epochs; AsymFlow is only 0.03 FID behind the x0 baseline.Source: §A.5 · pp. 23–24

A.6 RADIALLY-AVERAGED POWER SPECTRUM

  • The radially averaged power spectrum measures generated-image frequency content by averaging power across orientations at each radial frequency.The procedure converts images to luma, removes the mean, applies a Hann window, and averages profiles over 4,000 images.Source: §A.6 · p. 24

A.7 REFINERS

  • On ImageNet generations, adding ConvRefiner post-training to MFLOW-H+LPIPS produced overall improvement in texture details and average power gain in the high-frequency (0.5–0.7) region.For near-clean batches (t ∼ U(0.9, 0.98)), an additional GAN loss is applied with weight 0.02; flow-matching and LPIPS losses are downweighted by a factor of 0.25.Source: §A.7 · p. 24, Figure 22 · p. 32

A.8 LIMITATIONS.

  • MFlow's improvements over strong fixed parameterizations are sometimes modest, and lower diffusion loss does not consistently translate into better FID.The authors note that the cause of this mismatch remains unclear and identify staged or decoupled gain optimization as future work.Source: §A.8 · p. 25
  • ConvRefiner closes much of the high-frequency spectral gap with little or no FID change, so spectral fidelity and FID are complementary measures.Texture-sensitive metrics and human evaluation are needed to assess how strongly recovered high-frequency content translates to perceived image quality.Source: §A.8 · p. 25

Related papers

Loading 2610.23854v1…