🤖 machine learning

Joint Enhancement and Classification using Coupled Diffusion Models of Signals and Logits

This paper proposes a novel, domain-agnostic framework that jointly enhances signals and classifies them using coupled diffusion models operating on both input signals and classifier logits, enabling mutual guidance to achieve superior robustness in noisy environments without retraining the classifier.

Original authors: Gilad Nurko, Roi Benita, Yehoshua Dissen, Tomohiro Nakatani, Marc Delcroix, Shoko Araki, Joseph Keshet

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Gilad Nurko, Roi Benita, Yehoshua Dissen, Tomohiro Nakatani, Marc Delcroix, Shoko Araki, Joseph Keshet

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 recognize a friend's face in a crowded, foggy room. The fog is so thick that their features are blurry and distorted.

The Old Way (The "Clean First" Approach)
Traditionally, if you wanted to identify your friend, you would first hire a "fog-remover" to clear the air. You'd wait for the room to become crystal clear, and then you would look at your friend to guess who they are.

The Problem: The fog-remover doesn't know who you are looking for. They just try to make the image look "nice" or "sharp" based on general rules. Sometimes, in trying to make the image look good, they accidentally smooth out a key feature (like a unique scar or a specific hat) that is actually crucial for identifying your friend. They make the picture pretty, but they might ruin the clues needed for the ID.

The New Way (The "Coupled Diffusion" Approach)
This paper proposes a smarter team-up. Instead of working in separate steps, they work together simultaneously. Imagine you have two experts standing side-by-side:

The Image Restorer: Someone trying to clear up the blurry photo.
The Guessing Expert: Someone trying to figure out who the person is, even with the blur.

How They Help Each Other (The "Mutual Guidance")
Instead of waiting for the photo to be perfect before guessing, they talk to each other constantly:

The Guessing Expert says: "Hey, I think that blurry blob is a nose! If you sharpen that specific spot, it will look more like a nose."
The Image Restorer says: "Okay, I'll focus my cleaning power right there."
The Image Restorer says: "I think this blurry shape is a hat. Does that help you guess who it is?"
The Guessing Expert says: "Yes! If it's a hat, it's definitely my friend Bob, not Alice. Now that I know it's Bob, I can tell you exactly what his face should look like."

They keep doing this back-and-forth, refining the picture and the guess at the same time. The guess helps clean the picture, and the cleaned picture helps make a better guess.

The Three Ways They Can Talk
The paper tests three different ways for these two experts to coordinate their conversation:

Parallel (The "Fast Chat"): They talk every single second. As soon as the image gets slightly clearer, the guesser updates their thought, and the image gets slightly clearer again. It's fast and efficient, like a rapid-fire conversation.
Alternating (The "Turn-Taking"): The image restorer works alone until they have a "good enough" version of the photo. Then, they hand it to the guesser, who works alone to make a final guess. Then they swap again. It's like taking turns, which is very stable but takes longer.
Nested (The "Deep Dive"): Every time the guesser makes a tiny update, the image restorer goes into a "deep dive" to make sure the image is perfect before the guesser moves on. This is the most careful and accurate method, but it takes the most time and computing power.

The Results
The researchers tested this on two things:

Pictures: They took photos of numbers (like "8" or "2") and covered them with static noise. The old way often guessed the wrong number because the noise made the lines look weird. The new "team-up" method looked at the shape of the number (the "logits" or the guess) and used that to fix the missing lines in the picture, correctly identifying the number even when it was very noisy.
Speech: They tested it on spoken words (like "stop" or "go") mixed with loud background noise. The old way would clean the audio but sometimes remove the specific sound frequencies needed to tell "stop" from "top." The new method used the meaning of the word to guide the cleaning, resulting in much clearer speech recognition.

The Big Takeaway
The most important part of this paper is that they did not have to retrain the "Guessing Expert" (the classifier). They kept the expert exactly as they found them. They just added a new "Image Restorer" that learned to talk to the expert. This means you can take any powerful, pre-trained AI system and make it much more robust against noise without having to rebuild the whole thing from scratch.

In short: Instead of cleaning a mess and then solving the puzzle, they clean the puzzle pieces while solving the puzzle, using the solution to guide the cleaning.

Technical Summary: Joint Enhancement and Classification using Coupled Diffusion Models of Signals and Logits

1. Problem Statement

Robust classification in noisy environments remains a fundamental challenge in machine learning. Standard approaches typically treat signal enhancement and classification as separate, sequential stages: first enhancing the signal (e.g., denoising an image or speech waveform) and then applying a pre-trained classifier. This decoupling introduces a significant misalignment between the objectives of the two stages. Enhancement models are usually optimized for perceptual quality or signal-level reconstruction metrics (e.g., MSE, SNR, PESQ), which do not necessarily correlate with classification accuracy. In fact, aggressive denoising can introduce processing artifacts that are perceptually subtle but highly disruptive to downstream classifiers.

Alternative remedies, such as fine-tuning classifiers on noisy data, often induce catastrophic forgetting or distort pre-trained robust representations. Previous attempts to model the joint distribution of inputs and labels using Energy-Based Models (EBMs) suffer from training instabilities and require retraining the classifier itself. Furthermore, recent methods like Classification and Regression Diffusion (CARD) attempt to denoise classifier logits but fail to generalize to noisy inputs because the logit denoiser relies solely on the corrupted signal, losing valuable information present in the underlying clean input.

2. Methodology

The authors propose a general, domain-agnostic framework that integrates two interacting diffusion models: one operating on the input signal ( $x$ ) and the other on the classifier's output logits ( $y$ ). Crucially, this framework requires no retraining or fine-tuning of the underlying classifier ( $f_\phi$ ), which remains frozen.

Core Concept: Mutual Guidance

The central innovation is mutual guidance, a process where the two diffusion models inform one another iteratively:

The signal diffusion model refines the input signal, providing accurate features to the logit model.
The evolving class logits guide the signal reconstruction toward discriminative regions of the data manifold, steering the enhancement process toward semantically meaningful solutions rather than purely perceptual ones.

This interaction is analogous to alternating minimization procedures, where latent signal and latent logit estimates are dynamically refined through mutual updates.

Three Interaction Strategies

The authors formalize three strategies to model the joint distribution of the input and the logit, offering trade-offs between computational efficiency and the granularity of mutual guidance:

Parallel Strategy: Uses an interleaved scheduler where signal and logit updates alternate at every timestep ( $t$ ). Each update is conditioned on the most recent estimate of the other process. This allows for continuous, step-level bidirectional feedback and is the most computationally efficient of the coupled methods.
Alternating Strategy: Uses a block scheduler where a full diffusion trajectory is executed for the signal, followed by a full trajectory for the logits. Each process is guided by the fully denoised point estimate from the previous iteration. This aligns with standard diffusion pipelines by relying on clean estimates rather than intermediate states, offering stability at the cost of slower convergence.
Nested Strategy: Employs a hierarchical scheduler where a single outer loop over the logits triggers a complete inner diffusion loop over the signal at each logit timestep. This ensures that logit updates are conditioned on high-fidelity signal estimates, avoiding reliance on potentially inaccurate intermediate predictions, though it incurs higher computational complexity.

Training and Inference

The framework employs two coupled diffusion processes (modeled as Denoising Diffusion Probabilistic Models or SDE-based formulations). During training, the models minimize a joint noise-prediction objective. Crucially, cross-conditioning is simulated on-the-fly: intermediate estimates of the clean signal and logits are generated to condition the respective denoisers, aligning the training distribution with the inference dynamics.

3. Key Contributions

Unified Framework: A novel approach that bridges signal enhancement and generative classification by treating both as coupled diffusion processes without modifying the classifier weights.
Mutual Guidance Mechanism: A method where signal reconstruction and label prediction iteratively refine each other, overcoming the limitations of sequential "enhance-then-classify" pipelines and static conditional generation.
Strategy Flexibility: The introduction of Parallel, Alternating, and Nested strategies to balance computational cost and coupling strength.
Domain Agnosticism: The framework is applicable to high-dimensional data across modalities, demonstrated here on both image classification and Automatic Speech Recognition (ASR).

4. Experimental Results

The authors evaluated the framework on image classification (MNIST, CIFAR-10, CIFAR-100, ImageNet32-100) and ASR (Google Speech Commands, EARS-Reverb, EARS-WHAM) using a frozen Whisper-base model.

Image Classification: The coupled strategies consistently outperformed baselines (Noisy, Enhanced, URIE, and CARD) across various noise types (Gaussian noise, Gaussian blur).
- On CIFAR-10 and CIFAR-100, the Alternating and Nested strategies achieved the strongest robustness.
- The Parallel strategy scaled more favorably to complex data like ImageNet32-100.
- Qualitative analysis showed that coupled strategies could reconstruct missing semantic features (e.g., a digit stroke) that standard enhancement baselines failed to recover, leading to correct classification.
Speech Recognition: Coupling signal enhancement with evolving semantic logits consistently improved Word Error Rate (WER) compared to baselines.
- The Nested strategy yielded the lowest WER across all datasets, with gains increasing as acoustic complexity or task difficulty increased.
- The method demonstrated the potential to improve strong, pre-trained systems (like Whisper) on noisy speech benchmarks without retraining the ASR model.
Efficiency: While the strategies incur higher Neural Function Evaluations (NFE) than single-modality diffusion (ranging from 2x to 10x), the actual runtime increase is mitigated because logit diffusion operates in a lower-dimensional space and is computationally cheaper than signal diffusion.

5. Significance and Claims

The paper claims that enabling interaction between signal and label diffusion processes substantially improves robustness over both enhancement-based baselines and static conditional generation methods. The authors position their work as a shift in paradigm: generative models are not merely used for data synthesis or enhancement but serve as adaptive inference engines that reshape inputs in service of downstream decision-making.

The significance lies in the ability to achieve robust classification under diverse noise conditions without the computational and data costs associated with retraining large foundation models or the performance degradation caused by the misalignment of sequential enhancement and classification objectives. The proposed framework offers a flexible, domain-agnostic solution that leverages the semantic information in classifier outputs to guide signal reconstruction, effectively closing the gap between perceptual quality and task performance.