CrossFlow: One-Step Generation Across Latent and Pixel Spaces

Technical Summary: CrossFlow: One-Step Generation Across Latent and Pixel Spaces

1. Problem Statement

Current state-of-the-art image generation systems, particularly Latent Diffusion Models (LDMs) and flow-matching methods, operate by defining the prior, probability path, and prediction target within a single representation space. While LDMs improve computational efficiency by performing the generative trajectory in a compressed autoencoder latent space, they rely on a two-stage pipeline: iterative denoising in the latent space followed by a separate decoding step to map latents back to pixels.

This architecture introduces a fundamental mismatch:

• Optimization Mismatch: The generator is optimized to predict latent displacements, while the final image quality depends on a separately trained decoder.
• Distribution Shift: At inference, the decoder receives generated latents which may differ distributionally from the clean latents produced by the encoder during training. This shift can amplify errors.
• Supervision Limitation: The generator cannot be directly trained with image-level perceptual or adversarial losses unless its output passes through the separate decoder, weakening direct pixel-space supervision.

The paper asks whether the noise space and the output space must be identical, proposing a framework that preserves the efficiency of latent representations while predicting the image itself directly.

2. Methodology: CrossFlow

CrossFlow is a cross-space flow formulation that maps noisy latent inputs directly to pixel-space images in a single function evaluation. It resolves the space mismatch by deriving objectives where the latent-space velocity is eliminated from the supervised target.

Core Technical Formulation:

• Unified Framework: The authors define a unified target function $Z$ for single-step generative paradigms. Standard methods rely on latent-space velocity fields, which prevents cross-space generation.
• Design Criteria: To enable cross-space supervision, the authors establish two criteria:
  1. Eliminating Latent-Velocity Dependencies: The objective must be structurally decoupled from latent-space velocities. The network output should not be a latent displacement vector.
  2. Endpoint Boundary Conditions: The framework must yield a valid image at the destination boundary ($r=0$).
• The CrossFlow Objective: By satisfying these criteria, the authors derive a general loss function (Eq. 10) where the network $F_\theta$ predicts a pixel-space image $x_0$ rather than a latent velocity.
  • The latent trajectory (defined by $\alpha(t)$ and $\beta(t)$) still dictates how clean latents are corrupted to create inputs.
  • The supervised target is the clean pixel image.
  • The specific objective used in experiments (Eq. 11) involves a time derivative term computed via Jacobian-Vector Product (JVP) along the latent trajectory, ensuring the model learns the correct flow dynamics without predicting latent velocities.

Implementation Details:

• Architecture: The model uses a fixed, pretrained image encoder (VA-VAE) to map images to latents and a Vision Transformer (ViT) generator ($F_\theta$) to map noisy latents directly to pixels.
• Training Pipeline: A two-stage approach is used to ensure stability:
  1. Train the encoder $E$ using standard autoencoder objectives (reconstruction + adversarial).
  2. Freeze $E$ and train the generator $F_\theta$ from scratch using the CrossFlow objective.
• Loss Function: The total objective combines the analytical CrossFlow loss with pixel-space auxiliary losses:
  • L1 Reconstruction Loss: Pixel-wise alignment.
  • Perceptual Loss: Structural similarity.
  • Adversarial (GAN) Loss: High-frequency fidelity.
  • Note: The analytical weighting factor from the flow objective is omitted during perceptual loss computation to prevent training instability.

3. Key Contributions

1. Cross-Space Framework: Introduction of CrossFlow, a one-step generation framework where the input prior (latent) and output image (pixel) reside in different representation spaces.
2. Velocity-Free One-Step Objective: A novel derivation that allows a model to act as both a one-step latent-to-pixel generator and a decoder replacement, eliminating the need for a separate decoder at inference.
3. Direct Pixel-Space Supervision: The ability to train the generator directly with perceptual and adversarial losses on the generated image, bypassing the intermediate decoder stage.
4. Empirical Validation: Demonstration that a single model can achieve competitive FID scores on ImageNet-1k (256×256) with a single function evaluation (1-NFE) while outperforming standard one-step latent methods and matching multi-step latent diffusion baselines.

4. Experimental Results

Experiments were conducted on ImageNet-1k at 256×256 resolution using class-conditional generation.

• Main Performance:

  • CrossFlow-XL (461M parameters) achieves an FID of 1.62 with 1-NFE.
  • This is competitive with multi-step latent diffusion baselines like DiT-XL/2 (2.27 FID with 250 steps) and SiT-XL/2 (2.06 FID with 250 steps).
  • It significantly outperforms other one-step latent methods (e.g., iMF-XL/2 at 1.72 FID) and avoids the need for a separate decoder.
  • CrossFlow-B (95M) and CrossFlow-L (315M) achieve FIDs of 2.54 and 1.69, respectively, demonstrating scalability.

• Ablation Studies:

  • Encoder Choice: The VA-VAE encoder (32 channels) yielded the best results (2.54 FID), outperforming SD-VAE (4.08 FID) and FLUX.2 (3.81 FID), suggesting latent representation structure is critical.
  • Loss Components: Removing the perceptual loss caused a severe FID drop to 30.21, and removing the GAN loss dropped it to 8.86. The CrossFlow flow objective provides the structural signal, while perceptual and adversarial losses are vital for high-frequency fidelity.
  • Decoder Replacement: When used as a decoder for latents generated by a LightningDiT-B/1, CrossFlow consistently improved FID over the standard VA-VAE decoder throughout training.

5. Significance and Claims

The paper claims that CrossFlow successfully bridges the gap between the efficiency of latent representations and the fidelity of direct pixel-space supervision. By unifying the generator and decoder into a single model trained on noisy latent inputs and pixel-space targets, CrossFlow:

• Removes the explicit decoder stage at inference.
• Exposes the model during training to the noisy/imperfect latents that occur during generation, rather than just clean encoder outputs.
• Demonstrates that cross-space flow objectives can combine the computational benefits of latent diffusion with the quality improvements of direct pixel-level adversarial and perceptual training.

The authors maintain a modest tone, noting that while the latent representation and auxiliary losses are important, the core contribution is the architectural and objective shift that enables one-step, cross-space generation without requiring separate decoder modules.