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

Conference: ECCV 2026
Paper: ECCV Official
Area: Model Compression
Keywords: diffusion distillation, one-step generation, cross-space distillation, latent representation alignment, attention fidelity

TL;DR

Proposes the Cross-Space Distillation framework with a lightweight Bridge module that combines a frozen student VAE decoder prefix as a spatial prior with a SwinIR projector, overcoming latent resolution and VAE mismatches to enable compact SD 1.5 to achieve teacher-level generation quality in a single step.

Background & Motivation

Recent state-of-the-art text-to-image diffusion and flow matching models, such as Stable Diffusion 3.5 and FLUX.2, have achieved unprecedented visual fidelity at ultra-high resolutions. However, their immense parameter scale (e.g., 2.5B-4.0B) and iterative multi-step sampling trajectories (e.g., 50 network function evaluations) incur prohibitive computational and memory footprints, severely constraining deployment on resource-limited edge devices. To accelerate sampling, distribution-based distillation recipes—particularly variational score distillation (VSD) and adversarial distillation paradigms like DMD2 and ADD—have emerged as the de facto approach for compressing multi-step teachers into single-step generators. Yet, these methods universally rely on an implicit Shared-Space constraint, assuming that the teacher and student networks share identical latent spatial resolutions and VAE representations.

This assumption creates an impassable barrier when attempting to upgrade lightweight legacy models with modern foundation teachers: modern high-capacity teachers operate at 1024×1024 pixel resolution using high-channel, distinct VAE representations (e.g., 16 or 32 channels), whereas compact, deployment-friendly student backbones with rich ecosystem support (such as SD 1.5 and SD 2.1) are inherently locked to 512×512 resolution and legacy 4-channel VAEs. Because their latent spatial grids and latent manifolds do not match, standard distillation objectives—whether score-difference regression or latent adversarial discriminators—cannot be directly evaluated. Alternative paths that prune large teacher backbones or design bespoke lightweight architectures (such as OBS-Diff or SnapGen) suffer from steep sparsity performance cliffs, demand exhaustive architecture-specific tuning from scratch, and completely forfeit compatibility with the vast existing open-source ecosystem.

Given that the compact student architecture itself is intrinsically fast and ecosystem-compatible, the true bottleneck is the representational gap in the latent space. Core idea: formalize the Cross-Space Distillation regime and introduce Bridge, a lightweight module (~5M parameters) that repurposes the student's frozen VAE decoder prefix as a spatial upsampling scaffold coupled with a compact SwinIR projection head and a reverse-KL attention fidelity loss, precisely mapping student latents into the teacher latent space without modifying the student backbone to unlock direct one-step distribution distillation.

Method

Cross-Space Distillation introduces a plug-and-play Bridge interface \(B_\phi\) to decouple generator backbone inference from the teacher supervision representation space. During training, the two-stage Bridge maps low-resolution student latents to high-resolution teacher-compatible latents under dual supervision (latent reconstruction and attention fidelity); once trained and frozen, the Bridge projects student predictions into the teacher space to drive standard one-step distribution distillation (such as DMD2/VSD and adversarial losses), updating the student denoiser parameters; during inference, the student can either retain its native one-step generation at 512×512 or feed through the Bridge and teacher VAE decoder for training-free resolution upgrading.

Overall Architecture

The system consists of three coordinated pipelines: offline Bridge alignment training, Bridge-enabled one-step cross-space distillation, and inference-time resolution upgrading and post-training merging.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    subgraph S1["Bridge Alignment Training"]
        direction TB
        A["Image Sample x"] --> B["Dual VAE Encoders<br/>Extract Paired Latents (zS, zT)"]
        B --> C["Frozen Decoder Spatial Prior<br/>D_S^(1) Upsamples to H×W"]
        C --> D["SwinIR Projector Head g_ϕ<br/>Aligns Channels & Semantics to z_hat_T"]
        D --> E["Dual Objective Supervision<br/>L1 Latent Rec + Reverse-KL Attention Fidelity"]
    end
    subgraph S2["One-Step Cross-Space Distillation"]
        direction TB
        F["Gaussian Noise ξ + Text Prompt c"] --> G["One-Step Student Backbone S_θ<br/>Predicts Native Latent zS"]
        G --> H["Frozen Bridge Projection<br/>z_hat_T = B_ϕ(zS)"]
        H --> I["Mechanism-Agnostic Reparameterization<br/>Unifies Output as Clean x0"]
        I --> J["Teacher Distribution Supervision<br/>VSD Score Difference + Discriminator"]
        J -->|Back-Propagation Update| G
    end
    subgraph S3["Inference & Post-Training Extensions"]
        direction TB
        K["Native Student Inference<br/>Native 512×512 Ultra-Fast Output"]
        G -.->|Optional Direct| K
        H -.->|Optional Upgrade| L["Teacher VAE Decoder D_T<br/>1024×1024 High-Resolution Synthesis"]
        G -.->|Multiple Distilled Checkpoints| M["Parameter Averaging Model Merging<br/>Zero-Overhead Capability Consolidation"]
    end

Key Designs

1. Decoupled Two-Stage Bridge Architecture: Reconciling Resolution and Latent Manifolds Directly training a neural network from scratch to map representations across mismatched grids forces the model to learn spatial interpolation and texture upsampling simultaneously, leading to severe high-frequency artifacts or training instability. The authors observe that the student model's pretrained VAE decoder already contains rich spatial decoding and upsampling priors learned from extensive data. The Bridge is explicitly factorized into two stages: the first stage freezes the initial \(n\) decoding blocks of the student VAE decoder \(D_S^{(n)}\) (empirically \(n=1\)), upsampling the student latent tensor from \(h \times w\) (e.g., \(64 \times 64\)) directly to the teacher's spatial resolution \(H \times W\) (e.g., \(128 \times 128\)) to produce an intermediate scaffold feature \(f_{\text{prior}} \in \mathbb{R}^{H \times W \times C_{\text{prior}}}\) without introducing any learnable parameters; the second stage utilizes a compact SwinIR network (~5M parameters) as a learnable projection head \(g_\phi\) that maps \(f_{\text{prior}}\) from the decoder feature manifold into the teacher-compatible latent representation \(\hat{z}_T = g_\phi(f_{\text{prior}}) \in \mathbb{R}^{H \times W \times C_T}\). Decoupling geometric grid expansion from semantic channel transformation significantly lowers optimization complexity and preserves feature fidelity.

2. Reverse-KL Attention Fidelity Objective: Global Semantic Alignment in Teacher Space While latent-space regression is necessary, small numerical deviations in latent space often trigger noticeable semantic distortions or anatomical failure modes upon final image decoding. Traditional pixel-level supervision or additional adversarial tuning remains hyperparameter-sensitive and unstable. To enforce globally aware semantic consistency, the authors introduce an attention fidelity loss computed directly inside the frozen teacher denoiser: internal self-attention maps \(P^l(z; t, c) \in [0, 1]^{N \times N}\) are extracted across designated layers. Because self-attention captures long-range dependencies, representational discrepancies between \(\hat{z}_T\) and the ground-truth teacher latent \(z_T\) are sharply amplified in the teacher's attention responses. Drawing inspiration from MiniLLM, a reverse-KL divergence objective is employed to focus optimization on the dominant attention mass, offering substantially better stability than forward KL: $$ \mathcal{L}{\text{attn}} = \sum_T; t, c) \,||\, P^l(z_T; t, c) \right) $$ where }} \tau^2 \text{KL}\left( P^l(\hat{z\(\tau = 3.0\) denotes the temperature parameter for numerical stability, evaluated under noised state \(t=1\). This constraint ensures that the projected latent activates identical long-range contextual attention patterns within the teacher backbone.

3. Mechanism-Agnostic Reparameterization and Multi-Teacher Parameter Merging: Cross-Paradigm Synthesis Beyond spatial and VAE mismatch, modern teachers (often Rectified Flow velocity-prediction DiT/MMDiT models) and compact students (often Gaussian noise-prediction UNet models) differ fundamentally in architecture and mathematical dynamics. Avoiding fragile layer-wise feature matching, the method reformulates student predictions back into clean data estimations \(\hat{x}_0\) using closed-form endpoint relations, regardless of whether continuous-time flow matching or discrete diffusion formulations are used. The reparameterized state is passed through Bridge to receive standard distribution-level supervision (VSD and discriminator losses). Crucially, because all distilled students retain the identical, untouched student backbone parameterization, multiple student checkpoints distilled from distinct teachers can be merged post-training via simple parameter averaging: $$ \theta_{\text{merged}} = \frac{1}{|S|} \sum_{i=1}^{|S|} \theta_i $$ This lightweight post-training procedure incurs zero additional inference or memory cost while combining the visual strengths of diverse teacher families (SDXL, Kolors, PixArt-Σ, SD 3.5, and FLUX.2), frequently outperforming individual teacher-distilled checkpoints.

Loss & Training

The Bridge module is trained with the joint objective: $$ \mathcal{L}{\text{final}} = \alpha \mathcal{L} $$ where the latent reconstruction loss uses }} + \beta \mathcal{L}_{\text{attn}\(\mathcal{L}_{\text{rec}} = \|z_T - \hat{z}_T\|_1\). The \(\ell_1\) norm exhibits superior optimization stability and edge preservation in latent manifolds over the conventional \(\ell_2\) norm. Weighting coefficients are set to \(\alpha = \beta = 1\). Rather than collecting massive real-world datasets, the training pipeline utilizes a static synthetic dataset of ~2M image-text pairs generated by sampling prompts from JourneyDB and LAION and rendering them with the respective teacher models. After Bridge training, the student denoiser is distilled following the DMD2 protocol for 20 hours on 8× NVIDIA H100 80GB GPUs using the AdamW optimizer.

Key Experimental Results

Main Results

The framework is evaluated across two representative compact students (SD 1.5 and SD 2.1) distilled from five state-of-the-art 1024-resolution teachers: SDXL, Kolors, PixArt-Σ-1024, FLUX.2-klein-4B, and SD 3.5 Medium. Evaluated metrics include human preference benchmarks (HPSv3, HPSv2), prompt adherence and alignment (ImageReward, DPG Bench), and multidimensional quality (MPS).

Student / Teacher Setting	Backbone Params (B)	Inference NFE	HPSv3 ↑	HPSv2 ↑	ImageReward ↑	MPS ↑	DPG Bench ↑
Multi-Step Teacher Upper Bounds
SDXL (Large U-Net)	2.57	50	9.25	28.36	0.66	13.72	74.00
Kolors (Large U-Net)	2.57	50	10.59	30.94	0.88	14.12	76.52
PixArt-Σ-1024 (DiT)	0.61	20	9.62	30.39	0.92	14.14	80.00
FLUX.2-klein-4B (MM-DiT)	4.00	50	10.05	28.90	0.80	14.00	83.20
SD 3.5 Medium (MM-DiT)	2.50	50	10.86	30.02	0.96	14.13	84.50
SD 1.5 Student Family (Native 512×512)
SD 1.5 DMD2 (Initialization Baseline)	0.86	1	5.37	21.90	-0.29	10.42	59.85
+ SDXL Distillation	0.86	1	9.04	27.37	0.30	11.93	63.00
+ Kolors Distillation	0.86	1	9.33	27.38	0.34	11.66	64.46
+ PixArt-Σ-1024 Distillation	0.86	1	8.65	26.43	0.30	11.94	63.43
+ FLUX.2-klein-4B Distillation	0.86	1	9.49	26.64	0.40	12.35	64.06
+ SD 3.5 Medium Distillation	0.86	1	9.42	28.30	0.62	12.96	65.75
Merged (Average of All 5 Teachers)	0.86	1	10.53	29.07	0.65	12.62	66.67
SD 2.1 Student Family (Native 512×512)
SD 2.1 SiD (Initialization Baseline)	0.86	1	6.42	23.74	0.10	11.29	61.33
+ SDXL Distillation	0.86	1	8.40	26.73	0.23	11.88	67.03
+ Kolors Distillation	0.86	1	8.52	26.57	0.20	11.91	68.34
+ PixArt-Σ-1024 Distillation	0.86	1	8.42	28.29	0.44	12.20	68.51
+ FLUX.2-klein-4B Distillation	0.86	1	8.74	28.31	0.33	12.03	66.00
+ SD 3.5 Medium Distillation	0.86	1	8.64	29.11	0.45	12.92	68.52
Merged (Average of All 5 Teachers)	0.86	1	9.75	30.00	0.74	12.60	68.50

Ablation Study

1. Comparison with Structural Network Pruning

Benchmarking Bridge-enabled distillation against direct network pruning using OBS-Diff to prune SD 3.5 Medium by 30% sparsity (compressing 2.5B to 1.5B parameters), followed by distillation:

Compression Setting	Params (B)	Inference NFE	HPSv3 ↑	HPSv2 ↑	ImageReward ↑	MPS ↑	DPG Bench ↑
OBS-Diff Pruning (30% Sparsity)	1.50	1	2.76	21.94	-0.36	9.54	55.24
Ours: Bridge Distillation (SD 1.5)	0.86	1	9.42	28.30	0.62	12.96	65.75

2. Student Checkpoint Initialization Sensitivity

Distilling SD 1.5 with SD 3.5 Medium teacher under one-step DMD2 initialization versus multi-step base SD 1.5 initialization:

Initialization Setting	Training Iterations	HPSv3 ↑	HPSv2 ↑	ImageReward ↑	MPS ↑	DPG Bench ↑
Multi-Step SD 1.5 Weights Init	9k	9.22	28.44	0.66	13.08	66.89
SD 1.5 DMD2 One-Step Checkpoint Init	3k	9.42	28.30	0.62	12.96	65.75

Key Findings

• Cross-Space Distillation unlocks generational quality leaps: Under SD 3.5 Medium teacher guidance, SD 1.5 HPSv3 scores increase from 5.37 to 9.42 (+75.4%), while ImageReward jumps from -0.29 to 0.62. Qualitative samples show complete elimination of texture blurriness and anatomical artifacts common in legacy one-step generators.
• Model merging produces super-teacher performance: Merging five teacher-specific distilled SD 1.5 student checkpoints yields an HPSv3 score of 10.53 and HPSv2 of 29.07, surpassing the native multi-step performance of SDXL (9.25), PixArt-Σ (9.62), and FLUX.2-klein-4B (10.05) with zero added inference overhead.
• Structural pruning leads to catastrophic quality degradation: Even with 1.5B parameters (nearly double SD 1.5's size), the 30% pruned baseline suffers a severe collapse in HPSv3 to 2.76, confirming that pruning corrupts delicate internal attention dynamics.
• Pretrained one-step initialization accelerates convergence 3×: Initializing from a one-step checkpoint (DMD2) requires only 3k iterations compared to 9k iterations for multi-step initializations, saving 66.7% of training compute.

Highlights & Insights

• Repurposing VAE decoder prefix as spatial scaffold: Elegantly harnesses the student model's native spatial decoding capability, bypassing the instability and artifacts of learning latent upsamplers from scratch.
• Reverse-KL attention fidelity aligns latent topologies: Supervising teacher self-attention maps via reverse-KL divergence captures critical long-range dependencies, resolving subtle latent manifold misalignments that cause downstream decoding failures.
• Unified model consolidation across heterogeneous models: By distilling fundamentally incompatible teacher models (FLUX, SD 3.5, Kolors) into an identical lightweight student space, the framework enables zero-cost capability unification through parameter averaging.

Limitations & Future Work

• Reliance on teacher synthetic data distribution: Bridge pretraining relies on ~2M teacher-generated synthetic samples; domain or prompt coverage limitations in this corpus may introduce representational blind spots for atypical prompts.
• Inference-time resolution upgrading overhead: While Bridge can cascade into the teacher VAE decoder to output 1024×1024 images, the teacher decoder's latency and memory demands partially compromise the ultralight edge efficiency profile.
• Future directions: Investigating streaming online alignment without pre-generated synthetic datasets, and extending cross-space distillation to video diffusion foundation models (e.g., distilling Sora-like DiTs into compact spatial-temporal UNets).

Related Work & Insights

• vs DMD / DMD2 (Distribution Matching Distillation): DMD2 pioneered distribution matching for one-step generation but strictly required teacher and student to share identical VAEs and resolutions. This work removes the Shared-Space constraint via Bridge, scaling distribution matching across model generations.
• vs OBS-Diff / SnapFusion (Pruning and Structural Compression): Pruning methods compress large teachers directly but hit severe performance drop-offs and break downstream ecosystem tooling (such as LoRA and ControlNet). This work demonstrates that distilling into mature, compact baselines yields strictly superior fidelity at lower parameter counts.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Formulates the overlooked Cross-Space Distillation regime and introduces an effective, principled Bridge module.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Comprehensive evaluations across five state-of-the-art teachers, two student families, and varied architectures (DiT, MMDiT, UNet).
• Writing Quality: ⭐⭐⭐⭐⭐ Clear motivation, rigorous mathematical formulation, and thorough quantitative benchmarking.
• Value: ⭐⭐⭐⭐⭐ Revitalizes widely deployed lightweight models and introduces an accessible pathway for heterogeneous model capability merging.