arxiv:2610.36192

Beyond the Estimate: Wiener-Informed Decoding for Easier Training of Pixel-Space Flow Models

Published on Oct 26

· Submitted by

Ash Wang on Oct 28

National University of Singapore

Upvote

Authors:

Ziqiao Wang ,

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.

Generated by thinkingmachines/Inkling-Small

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.

View arXiv page View PDF GitHub 9 Add to collection

Community

ziqiaow

Paper author Paper submitter about 22 hours ago

•

edited about 22 hours ago

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.