Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

Technical Summary: Cross-Space Distillation

Problem Statement
Modern one-step diffusion models achieve high-quality generation through distribution-based timestep distillation. However, these methods rely on a restrictive Shared-Space constraint: the Teacher and Student must inhabit the same latent space, sharing identical latent resolution and VAE parameterization. This assumption prevents knowledge transfer from modern, high-capacity Teachers (e.g., Stable Diffusion 3.5, Flux, Kolors) operating at high resolutions (e.g., 1024×1024) and utilizing advanced VAEs, to compact, deployment-friendly Students (e.g., Stable Diffusion 1.5) operating at lower resolutions (e.g., 512×512) with different VAEs.

Standard distillation objectives fail in this setting because the latent tensors ($z_S$ and $z_T$) are inherently incompatible due to:

1. Cross-Resolution: Mismatched spatial grid shapes (e.g., $h \times w$ vs. $H \times W$).
2. Cross-VAE: Different autoencoder parameterizations resulting in misaligned latent spaces.
3. Cross-Architecture/Mechanism: Differences between UNet and DiT backbones, or between standard diffusion and flow matching formulations.

Existing alternatives, such as network pruning or building new compact backbones from scratch, face challenges including hard sparsity thresholds, extensive retraining requirements, and the loss of ecosystem compatibility.

Methodology: Cross-Space Distillation and the Bridge
The authors formalize this setting as Cross-Space Distillation and introduce Bridge ($B_\phi$), a lightweight module designed to map Student latents into the Teacher's latent space without modifying the Student backbone.

• Architecture: The Bridge is factorized into two stages to handle the mismatch efficiently:

  1. Spatial Prior (Frozen): The first $n$ decoding blocks of the Student's VAE decoder ($D_S^{(n)}$) are kept frozen. These blocks act as a pre-trained spatial upsampler, expanding the Student latent $z_S$ from low resolution to the Teacher's high resolution ($H \times W$). This leverages the Student's existing spatial knowledge without adding trainable parameters for upsampling.
  2. Learnable Projector ($g_\phi$): A compact, learnable network (implemented as SwinIR) maps the upsampled features from the Student's feature space to the Teacher's latent representation space, outputting a Teacher-compatible latent $\hat{z}_T$.

• Training Objectives: The Bridge is trained using a synthetic dataset of paired latents $(z_S, z_T)$ derived from the same image encoded by both VAEs. The optimization minimizes a weighted sum of two losses:

  1. Latent Reconstruction ($L_{rec}$): An $\ell_1$ loss between the predicted latent $\hat{z}_T$ and the ground-truth Teacher latent $z_T$. The authors note $\ell_1$ provides better stability and structure preservation than $\ell_2$ in this context.
  2. Attention Fidelity ($L_{attn}$): A reverse-KL divergence loss that aligns the self-attention distributions of the frozen Teacher denoiser when fed $\hat{z}_T$ versus $z_T$. This ensures the mapped latent induces the same internal structural responses (long-range dependencies) as the true Teacher latent, preserving fine-grained details that pixel-wise losses might miss.

• Distillation and Inference: Once trained, the Bridge is attached to the Student. During distillation, the Student's output is passed through the Bridge to generate $\hat{z}_T$, allowing standard distribution-based objectives (e.g., Variational Score Distillation, Adversarial Distillation) to be applied directly against the Teacher. Crucially, the Student backbone remains unchanged.

Key Contributions

1. Formalization of Cross-Space Distillation: The paper defines a new regime where distillation occurs across mismatched latent resolutions and VAE spaces, addressing a critical gap in transferring knowledge from modern high-resolution models to compact legacy models.
2. The Bridge Module: A lightweight, modular interface that enables standard distillation objectives without altering the Student architecture. It combines a frozen spatial prior with a learnable projector and utilizes an Attention Fidelity objective for stable alignment.
3. Practical Effectiveness: The method enables compact Students (e.g., SD 1.5) to inherit the visual priors of diverse, modern Teachers (SDXL, Kolors, PixArt-Σ, Flux, SD 3.5) while maintaining one-step inference, low latency, and broad ecosystem compatibility.
4. Model Merging Capability: Because the method distills diverse Teachers into Students with the same backbone and latent space, the resulting checkpoints can be trivially merged via parameter averaging to create a single model that aggregates knowledge from multiple Teacher sources.

Results
Experiments demonstrate substantial improvements in generation quality for compact one-step Students:

• Quantitative Gains: Distilling SD 3.5 Medium (Teacher) into SD 1.5 (Student) via Bridge improves the HPSv3 score from 5.37 to 9.42 and ImageReward from -0.29 to 0.62. Merging checkpoints from five different Teachers further boosts SD 1.5 HPSv3 to 10.53.
• Qualitative Improvements: Visual comparisons show that Bridge-enabled students produce images with clearer textures, stronger structural consistency, and fewer artifacts compared to initialized baselines or pruned models.
• Efficiency: The Bridge adds negligible overhead (approx. 5M parameters) and allows for one-step inference. It also enables inference-time resolution upgrades, where a 512×512 Student sample can be mapped to 1024×1024 via the Bridge and decoded with the Teacher's decoder without retraining the Student.
• Comparison to Pruning: Bridge-enabled distillation outperforms network pruning (e.g., pruning SD 3.5 to 30% sparsity), which tends to produce blur and artifacts.

Significance
The paper claims that Cross-Space Distillation unlocks the potential to transfer the generative capabilities of massive, modern Teachers into efficient, deployable backbones without the cost of training new architectures from scratch. By treating the alignment of heterogeneous latent representations as a reusable interface problem, the work enables:

• Ecosystem Compatibility: Compact models like SD 1.5 can be upgraded with state-of-the-art quality while retaining compatibility with existing tools and hardware.
• Knowledge Unification: The ability to merge knowledge from multiple, architecturally distinct Teachers (e.g., Flux and SD 3.5) into a single, lightweight Student model.
• Modularity: The Bridge serves as a plug-and-play component that can be reused for both distillation and inference-time resolution enhancement, decoupling the Student's generation capability from its latent representation constraints.

The authors position this work as a step toward broader exploration of distillation across heterogeneous latent spaces, applicable beyond text-to-image to video and multimodal generation.