Beyond the Estimate: Wiener-Informed Decoding for Easier Training of Pixel-Space Flow Models
Abstract
Wiener-informed Decoding accelerates pixel-space velocity model training and improves image generation by integrating statistical predictions into decoder computation without altering native blocks.
Learning velocity fields directly in pixel space is difficult and slow. We show that coarse second-order image statistics can ease this task: frozen Wiener estimates improve velocity learning, with gains extending beyond the linear predictions they explicitly provide. Moreover, we observe similar but complementary gains from frozen bases and pixel decoding. These observations motivate Wiener-informed Decoding (WiD), which uses statistical knowledge both for direct prediction and to organize decoder computation without modifying native blocks. WiD combines a trainable Wiener prediction path with coefficient interfaces that inject features into the native decoder and return its predicted coefficient corrections, with these components initialized with data statistics. All introduced parameters, including the spectrum and coordinate transforms, are trained jointly under the original pixel-space velocity objective. Across model scales and decoder architectures, WiD accelerates training and improves generation quality with modest parameter overhead. It also yields smoother gradient trajectories and remains effective with representation alignment. WiD improves class-conditional and text-to-image generation across different resolutions. On ImageNet 256 x 256, PixelDiT-XL with WiD reaches FID 1.53 after 500 epochs, achieving comparable quality to PixelDiT with 37.5% fewer training epochs. Statistical denoisers contribute beyond the estimate: their computational structure can also ease pixel-space training.
Community
Reading papers like JiT and PixelDiT got me thinking about why learning velocity directly in pixel space is so difficult. Some approaches tackle this by choosing an easier prediction target, such as x-prediction; others design architectures that separate semantic modeling from dense pixel prediction.
At first, my question was quite simple: what if we could first make a coarse estimate of the velocity from the noisy sample, and then let the model learn the remaining correction? Could that make the learning problem easier?
Inspired by Wiener filtering, we can construct an analytic estimator that gives us such a velocity prediction using the second-order statistics from data. In our experiments, we observed that simply adding this fixed prediction and letting the model learn the residual already led to improvements. Its benefits were similar to those of pixel decoders, but the two were not entirely redundant.
A frozen statistical estimator has limited flexibility, though. And when added to an architecture that already has a pixel decoder, its extra benefits tend to diminish as training progresses. So I wondered: why not make the estimator part of the architecture itself?
The pixel decoder turns out to provide a natural interface for this. Rather than using only the estimator’s prediction, we also use the way it organizes computation. We introduce modules and pathways within the decoder that interact with its existing structure, with the new parameters trained jointly with the model, and use the statistics to initialize these pathways.
We call this Wiener-informed Decoding (WiD). We’ve found it effective for both class-conditional and text-to-image generation, and across different resolutions. In our experiments, we estimate the statistics from around 50,000 images, which takes about only one minute with GPU acceleration. The method also adds relatively few parameters.
Statistical denoisers have a long history, but I think they are still able to inspire us in modern denoising models. In this first step of our exploration, we’ve demonstrated the idea in architectures with pixel decoders. In the future, we’d like to explore how it might apply to other architectures and generation tasks.
Get this paper in your agent:
hf papers read 2610.36192 Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash Models citing this paper 1
Datasets citing this paper 0
No dataset linking this paper
Spaces citing this paper 0
No Space linking this paper
Collections including this paper 0
No Collection including this paper