Skip to content

Pixel Space Is Back, and the Interesting Part Is the Loss

#diffusion-models #pixel-space #text-to-image #representation-learning #dit #image-editing

Pixel Space Is Back, and the Interesting Part Is the Loss ​

Every text-to-image model you've used in the last three years runs diffusion inside a compressed latent space. An autoencoder squeezes the image 8x, the diffusion model learns in that smaller space, a decoder puts it back. It saves a lot of FLOPs per step, and it costs you a second model to train, tune and ship, plus a reconstruction ceiling you can never diffuse past.

A cluster of papers now argues the compression step is optional, and in some settings actively harmful. JiT showed a plain diffusion transformer can work on raw pixels if you swap velocity prediction for clean-image prediction. The follow-up questions are the ones that matter in practice: how should you parameterize the target, what signal should you supervise with, and what do the resulting images still get wrong? (MFlow, LILIE, MOSS-Image, FlashReFocus, DiffPT)

The latent bottleneck is a compute decision, not a quality one ​

Nobody chose latent diffusion because latents looked better. It won on cost. Pixel space forces the transformer to spend tokens on detail the autoencoder would have thrown away, so you either pay for a long sequence or use big patches and lose the detail anyway.

The autoencoder also introduces a hard ceiling. Whatever the decoder can't reconstruct, the diffusion model can't generate. That's tolerable for natural photos, where the decoder's failure modes are rare. It's a problem for dense text, thin strokes, and anything where a misplaced glyph is the difference between a usable image and a corrupt one. Which is exactly the regime MOSS-Image targets.

So the interesting shift isn't "pixels beat latents." It's that the bottleneck moved. Once you remove the autoencoder, the next question is what you're predicting and what you're measuring against.

Predicting the clean image vs predicting the velocity ​

In flow matching the two targets are algebraically equivalent: given the noisy sample, velocity and clean image determine each other. They are not numerically equivalent. The relation carries a factor of the timestep, so an L2 loss on one target weights errors at low noise levels very differently from an L2 loss on the other. Same model, same data, different gradient landscape.

JiT's contribution was picking clean-image prediction, and it made vanilla pixel-space DiTs train. MFlow's contribution is pointing out that this was a global constant chosen once, and it probably shouldn't be. MFlow learns an adaptive parameterization that interpolates between clean-image and velocity prediction along the principal directions of the data, trained end to end with the diffusion objective.

The result that should change your training config: the preferred parameterization shifts systematically with data representation and model capacity, and larger models lean toward velocity prediction. My read is that big models can absorb the steep low-noise weighting that velocity prediction imposes, while smaller ones do better with the flatter clean-image target. Either way, the practical consequence is blunt. A configuration tuned at 300M parameters is not a configuration at 3B.

Quick Take: The parameterization of the diffusion target is a hyperparameter that most pixel-space training runs set once and never revisit, and MFlow shows it should be re-tuned every time you change model size.

Five papers side by side ​

PaperWhat it changesHeadline resultWhere it pays off
MFlowLearns a per-direction mix of clean-image and velocity predictionCompetitive or better generation on ImageNet, UCF-101 video, and text-to-imageScaling pixel DiTs past the size where JiT's default holds
ConvRefinerPost-hoc convolutional refiner, adversarial loss, backbone frozenRestores high-frequency statistics while preserving FIDFixing the waxy look without retraining the base model
LILIEInvertible flow encoder plus a VLM gives a language-informed lossgFD drops 66.84 to 53.01 on GPICSemantic supervision without an image tokenizer
MOSS-Image5B single-stream DiT trained from scratch, plus an agentic prompt enhancer96.3% average text coverage; top open-weight model on LLTBDense bilingual text and layout-heavy images
FlashReFocusOne diffusion pass, then a lightweight defocus renderer462x faster per edit than GenRefocusInteractive depth-of-field editing
DiffPTDiffusion as pretraining for a classifier backboneBeats self-supervised baselines on medical imaging and time seriesSmall domain-specific datasets

Key numbers

  • 96.3% average text coverage for MOSS-Image on its long-form text benchmark, the number that separates "can render a word" from "can render a paragraph."
  • 66.84 to 53.01 gFD for JiT models trained with a language-informed loss instead of pixel loss, a 14-point move with no architecture change.
  • 462x faster per edit for FlashReFocus, which turns a multi-minute edit into roughly half a second.
  • 5B parameters for the MOSS-Image backbone, about 10GB in bf16, so single-GPU inference without model parallelism.
  • 89 image pairs across three cameras in the 3CReal benchmark built to test zero-shot refocusing on unseen hardware.

The high-frequency deficit nobody puts in the abstract ​

Here's the finding I keep coming back to. Vanilla pixel diffusion transformers, regardless of which parameterization you pick, retain a persistent high-frequency deficit relative to real images. The generated images are missing fine-grained energy in the spectrum. They're smooth in a way that photons aren't.

FID won't tell you this. FID pushes both distributions through Inception features, which are tuned for semantic content, so a model can post a strong FID while its 100% crops look like they were rendered on slightly damp paper. If you've shipped a diffusion model and gotten the "it looks AI-generated but I can't say why" feedback, this is likely the why.

ConvRefiner is a small convolutional module trained after the fact with an adversarial loss while the diffusion backbone stays frozen. It restores the high-frequency statistics, and the paper reports FID is preserved rather than sacrificed. The training cost is small because you're not touching the backbone.

I like this result because it's cheap and it's honest about where the cost sits. Pixel-space models have a known texture gap; a discriminator is a reasonable tool for closing a gap that a reconstruction loss can't see.

Language-informed supervision instead of pixel loss ​

Pixel loss tells the model nothing about meaning. Two images with identical per-pixel error can be a legible street sign and a smear of glyph soup. LILIE attacks this by training an invertible image transform, a normalizing flow transformer, jointly with a vision-language model through next-token prediction, such that the VLM can recover captions from noised features. Then that encoder becomes the loss function for a text-to-image model.

The invertibility is the whole trick. A VAE encoder or a CLIP encoder would throw information away, and you'd be back to supervising through a lossy projection. A normalizing flow is lossless, so the supervision is semantic without introducing a new bottleneck.

Training JiT models with the LILIE loss on GPIC dropped gFD from 66.84 to 53.01. That's a 14-point improvement from the training signal alone, with the same architecture and the same data. Most teams would spend 4x the GPU hours chasing that.

The obvious cost: you now run an invertible encoder and a VLM forward pass on every training step, on noised features. That's real overhead, and it's the reason this approach lands in research code before it lands in production pipelines.

MOSS-Image and the uncomfortable question of what you're comparing ​

MOSS-Image is a 5B single-stream diffusion transformer trained from scratch on a recipe built for bilingual long-form text rendering and precise layout control. Around it sits an intention-aware agentic prompt enhancement module that turns a vague user request into a detailed content and layout specification. Backbone plus agent, that's the system.

The numbers: 96.3% average text coverage on LongLongTextBench, a benchmark the authors introduce for English and Chinese long-form rendering. State of the art among open-weight models on that benchmark. On RealLayoutArena, an end-to-end layout-rich generation benchmark, the full system ranks second. The backbone alone, evaluated with shared enhanced prompts, ranks first among generators.

That rank flip is the most useful thing in the report. When the backbone is evaluated with identical enhanced prompts it wins. When the whole system is evaluated end to end, it doesn't. The agent is doing measurable work, and the parts of the gap that aren't the model are the parts most comparisons quietly ignore.

If you're benchmarking text-to-image systems, ask what's in the number. A prompt enhancer, a best-of-N sampler, and a reranker can all live inside "our model," and none of them will transfer when you swap the backbone.

Five billion parameters also means about 10GB of weights in bf16. That fits on a single 24GB card for inference at moderate resolution, which is the difference between an open-weight model you can actually experiment with and one you can only read about.

FlashReFocus: split the problem, then stop re-running diffusion ​

Refocusing a photo after capture is interactive by nature. You poke at different focal points and blur strengths, dozens of times, on one image. Generative approaches rerun a diffusion model for every poke, which is the wrong shape of compute for the interaction.

FlashReFocus observes that refocusing is two problems, not one. Recovering the all-in-focus image is ill-posed and needs a generative prior. Rendering blur is mostly determined by depth and focal plane, which is geometry. So run diffusion once per photo with a one-step deblurrer and a decoder refiner, then render every edit with a lightweight network conditioned on an explicit defocus map. Defocus-bias attention lets pixels at similar defocus share information without interference from pixels at other depths. In-focus content stays untouched, and editing the map gives you tilt-shift or two focal planes.

MethodDiffusion passes per photoFirst refocusEach further edit
GenRefocusOne per editBaselineBaseline
FlashReFocusOne total148.3x faster than GenRefocus~0.5 s, 462x faster

Half a second per edit is the threshold where a slider feels like a slider instead of a submit button. The paper also reports the best perceptual fidelity for zero-shot bokeh rendering on EBB! and on 3CReal, a 89-pair, three-camera benchmark built to test generalization to unseen cameras. Zero-shot matters here because depth-of-field behavior is camera-specific, and a method that only works on the training camera is a demo, not a tool.

The general lesson generalizes past refocusing: if a sub-problem is well-posed, don't spend a diffusion model on it.

DiffPT: diffusion as pretraining, and where it actually wins ​

Before AlexNet, people pretrained with autoencoders and denoising autoencoders, largely because their datasets were tiny. DiffPT revives that idea with a modern diffusion objective and asks whether it beats the self-supervised methods we reach for by default.

The answer is modality-dependent, and the paper says so plainly. Contrastive pretraining wins on CIFAR-10 and Tiny ImageNet in low-data regimes. DiffPT wins on medical imaging and time series across a range of low-data settings.

SettingBetter pretrainingWhat I'd do
CIFAR-10, low dataContrastive SSLUse the SSL baseline, don't overthink it
Tiny ImageNet, low dataContrastive SSLSame
Medical imaging, low dataDiffPTTry diffusion pretraining first
Time series, low dataDiffPTTry diffusion pretraining first

The pattern makes sense. Contrastive objectives need semantically distinct samples to push apart and pull together. A dataset of 2,000 CT slices has plenty of variation but little semantic diversity, and augmentation invariances that hold for photos often don't hold for scans. Diffusion pretraining just learns the data distribution, which is a weaker assumption.

If you're working with a small domain-specific dataset, run both. Pretraining is cheap relative to the fine-tuning that follows, and the default choice isn't the right one as often as people assume.

Common Pitfalls ​

  1. Inheriting pixel-space hyperparameters across model scales. MFlow shows the preferred parameterization shifts as capacity grows. If you take a recipe tuned on a small DiT and apply it to a 3B pixel model, you're running a config that was never evaluated at that size, and the loss curve won't tell you.

  2. Judging a pixel-space generator by FID alone. The high-frequency deficit survives every parameterization choice and barely moves FID. Crop at 100%, look at the spectrum, or run a small paired comparison. A single number from Inception feature space hides exactly the artifact your users notice.

  3. Rerunning the whole pipeline for every interactive edit. If the user is going to move a slider 40 times, cache the expensive generative step and make the per-edit path a small network. FlashReFocus got 462x from this split, and most editing UIs leave it on the table.

  4. Comparing a system when you mean a backbone. The MOSS-Image ranking flips depending on whether the prompt-enhancement agent is inside the number. Before you conclude a model is better, check whether the eval included an enhancer, a reranker, or best-of-N sampling.

  5. Defaulting to contrastive SSL on small domain datasets. DiffPT loses on natural-image benchmarks and wins on medical imaging and time series. The modality, not the dataset size, is the deciding factor.

One thing to remember ​

Every paper here measures something that a standard FID number would have missed: a spectral deficit, a rank flip caused by an agent, a 462x latency gap hiding behind equal output quality, a gFD move that never shows up in pixel error. The metric you optimize is narrower than the thing you ship.

The Bottom Line ​

If you're training a pixel-space DiT above roughly a billion parameters, stop inheriting the clean-image prediction default. MFlow's result is that parameterization preference shifts with capacity, and larger models want more velocity prediction. Re-run the comparison at your actual scale, it's one of the cheapest experiments you'll do all quarter.

If your FID is fine but your crops look plasticky, bolt on a post-hoc adversarial refiner instead of retraining. ConvRefiner trains against a frozen backbone with a small amount of extra compute and restores high-frequency statistics without moving FID. Retraining the base model to fix a texture gap is a much more expensive way to solve the same problem.

If you're building anything interactive, separate the ill-posed step from the cheap step. Restore once with diffusion, render every edit with a light network conditioned on explicit geometry. That split is worth 462x, and it applies well beyond depth of field.

Sources ​