跳转至

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

会议: ECCV 2026
论文: ECCV 官方
领域: 模型压缩
关键词: 扩散模型蒸馏, 单步生成, 跨空间蒸馏, 潜在表征对齐, 注意力保真度

一句话总结

提出跨空间蒸馏框架与轻量级 Bridge 接口,冻结学生 VAE 解码器前缀作空间先验并结合 SwinIR 投影器,突破师生潜在分辨率与 VAE 不匹配的限制,使轻量 SD 1.5 单步生成达到现代 1024 分辨率大教师水平。

研究背景与动机

近年来,Stable Diffusion 3.5、FLUX.2 等现代扩散与流匹配模型在超高分辨率生成上取得了突破性画质,但其庞大的参数量(如 2.5B-4.0B)与多步反向采样(如 50 步 NFE)带来了极其高昂的计算和显存开销,难以在端侧消费级硬件上实时部署。为了实现极速生成,变分分数蒸馏(VSD)与对抗式蒸馏(如 DMD2、ADD)等基于分布匹配的技术已成为将多步教师压缩为单步生成器的核心范式。然而,这类方法普遍依赖一个隐式假设——“共享潜在空间(Shared-Space)约束”,即假定教师网络与学生网络共享完全相同的潜在网格分辨率与 VAE 编解码器表征。

这一假设在面对跨代模型时构成了不可逾越的壁垒:现代大容量教师模型普遍在 1024×1024 图像空间以及更高维通道、不同结构的 VAE 潜在空间(如 16 或 32 通道)下运行,而轻量、端侧友好且具备广泛插件生态的存量骨干(如 SD 1.5、SD 2.1)则严格绑定在 512×512 分辨率与旧版 4 通道 VAE 上。由于潜在网格空间维度与潜在表征流形完全不匹配,标准的单步蒸馏目标(无论是分数匹配还是隐空间对抗鉴别器)根本无法直接计算。另一条常见的直接剪枝或重新设计微型架构路线(如 OBS-Diff、SnapGen),不仅面临严重的稀疏度质量悬崖,更需要极其高昂的特定调优与从头重训成本,且彻底丢失了 SD 1.5 既有的庞大生态兼容性。

既然轻量骨干的架构本身具备极高的推理效率与生态价值,阻碍其学习先进大教师知识的核心瓶颈仅在于潜在特征的异构性。核心 idea:形式化提出“跨空间蒸馏”范式,设计参数仅约 5M 的轻量级 Bridge 映射模块,以冻结的学生 VAE 解码器早期块提供空间升采样先验,配合紧凑的 SwinIR 投影器和逆向 KL 注意力保真度损失,在不改变学生骨干网络的前提下将单步学生潜在特征精确对齐至现代大教师潜在空间,解锁异构师生间的直接单步分布蒸馏。

方法详解

跨空间蒸馏(Cross-Space Distillation)通过引入即插即用的 Bridge 模块 \(B_\phi\),彻底解耦了生成骨干推理与教师监督表征空间之间的强绑定关系。在训练阶段,首先利用两阶段特征投影器将低分辨率学生潜变量映射为高分辨率教师兼容潜变量,并以潜在重建与注意力保真度进行双重对齐;完成 Bridge 训练后,保持 Bridge 参数冻结,直接将学生生成的潜在表征经 Bridge 投射到教师潜在空间,驱动标准 DMD2/VSD 及对抗判别器计算损失,更新学生骨干参数;在推理阶段,学生既可保持原生单步极速输出,亦可级联 Bridge 与教师解码器实现免训练分辨率无损升阶。

整体框架

整体框架由 Bridge 离线表征对齐训练、基于 Bridge 的单步学生端到端跨空间蒸馏、以及可选的推理期分辨率升级三部分构成。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    subgraph S1["Bridge 表征对齐训练"]
        direction TB
        A["图像样本 x"] --> B["双 VAE 编码<br/>提取配对潜变量 (zS, zT)"]
        B --> C["冻结解码器空间先验<br/>D_S^(1) 空间升采样到 H×W"]
        C --> D["SwinIR 投影头 g_ϕ<br/>对齐通道与语义至 z_hat_T"]
        D --> E["双重监督优化<br/>L1 潜在重建 + 逆向 KL 注意力保真度"]
    end
    subgraph S2["单步跨空间知识蒸馏"]
        direction TB
        F["随机高斯噪声 ξ + 文本 c"] --> G["单步学生骨干 S_θ<br/>预测原生潜变量 zS"]
        G --> H["冻结 Bridge 映射<br/>z_hat_T = B_ϕ(zS)"]
        H --> I["机制解耦重参数化<br/>统一转换为 x0 预测"]
        I --> J["大教师分布监督<br/>VSD 分数差异 + 对抗鉴别器"]
        J -->|反向传播更新| G
    end
    subgraph S3["推理期与后处理扩展"]
        direction TB
        K["单步学生原生推理<br/>原生分辨率 512×512 极速输出"]
        G -.->|可选直通| K
        H -.->|可选升级| L["教师 VAE 解码器 D_T<br/>超分辨率 1024×1024 输出"]
        G -.->|多教师学生检查点| M["多教师模型参数平均<br/>无额外推理开销的能力融合"]
    end

关键设计

1. 空间先验与投影头解耦的两阶段 Bridge 架构:无痛解决分辨率与通道异构 传统跨网格表征映射若直接从头学习卷积或转置卷积上采样,网络必须耗费大量容量去从零拟合复杂的空间插值与纹理重建规律,极易产生高频伪影或训练发散。本文敏锐观察到学生模型自有的预训练 VAE 解码器早已经过海量图像训练,其浅层解码块天生蕴含强大的隐空间空间上采样先验。作者将 Bridge 显式分解为两个阶段:第一阶段直接截取并冻结学生 VAE 解码器的前 \(n\) 个解码块 \(D_S^{(n)}\)(实验中直接取第一块 \(n=1\)),将潜在特征图的空间尺寸从学生端的 \(h \times w\)(如 \(64 \times 64\))直接升采样到匹配教师维度的 \(H \times W\)(如 \(128 \times 128\)),输出中间先验特征 \(f_{\text{prior}} \in \mathbb{R}^{H \times W \times C_{\text{prior}}}\),且不引入任何新参数;第二阶段则采用仅约 5M 参数的轻量 SwinIR 网络作为可学习投影头 \(g_\phi\),专职负责将 \(f_{\text{prior}}\) 的通道维度与语义流形投影为教师兼容潜变量 \(\hat{z}_T = g_\phi(f_{\text{prior}}) \in \mathbb{R}^{H \times W \times C_T}\)。该设计清晰地将“空间几何形变”与“流形语义变换”解耦,极大降低了学习难度并确保了特征重建的精准度。

2. 逆向 KL 散度注意力保真度损失:在全局语义空间锚定教师注意力响应 仅依靠潜在空间数值上的简单距离回归虽然能恢复主体结构,但微小的潜变量扰动在经过非线性解码后常放大为显著的局部语义畸变或解剖结构崩溃。常规的像素级重建或额外判别器微调极不稳定且超参敏感。为此,本文直接在冻结的教师去噪骨干网络内部引入注意力保真度约束:从教师网络指定层抽取自注意力图 \(P^l(z; t, c) \in [0, 1]^{N \times N}\)。自注意力天然捕获全局长程依赖,任何潜在表征的语义瑕疵都会在教师自注意力响应中被成倍放大。作者采用 MiniLLM 中的逆向 KL 散度(reverse-KL divergence),使其高度聚焦于注意力质量的主导分布,有效避免正向 KL 散度容易导致的平滑模糊。注意力保真度损失公式为: $$ \mathcal{L}{\text{attn}} = \sum_T; t, c) \,||\, P^l(z_T; t, c) \right) $$ 其中 }} \tau^2 \text{KL}\left( P^l(\hat{z\(\tau\) 为数值稳定性温度系数(实验设为 \(\tau = 3.0\)),加噪水平设定为 \(t=1\)。该损失迫使 Bridge 生成的潜变量不仅在像素/潜在空间数值接近,更能完全激发与真实教师潜变量一致的长程拓扑注意特征。

3. 机制无关的端点重参数化与多教师模型融合:彻底打通异构范式与能力集成 在完成跨空间对齐后,学生骨干(如 DDPM 噪声预测、UNet 架构)与教师模型(如 Rectified Flow 流匹配速度预测、DiT/MMDiT 架构)之间仍存在显著的机制代沟。如果采用层间特征对齐(layer-wise matching),将不可避免地受制于骨干结构的层数与连接限制。作者设计了完全机制无关的端点输出形式:无论底层是连续流还是反向随机扩散,在加噪路径端点处均存在确定性的解析映射,统一将学生输出反解并重参数化为一致的清洁潜变量估计 \(\hat{x}_0\)。通过将其输入 Bridge 映射后再注入教师 VSD 与对抗损失,实现了对 DiT、MMDiT 及 Large UNet 教师的无缝透明监督。更关键的是,由于所有蒸馏后的单步学生严格保留了原始完全相同的轻量骨干结构(如 SD 1.5),多教师蒸馏出的多个学生检查点集合可以通过简单的参数直接线性平均(Parameter Averaging)进行免训练模型融合: $$ \theta_{\text{merged}} = \frac{1}{|S|} \sum_{i=1}^{|S|} \theta_i $$ 该设计以零额外推理与存储成本,将 SDXL、Kolors、PixArt-Σ、SD 3.5 和 FLUX.2 五大现代教师的优势集于一身,甚至在多项评估中全面超越单教师蒸馏模型。

损失函数 / 训练策略

Bridge 模块本身的训练采用多任务联合损失: $$ \mathcal{L}{\text{final}} = \alpha \mathcal{L} $$ 其中潜在重建损失采用 }} + \beta \mathcal{L}_{\text{attn}\(\mathcal{L}_{\text{rec}} = \|z_T - \hat{z}_T\|_1\)。实验发现 \(\ell_1\) 范数相比常规 \(\ell_2\) 范数在潜在空间优化上具备显著更优的稳定性和边缘保持能力。超参数固定为 \(\alpha = 1, \beta = 1\)。 训练数据完全不依赖大规模真实图像数据集,而是从 JourneyDB 与 LAION 中采样文本 Prompt,由对应的教师模型离线生成约 200 万张高质合成图像构建静态训练集。利用该合成数据集先独立训练 Bridge 模块;在随后的单步学生蒸馏中,沿用 DMD2 训练范式,在 8 张 NVIDIA H100 80GB GPU 上使用 AdamW 优化器训练 20 小时,同时复用该 2M 合成集计算辅助对抗目标。

实验关键数据

主实验

论文在 SD 1.5 和 SD 2.1 两个经典的轻量级学生骨干上验证了跨空间蒸馏的卓越有效性,蒸馏教师涵盖 SDXL、Kolors、PixArt-Σ-1024、FLUX.2-klein-4B 以及 SD 3.5 Medium 五款代表性 1024 分辨率现代模型。评估指标覆盖人类偏好评分(HPSv3、HPSv2)、图文对齐质量(ImageReward、DPG Bench)及综合多维质量(MPS)。

学生模型 / 教师模型配置 骨干参数量 (B) 推理步数 (NFE) HPSv3 ↑ HPSv2 ↑ ImageReward ↑ MPS ↑ DPG Bench ↑
多步大教师上限参考
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 学生家族 (原生 512×512)
SD 1.5 DMD2 (初始化基线) 0.86 1 5.37 21.90 -0.29 10.42 59.85
+ SDXL 教师蒸馏 0.86 1 9.04 27.37 0.30 11.93 63.00
+ Kolors 教师蒸馏 0.86 1 9.33 27.38 0.34 11.66 64.46
+ PixArt-Σ-1024 教师蒸馏 0.86 1 8.65 26.43 0.30 11.94 63.43
+ FLUX.2-klein-4B 教师蒸馏 0.86 1 9.49 26.64 0.40 12.35 64.06
+ SD 3.5 Medium 教师蒸馏 0.86 1 9.42 28.30 0.62 12.96 65.75
Merged (五大教师参数均值融合) 0.86 1 10.53 29.07 0.65 12.62 66.67
SD 2.1 学生家族 (原生 512×512)
SD 2.1 SiD (初始化基线) 0.86 1 6.42 23.74 0.10 11.29 61.33
+ SDXL 教师蒸馏 0.86 1 8.40 26.73 0.23 11.88 67.03
+ Kolors 教师蒸馏 0.86 1 8.52 26.57 0.20 11.91 68.34
+ PixArt-Σ-1024 教师蒸馏 0.86 1 8.42 28.29 0.44 12.20 68.51
+ FLUX.2-klein-4B 教师蒸馏 0.86 1 8.74 28.31 0.33 12.03 66.00
+ SD 3.5 Medium 教师蒸馏 0.86 1 8.64 29.11 0.45 12.92 68.52
Merged (五大教师参数均值融合) 0.86 1 9.75 30.00 0.74 12.60 68.50

消融实验

1. 结构剪枝对比实验

对比采用最先进的扩散模型剪枝方案 OBS-Diff 将 SD 3.5 Medium 剪枝 30% 稀疏度(参数从 2.5B 降至 1.5B)后进行单步蒸馏,与本文 Bridge 赋能的 SD 1.5 极轻量学生(0.86B)的生成性能表现:

压缩/蒸馏设定 模型参数量 (B) 单步推理 NFE HPSv3 ↑ HPSv2 ↑ ImageReward ↑ MPS ↑ DPG Bench ↑
OBS-Diff 剪枝 30% 初始化 + 蒸馏 1.50 1 2.76 21.94 -0.36 9.54 55.24
本文 Bridge 跨空间蒸馏 (SD 1.5) 0.86 1 9.42 28.30 0.62 12.96 65.75

2. 学生模型初始化权重敏感度消融

以 SD 3.5 Medium 为教师模型,分析单步 DMD2 检查点初始化与标准多步 SD 1.5 预训练权重初始化的收敛与质量对比:

学生初始化方案 训练收敛步数 (Iterations) HPSv3 ↑ HPSv2 ↑ ImageReward ↑ MPS ↑ DPG Bench ↑
原始多步 SD 1.5 权重初始化 9k 9.22 28.44 0.66 13.08 66.89
SD 1.5 DMD2 单步检查点初始化 3k 9.42 28.30 0.62 12.96 65.75

关键发现

  • 跨空间蒸馏带来跨代级画质跨越:以 SD 1.5 为例,引入 SD 3.5 Medium 教师指导后,HPSv3 由 5.37 剧增至 9.42(+75.4%),ImageReward 由负值 -0.29 跃升为 0.62。定性对比显示,初始单步学生常伴随的纹理模糊、解剖结构畸变和低对比度问题被完全根除,高频质感与构图审美直逼 50 步大教师。
  • 参数平均融合形成“超教师”学生:由于所有经过不同教师蒸馏的学生处于完全相同的参数空间,直接线性平均生成的 Merged Checkpoint 在 SD 1.5 上取得了 10.53 的 HPSv3 与 29.07 的 HPSv2,甚至超越了 SDXL (9.25)、PixArt-Σ (9.62) 以及 FLUX.2-klein-4B (10.05) 原生多步教师自身的表现,证明不同先进大模型的视觉先验具有极强的互补性。
  • 剪枝路线存在不可逆的质量崩塌:剪枝基线即便保留高达 1.5B 参数(几乎为 SD 1.5 的两倍),其 HPSv3 仍断崖式暴跌至 2.76。这证实直接稀疏化大模型权重会破坏注意力拓扑结构的完整性,而借助 Bridge 复用高度成熟小模型的紧凑表征,是性价比显著更高的模型轻量化路径。
  • 单步检查点初始化提升 3 倍训练效率:从已具备单步生成能力的 DMD2 权重起步,仅需 3k 迭代即可达到最优指标,比从原生多步权重冷启动(9k 迭代)节省了 66.7% 的算力消耗。

亮点与洞察

  • 复用 VAE 解码器早期块作为空间先验:极其巧妙地借力了学生网络自带的预训练解码器浅层权重完成尺度跨越,避免了从零训练潜在升采样层引入的结构伪影与过拟合风险。
  • 逆向 KL 注意力保真度拉齐隐空间拓扑:在特征回归之外将教师自注意力分布的逆向 KL 散度作为正则项,抓住了扩散模型表征长程语义一致性的命脉,解决了潜在空间局部微差导致全局语义破损的痛点。
  • 均值融合打通异构大模型联合沉淀路径:将 FLUX、SD 3.5、Kolors 等结构迥异、潜在空间完全不兼容的前沿模型,通过各自向同构小模型的单步蒸馏,最终在统一权重空间中实现了“零成本能力大一统”。

局限与展望

  • 教师生成合成数据的领域偏置依赖:Bridge 的预训练依赖教师模型离线生成的约 200 万张合成图,若 Prompt 分布或教师生成域受限,可能导致 Bridge 在冷门极端分布上的表征映射产生偏差。
  • 推理期分辨率升级受限于教师解码速度:虽然 Bridge 允许在推理期级联教师 VAE 解码器实现 1024 分辨率输出,但教师解码器本身的显存与计算开销相对较大,削弱了部分在极致端侧场景下的纯单步效率优势。
  • 未来方向:探索无须离线预先合成数据的流式自适应 Bridge 训练,以及将跨空间蒸馏推广至视频扩散大模型(如 Sora 类架构向轻量时空 UNet 迁移)。

相关工作与启发

  • vs DMD / DMD2 (One-step Diffusion with Distribution Matching): DMD2 奠定了基于假教师与真教师分布匹配的高质量单步蒸馏基石,但其本质被严格限制在“同潜在空间、同 VAE、同分辨率”的自蒸馏范围内。本文通过 Bridge 解耦了表征空间,将 DMD2 范式成功拓展至跨代、跨分辨率、跨架构的现代大教师。
  • vs OBS-Diff / SnapFusion (Structural Pruning & Architecture Search): 剪枝与架构搜索旨在压缩大网络本身,但面临严重的非线性性能骤降且失去下游 LoRA 与 ControlNet 繁荣生态支持。本文证明通过跨空间蒸馏复用成熟轻量基座(如 SD 1.5),在相同推理算力下显著战胜结构剪枝模型。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次系统形式化跨空间单步扩散蒸馏问题,提出的解耦 Bridge 结构与注意力对齐策略新颖且通用。
  • 实验充分度: ⭐⭐⭐⭐⭐ 覆盖 5 种现代大教师、2 种轻量学生,横跨 DiT/UNet/Flow Matching 多种范式,指标详实全面。
  • 写作质量: ⭐⭐⭐⭐⭐ 动机阐述深刻,架构图解与形式化定义清晰严密,论证有力。
  • 价值: ⭐⭐⭐⭐⭐ 彻底盘活了生态庞大的存量轻量级模型,为端侧单步生成赋能与多前沿模型能力融合开辟了极具实用性的技术路线。