WaiT for the Signal: Simple Frequency-Aware Flow-Matching

AI summary

Overview

• Research area: Generative computer vision, specifically pixel-space flow-matching diffusion models for high-resolution image and video synthesis.
• Technical level: Advanced. The paper assumes familiarity with flow matching, diffusion sampling, wavelet transforms, and generative evaluation metrics (FID, FVD, CLIPScore, GenEval, DPG).
• Scope: The authors introduce WaiT (Wavelet-aware image Transformer), a modification to the noise schedule of a pixel-space JiT model that delays denoising of high-frequency wavelet bands until coarse structure has emerged, plus a three-axis evaluation protocol for native-resolution image quality.

What This Paper Is About

Standard flow-matching image generators treat all spatial frequencies the same, even though high-frequency image detail is swamped by noise much earlier than coarse structure. This wastes model capacity and compute on frequency bands that carry no learnable signal during the early part of the sampling trajectory. The authors fix this by splitting each image into wavelet bands and giving the high-frequency bands a delayed noise schedule so they "wait" until coarse structure has formed, then join a joint refinement stage.

Key Contributions

1. A frequency-aware noise schedule. WaiT decouples the flow-matching time variables for the low-frequency (LF) and high-frequency (HF) wavelet bands, with the HF band following the accelerated schedule t_HF = max(0, (t_LF − t*)/(1 − t*)). High frequencies stay as pure unit-variance noise until the crossover point t*, at which fresh noise is injected that exactly matches the training distribution, avoiding the train–test mismatch of discontinuous cascades.

2. A three-axis evaluation protocol. Because standard FID downsamples to 299×299 and discards fine detail, the authors add 5-crop FID (5cFID) over five native-resolution crops (four corners plus center) for local detail, and high-frequency Fréchet Wavelet Distance (hFWD) with the DC wavelet packet excluded for pure texture fidelity. They validate the protocol against ~1.13M human preference judgments on the PIPAL benchmark (Pearson |PLCC| up to 0.71, Spearman |SRCC| up to 0.66, Holm–Bonferroni-corrected p < 4×10⁻⁴).

3. Pareto-optimal results on ImageNet 512×512 and OpenImages. WaiT is Pareto-optimal across all three axes among pixel-space models at both OpenImages 512×512 and 1024×1024, and reduces sampling compute by up to 50%.

4. Transfer of the same recipe to text-to-image and video. No algorithmic modification is needed for temporal wavelets: the method reaches a state-of-the-art FVD of 0.84 on Kinetics-600, and improves over JiT on almost all text-to-image metrics with up to 3× higher throughput.

Main Findings

• ImageNet 512×512 headline numbers: WaiT-H/16 achieves FID 1.43, 5cFID 1.63, and hFWD 0.67 at 397 GFLOPs. The largest 2B-parameter model, WaiT-G/16, sets a new state-of-the-art pixel-space FID of 1.30, with 5cFID 1.45 and hFWD 0.59 at 822 GFLOPs. The paper reports improving FID by 23% and 5-crop FID by 56% over compute-matched JiT baselines.

• Compute savings: Against the JiT-H/16 baseline (810 GFLOPs), WaiT-H/16 uses 397 GFLOPs while also improving FID (1.43 vs. 1.70). Video results show roughly 22% compute reduction on Taichi-HD and about 30% on Kinetics-600.

• The frequency asymmetry is real and measurable: Mutual-information analysis shows fine (HF) wavelet bands lose mutual information with clean images far earlier than coarse (LF) bands, and the same asymmetry persists in the backward process of a trained JiT model. The crossover t* = 0.25 marks where fine-band mutual information drops below 0.01 nats.

• Cumulative ablation on ImageNet 256×256 with WaiT-B/16: Starting at FID 7.51 for a naïve two-stage JiT, adding the delayed linear schedule gives 5.43 (−2.08), global coarse training over [0,1] gives 3.81 (−1.62), and the optimal t* = 0.25 gives 3.57 (−0.24). The final model halves FID while saving 50% GFLOPs.

• Texture fidelity is where WaiT is strongest: On hFWD it outperforms both state-of-the-art pixel- and latent-space methods. On OpenImages 1024, JiT's hFWD plateaus around 0.5 already at B/32, while WaiT pushes hFWD well below 0.3 and matches JiT's 5-crop FID at roughly 2× lower compute.

• Text-to-image at 1024 resolution: With the SigLIP + Qwen3-VL + T5-v1.1-XXL pipeline, WaiT-H/32 scores FID 5.15, 5cFID 8.40, hFWD 3.30, CLIP 26.75, GenEval 0.491, DPG 0.823 at 0.29 imgs/s versus JiT-H/32 at FID 5.35, 5cFID 8.90, hFWD 3.40, CLIP 26.71, GenEval 0.442, DPG 0.824 at 0.10 imgs/s. With MetaCLIP + Llama 3.2 Vision + Llama 3.1, WaiT reaches FID 4.78, 5cFID 6.53, hFWD 2.31, CLIP 39.43, GenEval 0.460, DPG 0.805 at 0.28 imgs/s versus JiT at FID 5.24, 5cFID 6.62, hFWD 2.31, CLIP 39.10, GenEval 0.419, DPG 0.798 at 0.12 imgs/s.

• Video: On Kinetics-600 at 128×128 (5→16 frames), WaiT-XL/8 reaches FVD 0.84 at 1,110 ×10² GFLOPs, versus JiT-XL/8 at 0.89 at 1,580 ×10² GFLOPs and Unified Latents (M) at 1.3 at 700 ×10² GFLOPs. At 256×256, WaiT-XL/32 reaches FVD 1.20 at 319 versus JiT-XL/32 at 1.32 at 371. On Taichi-HD (128×128), WaiT-B/8 improves unconditional FVD from 28.92 to 27.13 and conditional FVD from 19.92 to 19.75 (LPIPS 0.292 vs. 0.297).

• Multi-level transitions save more: On OpenImages-1M at 1024×1024 (L/32), moving from 2 levels (t* = 0.25: FID 9.74, 5cFID 21.71, hFWD 0.918, 18.5% savings) to 3 levels (t* = {0.15, 0.3}: FID 10.11, 5cFID 16.03, hFWD 0.867, 33.0% savings) improves both high-frequency metrics at a marginal FID cost.

• The recipe also works in latent space: Applied unchanged to DDT/XL-2 on ImageNet 512×512, WaiT+DDT trained from scratch for only 500k steps reaches FID 1.33, 5cFID 1.77, hFWD 1.20 at 329 GFLOPs, versus the DDT baseline's FID 1.28, 5cFID 1.98, hFWD 1.16 at 525 GFLOPs (1.28M + 500k training steps).

• Minimal architectural footprint: The only architectural addition is a resolution embedding (sinusoidal embedding plus MLP) added to the existing time and class conditioning, following PixelFlow. The wavelet transform is lossless, parameter-free, and requires no training. Haar wavelets are used throughout; alternative wavelet families are ablated in Section D.1 and the method is reported robust to this choice.

Methodology in Plain English

The authors start from the observation that images have a natural frequency hierarchy: coarse shapes carry meaningful signal very early in a denoising trajectory, while fine textures are indistinguishable from noise. Instead of building a complicated pyramid architecture, they make one change — the noise schedule.

They apply a single-level Discrete Wavelet Transform (DWT) to each clean image, splitting it into a low-frequency (LF) approximation band of size D/4 and three high-frequency (HF) detail bands totaling 3D/4. The LF band is normalized by a scalar S_LF (the 95th percentile of absolute LF coefficients over the training set, e.g. 1.94 for ImageNet 256); HF bands are left unnormalized to preserve their natural sparsity.

During training, the LF band follows the standard flow-matching schedule over [0,1], while the HF bands follow a lagged schedule that keeps them at pure noise until t*. Both phases are supervised with a v-loss reweighted band-wise against wavelet targets rather than raw pixels.

During sampling, generation happens in two phases. Phase 0 runs the ODE for t from 0 to t*, operating only on the LF band, which means 4× fewer tokens (for a single-level DWT) and therefore major compute savings. At t*, the LF band is un-normalized, fresh HF noise is injected at the exact variance the schedule prescribes, the bands are recombined via the inverse DWT, and Phase 1 continues the ODE for t from t* to 1 with joint refinement. Because the injected noise distribution matches the training distribution precisely at t*, there is no train–test mismatch.

Step allocation is controlled by two knobs: a timestep shift α that warps the uniform t-grid to concentrate steps near t = 0, and a Phase 0 multiplier m that scales how many steps fall below t*. Larger n in Phase 0 means more steps at the 4× cheaper token count.

For their evaluation, they use ImageNet-1k at 256 and 512, plus natively curated 512 and 1024 resolution subsets of OpenImages V6 — because only 5.3% of ImageNet images have a shortest side ≥ 512, upsampling would corrupt the exact high-frequency statistics they want to measure.

Why This Matters

• Research impact: The paper argues that a single noise-schedule change, with no architectural redesign, can produce Pareto-optimal compute–quality trade-offs across three quality axes. It also provides evidence that standard FID is insufficient for high-resolution generation, since its 299×299 downsampling discards precisely the detail that distinguishes high-fidelity models. It highlights complementary weaknesses in both pixel-space and latent-space methods that FID alone cannot detect.

• Real-world applications:

  • High-resolution image synthesis for design and media, where fidelity at 512×512, 1024×1024 and beyond matters.
  • Text-to-image generation at 1024 resolution with up to 3× higher throughput, reducing inference cost per image.
  • Video generation for content creation, with state-of-the-art FVD and roughly 30% less compute.
  • Compute-constrained deployment, where halving sampling GFLOPs at equal or better quality directly reduces serving costs.

• Industry relevance: WaiT plugs into an existing open-source pixel-space backbone (JiT) with minimal changes and is reported to transfer out of the box to latent models like DDT and to video. That makes the recipe attractive for teams already running diffusion infrastructure who want quality and cost improvements without retraining a new architecture from scratch.

Future Directions

• Deeper coarse-phase compression: The current formulation uses a single DWT level (except the three-level experiment), and the authors suggest exploring a two- or three-level low-frequency band to shrink the coarse token grid drastically before injecting all high-frequency levels at once or progressively. Whether extreme compression retains enough signal for fine-phase recovery is an open question for quality–compute frontiers.

• Integration with state-of-the-art architectures and optimizations: The text-to-image and video experiments deliberately use basic baseline configurations to isolate the method's effect. How standard architectural and optimization improvements compound with WaiT is left for future work.

• High-resolution video evaluation: The video experiments align with current pixel-space benchmarks at 128×128 (Kinetics-600, Taichi-HD), though models were also trained at 256×256. Because standard FVD cannot capture high-frequency textures, the full three-axis protocol could not be applied to video; the authors argue that natively high-resolution video datasets are needed to verify that image-domain benefits transfer.

• Tuning the sampling schedule for video: The video experiments reused image-optimized shift parameters α and m. The authors expect minor targeted tuning could yield further substantial GFLOP savings at minimal quality cost.

Target Audience

Generative-model researchers and engineers working on diffusion and flow matching, particularly those focused on high-resolution image or video synthesis and on inference-compute efficiency. It is also useful for practitioners who need better evaluation than FID alone for high-resolution outputs, and for teams interested in low-overhead modifications that can be layered onto an existing pixel-space or latent-space backbone rather than requiring a new architecture.

Authors’ abstract

As image generation models scale to ever higher resolutions, global coherence, local detail, and texture fidelity become critical axes for generation quality. However, standard flow matching treats all spatial frequencies uniformly, ignoring the natural frequency hierarchy where high-frequency bands become indistinguishable from pure noise far earlier than coarse structures. We introduce WaiT, a Wavelet-aware image Transformer that decomposes generation into coarse and fine bands via lossless wavelets. True to its name, the high-frequency bands wait for the signal: staying pure noise until coarse structure has emerged, then joining the flow for joint refinement. Since standard FID discards fine-grained detail through aggressive downsampling, we introduce a more stringent three-axis evaluation protocol to assess quality at native resolution. On ImageNet 512x512, WaiT achieves a pixel-space FID of 1.43 and is Pareto-optimal across all three axes, reducing sampling compute by up to 50%. With our largest 2B model, we set a new state-of-the-art FID of 1.3 for pixel-space models on ImageNet 512 resolution. Our formulation outperforms even the strongest latent-space models on texture fidelity, and scales seamlessly to high-resolution OpenImages and to video generation, achieving a state-of-the-art FVD of 0.84 on Kinetics-600 with no algorithmic modifications.