CoD: A Diffusion Foundation Model for Image Compression

AI summary

Overview

Research area: Generative image compression, specifically replacing text-to-image diffusion foundation models (Stable Diffusion) with a diffusion foundation model trained from scratch for compression.

Technical level: Intermediate. The paper assumes familiarity with denoising diffusion, rectified flow, rate-distortion-perception theory, and neural codec architecture, though the core argument is intuitive.

Scope: The paper introduces CoD, a compression-oriented diffusion foundation model trained on open image-only datasets, and evaluates it as a drop-in replacement for Stable Diffusion inside downstream diffusion codecs such as DiffC.

What This Paper Is About

Diffusion-based image codecs have generally been built on top of large text-to-image models like Stable Diffusion, inheriting text conditioning that is not designed for compression. Prior work (PerCo, later studies, and the zero-shot framework DiffC) shows that text conditions can actually hurt compression performance, especially at low bitrates, and that Stable Diffusion is not an ideal foundation model for codecs. The paper's goal is to train a diffusion foundation model from scratch that is explicitly oriented toward compression, so that compression and generation are optimized jointly end-to-end, with a native image encoder producing compact image tokens rather than text captions as the conditioning signal.

Key Contributions

1. The first compression-oriented diffusion foundation model (CoD): A model designed not as a fixed codec but as a general foundation model that can replace Stable Diffusion inside diverse diffusion-based compression algorithms, such as DiffC.

2. Low-cost, reproducible training: CoD trains in roughly 20 A100 GPU days versus about 6,250 A100 GPU days for Stable Diffusion v1.5 (about 300x faster, roughly 0.3% of the cost), using only open image datasets (ImageNet-21K, OpenImages, SA-1B, 22M images total) and no text-image pairs.

3. High compression performance at ultra-low bitrates: Replacing Stable Diffusion with CoD in DiffC yields state-of-the-art results, especially at ultra-low bitrates such as 0.0039 bpp.

4. New insights into diffusion codecs: A scaling-law study showing compression performance improves consistently with model size, a demonstration that a 49M-parameter CoD outperforms the 181M-parameter GAN-based MS-ILLM, and a comprehensive pixel-space versus latent-space comparison.

Main Findings

• Text conditioning is suboptimal for compression: The paper frames a text captioner plus Stable Diffusion as a compression encoder and decoder pair, and argues text has two drawbacks: human text describes fine spatial and texture detail inefficiently, and discrete text vocabularies are non-differentiable, preventing joint rate-distortion optimization.

• CoD enables joint distortion and perception optimization: Because CoD's condition encoder, entropy bottleneck, and diffusion decoder are trained together, unified rectified flow training with some samples at t = 0 directly minimizes one-step reconstruction distortion while multi-step sampling promotes perceptual quality.

• Training cost and reproducibility: Roughly 20 A100 GPU days (4 NVIDIA A100 GPUs for about 5 days) versus about 6,250 A100 GPU days for Stable Diffusion v1.5. Training is fully self-supervised on open image-only datasets totaling 22M images, and progresses from 256x256 for 400k steps (batch size 128) to 512x512 for 150k steps (batch size 64).

• Scaling law: Increasing pixel-space CoD parameters consistently improves compression performance. CoD achieves better reconstruction quality than the GAN-based MS-ILLM (181M parameters) even with only 49M parameters, indicating the gain comes from algorithmic improvements rather than model size. Most diffusion codecs are otherwise constrained near Stable Diffusion's approximately 860M parameters.

• Zero-shot distortion-perception control: At 0.0039 bpp, pixel-space CoD attains its best perceptual quality at 25 sampling steps and gains 3.4 dB PSNR (from 16.2 dB to 19.6 dB) when reduced to a single step, with intermediate step counts interpolating smoothly. PerCo (SD), which lacks this training formulation, does not show such controllable improvement.

• Pixel-space diffusion escapes the VAE bottleneck: DiffC on Stable Diffusion cannot surpass the SD-VAE reconstruction limit (PSNR 25.77 dB, LPIPS 0.069) and has a bitrate ceiling around 0.26 bpp; latent-space CoD remains bounded by VAE quality even up to 0.60 bpp. Pixel-space CoD scales beyond 1.0 bpp, surpasses latent diffusion models in LPIPS around 0.3 bpp, and reaches near-lossless about 47 dB PSNR at 4 bpp.

• VTM-level PSNR with better perceptual quality: Pixel-space CoD achieves BPP-PSNR comparable to VTM with a BD-Rate of -2.1% using VTM as anchor, while delivering better perceptual quality. HiFiC, MS-ILLM, and CDC (rho = 0.9) instead trade PSNR for perception, and TACO's removal of the GAN loss does not reach an optimal trade-off. Note that HiFiC, MS-ILLM, CDC, and TACO are optimized with LPIPS while CoD is not.

• Latent-space results: DiffC with latent-space CoD achieves the best reconstruction quality at lower bitrates (below 0.02 bpp); the gap narrows at higher bitrates as performance approaches the SD-VAE limit, and at some points (e.g., FID at 0.02 bpp on CLIC) CoD is slightly worse than DiffC due to smaller training scale.

• Extreme 64-bit compression (appendix): Compressing a 512x512 image to 64 bits (0.00024 bpp) with a 1/128 downsample (4x4 patches) and codebook size 16, latent-space 64-bit CoD reconstructs correct semantics on Kodak, while Stable Diffusion-based DDCM and DiffC fail to reconstruct correct semantics even at four times the bit cost. CoD-based DiffC reaches a FID of 70 using less than 10% of the bits of prior codecs.

• One-step distillation (appendix): A distilled one-step CoD is competitive with OneDC and StableCodec at ultra-low bitrates without relying on any pretrained Stable Diffusion components. It processes a 512x512 image in 25.2 ms, faster than a single Stable Diffusion step as used in OneDC and StableCodec.

• CoD as perceptual supervision (appendix): Finetuning only the MS-ILLM decoder with a CoD-based DMD loss for 2K steps (20K steps total including fake score learning) at BPP 0.011 improves LPIPS from 0.403 to 0.376, DISTS from 0.271 to 0.248, and FID from 92.5 to 80.92, at a PSNR cost from 21.43 dB to 21.08 dB.

Methodology in Plain English

CoD keeps the standard compression shape but swaps the components. An image encoder (residual blocks and attention layers) compresses the input to 1/32 resolution, an entropy bottleneck quantizes it into image tokens, and a condition decoder reconstructs intermediate conditions at 1/16 resolution. A diffusion model then denoises Gaussian noise conditioned on those conditions to reconstruct the image.

The bottleneck is deliberately tight: vector quantization with a codebook of 16 entries gives 4 bits per 32x32 patch, which equals 0.0039 bpp. This forces the encoder to transmit only essential semantic information and leaves the diffusion model responsible for realistic detail. The diffusion module decouples a DiT backbone from a DDT head so the heavy backbone runs at 1/16 resolution, making pixel-space CoD similar in computational complexity to latent-space CoD. In latent space, VAE latents at 1/8 resolution are downsampled with a 2x2 patch embedding; in pixel space, a 16x16 patch embedding maps noised images to 1/16 resolution.

Training uses rectified flow, which learns a velocity field rather than a score and gives a clearer geometric view and more stable training. The key idea is unified training: the rectified flow loss at t = 0 reduces to MSE between the image and the one-step estimate, which is exactly a distortion term. So during training, a portion of samples are taken at t = 0 and the rest across t in [0, 1], jointly optimizing distortion and perception. The authors note color shifts appear without this t = 0 term, and that DDPM-based codecs like PerCo implicitly include the fully noised state as a special case.

The overall objective combines the rectified flow loss, a REPA representation alignment loss using DINOv2, a codebook commitment loss, and an auxiliary loss that reconstructs both pixels and DINOv2 features from the condition, with weights of 0.5, 0.25, and 1.0. Timesteps follow a log-normal distribution. Sampling uses an Adam-like second-order solver with 25 steps and classifier-free guidance at scales of 3.0 (pixel) and 1.25 (latent). Evaluation uses DiffC on Kodak, CLIC2020, and Div2K at 512x512, reporting PSNR, LPIPS, DISTS, and FID, with FID computed on extracted patches (64x64 for Kodak, 128x128 for CLIC2020).

Why This Matters

The paper reframes a foundation-model question for a specific domain: instead of adapting a general-purpose text-to-image model to compression, it trains a purpose-built model that is cheaper, reproducible, and outperforms the general model it replaces. It also provides evidence that a diffusion codec can reach VTM-level PSNR while beating GAN-based perceptual codecs, which argues against the assumption that perceptual quality must come at a large distortion cost. Releasing training code and relying only on open datasets lowers the barrier for other groups to push diffusion codec research beyond Stable Diffusion's roughly 860M-parameter ceiling.

Real-world applications (implications of image compression foundation models rather than claims measured in this paper):

• Bandwidth-constrained image transmission, where a 512x512 image can be represented in as little as 64 bits while preserving semantics, and where bitrate can be traded against sampling steps at decode time.
• Photo storage and cloud services that need both high fidelity and plausible perceptual quality across a wide bitrate range, rather than being capped by a VAE's reconstruction limit.
• Generative restoration, super-resolution, or perceptual post-processing, since CoD can act as a perceptual supervision signal (demonstrated by finetuning the MS-ILLM decoder).
• Latency-sensitive deployments, supported by the distilled one-step CoD running at 25.2 ms for a 512x512 image.

Industry relevance centers on the cost of producing a domain foundation model (about 20 A100 GPU days versus about 6,250 for Stable Diffusion v1.5), the use of entirely open datasets, and the ability to plug the model into existing codec frameworks as a drop-in replacement for Stable Diffusion.

Future Directions

• High-resolution training: CoD is not yet trained for high-resolution inputs. Extending to resolutions such as 2K is described as technically feasible but computationally expensive, and is left for future work.
• Real-time coding: Like all diffusion codecs, CoD does not meet real-time requirements, though the authors see diffusion distillation as a path toward it.
• Better codebooks and entropy coding: The one-step variant currently uses fixed-length coding for codebook indices; the paper suggests gains from an entropy model or latent compression.
• Scaling data and training: Some metrics (for example FID at 0.02 bpp on CLIC) lag behind DiffC because of smaller training scale, which the authors expect larger data and training to close.
• Broader downstream integration: The paper positions CoD as a backbone for future diffusion codec architectures, training schemes, and applications such as perceptual supervision and joint finetuning of entire codec networks.

Target Audience

Researchers and engineers working on learned image compression, diffusion models, or generative codecs; practitioners evaluating whether to replace Stable Diffusion with a domain-specific foundation model; and readers interested in scaling laws, the pixel-space versus latent-space trade-off, and the rate-distortion-perception trade-off. The paper is most useful to those with some background in diffusion and neural compression, but its central argument and headline numbers are accessible to a broader technical audience.

Authors’ abstract

Existing diffusion codecs typically build on text-to-image diffusion foundation models like Stable Diffusion. However, text conditioning is suboptimal from a compression perspective, hindering the potential of downstream diffusion codecs, particularly at ultra-low bitrates. To address it, we introduce \textbf{CoD}, the first \textbf{Co}mpression-oriented \textbf{D}iffusion foundation model, trained from scratch to enable end-to-end optimization of both compression and generation. CoD is not a fixed codec but a general foundation model designed for various diffusion-based codecs. It offers several advantages: \textbf{High compression efficiency}, replacing Stable Diffusion with CoD in downstream codecs like DiffC achieves SOTA results, especially at ultra-low bitrates (e.g., 0.0039 bpp); \textbf{Low-cost and reproducible training}, 300$\times$ faster training than Stable Diffusion ($\sim$ 20 vs. $\sim$ 6,250 A100 GPU days) on entirely open image-only datasets; \textbf{Providing new insights}, e.g., We find pixel-space diffusion can achieve VTM-level PSNR with high perceptual quality and can outperform GAN-based codecs using fewer parameters. We hope CoD lays the foundation for future diffusion codec research. Codes are released at https://github.com/microsoft/GenCodec/tree/main/CoD.