adu2021 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.
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.
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.
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.
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.