Table of Contents
Fetching ...

DiT-IC: Aligned Diffusion Transformer for Efficient Image Compression

Junqi Shi, Ming Lu, Xingchen Li, Anle Ke, Ruiqi Zhang, Zhan Ma

Abstract

Diffusion-based image compression has recently shown outstanding perceptual fidelity, yet its practicality is hindered by prohibitive sampling overhead and high memory usage. Most existing diffusion codecs employ U-Net architectures, where hierarchical downsampling forces diffusion to operate in shallow latent spaces (typically with only 8x spatial downscaling), resulting in excessive computation. In contrast, conventional VAE-based codecs work in much deeper latent domains (16x - 64x downscaled), motivating a key question: Can diffusion operate effectively in such compact latent spaces without compromising reconstruction quality? To address this, we introduce DiT-IC, an Aligned Diffusion Transformer for Image Compression, which replaces the U-Net with a Diffusion Transformer capable of performing diffusion in latent space entirely at 32x downscaled resolution. DiT-IC adapts a pretrained text-to-image multi-step DiT into a single-step reconstruction model through three key alignment mechanisms: (1) a variance-guided reconstruction flow that adapts denoising strength to latent uncertainty for efficient reconstruction; (2) a self-distillation alignment that enforces consistency with encoder-defined latent geometry to enable one-step diffusion; and (3) a latent-conditioned guidance that replaces text prompts with semantically aligned latent conditions, enabling text-free inference. With these designs, DiT-IC achieves state-of-the-art perceptual quality while offering up to 30x faster decoding and drastically lower memory usage than existing diffusion-based codecs. Remarkably, it can reconstruct 2048x2048 images on a 16 GB laptop GPU.

DiT-IC: Aligned Diffusion Transformer for Efficient Image Compression

Abstract

Diffusion-based image compression has recently shown outstanding perceptual fidelity, yet its practicality is hindered by prohibitive sampling overhead and high memory usage. Most existing diffusion codecs employ U-Net architectures, where hierarchical downsampling forces diffusion to operate in shallow latent spaces (typically with only 8x spatial downscaling), resulting in excessive computation. In contrast, conventional VAE-based codecs work in much deeper latent domains (16x - 64x downscaled), motivating a key question: Can diffusion operate effectively in such compact latent spaces without compromising reconstruction quality? To address this, we introduce DiT-IC, an Aligned Diffusion Transformer for Image Compression, which replaces the U-Net with a Diffusion Transformer capable of performing diffusion in latent space entirely at 32x downscaled resolution. DiT-IC adapts a pretrained text-to-image multi-step DiT into a single-step reconstruction model through three key alignment mechanisms: (1) a variance-guided reconstruction flow that adapts denoising strength to latent uncertainty for efficient reconstruction; (2) a self-distillation alignment that enforces consistency with encoder-defined latent geometry to enable one-step diffusion; and (3) a latent-conditioned guidance that replaces text prompts with semantically aligned latent conditions, enabling text-free inference. With these designs, DiT-IC achieves state-of-the-art perceptual quality while offering up to 30x faster decoding and drastically lower memory usage than existing diffusion-based codecs. Remarkably, it can reconstruct 2048x2048 images on a 16 GB laptop GPU.
Paper Structure (18 sections, 7 equations, 19 figures, 4 tables)

This paper contains 18 sections, 7 equations, 19 figures, 4 tables.

Figures (19)

  • Figure 1: Overview of reconstructed results and efficiency of our proposed DiT-IC.
  • Figure 2: Architectural comparison. The left panel illustrates the overall diffusion-based image compression framework. U-Net-based diffusers perform multi-stage downsampling, while DiTs maintain a constant spatial resolution throughout the denoising process, making them naturally compatible with deeply compressed latent inputs.
  • Figure 3: Overview of the proposed DiT-IC framework. Following StableCodec zhang2025stablecodec, we adopt ELIC he2022elic as our auxiliary encoder.
  • Figure 4: Variance-Guided Flow Matching. Unlike standard diffusion that starts from Gaussian noise, compression reconstruction begins from a quantized latent ${\mathbf{y}}_t$ containing structured noise. The local variance $\sigma(\mathbf{y}_t)$ measures spatial uncertainty, which we map to pseudo-timesteps $t = \mathcal{F}(\sigma)$ for spatially adaptive one-step flow matching.
  • Figure 5: Ablation study of variance-guided reconstruction flow.
  • ...and 14 more figures