⚡ electrical engineering

Diffusion-OAMP for Joint Image Compression and Wireless Transmission

The paper proposes Diffusion-OAMP, a training-free framework that integrates a pre-trained diffusion model into the OAMP algorithm to effectively solve the joint image compression and wireless transmission problem by leveraging generative priors for reconstruction.

Original authors: Wentao Hou, Yimin Bai, Zelei Luo, Jiadong Hong, Lei Liu

Published 2026-05-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Wentao Hou, Yimin Bai, Zelei Luo, Jiadong Hong, Lei Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to send a high-definition photo of a friend across a very noisy, crowded room. You want to shout the details of the photo to someone on the other side, but the room is full of static, and you can only whisper a few words at a time. This is the real-world problem the paper tackles: how to compress an image, send it over a wireless connection, and get it back perfectly clear, even when the signal is messy.

Here is how the authors solved this, explained simply:

1. The Problem: A Broken Puzzle

Usually, when we send images, we compress them (shrink them) and then send them. If the wireless signal is bad (like static on an old radio), the image comes back blurry or distorted.

The Old Way: Engineers used to guess what the missing pieces of the puzzle looked like using simple math rules (like "this part is probably blue because the sky is blue"). But these rules are too simple for complex photos of faces or landscapes.
The New Way: The authors realized that if we could teach a computer to "dream" of what a perfect face looks like, it could help fill in the missing puzzle pieces much better.

2. The Solution: "Diffusion-OAMP"

The paper introduces a new system called Diffusion-OAMP. Think of this system as a two-person team working together to fix the broken image:

Team Member A (The Linear Estimator): This is the "Math Guy." He looks at the noisy signal and the known rules of the wireless channel. He says, "Based on the noise, the image probably looks like this rough sketch." He is good at handling the math of the transmission but bad at making the picture look realistic.
Team Member B (The Diffusion Model): This is the "Art Expert." This is a pre-trained AI that has seen millions of photos. It knows exactly what a human face, a tree, or a building should look like. It doesn't need to be retrained; it just brings its knowledge to the table.

3. How They Work Together (The Dance)

The magic happens in how these two talk to each other. The paper describes a loop that happens over and over again:

The Rough Draft: The "Math Guy" takes the noisy signal and makes a rough guess.
The Translation: He translates this guess into a format the "Art Expert" can understand.
The Cleanup: The "Art Expert" looks at the rough draft and says, "Okay, this looks like a face, but the nose is weird. Let me fix it based on what I know about faces." It cleans up the image, removing the static and filling in missing details.
The Feedback: The "Art Expert" sends the cleaned-up version back to the "Math Guy."
The Check: The "Math Guy" checks this new version against the original noisy signal to make sure they didn't change the picture too much. If it looks good, they keep it. If not, they adjust and try again.

They repeat this "dance" a few times (usually just 3 times!) until the image is crystal clear.

4. The Secret Sauce: "SNR Matching"

One of the clever tricks in the paper is how they decide how much help the "Art Expert" needs at each step.

Imagine the "Art Expert" is a sculptor. If the clay is very muddy (high noise), the sculptor needs to do a lot of heavy lifting. If the clay is almost clean (low noise), the sculptor just needs to do a little polishing.
The system automatically measures how "muddy" the current guess is and tells the AI exactly how much to "denoise" it. This ensures the AI doesn't accidentally invent fake features (like adding a third eye to a face) just because it's trying too hard.

5. The Results

The authors tested this on a dataset of celebrity faces (CelebA) under various conditions:

Heavy Compression: Sending very small amounts of data.
Bad Channels: Simulating a very noisy wireless environment (like a busy city street).

The findings were:

Better Quality: Their method produced sharper, more realistic faces than older methods (like OAMP+BM3D or DDRM).
Robustness: Even when the signal was terrible or the image was compressed heavily, their system didn't break; it kept the important details (like facial features) intact.
Speed: The system converged (finished the job) very quickly, usually in just 3 rounds of the "dance."

Summary

In short, the paper proposes a smart receiver that doesn't just try to mathematically reverse the noise. Instead, it uses a pre-trained AI artist to "hallucinate" the missing details in a realistic way, while a mathematical guard ensures the hallucinations stay true to the original signal. The result is a clearer, more reliable way to send images over wireless networks, even when the connection is bad.

1. Problem Statement

The paper addresses the challenge of Joint Image Compression and Wireless Transmission (JCWT). Unlike generic image restoration, JCWT involves recovering a source image ( $s$ ) from a received signal ( $y$ ) that has undergone two distinct linear transformations:

Compression and Modulation: The source image is compressed and modulated into a transmitted signal ( $x$ ) via an operator $F$ .
Wireless Transmission: The signal passes through a wireless channel ( $A$ ) corrupted by Additive White Gaussian Noise (AWGN, $n$ ).

The system is modeled as a linear inverse problem:
$y = A(Fs) + n = Ax + n$
where $A$ represents the channel matrix and $F$ represents the joint compression/modulation operator.

Key Challenges:

Ill-Conditioning: The combined operator $AF$ often results in a severely ill-conditioned system, especially in fading channels (e.g., 3GPP TDL-A) and low compression ratios.
Prior Mismatch: Traditional recovery algorithms like Orthogonal Approximate Message Passing (OAMP) rely on analytical priors (e.g., Bernoulli-Gaussian) or hand-crafted denoisers (e.g., BM3D). These fail to capture the complex, high-dimensional distribution of natural images, leading to sub-optimal performance when the assumed prior mismatches the true image distribution.

2. Methodology: Diffusion-OAMP

The authors propose Diffusion-OAMP, a training-free reconstruction framework that integrates pre-trained generative diffusion models into the OAMP algorithm. The core idea is to use the diffusion model as a powerful Non-Linear Estimator (NLE) within the OAMP iterative loop.

A. System Architecture

The framework operates in the transmission domain ( $x$ ) but leverages the source domain ( $s$ ) for generative priors:

Linear Estimator (LE): An LMMSE estimator processes the noisy observation $y$ to produce a posterior estimate of $x$ .
Orthogonalization: The output is orthogonalized to decouple the linear estimator from the non-linear estimator, producing "extrinsic" pseudo-observations.
Inverse Random Multiplexing (RM): The estimate is mapped from the transmission domain back to the source image domain ( $s$ ) using the inverse of the compression operator ( $F^{-1}$ ).
SNR Matching: The algorithm calculates the effective noise variance of the current estimate and maps it to a specific diffusion time step ( $t^*$ ). This ensures the diffusion model operates at the correct noise level corresponding to the current iteration's uncertainty.
Non-Linear Estimator (NLE): A pre-trained diffusion model (either DDIM or Flow Matching) acts as a denoiser, projecting the noisy image estimate onto the learned data manifold.
MMSE Correction: The denoised image is remodulated ( $F$ ), and a correction step updates the prior for the next OAMP iteration using Stein's Unbiased Risk Estimator (SURE) to ensure orthogonality.

B. Key Technical Components

Random Multiplexing (RM): The compression operator $F$ combines a random sign matrix, a DCT matrix, a permutation matrix, and a row selector. This design ensures the equivalent channel matrix is right-unitarily invariant, a requirement for OAMP's state evolution analysis.
SNR-Matching Rule: The diffusion time step $t^*$ $t^{*}$ is dynamically selected based on the OAMP error variance ( $v_{post}^{orth}$ $v_{p os t}^{or t h}$ ):
- For Flow Matching: $t^* = \frac{1}{1 + \sqrt{v_{post}^{orth}}}$
- For DDIM: $t^* = \frac{1}{1 + v_{post}^{orth}}$
- As iterations progress and error variance decreases, $t^*$ moves toward the "clean" end of the diffusion trajectory.
Training-Free: The framework does not require retraining the diffusion models; it utilizes existing pre-trained models (e.g., on CelebA) as plug-and-play components.

3. Key Contributions

Novel Framework: Introduction of Diffusion-OAMP, the first framework to embed diffusion models specifically into the OAMP algorithm for the joint task of image compression and wireless transmission.
Prior Flexibility: The method is compatible with multiple generative priors. The authors instantiate it with both DDIM (Denoising Diffusion Implicit Models) and Flow Matching, demonstrating versatility.
SNR-Matching Mechanism: A novel rule to align the internal noise variance of the OAMP algorithm with the noise schedule of the diffusion model, enabling seamless integration without fine-tuning.
Robustness to Channel Distortion: The framework effectively handles the ill-conditioning caused by realistic wireless fading channels (3GPP TDL-A), where traditional methods often fail.

4. Experimental Results

The authors evaluated the framework on the CelebA dataset (256x256) under varying compression ratios ( $\beta \in \{0.1, 0.4, 0.7\}$ ) and noise levels ( $\sigma \in \{0.05, 0.5\}$ ).

Pure Image Compression (AWGN):
- Diffusion-OAMP (DDIM) achieved the highest PSNR, particularly in high-noise regimes ( $\sigma=0.5$ ), outperforming OAMP+BM3D and DDRM.
- Flow Matching showed superior SSIM (structural similarity) in some cases and faster convergence (reaching near-optimal PSNR in ~6 NFEs).
Joint Compression & Wireless Transmission (Fading Channels):
- In realistic 3GPP TDL-A fading channels, the observation matrix becomes severely ill-conditioned.
- Diffusion-OAMP maintained stable reconstruction, significantly outperforming baselines. For example, at low compression ( $\beta=0.1$ ) and high noise ( $\sigma=0.5$ ), Diffusion-OAMP (DDIM) achieved 21.95 dB PSNR, whereas OAMP+BM3D dropped to 16.20 dB.
- Visual results showed superior preservation of facial details and identity, avoiding the blurring and artifacts seen in baseline methods.
Convergence:
- The algorithm converges rapidly, stabilizing within 3 OAMP iterations across all tested scenarios.

5. Significance

Bridging Generative AI and Communications: This work demonstrates a practical application of large-scale generative models (Diffusion/Flow Matching) in physical layer communications, moving beyond simple image restoration to complex joint source-channel coding problems.
Handling Ill-Posed Problems: By using a learned generative prior, the framework regularizes the inverse problem more effectively than hand-crafted priors, making it robust against severe channel distortions and low-bandwidth constraints.
Efficiency: The "training-free" nature of the approach means it can leverage existing, state-of-the-art generative models without the computational cost of retraining them for specific channel conditions.
Future Direction: It opens a new pathway for "Semantic Communication" or "Deep Learning-based Physical Layer" design, where advanced generative models serve as the core receiver logic for next-generation intelligent communication systems.