PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ziqiao Wang, Yang You

arXiv:2610.36192v1cs.CVcs.AI

TL;DR

Pixel-space flow models avoid a separate pretrained autoencoder, but learning velocity fields from noisy, high-dimensional pixels remains slow and challenging. The paper introduces WiD, combining a trainable Wiener prediction path with coefficient interfaces in the pixel decoder. Across model scales and decoder architectures, WiD improves training efficiency and stability and generation quality with modest overhead.

  • Problem

    Pixel-space flow models remove the need for a separate pretrained autoencoder, but learning natural image structure from noisy, high-dimensional pixels remains slow and challenging.

    Source: §1 · p. 1

  • Method

    WiD combines a trainable statistical prediction path with coefficient interfaces that connect it to the native pixel decoder.

    Source: §1 · p. 3

  • Results

    WiD improves training efficiency and stability while improving generation quality across model scales and decoder architectures with modest parameter overhead.

    Source: §1 · p. 3, §7 · p. 10

  • Takeaways & Limitations

    The findings suggest statistical denoisers can contribute beyond their estimates through computational structure that eases pixel-space training.

    Source: §7 · p. 10

    The work mainly focuses on pixel-space image generation; integrating statistical denoisers into other generative tasks remains future exploration.

    Source: C LIMITATIONS AND FUTURE WORK · p. 19

Abstract

from arXiv · show

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.

1 INTRODUCTION

Pixel-space flow models avoid a separate pretrained autoencoder but still learn natural image structure from noisy, high-dimensional pixels, making optimization slow and challenging. This section motivates Wiener-informed Decoding (WiD): coarse statistical estimates ease velocity learning, and their computation can be integrated into pixel decoding; WiD improves efficiency, stability, and generation quality across decoder architectures.

1 INTRODUCTION

  • Pixel-space flow models avoid a separate pretrained autoencoder, but learning natural image structure directly from noisy, high-dimensional pixels makes optimization slow and challenging.Recent approaches address the challenge through prediction-space parameterizations or architectures that separate contextual modeling from dense pixel prediction.Source: §1 · p. 1
  • Natural-image spatial and color correlations differ from isotropic added Gaussian noise, yielding a closed-form linear Wiener denoiser and motivating its use to ease velocity learning.The estimator is a second-order approximation and cannot capture the full complexity of natural images.Source: §1 · p. 1
  • Frozen DCT and patch-PCA Wiener bases rescue the struggling plain DiT and bring generation quality into the range of contemporary approaches.The bases supply an analytic velocity estimate, while the DiT learns a residual correction under the velocity objective.Source: §1 · pp. 1–2
  • Wiener-basis gains extend beyond the fitted linear response, motivating use of the estimator’s computation to organize denoiser computation.For DiT, diagnostics show much of the velocity-prediction error reduction lies beyond the fitted linear component.Source: §1 · p. 2
  • Pixel decoding and frozen Wiener bases provide similar but complementary gains, and input-perturbation tests show closer Wiener alignment in the direct pixel pathway than in the semantic pathway.Pixel decoders map noisy pixels to velocity, the same input-output task addressed by the Wiener estimator, while retaining the semantic backbone.Source: §1 · p. 2
  • WiD initializes a direct path from patch-PCA statistics and adds noisy-patch coordinates to decoder embeddings, with the decoder predicting coefficient-space corrections alongside pixel-space velocity.The Wiener-branch and pixel-head predictions are summed and supervised by the original objective; introduced parameters are trained jointly.Source: §1 · p. 2
  • WiD combines a trainable statistical prediction path with coefficient-based feature processing and output correction, improving stability, efficiency, and generation quality across pixel-decoder architectures with modest overhead.The gains persist further into training than frozen Wiener bases, and WiD is validated for class-conditional and text-to-image generation across resolutions.Source: §1 · pp. 2–3

2 PRELIMINARIES

The preliminaries derive an analytic Wiener velocity estimate from a Gaussian approximation to the data and describe the shared interface used by pixel-decoder architectures. For natural images, the estimate is a second-order statistical approximation; pixel decoders map noisy pixels and contextual features to velocity.

2.1 LINEAR WIENER ESTIMATION FOR FLOW MATCHING

  • For linear flow matching, the observation interpolates between data and Gaussian noise, and the velocity is their difference.The time parameter t runs from noise to data.Source: §2.1 · p. 3
  • Under the Gaussian prior, the analytic Wiener velocity estimate is exact; for natural images, it is a second-order statistical approximation to conditional velocity.The practical linear estimator omits an additive offset and uses the uncentered second moment; with exact second moments, its linear map minimizes expected squared velocity error among zero-offset linear estimators.Source: §2.1 · p. 3

2.2 PIXEL DECODER INTERFACES

  • The pixel-decoder interface computes contextual features from noisy pixels, time, and conditioning, then maps noisy pixels and those features to a velocity prediction.The contextual backbone is B, its conditioning features are s, and D denotes the pixel decoder.Source: §2.2 · p. 3

3 LEARNING WITH ANALYTIC WIENER BASES

This section tests how structured Wiener references affect pixel-space velocity prediction and whether their benefits extend beyond fitted linear estimates. It finds residual improvements and complementary gains with pixel decoding, motivating statistical integration into the decoder.

3.3 PIXEL DECODING AND STATISTICAL PREDICTION

  • Adding a frozen PCA base to PixelDiT further reduces both fitted linear and remaining error, with these additional gains narrowing over training.Figure 2b compares PixelDiT and PixelDiT+PCA using patch-affine fits.Source: Figure 2 · p. 4, §3.3 · p. 5
  • The direct pixel pathway is more aligned with Wiener responses than the contextual pathway, and this alignment gap widens over training.The JVP comparison removes each response’s component parallel to the perturbation and averages cosine similarities across held-out images and time points.Source: Figure 2 · p. 4, §3.3 · p. 5

3.1 STRUCTURED STATISTICAL REFERENCES

  • For structured references, 51,200 images supply uncentered second moments; eigendecomposition defines U and Λ, and Figure 2 compares frozen PCA/DCT bases.The statistics use the same center-crop and scaling as training preprocessing, without other augmentation.Source: §3.1 · pp. 3–4, Figure 2 · p. 4

3.2 BEYOND THE ADDITIVE LINEAR PREDICTION

  • After removing fitted affine responses, the remaining prediction contributes most of the Wiener-base improvement, while linear-error reduction shrinks with training.The comparison uses the official PixelDiT-XL checkpoint at 320 epochs as reference; Figure 2a reports reductions averaged over five time points.Source: §3.2 · p. 4, Figure 2 · p. 4

3.3 PIXEL DECODING AND STATISTICAL PREDICTION

  • The complementary gains and pathway alignment motivate integrating the estimator’s prediction path and coordinate transforms into pixel decoding as trainable components.The authors aim for gains that persist as learning proceeds.Source: §3.3 · p. 5

4 WIENER-INFORMED DECODING

WiD integrates a trainable Wiener prediction path with coefficient-based interactions in the native pixel decoder. It starts from a patch PCA Wiener estimate and improves training stability and generation quality across decoder architectures, with gains also present alongside REPA.

4.1 ARCHITECTURE

  • WiD combines a trainable Wiener prediction path with coefficient features for the native decoder and coefficient corrections returned by that decoder.The statistical prediction and native pixel prediction are supervised together under the original velocity objective.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 INITIALIZATION

  • At initialization, the total prediction equals the patch PCA Wiener estimate while the native decoder’s initial input features and block initialization are preserved.The analysis and synthesis transforms are trained independently, without orthogonality or inverse constraints.Source: §4.2 · p. 6

4.3 TRAINING OBSERVATIONS

  • WiD improves generation quality with REPA in PixelDiT and over native DeCo under DeCo’s training recipe, while producing smoother gradient trajectories with fewer large excursions.The comparisons use B-size models trained for 40 epochs on ImageNet 256×256.Source: §4.3 · pp. 6–7, Figure 4 · p. 6

5 EXPERIMENTS

The experiments evaluate WiD on class-conditional and text-to-image generation, then ablate its prediction path, coefficient interfaces, and statistical initialization. WiD improves generation quality and reaches comparable ImageNet quality with fewer training epochs.

5.1 CLASS-TO-IMAGE GENERATION

  • FID 1.53 after 500 epochs lets PixelDiT-XL+WiD save 37.5% of PixelDiT’s training epochs; at 512×512, it reaches FID 1.64 after 50 fine-tuning epochs, outperforming native PixelDiT-XL/16 at 320+530 epochs.The 256×256 results also report FID 2.09 at 80 epochs versus 2.36 for native PixelDiT at the same duration, and FID 1.61 at 260 epochs matching native PixelDiT-XL at 320 epochs.Source: §5.1 · p. 8, Table 2 · p. 8
  • WiD’s PixelDiT improvement requires around 1% more parameters, increasing model size from 797M to 805M.Together with the B-size comparisons, the authors report consistent improvement and acceleration across model scales.Source: §5.1 · p. 8

5.2 TEXT-TO-IMAGE GENERATION

  • At 80K pretraining steps, WiD improves DeCo-XXL/16 on both GenEval and DPG-Bench; after full training, it improves generation across resolutions with 14M additional parameters, about 1.3% of the 1.1B-parameter baseline.The 512×512 comparison is reported in Table 4, where † marks the authors’ test of the public DeCo checkpoint using its own evaluation protocol.Source: Table 3 · p. 8, §5.2 · p. 9, Table 4 · p. 9

5.3 ABLATIONS

  • Removing the Wiener prediction path worsens FID from 41.71 to 49.54, while removing coefficient interfaces raises FID from 41.71 to 50.67.The ablations also report that random initialization of Uenc and Udec and setting all spectral values λ to 1 support the benefit of statistical initialization; removing Wiener gain still outperforms native PixelDiT-B/16.Source: §5.3 · p. 9

6 RELATED WORK

Prior work uses Wiener statistics in restoration and to study diffusion denoisers, while pixel-space flow methods address learning difficulty through prediction parameterizations or decoder designs. WiD applies statistical structure to velocity prediction and decoder computation.

6.1 STATISTICAL DENOISING

  • Wiener-based restoration methods include BLS-GSM and NL-Bayes, while BM3D organizes computation through coefficient representations, empirical Wiener filtering, and inverse transforms.Deep Wiener Deconvolution integrates Wiener operations into a learnable feature-processing pipeline for deblurring.Source: §6.1 · pp. 9–10
  • Diffusion studies use statistical estimates to analyze learned denoisers: linear fits reveal Gaussian structure, Gaussian score approximations capture moderate- and high-noise behavior, and data statistics can explain locality patterns.These connections motivate exploring statistical priors in generative learning.Source: §6.1 · p. 10
  • WiD uses statistics to initialize a prediction path and integrates statistical operations and coordinates into the decoder; ablations report lower generation quality when its prediction, coefficient interface, or statistical initialization is removed.Source: Table 5 · p. 9, §6.1 · p. 10

6.2 PIXEL-SPACE FLOW MODELS

  • Pixel-space flow methods also change prediction parameterization: JiT predicts clean images, AsymFlow combines full-rank data and low-rank noise prediction, and k-Diff learns a scalar controlling the predicted data–noise mixture.WiD addresses the learning difficulty while retaining velocity prediction.Source: §6.2 · p. 10
  • Pixel-decoding approaches separate contextual modeling from dense pixel prediction, and WiD adds statistical interfaces around the existing PixelDiT and DeCo decoder computation.Source: §6.2 · p. 10
  • EDM and Preconditioned Flow Matching use normalization or reversible transformations to address conditioning, whereas WiD focuses on decoder design with a trainable Wiener path interacting with the native decoder.Source: §6.2 · p. 10

7 CONCLUSION

WiD combines a trainable statistical prediction path with coefficient interfaces connecting it to the native pixel decoder, improving training and generation across scales and decoder architectures. The appendix describes how these interfaces preserve native decoder behavior at initialization and output.

7 CONCLUSION

  • WiD combines a trainable statistical prediction path with coefficient interfaces to the native pixel decoder, improving training efficiency and stability and generation quality across model scales and decoder architectures.The reported applications include class-conditional and text-to-image generation with modest parameter overhead.Source: §7 · p. 10

A.1 DECODER STRUCTURES IN PIXELDIT AND DECO

  • PixelDiT and DeCo decoders both process noisy input and backbone context to produce velocity, using pixel-wise modulation before projection; PixelDiT uses MHSA and FFN blocks, while DeCo mainly uses MLP layers.Source: §A.1 · p. 15, Figure 5 · p. 15
  • For DeCo, WiD’s initialized injected coefficients make the summed features match the original embedding before native decoder blocks, while preserving native positional encoding.Source: §A.1 · p. 15
  • Both decoders zero-initialize coefficient-correction and native pixel-output projections; output normalization follows each native decoder, using LayerNorm for DeCo and RMSNorm for PixelDiT’s coefficient head.Source: §A.1 · p. 15

A.2 ADDITIONAL IMPLEMENTATION DETAILS

The appendix provides model configuration references, native training and inference settings, and reproducibility procedures for the experiments.

A.2 ADDITIONAL IMPLEMENTATION DETAILS

  • Table 6 lists settings for five model configurations, with parameter counts excluding WiD, REPA projectors, and external encoders.Source: Table 6 · p. 16
  • Frozen Wiener bases and WiD use the native denoising networks’ original training and inference settings; the DiT-B/16 baseline uses PixelDiT-B’s 12-block backbone and a conditioned patch-wise velocity head.PixelDiT-XL/16+WiD uses timeshift 1.0 on ImageNet 256×256 and 3.0 for training and 2.0 for sampling on ImageNet 512×512.Source: §A.2 · p. 16
  • The authors report fixing random seeds for model initialization and patch-statistics sampling, saving and restoring each rank’s RNG state, and releasing code, checkpoints, logs, and statistical priors.The reported runs include PixelDiT-XL/16+WiD and DeCo-XXL/16+WiD on 8 NVIDIA H200 GPUs, with the last 180 PixelDiT-XL epochs on 8 NVIDIA B200 GPUs.Source: §A.2 · p. 16

A.3 ADDITIONAL EVALUATION DETAILS

Additional evaluations report results under native sampling and evaluation protocols, along with WiD’s parameter, computational, and inference-throughput overhead. The accompanying diagnostics section focuses on frozen Wiener bases and native pixel decoding before adaptive WiD.

A.3 ADDITIONAL EVALUATION DETAILS

  • At 160 and 200 training epochs, WiD outperforms native PixelDiT-XL/16, which reports FIDs of 1.97 and 1.94, respectively.The comparison follows each original backbone’s native sampling and evaluation protocol.Source: §A.3 · p. 17
  • WiD raises PixelDiT-XL/16 parameters from 797.38M to 804.86M and analytical forward cost from 311.19 to approximately 315.32 GFLOPs.These correspond to increases of 0.94% and 1.33%, respectively.Source: §A.3 · p. 17
  • At batch size four, WiD reduces PixelDiT-XL forward throughput from 238.8 to 227.6 samples/s and reproduced DeCo-XXL/16 throughput from 260.2 to 245.1 samples/s.Measurements are denoiser-only inference on one NVIDIA H200 using EMA weights and BF16.Source: §A.3 · p. 17
  • With a 25-step Adams-2nd sampler and CFG 4.0, DeCo-XXL/16+WiD generates 7.04 images/s versus 7.50 images/s for the baseline.This is batch-four text-to-image generation throughput.Source: §A.3 · p. 17

B ADDITIONAL DIAGNOSTIC DETAILS

  • The additional diagnostics examine frozen Wiener bases and native pixel decoding before adaptive WiD.The 256×256 ImageNet evaluation samples 50K images with FlowDPMSolver-100 and the official PixelDiT-XL/16 seed 1600.Source: §A.3 · p. 17, B ADDITIONAL DIAGNOSTIC DETAILS · p. 18

B.1 COMPARISON OF PREDICTION ERRORS

Frozen Wiener bases and native pixel decoding improve remaining velocity-prediction errors beyond their fitted affine responses. Pathway comparisons generally find stronger Wiener alignment in the direct pixel pathway, with an exception at the lowest evaluated noise time.

B.1 COMPARISON OF PREDICTION ERRORS

  • Diagnostics fit affine maps to each model’s full velocity output and subtract fitted responses to assess remaining predictions against a well-trained reference.PCA uses a shared patch-affine fit; DCT uses independent RGB-affine fits at each frequency, and the fitted families include their corresponding frozen bases.Source: §B.1 · p. 18
  • Both frozen bases and native pixel decoding reduce remaining and total prediction errors relative to plain DiT, even when fitted linear errors are higher.At 40 epochs and t = 0.9, DiT+DCT and native PixelDiT have higher linear errors but lower remaining and total errors; Table 8 reports remaining-error reductions over time.Source: Table 8 · p. 18, §B.1 · p. 19
  • Adding a frozen PCA base to native PixelDiT yields positive remaining-error gains at all five evaluated times through 40 epochs.The gains shrink overall with training, consistent with decreasing FID gains in Figure 1c.Source: §B.1 · p. 19

B.2 PATHWAY RESPONSES TO INPUT PERTURBATIONS

  • Pathway tests use velocity-output Jacobian-vector products while holding contextual features or the decoder’s direct noisy input fixed, then compare alignment after removing the isotropic perturbation component.Cosine similarities are computed per image, with the reported difference favoring the pixel pathway when positive.Source: §B.2 · p. 19
  • Averaged over all five evaluated times, the direct pixel pathway has higher direction-removed cosine than the contextual pathway for each Wiener reference; at t = 0.1, contextual alignment is generally closer.At t = 0.1, contextual alignment is closer to Wiener bases for DCT at every checkpoint and for PCA at 10–30 epochs.Source: §B.2 · p. 19, Table 9 · p. 19

C LIMITATIONS AND FUTURE WORK

The paper’s stated focus is easing velocity learning for pixel-space image generation, while identifying integration into other generative tasks and architectures as future work. Additional visualizations show class-conditional samples and text-to-image results at 256×256 and 512×512.

C LIMITATIONS AND FUTURE WORK

  • The work mainly focuses on easing velocity learning for pixel-space image generation; latent, language, audio, and video generation are proposed as directions for exploration.The authors also plan to investigate Wiener-informed structures in diffusion transformers and develop a flow-matching foundation architecture.Source: C LIMITATIONS AND FUTURE WORK · p. 19

D ADDITIONAL VISUALIZATIONS

  • Additional visualizations include 256×256 PixelDiT-XL/16+WiD class samples and 512×512 PixelDiT-XL and DeCo-XXL/16+WiD generations.The PixelDiT samples use classifier-free guidance; the 512×512 PixelDiT samples use 260 + 50 epochs.Source: Figure 6 · p. 20, Figure 7 · p. 21, Figure 8 · p. 22
Loading 2610.36192v1…