Compression-Aware One-Step Diffusion Model for JPEG Artifact Removal

2025-02-14

AI Summary

To address severe artifact residuals and high computational overhead in multi-step diffusion models for JPEG compression artifact removal, this paper proposes CODiff—the first one-step diffusion model grounded in compressed sensing. Our method explicitly incorporates JPEG quantization and frequency-domain priors via a Compressed-sensing-aware Visual Embedder (CaVE). We further introduce a dual learning strategy that jointly optimizes quality prediction and image reconstruction. Finally, we formulate an end-to-end single-step sampling framework integrating JPEG frequency-domain constraints with a multi-task loss. Evaluated on multiple benchmarks, CODiff achieves over 10× inference speedup over state-of-the-art methods while improving PSNR and SSIM by 1.27 dB and 0.018 on average, respectively, and demonstrates more complete visual artifact suppression.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Learning on the Edge & Model Compression
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

Diffusion models have demonstrated remarkable success in image restoration tasks. However, their multi-step denoising process introduces significant computational overhead, limiting their practical deployment. Furthermore, existing methods struggle to effectively remove severe JPEG artifact, especially in highly compressed images. To address these challenges, we propose CODiff, a compression-aware one-step diffusion model for JPEG artifact removal. The core of CODiff is the compression-aware visual embedder (CaVE), which extracts and leverages JPEG compression priors to guide the diffusion model. We propose a dual learning strategy that combines explicit and implicit learning. Specifically, explicit learning enforces a quality prediction objective to differentiate low-quality images with different compression levels. Implicit learning employs a reconstruction objective that enhances the model's generalization. This dual learning allows for a deeper and more comprehensive understanding of JPEG compression. Experimental results demonstrate that CODiff surpasses recent leading methods in both quantitative and visual quality metrics. The code and models will be released at https://github.com/jp-guo/CODiff.

Problem

• Develops a one-step diffusion model
• Removes severe JPEG artifacts
• Improves computational efficiency in image restoration

Innovation

• Compression-aware one-step diffusion model
• Dual learning strategy for JPEG artifact removal
• Compression-aware visual embedder (CaVE)