DiT-IC: Aligned Diffusion Transformer for Efficient Image Compression

会议: CVPR 2025
arXiv: 2603.13162
代码: https://njuvision.github.io/DiT-IC/
领域: 图像压缩 / 扩散模型
关键词: Diffusion Transformer, 图像压缩, 单步推理, 方差引导, 自蒸馏

一句话总结

DiT-IC 将预训练 T2I 扩散 Transformer 适配为单步图像压缩重建模型，在 32x 下采样的深层潜空间工作，通过方差引导重建流、自蒸馏对齐和潜变量条件引导三种对齐机制，实现 SOTA 感知质量且解码比现有扩散 codec 快 30 倍。

研究背景与动机

1. 领域现状：基于扩散的图像压缩感知质量突出，但采样开销（4-50步）和高内存限制实用性。
2. 现有痛点：现有扩散 codec 基于 U-Net，被迫在 8x 浅层潜空间工作。传统 VAE codec 在 16x-64x。
3. 核心矛盾：压缩重建从结构化潜变量开始（接近数据流形），多步去噪可能冗余；但直接微调生成模型导致流形不匹配。
4. 本文要解决什么：让扩散在 32x 深层潜空间单步推理。
5. 切入角度：三种"对齐"机制将预训练 T2I DiT（SANA）适配为压缩重建模型。
6. 核心idea一句话：生成到重建的三重对齐——方差引导去噪、自蒸馏多步到单步、潜条件替代文本。

方法详解

整体框架

ELIC 编码器 + SANA DiT 重建器。编码器产生量化潜变量，DiT 在 32x 空间做单步方差引导 flow matching。LoRA 高效微调。

关键设计

1. 方差引导重建流
2. 做什么：将多步扩散折叠为单步自适应变换
3. 核心思路：利用编码器方差作空间不确定性度量，映射为逐像素伪时步。高方差区域更强去噪

4. 设计动机：压缩噪声空间异质，全局单一时步不够

5. 自蒸馏对齐

6. 做什么：多步扩散行为蒸馏到单步
7. 核心思路：冻结编码器，DiT 输出对齐编码器潜变量（余弦对齐 + margin）

8. 设计动机：深层潜空间无现成多步教师

9. 潜条件引导

10. 做什么：用压缩表示替代文本条件
11. 核心思路：轻量投影映射到预训练文本编码器同一嵌入空间；CLIP 式对比对齐
12. 设计动机：文本对细粒度空间不足，且需重型编码器

损失函数 / 训练策略

两阶段 IBP：Stage1 100K iter(256)，Stage2 60K iter(512)。LoRA rank VAE=32, DiT=64。

实验关键数据

主实验

方法	步数	延迟	LPIPS BD-rate	DISTS BD-rate
StableCodec	1	0.34s	-79.19%	-83.95%
ResULIC	4	0.83s	-62.27%	-65.64%
OSCAR	1	0.32s	-19.04%	-58.38%
DiT-IC	1	0.15s	-83.65%	-87.88%

消融实验

配置	LPIPS	DISTS
完整	0.00%	0.00%
从零训练	+22.00%	+32.45%
Full finetuning	+7.95%	+8.05%

关键发现

• SOTA 感知质量，解码 0.15s 比 DiffEIC 12.4s 快 80x
• 可在 16GB 笔记本 GPU 重建 2048x2048 图像
• 预训练初始化极关键，从零训练 DISTS -32.45%
• LoRA 32/64 优于 full finetuning

亮点与洞察

• 首次将 DiT 用于图像压缩，证明深层潜空间可行性
• 三种对齐机制层层递进
• 单步+深层潜空间=极低延迟和内存

局限性 / 可改进方向

• 极低码率潜变量信息不足
• 训练数据仅 150K
• 未与 OneDC 等同期工作对比

相关工作与启发

• StableCodec 开创单步扩散压缩但受限于 U-Net 架构，无法在深层潜空间工作
• SANA 的线性注意力 DiT 为高效扩散提供了架构基础，本文首次将其适配到压缩任务
• ResULIC 在 32x 潜空间做 4 步扩散，本文证明了单步即可在同一空间达到更好效果
• 自蒸馏的思路（用编码器自身做目标）比外部教师蒸馏更优雅，可推广到其他重建任务
• VA-VAE 的特征对齐范式启发了本文的自蒸馏设计

评分

• 新颖性: ⭐⭐⭐⭐ 首次 DiT 用于压缩，三种对齐各有独到
• 实验充分度: ⭐⭐⭐⭐ 三测试集+多基线+详尽消融
• 写作质量: ⭐⭐⭐⭐ 架构图清晰，消融逐步论证
• 价值: ⭐⭐⭐⭐⭐ 30x加速+SOTA质量=扩散压缩实用化关键一步