Sodec Diffusion Compression

---
name: sodec-diffusion-compression
title: SODEC - Steering One-Step Diffusion for Fast Image Compression
version: 0.0.2
engine: skillxiv-v0.0.2-claude-opus-4.6
license: MIT
url: https://arxiv.org/abs/2508.04979
keywords: [image-compression, diffusion-models, generative-models, single-step-decoding]
description: "Replaces iterative diffusion with single-step decoding for image compression. Combines VAE latents with fidelity guidance and rate annealing training. Achieves 20× decoding speedup with improved perceptual quality."
---

# SODEC: Steering One-Step Diffusion for Fast Image Compression

## Core Concept

Traditional generative compression uses iterative diffusion refinement that requires dozens of denoising steps, making inference prohibitively slow. SODEC demonstrates that information-rich latent representations eliminate the need for iteration. By combining pre-trained VAE encodings with a fidelity guidance module and rate-annealing training strategy, the approach achieves single-step decoding with 20× speedup while maintaining or improving perceptual quality.

## Architecture Overview

- **VAE-Based Latent Generation**: High-information-density latent representation
- **Fidelity Guidance Module**: Keeps generated images faithful to originals
- **Single-Step Decoding**: Direct latent-to-image generation without iteration
- **Rate Annealing Training**: Progressive adjustment of compression-quality trade-off
- **Adaptive Rate Control**: Dynamic bit allocation based on image complexity

## Implementation Steps

### Step 1: Build Information-Rich Latent Space

Create VAE-based latent representation that captures sufficient information for reconstruction.

```python
import torch
import torch.nn as nn
from typing import Tuple

class InformationRichVAE(nn.Module):
    """
    VAE encoder that creates information-rich latents for compression.
    """

    def __init__(self, in_channels=3, latent_dim=16, hidden_dim=128):
        super().__init__()
        self.latent_dim = latent_dim

        # Encoder: image -> latent
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, hidden_dim, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim, hidden_dim * 2, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(hidden_dim * 2, hidden_dim * 4, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
            nn.Flatten()
        )

        self.fc_mu = nn.Linear(hidden_dim * 4 * 16, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim * 4 * 16, latent_dim)

        # Decoder: latent -> image
        self.fc_decode = nn.Linear(latent_dim, hidden_dim * 4 * 4 * 4)

        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(hidden_dim * 4, hidden_dim * 2, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(hidden_dim * 2, hidden_dim, 4, stride=2, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(hidden_dim, in_channels, 4, stride=2, padding=1),
            nn.Tanh()
        )

    def encode(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Encode image to latent distribution.

        Args:
            x: Image tensor [batch, channels, height, width]

        Returns:
            (mu, logvar) latent distribution parameters
        """
        h = self.encoder(x)
        mu = self.fc_mu(h)
        logvar = self.fc_logvar(h)

        return mu, logvar

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        """Sample from latent distribution."""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        """
        Decode latent to image.

        Args:
            z: Latent code [batch, latent_dim]

        Returns:
            Reconstructed image [batch, channels, height, width]
        """
        h = self.fc_decode(z)
        h = h.view(h.shape[0], -1, 4, 4)
        x_recon = self.decoder(h)

        return x_recon

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Full forward pass."""
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        x_recon = self.decode(z)

        return x_recon, mu, logvar
```

### Step 2: Implement Fidelity Guidance Module

Create module that ensures generated images remain faithful to original.

```python
class FidelityGuidanceModule(nn.Module):
    """
    Steers generation toward fidelity to original image.
    """

    def __init__(self, hidden_dim=128):
        super().__init__()

        # Learned guidance function
        self.guidance_net = nn.Sequential(
            nn.Linear(16, hidden_dim),  # latent_dim
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 3),  # Guidance for RGB channels
        )

    def compute_guidance(
        self,
        latent: torch.Tensor,
        generated_image: torch.Tensor,
        original_image: torch.Tensor,
        guidance_strength: float = 1.0
    ) -> torch.Tensor:
        """
        Compute fidelity guidance signal.

        Args:
            latent: Latent code
            generated_image: Generated image from latent
            original_image: Original image to stay faithful to
            guidance_strength: Strength of fidelity guidance

        Returns:
            Guidance tensor to apply to generation
        """
        # Compute pixel-level difference
        pixel_diff = (generated_image - original_image).abs().mean(dim=(2, 3))  # [batch, channels]

        # Network-based guidance
        net_guidance = self.guidance_net(latent)

        # Combine pixel and network guidance
        guidance = pixel_diff + guidance_strength * net_guidance

        return guidance

    def apply_guidance(
        self,
        generated_image: torch.Tensor,
        guidance: torch.Tensor,
        learning_rate: float = 0.01
    ) -> torch.Tensor:
        """
        Apply guidance to adjust generation.

        Args:
            generated_image: Current generated image
            guidance: Guidance signal
            learning_rate: Strength of adjustment

        Returns:
            Adjusted image
        """
        # Gradient-based adjustment toward fidelity
        guided_image = generated_image - learning_rate * guidance.view(-1, 3, 1, 1)

        return guided_image
```

### Step 3: Implement Single-Step Decoder

Create decoder that generates image in one step from latent.

```python
class SingleStepDecoder(nn.Module):
    """
    Direct latent-to-image decoding without iterative refinement.
    """

    def __init__(self, latent_dim=16, hidden_dim=256):
        super().__init__()

        # Direct mapping from latent to image
        self.latent_to_image = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim * 8),
            nn.ReLU(),
            nn.Linear(hidden_dim * 8, hidden_dim * 4),
            nn.ReLU(),
            nn.Linear(hidden_dim * 4, 3 * 256 * 256),  # 256x256 RGB image
            nn.Tanh()
        )

        # Confidence estimation network
        self.confidence_net = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )

    def forward(self, latent: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Generate image from latent in single step.

        Args:
            latent: Latent code [batch, latent_dim]

        Returns:
            (generated_image, confidence)
        """
        # Direct generation
        image_flat = self.latent_to_image(latent)
        image = image_flat.view(-1, 3, 256, 256)

        # Confidence score
        confidence = self.confidence_net(latent)

        return image, confidence

    def decode_batch(
        self,
        latents: torch.Tensor,
        apply_guidance: bool = False,
        guidance_module: FidelityGuidanceModule = None,
        original_images: torch.Tensor = None
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Decode batch of latents with optional guidance.

        Args:
            latents: Batch of latent codes
            apply_guidance: Whether to apply fidelity guidance
            guidance_module: Guidance module if applying guidance
            original_images: Original images for guidance

        Returns:
            (generated_images, confidences)
        """
        images, confidences = self.forward(latents)

        if apply_guidance and guidance_module is not None:
            for _ in range(3):  # 3 refinement steps with guidance
                guidance = guidance_module.compute_guidance(
                    latents,
                    images,
                    original_images,
                    guidance_strength=0.1
                )

                images = guidance_module.apply_guidance(images, guidance)

        return images, confidences
```

### Step 4: Implement Rate Annealing Training

Create training strategy that progressively tightens compression constraints.

```python
class RateAnnealingTrainer:
    """
    Train with rate annealing: progressively increase compression.
    """

    def __init__(self, vae, decoder, guidance_module):
        self.vae = vae
        self.decoder = decoder
        self.guidance = guidance_module

    def train_step(
        self,
        images: torch.Tensor,
        current_rate: float,
        target_rate: float,
        num_steps: int
    ):
        """
        Single training step with rate annealing.

        Args:
            images: Batch of images
            current_rate: Current compression rate (bits/pixel)
            target_rate: Target compression rate
            num_steps: Current step number

        Returns:
            Loss value
        """
        # Encode images
        mu, logvar = self.vae.encode(images)
        latent = self.vae.reparameterize(mu, logvar)

        # Decode
        decoded_images, confidence = self.decoder(latent)

        # Apply guidance
        guided_images = self.guidance.compute_guidance(latent, decoded_images, images)

        # Reconstruction loss
        recon_loss = torch.mean((decoded_images - images) ** 2)

        # Rate loss: encourage lower bitrate as training progresses
        # Current rate decreases as we train
        rate_schedule = target_rate + (current_rate - target_rate) * (1.0 - num_steps / 10000)

        # KL divergence for compression
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        kl_loss = kl_loss / images.shape[0]

        # Annealed loss: increase rate penalty over time
        annealing_factor = min(1.0, num_steps / 5000)
        rate_loss = annealing_factor * kl_loss

        # Total loss
        total_loss = recon_loss + 0.1 * rate_loss

        return total_loss

    def train_epoch(
        self,
        dataloader,
        num_epochs: int = 10,
        target_rate: float = 0.5
    ):
        """
        Train for one epoch with rate annealing.

        Args:
            dataloader: Training data
            num_epochs: Number of epochs
            target_rate: Target compression rate
        """
        optimizer = torch.optim.Adam(
            list(self.vae.parameters()) +
            list(self.decoder.parameters()) +
            list(self.guidance.parameters()),
            lr=1e-4
        )

        total_steps = 0

        for epoch in range(num_epochs):
            epoch_loss = 0

            for batch_images in dataloader:
                # Get current rate based on progress
                current_rate = 2.0  # Start with 2 bits/pixel

                # Training step
                loss = self.train_step(
                    batch_images,
                    current_rate=current_rate,
                    target_rate=target_rate,
                    num_steps=total_steps
                )

                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

                epoch_loss += loss.item()
                total_steps += 1

            print(f"Epoch {epoch}: Loss={epoch_loss / len(dataloader):.4f}")
```

## Practical Guidance

### When to Use SODEC

- **Real-time image compression**: Latency-critical applications requiring single-step decoding
- **Edge deployment**: 20× speedup enables efficient compression on mobile/edge devices
- **High-bitrate applications**: Image compression where generation quality matters
- **Perceptual compression**: Generative approach preserves perceptual quality

### When NOT to Use SODEC

- **Maximum compression**: Iterative refinement can achieve slightly better compression
- **Lossless compression**: Generative approach inherently lossy
- **Legacy format compatibility**: Output is image, not traditional bitstream
- **Ultra-low compute**: Requires capable decoder, still heavier than traditional codecs

### Hyperparameter Recommendations

- **Latent dimension**: 12-24 bits effective per image
- **Guidance strength**: 0.05-0.2 balances fidelity and generation quality
- **Target rate**: 0.3-1.0 bits/pixel depending on quality requirements
- **Rate annealing schedule**: Linear from 2.0 to target_rate over 10k steps

### Key Insights

The critical insight is that information-rich VAE latents reduce the need for iterative refinement. By starting from a high-quality latent representation and applying single-step decoding with guidance, the approach maintains quality while achieving massive speedup. Rate annealing prevents early mode collapse during training.

## Reference

**Steering One-Step Diffusion for Fast Image Compression** (arXiv:2508.04979)

Replaces iterative diffusion decoding with single-step generation for image compression. Achieves 20× decoding speedup through VAE latents, fidelity guidance, and rate-annealing training while maintaining or improving perceptual quality.