PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Anh Nguyen, Ngan Nguyen, Duc Vu, Trung Dao, Viet Nguyen, Quan Dao, Kien Nguyen, Chi Tran, Phong Nguyen, Khoi Nguyen, Cuong Pham, Dimitris Metaxas, Vishal M. Patel, Anh Tran

arXiv:2606.32020v1cs.CV

TL;DR

One-step diffusion distillation assumes teacher and student share a latent space, leaving mismatched-resolution and VAE architectures without directly applicable objectives. Cross-Space Distillation introduces Bridge, a lightweight latent interface that aligns student latents to the teacher space; across teachers and metrics, it improves SD 1.5 from 5.37 to 9.42 HPSv3 while preserving one-step inference.

  • Problem

    Different latent resolutions and VAE designs make standard distillation objectives directly inapplicable between teacher and student representations.

  • Method

    Bridge maps student latents into the teacher space through a lightweight interface while preserving the student backbone.

  • Results

    5.37 to 9.42 HPSv3: Cross-Space Distillation improves one-step students across all evaluated teachers and five complementary metrics.

  • Takeaways & Limitations

    The approach transfers modern teacher knowledge to compact students while preserving one-step inference, the original backbone, and deployment compatibility.

  • Takeaways & Limitations

    The study is limited to the text-to-image setting, with extensions to video, editing, and richer conditional pipelines left for future work.

Abstract

from arXiv · show

Modern one-step diffusion models achieve impressive quality through distribution-based timestep distillation. Yet, they rely on a critical assumption: Teacher and Student must inhabit the same latent space. This Shared-Space constraint prevents knowledge transfer from modern high-capacity Teachers (e.g., SD 3.5 and Flux) into compact, deployment-friendly Students such as SD 1.5, whose latent resolution and VAE parameterization differ from the Teacher. We formalize this overlooked regime as Cross-Space Distillation, where Teacher and Student differ in both latent resolution and VAE space. To enable distillation under this mismatch, we introduce the Bridge, a lightweight latent interface that maps Student latents into the Teacher space without modifying the Student backbone. Bridge combines a frozen Student VAE decoder as a spatial prior with a compact learnable projector, and is trained with latent reconstruction and attention fidelity objectives for stable Teacher-space alignment. Across diverse modern Teachers, Bridge enables substantial gains for compact one-step Students; for example, it improves SD 1.5 from 5.4 to 9.4 HPSv3 while preserving one-step inference, low latency, and broad ecosystem compatibility. These results show that heterogeneous large Teachers can be distilled into efficient, deployable backbones through a lightweight latent-space interface.

1 Introduction

The introduction identifies latent resolution and VAE mismatches as the main barriers to distilling modern high-capacity teachers into compact students. It proposes Cross-Space Distillation and Bridge, a lightweight latent interface that aligns student latents with teacher space while preserving backbone compatibility.

  • Motivation: Modern high-fidelity teachers such as Stable Diffusion 3.5 and Flux incur substantial computational and memory costs because of large backbones and multi-step sampling.These costs limit deployment on consumer hardware.
  • Motivation: The proposed strategy reuses compact, widely supported models like Stable Diffusion 1.5 instead of building new backbones, leveraging pretrained weights and existing ecosystems.The approach is described as compact, model-agnostic, and easier to optimize.
  • Problem: Resolution and latent-space mismatches are the main barriers among four cross-model gaps, preventing standard distillation objectives from being applied directly.The four gaps are Cross-Resolution, Cross-VAE, Cross-Architecture, and Cross-Mechanism; architecture and mechanism gaps are comparatively minor.
  • Bridge: Bridge maps Student latents into Teacher space through a lightweight module without modifying the student architecture.It freezes the Student VAE for spatial alignment and uses an Attention Fidelity loss to preserve fine structure and high-quality reconstruction.
  • Results: Bridge substantially improves SD 1.5 through Cross-Resolution and Cross-VAE distillation while preserving backbone compatibility and enabling high-resolution synthesis upgrades.The module is presented as modular and capable of transferring the teacher’s visual priors.
  • Cross-Space Distillation: Cross-Space Distillation formalizes distillation between Teachers and Students that differ simultaneously in latent resolution and VAE space.Cross-Resolution involves different latent-grid shapes, while Cross-VAE involves differing autoencoders and unaligned latent spaces.

2 Preliminary

Diffusion and flow matching provide related continuous-time formulations for generative modeling, while one-step distillation matches a Student’s sample distribution to a multi-step Teacher. Distribution-based objectives commonly operate in latent space under a shared-representation assumption, motivating alignment methods for mismatched Teacher and Student spaces.

  • Diffusion: Diffusion models progressively transform data into Gaussianlike noise and learn reverse-time dynamics to recover the data distribution.Common reverse-process parameterizations include ϵ-prediction, x0-prediction, and v-prediction.
  • Flow matching: Flow matching trains a continuous normalizing flow by predicting a conditional velocity field along a chosen probability path.Related formulations include rectified flow and conditional/OT flow matching objectives.
  • Connection: Diffusion and flow matching share a unified continuous-time view, with endpoint transformations relating flow velocities to diffusion outputs.The work converts outputs to a consistent x0-estimate in a shared latent space.
  • One-step distillation: One-step distillation trains a fast Student to match the sample distribution of a multi-step Teacher using distribution-based objectives rather than trajectory supervision.Variational score distillation regresses the difference between Student and Teacher scores on the same noised sample, while adversarial distillation adds discriminator-based distribution matching.
  • Cross-space mismatch: Latent-space distillation commonly assumes Teacher and Student share the same representation, whereas Cross-Resolution and Cross-VAE mismatch requires aligning Student states to a Teacher-compatible space.This alignment enables standard VSD and adversarial objectives with minimal modification.

3 Distillation Enabled by the Bridge

The Bridge enables Cross-Space Distillation by aligning mismatched Student and Teacher latent spaces in spatial resolution and VAE representation. It combines a frozen Student decoder prefix with a learnable projection head and trains this interface using latent reconstruction and Teacher attention fidelity objectives.

  • Problem and formulation: Cross-Space Distillation addresses teacher–student mismatches in both latent spatial resolution and VAE parameterization through an explicit alignment mapping.Standard distillation objectives cannot directly supervise mismatched Teacher and Student latent variables.
  • Bridge usage: After training, the frozen Bridge attaches to the Student so standard one-step distillation losses and one-step inference can operate through the Teacher latent space.This produces an augmented Student without modifying the Student denoiser backbone.
  • Bridge architecture: The Bridge first expands Student latent resolution with frozen early Student VAE decoder blocks, then projects decoder features into the Teacher latent representation.The frozen decoder prefix supplies a pretrained spatial upsampling scaffold without additional trainable parameters, while the compact projector performs channel and semantic alignment.
  • Training objectives: The Bridge is optimized with complementary latent reconstruction and Teacher-based attention fidelity objectives.Latent reconstruction uses an ℓ1 loss, while attention fidelity matches the Teacher denoiser’s internal attention responses induced by Bridge and Teacher latents.
  • Training objectives: Attention fidelity uses reverse KL divergence over Teacher self-attention distributions, emphasizing dominant attention mass and providing a stable alignment signal.Attention captures long-range dependencies, so matching these responses helps preserve fine-grained semantic structure beyond latent regression alone.

4 Experiments

Experiments show that Bridge enables cross-architecture, cross-mechanism distillation from diverse 1024×1024 Teachers into 512×512 one-step Students. It improves quantitative and qualitative generation while preserving one-step inference, and supports checkpoint merging and resolution upgrades.

  • Experimental Setup: Students use one-step SD 1.5 DMD2 and SD 2.1 SiD-LSG backbones at 512 × 512, while Teachers include five 1024×1024 models spanning U-Net, DiT, and MMDiT architectures.The Teachers use VAE configurations with 4, 16, and 32 channels.
  • Bridge Configuration: Bridge is a lightweight alignment network with approximately 5M trainable parameters, combining a frozen first Student VAE decoding block with a SwinIR projection head.The frozen spatial prior expands Student latents to Teacher resolution, while the projection head produces Teacher-compatible latents.
  • Cross-Architecture and Cross-Mechanism Distillation: Bridge enables distillation across architectures and diffusion mechanisms, including from DiT-based flow-matching Teachers such as SD 3.5 and FLUX.2 to UNet-based noise-prediction SD 1.5.The framework explicitly addresses cross-resolution and cross-VAE constraints with a mechanism-agnostic formulation.
  • Quantitative Results: 5.37 to 9.42 on HPSv3 and −0.29 to 0.62 on ImageReward are the SD 1.5 gains when SD 3.5 Medium is used as Teacher.Improvements also occur on HPSv2, MPS, and DPG Bench, indicating gains across five complementary metrics.
  • Qualitative Results: Distilled outputs have clearer textures, stronger structural consistency, and more faithful local details than the initialized baseline, while preserving one-step generation.The baseline more often shows softness, broken fine structure, and occasional anatomical artifacts.
  • Model Merging: Merging distilled checkpoints improves the SD 1.5 HPSv3 score from 5.37 to 10.53 and HPSv2 from 21.90 to 29.07 through direct parameter averaging.Because distilled students share a backbone and latent space, capabilities from different Teachers can be combined without changing architecture, parameter count, or inference cost.

5 Related Work

Related work spans diffusion and flow-based generation, fast-sampling distillation, and the shared-space assumption that limits transferring modern Teachers to compact Students. This paper addresses that limitation with Cross-Space Distillation, using a lightweight alignment module to map Student representations into the Teacher latent space without modifying the Student architecture.

  • Diffusion and Flow: Diffusion models use stochastic corruption and reverse denoising, while later solvers and training formulations accelerate sampling and improve efficiency.Flow matching offers an alternative continuous-time formulation that directly learns the velocity field transporting noise to data.
  • Diffusion and Flow: Modern systems such as SD 3.5 and Flux use large backbones and high-resolution latents for strong fidelity and prompt alignment, but require multiple network evaluations.
  • Distillation for Fast Diffusion Sampling: Distribution-based distillation matches a Student generator to a pretrained Teacher’s sample distribution to reduce multi-step inference cost.Representative methods use score differences between Teacher and Student predictions, and some add adversarial supervision to improve perceptual quality.
  • Distillation for Fast Diffusion Sampling: Existing fast-distillation methods typically require Teacher and Student to share latent representation and spatial resolution, causing distilled models to inherit the Teacher’s scale or latent structure.
  • Distillation for Fast Diffusion Sampling: Cross-Space Distillation introduces a lightweight alignment module that maps Student representations into Teacher latent space, enabling standard objectives without changing the Student architecture.This lets compact backbones such as Stable Diffusion 1.5 inherit modern high-capacity models’ generative capability while preserving ecosystem compatibility.

6 Conclusion

The paper identifies shared latent representation as a restrictive assumption in distribution-based one-step distillation. Cross-Space Distillation and Bridge relax this constraint by mapping Student latents into Teacher space under latent-resolution and VAE mismatch.

  • Shared latent representation is a common but restrictive assumption in distribution-based one-step distillation.
  • Cross-Space Distillation relaxes the shared-space constraint between Teachers and Students.
  • Bridge (Bϕ) is a lightweight latent interface that maps Student latents into Teacher space despite latent-resolution and VAE mismatch.This makes standard one-step distillation objectives applicable under the mismatch.

Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion … B.1 Architectural Design

The supplementary results detail Bridge’s practical implementation and architectural design for Cross-Space Distillation. They show that a shared Bridge setup with a frozen spatial prior and SwinIR backbone provides strong latent-space reconstruction while remaining lightweight and deployable.

  • – Supplementary Materials –: The supplementary material reports that Bridge is lightweight, reusable across diverse teachers, efficient to train as a post-training module, and effective for compact one-step students.These claims summarize the supplementary implementation details, ablations, efficiency analyses, and qualitative results.
  • Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion: The supplementary material expands the main paper with implementation details, Bridge design ablations, efficiency analyses, and extended qualitative results.It frames these additions as support for reproducibility and practical evaluation.
  • A Implementation and Hyperparameter Details: Bridge uses one architecture across teachers, changing only the final projection layer to match each teacher’s latent dimensionality.The Bridge is trained from scratch without pretrained initialization.
  • B.1 Architectural Design: Under 5M trainable parameters and 10K training iterations, MLP performs worst, followed by UNet, while SwinIR improves further.The ablations compare Bridge architectures under the same training budget.
  • A Implementation and Hyperparameter Details: Both Bridge training and downstream distillation fit on a single 8×H100 (80GB) node, supporting a practical post-training recipe.The pipeline is presented as reusable post-training rather than full model pretraining.
  • Teachers: The reported setup covers Bridge training and downstream one-step distillation for students paired with different teachers.Table A1 organizes the implementation, hyperparameter, and compute settings across teacher models.
  • B.1 Architectural Design: SwinIR with the proposed Spatial Prior achieves the best overall reconstruction fidelity among the evaluated Bridge architectures.The result supports using a stronger image-restoration-style backbone for latent-space alignment.
  • B.1 Architectural Design: The frozen Spatial Prior provides an additional reconstruction-fidelity gain beyond the choice of backbone alone.This supports the architectural role of the Spatial Prior in latent-space alignment.

B.2 Training Objectives … B.6 Bridge Ablation Studies

The appendix finds that Attention Fidelity and the full Bridge design improve reconstruction fidelity and convergence, while Bridge remains a modest-overhead, frozen interface for cross-space distillation. Bridge also supports the practical 512 →1024 inference-time resolution upgrade, though larger or arbitrary output resolutions are outside scope.

  • B.2 Training Objectives: Attention Fidelity produces sharper and more faithful reconstructions than ℓ1 or ℓ1 + E-LatentLPIPS.ℓ1 preserves coarse layout but over-smooths outputs, while E-LatentLPIPS provides limited perceptual improvement.
  • B.3 Efficiency: Bridge adds only modest runtime and memory overhead to the base SD 1.5 pipeline, with overhead increasing with target latent dimensionality.Measurements average 500 iterations on a single H100 80GB GPU with batch size 1 and include SDXL, SD 3.5 Medium, and FLUX.2-klein-4B target spaces.
  • B.3 Efficiency: Bridge remains lightweight even for the 32-channel FLUX latent space, supporting its use as a latent-space interface rather than a second-stage generator.Kolors and PixArt-σ use the SDXL Bridge configuration because they share SDXL’s VAE.
  • B.4 Usage During Distillation: Jointly updating Bridge with the Student causes unstable optimization and occasional training collapse, so Bridge is frozen during downstream distillation.Freezing reduces memory usage and preserves Bridge as a pretrained interface rather than a second trainable generator component.
  • B.5 Inference-Time Resolution Upgrade: In the practical 512 →1024 setting, Bridge maps low-resolution Student latents to teacher-compatible high-resolution latents while preserving structure, texture, and global coherence.The upgraded latents can be decoded at 1024 × 1024 and closely match target reconstructions across diverse scenes.
  • B.5 Inference-Time Resolution Upgrade: Bridge’s resolution upgrade is scoped to mapping low-resolution Student latents onto the Teacher latent grid under Cross-Resolution and Cross-VAE mismatch.Larger output scales and arbitrary output resolutions are beyond the scope of the work.
  • B.6 Bridge Ablation Studies: The full Bridge with Spatial Prior plus Attention Fidelity achieves the best decoded-image fidelity, and Attention Fidelity accelerates convergence while increasing final SSIM.Ablations evaluate L1, PSNR, and SSIM between decoded Teacher-space outputs and target reconstructions under matched training settings.
  • B.6 Bridge Ablation Studies: Among Bridge architectures, MLP underperforms UNet and SwinIR, while adding the frozen Spatial Prior to SwinIR improves spatial modeling and alignment.The Spatial Prior expands Student latents to the Teacher grid before projection, allowing the projector to focus on feature and semantic alignment.

C Limitations

The study focuses on latent-resolution and VAE-space heterogeneity for compact one-step generation, using Bridge as a lightweight interface that leaves the Student backbone unchanged. Its formulation suggests broader applications across heterogeneous latent representations, while extensions beyond the studied setting remain future work.

  • Scope and design: Bridge addresses latent-resolution and VAE-space mismatch after conditioning, leaving the Student backbone unchanged.This keeps the method lightweight, modular, and easy to integrate with existing one-step backbones.
  • Operating point: The experiments target compact one-step image generation rather than reproducing the full capacity of large multi-step Teachers.The intended trade-off is to transfer as much Teacher knowledge as possible through a minimal additional module while preserving deployment advantages.
  • Broader scope: Cross-Space Distillation may extend beyond text-to-image systems to latent video models, image editing pipelines, and other multimodal generators.The paper frames alignment across heterogeneous latent representations as a reusable interface problem, but extending Bridge to video remains future work.

D Additional Qualitative Results … F Societal Impacts

The extended material presents additional SD 1.5 and SD 2.1 qualitative samples, situates the method among diffusion, distillation, representation-alignment, and efficient-modeling research, and discusses societal risks of image manipulation.

  • D Additional Qualitative Results: Additional uncurated samples are provided for the distilled one-step SD 1.5 and SD 2.1 models.The samples appear in Fig. A4 and Fig. A5, respectively.
  • E.1 Diffusion and Flow.: Diffusion research advanced from reverse denoising formulations to faster deterministic and solver-based sampling, while flow methods model transport paths between noise and data.DDIM, DPM-Solver, EDM, Flow Matching, Rectified Flow, and optimal-transport conditional flow matching are cited as representative developments.
  • E.1 Diffusion and Flow.: Stable Diffusion 3.5 and Flux achieve strong visual fidelity and prompt alignment with stronger backbones, larger data, and higher-capacity latent representations, but increase inference cost and model size.These costs make direct deployment difficult in resource-constrained settings.
  • E.2 Distillation for Fast Diffusion Sampling.: One-step and few-step distillation methods shorten sampling trajectories or match Teacher-induced distributions, while adversarial supervision improves sharpness and realism.Most such approaches assume Teacher and Student operate in the same latent space.
  • E.2 Distillation for Fast Diffusion Sampling.: Efficient-modeling methods compress diffusion systems through pruning, structural compression, and lightweight architectures, but pruning can degrade quality and redesigns often require Teacher-specific retraining.The proposed approach instead keeps the compact Student architecture unchanged while transferring knowledge across different latent spaces.
  • E.3 Representation Alignment and Internal Distillation.: Representation-alignment distillation transfers internal structure in addition to matching final outputs, and Cross-Space Distillation targets mismatches in both latent resolution and latent parameterization with a lightweight Bridge.The Bridge maps Student features into the Teacher space.
  • E.4 Comparison to Deep Compression Autoencoder Approaches.: The method distills a 1024×1024 Teacher into a 512 × 512 Student using existing pretrained weights, 2M synthetic images, and around 10K training iterations.It reuses standard Stable Diffusion UNet backbones instead of requiring SANA-style DiT architectures and specialized acceleration modules.
  • F Societal Impacts: The work promotes fast, accessible, and high-quality image generation while recognizing that image manipulation methods may be misused to create deceptive content.It emphasizes robust detection of AI-generated or manipulated media and responsible deployment.
Loading 2606.32020v1…